merge(main): sync latest main into security branch

This commit is contained in:
elky
2026-09-05 00:30:16 +08:00
40 changed files with 2513 additions and 119 deletions
@@ -14,6 +14,9 @@ use aether_data::repository::provider_catalog::InMemoryProviderCatalogReadReposi
use aether_data_contracts::repository::candidate_selection::{
StoredMinimalCandidateSelectionRow, StoredProviderModelMapping,
};
use aether_data_contracts::repository::candidates::{
RequestCandidateReadRepository, RequestCandidateStatus,
};
use aether_data_contracts::repository::provider_catalog::{
StoredProviderCatalogEndpoint, StoredProviderCatalogKey, StoredProviderCatalogProvider,
};
@@ -427,6 +430,112 @@ async fn gateway_stops_execution_runtime_stream_when_client_disconnects_impl() {
upstream_handle.abort();
}
#[test]
fn gateway_settles_stream_attempt_when_client_disconnects_before_first_byte() {
run_lifecycle_test(
"gateway_settles_stream_attempt_when_client_disconnects_before_first_byte",
gateway_settles_stream_attempt_when_client_disconnects_before_first_byte_impl,
);
}
async fn gateway_settles_stream_attempt_when_client_disconnects_before_first_byte_impl() {
// The execution runtime accepts the plan and then goes quiet, so the attempt
// is parked between its `pending` rows and the first upstream byte.
let execution_runtime = Router::new().route(
"/v1/execute/stream",
any(|_request: Request| async move {
tokio::time::sleep(std::time::Duration::from_secs(30)).await;
StatusCode::OK
}),
);
let (execution_runtime_url, execution_runtime_handle) = start_server(execution_runtime).await;
let auth_repository = Arc::new(InMemoryAuthApiKeySnapshotRepository::seed(vec![(
Some(hash_api_key("sk-client-openai-stream-precommit-disconnect")),
sample_local_openai_auth_snapshot(
"api-key-openai-lifecycle-local-1",
"user-openai-lifecycle-local-1",
),
)]));
let candidate_selection_repository =
Arc::new(InMemoryMinimalCandidateSelectionReadRepository::seed(vec![
sample_local_openai_candidate_row(),
]));
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![sample_local_openai_provider()],
vec![sample_local_openai_endpoint()],
vec![sample_local_openai_key()],
));
let request_candidate_repository = Arc::new(InMemoryRequestCandidateRepository::default());
let gateway = build_router_with_state(
build_state_with_execution_runtime_override(execution_runtime_url)
.with_data_state_for_tests(
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_url, gateway_handle) = start_server(gateway).await;
let request = reqwest::Client::new()
.post(format!("{gateway_url}/v1/chat/completions"))
.header(http::header::CONTENT_TYPE, "application/json")
.header(
http::header::AUTHORIZATION,
"Bearer sk-client-openai-stream-precommit-disconnect",
)
.header(
TRACE_ID_HEADER,
"trace-openai-chat-stream-precommit-disconnect-123",
)
.body("{\"model\":\"gpt-5\",\"messages\":[],\"stream\":true}")
.send();
// Drop the in-flight request the way a downstream client does when its own
// first-byte timeout fires, before any response header exists.
assert!(
tokio::time::timeout(std::time::Duration::from_millis(750), request)
.await
.is_err(),
"the execution runtime should not have answered before the client gave up"
);
let mut stored_candidates = Vec::new();
for _ in 0..200 {
stored_candidates = request_candidate_repository
.list_by_request_id("trace-openai-chat-stream-precommit-disconnect-123")
.await
.expect("request candidate trace should read");
if stored_candidates
.iter()
.any(|candidate| candidate.status == RequestCandidateStatus::Cancelled)
{
break;
}
tokio::time::sleep(std::time::Duration::from_millis(10)).await;
}
let cancelled = stored_candidates
.iter()
.find(|candidate| candidate.status == RequestCandidateStatus::Cancelled)
.unwrap_or_else(|| {
panic!("dropped stream attempt should settle as cancelled: {stored_candidates:?}")
});
assert_eq!(cancelled.status_code, Some(499));
assert_eq!(
cancelled.error_type.as_deref(),
Some("local_stream_attempt_cancelled")
);
assert!(cancelled.finished_at_unix_ms.is_some());
gateway_handle.abort();
execution_runtime_handle.abort();
}
#[test]
fn gateway_returns_error_body_when_prefetch_detects_embedded_stream_error() {
run_lifecycle_test(
@@ -1987,8 +1987,9 @@ fn admin_provider_oauth_quota_mod_stays_thin() {
"handlers/admin/provider/oauth/quota/antigravity.rs should import common quota helpers from shared.rs"
);
assert!(
quota_antigravity
.contains("use aether_provider_pool::build_antigravity_pool_quota_request;"),
quota_antigravity.contains("use aether_provider_pool::{")
&& quota_antigravity.contains("build_antigravity_pool_quota_request")
&& quota_antigravity.contains("build_antigravity_pool_quota_summary_request"),
"handlers/admin/provider/oauth/quota/antigravity.rs should delegate antigravity quota request construction to aether-provider-pool"
);
let quota_chatgpt_web = read_workspace_file(
@@ -1,9 +1,15 @@
use std::collections::BTreeMap;
use std::sync::{Arc, Mutex};
use aether_crypto::{encrypt_python_fernet_plaintext, DEVELOPMENT_ENCRYPTION_KEY};
use aether_crypto::{
decrypt_python_fernet_ciphertext, encrypt_python_fernet_plaintext, DEVELOPMENT_ENCRYPTION_KEY,
};
use aether_data::repository::global_models::InMemoryGlobalModelReadRepository;
use aether_data::repository::provider_catalog::InMemoryProviderCatalogReadRepository;
use aether_data::repository::proxy_nodes::InMemoryProxyNodeRepository;
use aether_data_contracts::repository::global_models::{
AdminProviderModelListQuery, GlobalModelReadRepository,
};
use aether_data_contracts::repository::provider_catalog::{
ProviderCatalogReadRepository, StoredProviderCatalogKey, StoredProviderCatalogProvider,
};
@@ -2308,6 +2314,12 @@ async fn gateway_refreshes_admin_provider_quota_locally_for_antigravity_with_tru
},
"gemini-2.5-pro": {
"displayName": "Gemini 2.5 Pro"
},
"gemini-3.7-flash-tiered": {
"displayName": "Gemini 3.7 Flash"
},
"chat_23310": {
"displayName": "Internal Chat"
}
}
}),
@@ -2409,6 +2421,7 @@ async fn gateway_refreshes_admin_provider_quota_locally_for_antigravity_with_tru
)],
vec![key],
));
let global_model_repository = Arc::new(InMemoryGlobalModelReadRepository::default());
let (upstream_url, upstream_handle) = start_server(upstream).await;
let (execution_runtime_url, execution_runtime_handle) = start_server(execution_runtime).await;
@@ -2418,6 +2431,7 @@ async fn gateway_refreshes_admin_provider_quota_locally_for_antigravity_with_tru
GatewayDataState::with_provider_catalog_repository_for_tests(
provider_catalog_repository.clone(),
)
.with_global_model_repository_for_tests(global_model_repository.clone())
.with_encryption_key_for_tests(DEVELOPMENT_ENCRYPTION_KEY),
),
);
@@ -2519,6 +2533,23 @@ async fn gateway_refreshes_admin_provider_quota_locally_for_antigravity_with_tru
.and_then(|value| value.get("remaining_fraction")),
Some(&json!(0.25))
);
let imported_provider_models = global_model_repository
.list_admin_provider_models(&AdminProviderModelListQuery {
provider_id: "provider-antigravity".to_string(),
is_active: None,
offset: 0,
limit: 100,
})
.await
.expect("imported Antigravity provider models should read");
let imported_model_names = imported_provider_models
.iter()
.map(|model| model.provider_model_name.as_str())
.collect::<std::collections::BTreeSet<_>>();
assert!(imported_model_names.contains("claude-sonnet-4"));
assert!(imported_model_names.contains("gemini-2.5-pro"));
assert!(imported_model_names.contains("gemini-3.7-flash-tiered"));
assert!(!imported_model_names.contains("chat_23310"));
assert_eq!(
reloaded[0]
.upstream_metadata
@@ -2,11 +2,9 @@ use std::sync::{Arc, Mutex};
use aether_data::repository::global_models::InMemoryGlobalModelReadRepository;
use aether_data::repository::provider_catalog::InMemoryProviderCatalogReadRepository;
use aether_data::repository::routing_profiles::InMemoryRoutingGroupRepository;
use aether_data_contracts::repository::global_models::{
AdminProviderModelListQuery, GlobalModelReadRepository,
};
use aether_data_contracts::repository::routing_profiles::StoredRoutingGroup;
use axum::body::Body;
use axum::routing::any;
use axum::{extract::Request, Router};
@@ -843,29 +841,6 @@ async fn gateway_handles_admin_global_model_routing_locally_with_trusted_admin_p
]),
);
let routing_group_repository = Arc::new(InMemoryRoutingGroupRepository::seed(
[StoredRoutingGroup {
id: "system-default".to_string(),
name: "system-default".to_string(),
description: Some("test system default routing strategy".to_string()),
enabled: true,
is_system_default: true,
sort_order: 0,
config_json: json!({
"default_policy": {
"priority_mode": "global_key",
"scheduling_mode": "fixed_order"
}
}),
version: 1,
created_at: 1,
updated_at: 1,
published_at: Some(1),
}],
std::iter::empty(),
std::iter::empty(),
));
let (upstream_url, upstream_handle) = start_server(upstream).await;
let gateway = build_router_with_state(
AppState::new()
@@ -875,7 +850,7 @@ async fn gateway_handles_admin_global_model_routing_locally_with_trusted_admin_p
provider_catalog_repository,
)
.with_global_model_repository_for_tests(global_model_repository)
.with_routing_group_repository_for_tests(routing_group_repository),
.with_system_default_routing_group_for_tests(),
),
);
let (gateway_url, gateway_handle) = start_server(gateway).await;
@@ -898,8 +873,8 @@ async fn gateway_handles_admin_global_model_routing_locally_with_trusted_admin_p
assert_eq!(payload["global_model_name"], "gpt-5");
assert_eq!(payload["display_name"], "GPT 5");
assert_eq!(payload["global_model_mappings"], json!(["gpt-5-upstream"]));
assert_eq!(payload["scheduling_mode"], "fixed_order");
assert_eq!(payload["priority_mode"], "global_key");
assert_eq!(payload["scheduling_mode"], "cache_affinity");
assert_eq!(payload["priority_mode"], "provider");
assert_eq!(payload["total_providers"], 2);
assert_eq!(payload["active_providers"], 2);
@@ -3912,6 +3912,175 @@ async fn gateway_completes_admin_provider_oauth_provider_locally_with_trusted_ad
upstream_handle.abort();
}
#[test]
fn gateway_names_new_antigravity_oauth_account_from_google_userinfo_email() {
run_admin_oauth_test(
"gateway_names_new_antigravity_oauth_account_from_google_userinfo_email",
gateway_names_new_antigravity_oauth_account_from_google_userinfo_email_impl,
);
}
async fn gateway_names_new_antigravity_oauth_account_from_google_userinfo_email_impl() {
let upstream_hits = Arc::new(Mutex::new(0usize));
let upstream_hits_clone = Arc::clone(&upstream_hits);
let upstream = Router::new().fallback(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 token_hits = Arc::new(Mutex::new(0usize));
let token_hits_clone = Arc::clone(&token_hits);
let user_info_hits = Arc::new(Mutex::new(0usize));
let user_info_hits_clone = Arc::clone(&user_info_hits);
let seen_user_info_authorization = Arc::new(Mutex::new(None::<String>));
let seen_user_info_authorization_clone = Arc::clone(&seen_user_info_authorization);
let google_server = Router::new()
.route(
"/oauth/token",
post(move || {
let token_hits_inner = Arc::clone(&token_hits_clone);
async move {
*token_hits_inner.lock().expect("mutex should lock") += 1;
Json(json!({
"access_token": "antigravity-access-token",
"refresh_token": "antigravity-refresh-token",
"token_type": "Bearer",
"expires_in": 3600,
"scope": "https://www.googleapis.com/auth/userinfo.email"
}))
}
}),
)
.route(
"/oauth/userinfo",
get(move |headers: HeaderMap| {
let user_info_hits_inner = Arc::clone(&user_info_hits_clone);
let seen_authorization_inner = Arc::clone(&seen_user_info_authorization_clone);
async move {
*user_info_hits_inner.lock().expect("mutex should lock") += 1;
*seen_authorization_inner.lock().expect("mutex should lock") = headers
.get(http::header::AUTHORIZATION)
.and_then(|value| value.to_str().ok())
.map(ToOwned::to_owned);
Json(json!({
"email": "[email protected]",
"verified_email": true,
"name": "Antigravity User"
}))
}
}),
);
let mut provider = sample_provider("provider-antigravity", "antigravity", 10);
provider.provider_type = "antigravity".to_string();
let endpoint = sample_endpoint(
"endpoint-antigravity",
"provider-antigravity",
"gemini:generate_content",
"https://daily-cloudcode-pa.googleapis.com",
);
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![provider],
vec![endpoint],
vec![],
));
let (upstream_url, upstream_handle) = start_server(upstream).await;
let (google_url, google_handle) = start_server(google_server).await;
let gateway = build_router_with_state(
AppState::new()
.expect("gateway should build")
.with_data_state_for_tests(
GatewayDataState::with_provider_catalog_repository_for_tests(
provider_catalog_repository.clone(),
)
.with_encryption_key_for_tests(DEVELOPMENT_ENCRYPTION_KEY),
)
.with_provider_oauth_state_entry_for_tests(
"nonce-antigravity-123",
json!({
"nonce": "nonce-antigravity-123",
"key_id": "",
"provider_id": "provider-antigravity",
"provider_type": "antigravity",
"pkce_verifier": "verifier-antigravity-123",
}),
)
.with_provider_oauth_token_url_for_tests(
"antigravity",
format!("{google_url}/oauth/token"),
)
.with_provider_oauth_token_url_for_tests(
"antigravity_user_info",
format!("{google_url}/oauth/userinfo"),
),
);
let (gateway_url, gateway_handle) = start_server(gateway).await;
let response = reqwest::Client::new()
.post(format!(
"{gateway_url}/api/admin/provider-oauth/providers/provider-antigravity/complete"
))
.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")
.json(&json!({
"callback_url": "http://localhost:51121/oauth2callback?code=antigravity-code-123&state=nonce-antigravity-123"
}))
.send()
.await
.expect("request should succeed");
let status = response.status();
let payload: Value = response.json().await.expect("json body should parse");
assert_eq!(status, StatusCode::OK, "payload={payload}");
assert_eq!(payload["provider_type"], "antigravity");
assert_eq!(payload["email"], "[email protected]");
assert_eq!(payload["replaced"], false);
assert_eq!(*token_hits.lock().expect("mutex should lock"), 1);
assert_eq!(*user_info_hits.lock().expect("mutex should lock"), 1);
assert_eq!(
seen_user_info_authorization
.lock()
.expect("mutex should lock")
.as_deref(),
Some("Bearer antigravity-access-token")
);
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
let key_id = payload["key_id"]
.as_str()
.expect("created key id should be returned")
.to_string();
let persisted_keys = provider_catalog_repository
.list_keys_by_ids(std::slice::from_ref(&key_id))
.await
.expect("created key should load");
let persisted = persisted_keys.first().expect("created key should exist");
assert_eq!(persisted.name, "[email protected]");
let decrypted_auth_config = decrypt_python_fernet_ciphertext(
DEVELOPMENT_ENCRYPTION_KEY,
persisted
.encrypted_auth_config
.as_deref()
.expect("auth config should be stored"),
)
.expect("auth config should decrypt");
let auth_config: Value =
serde_json::from_str(&decrypted_auth_config).expect("auth config json should parse");
assert_eq!(auth_config["email"], "[email protected]");
assert_eq!(auth_config["refresh_token"], "antigravity-refresh-token");
gateway_handle.abort();
google_handle.abort();
upstream_handle.abort();
drop(upstream_url);
}
#[test]
fn gateway_imports_admin_provider_oauth_refresh_token_locally_with_trusted_admin_principal() {
run_admin_oauth_test(
@@ -2184,7 +2184,7 @@ async fn gateway_handles_admin_system_config_locally_with_bearer_admin_session()
}),
);
let (upstream_url, upstream_handle) = start_server(upstream).await;
let (_upstream_url, upstream_handle) = start_server(upstream).await;
let state = AppState::new().expect("gateway should build");
let access_token = issue_test_admin_access_token(&state, "device-admin-config").await;
let gateway = build_router_with_state(state);
@@ -2208,6 +2208,44 @@ async fn gateway_handles_admin_system_config_locally_with_bearer_admin_session()
upstream_handle.abort();
}
#[tokio::test]
async fn gateway_rejects_removed_admin_system_provider_priority_mode_with_bearer_admin_session() {
let upstream_hits = Arc::new(Mutex::new(0usize));
let upstream_hits_clone = Arc::clone(&upstream_hits);
let upstream = Router::new().route(
"/api/admin/system/configs/provider_priority_mode",
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 (_upstream_url, upstream_handle) = start_server(upstream).await;
let state = AppState::new().expect("gateway should build");
let access_token = issue_test_admin_access_token(&state, "device-admin-config").await;
let gateway = build_router_with_state(state);
let (gateway_url, gateway_handle) = start_server(gateway).await;
let response = reqwest::Client::new()
.get(format!(
"{gateway_url}/api/admin/system/configs/provider_priority_mode"
))
.header("authorization", format!("Bearer {access_token}"))
.header("x-client-device-id", "device-admin-config")
.send()
.await
.expect("request should succeed");
assert_eq!(response.status(), StatusCode::NOT_FOUND);
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
gateway_handle.abort();
upstream_handle.abort();
}
#[tokio::test]
async fn gateway_sets_admin_system_config_locally_with_trusted_admin_principal() {
let upstream_hits = Arc::new(Mutex::new(0usize));
@@ -5,6 +5,7 @@ use super::{
InMemoryVideoTaskRepository, StoredAuthApiKeySnapshot, UpsertVideoTask, VideoTaskLookupKey,
VideoTaskReadRepository, VideoTaskStatus, VideoTaskWriteRepository, DEVELOPMENT_ENCRYPTION_KEY,
};
use crate::data::GatewayDataState;
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,
@@ -6,6 +6,7 @@ use super::{
InMemoryVideoTaskRepository, UpsertVideoTask, VideoTaskStatus, VideoTaskWriteRepository,
DEVELOPMENT_ENCRYPTION_KEY,
};
use crate::data::GatewayDataState;
use crate::tests::{
any, build_router, build_router_with_state, build_state_with_execution_runtime_override, json,
start_server, strip_sse_keepalive_comments, AppState, Arc, Body, HeaderValue, Json, Mutex,
@@ -1850,7 +1851,7 @@ async fn gateway_returns_internal_gateway_decision_sync_fallback_with_resolved_a
AppState::new()
.expect("gateway should build")
.with_data_state_for_tests(
crate::data::GatewayDataState::with_auth_api_key_reader_for_tests(auth_repository)
GatewayDataState::with_auth_api_key_reader_for_tests(auth_repository)
.with_system_default_routing_group_for_tests(),
),
);
+16
View File
@@ -27,6 +27,10 @@ use aether_data_contracts::repository::provider_catalog::{
StoredProviderCatalogEndpoint, StoredProviderCatalogKey, StoredProviderCatalogProvider,
};
use aether_runtime_state::{RedisClientConfig, RuntimeState};
use aether_scheduler_core::{
build_scheduler_affinity_cache_key_for_api_key_id_with_client_session_and_scope,
SchedulerAffinityScope,
};
use aether_test_support::ManagedRedisServer;
use sha2::{Digest, Sha256};
@@ -41,6 +45,18 @@ fn hash_api_key(value: &str) -> String {
format!("{:x}", hasher.finalize())
}
fn system_default_affinity_cache_key(api_key_id: &str, api_format: &str, model: &str) -> String {
let scope = SchedulerAffinityScope::new("system-default", Some(1));
build_scheduler_affinity_cache_key_for_api_key_id_with_client_session_and_scope(
api_key_id,
api_format,
model,
None,
Some(&scope),
)
.expect("system-default affinity cache key should build")
}
fn sample_auth_snapshot(
api_key_id: &str,
user_id: &str,