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

fix(codex): fence concurrent quota updates

Browse Source
This commit is contained in:
elky
2026-08-14 09:28:07 +08:00
parent f3a12c1008
commit 5b0c763086
65 changed files with 13009 additions and 738 deletions
Download Patch File Download Diff File Expand all files Collapse all files
apps/aether-gateway/src/tests/control/admin/endpoints/keys.rs
+99
View File
@@ -774,6 +774,105 @@ async fn generic_key_routes_reject_agent_identity_credential_writes() {
gateway_handle.abort();
}
#[tokio::test]
async fn generic_codex_key_credential_switch_rotates_generation_and_clears_quota() {
let mut existing_key = sample_key(
"key-codex-existing",
"provider-codex",
"openai:responses",
"old-oauth-access-token",
);
existing_key.auth_type = "oauth".to_string();
existing_key.encrypted_auth_config = Some(
encrypt_python_fernet_plaintext(
DEVELOPMENT_ENCRYPTION_KEY,
r#"{"provider_type":"codex","refresh_token":"old-refresh-token"}"#,
)
.expect("old auth config should encrypt"),
);
existing_key.upstream_metadata = Some(json!({
"codex": {
"credential_generation": "generation-before-switch",
"primary_used_percent": 80.0,
},
"unrelated": {"preserved": true},
}));
existing_key.status_snapshot = Some(json!({
"oauth": {"status": "valid"},
"quota": {"used_ratio": 0.8},
}));
let mut provider = sample_provider("provider-codex", "codex", 10);
provider.provider_type = "codex".to_string();
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![provider],
vec![],
vec![existing_key],
));
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),
),
);
let (gateway_url, gateway_handle) = start_server(gateway).await;
let response = reqwest::Client::new()
.put(format!(
"{gateway_url}/api/admin/endpoints/keys/key-codex-existing"
))
.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!({
"auth_type": "api_key",
"api_key": "new-codex-api-key"
}))
.send()
.await
.expect("credential switch should complete");
assert_eq!(response.status(), StatusCode::OK);
let reloaded = provider_catalog_repository
.list_keys_by_ids(&["key-codex-existing".to_string()])
.await
.expect("key should reload");
assert_eq!(reloaded.len(), 1);
let key = &reloaded[0];
assert_eq!(key.auth_type, "api_key");
let codex = key
.upstream_metadata
.as_ref()
.and_then(|metadata| metadata.get("codex"))
.and_then(serde_json::Value::as_object)
.expect("codex metadata should exist");
assert_eq!(codex.len(), 1, "unexpected Codex metadata: {codex:?}");
assert_ne!(
codex
.get(aether_admin::provider::quota::CODEX_CREDENTIAL_GENERATION_KEY)
.and_then(serde_json::Value::as_str),
Some("generation-before-switch")
);
assert_eq!(
key.upstream_metadata
.as_ref()
.and_then(|metadata| metadata.pointer("/unrelated/preserved")),
Some(&json!(true))
);
assert_eq!(
key.status_snapshot
.as_ref()
.and_then(|snapshot| snapshot.get("quota")),
Some(&serde_json::Value::Null)
);
gateway_handle.abort();
}
#[tokio::test]
async fn provider_key_concurrent_limit_create_and_list_responses() {
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
apps/aether-gateway/src/tests/control/admin/endpoints/quota.rs
+473 -1
View File
@@ -1,7 +1,9 @@
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::provider_catalog::InMemoryProviderCatalogReadRepository;
use aether_data::repository::proxy_nodes::InMemoryProxyNodeRepository;
use aether_data_contracts::repository::provider_catalog::{
@@ -121,6 +123,7 @@ async fn gateway_refreshes_admin_provider_quota_locally_for_codex_with_trusted_a
"1900500000".to_string(),
),
]),
response_observation: None,
body: Some(aether_contracts::ResponseBody {
json_body: Some(json!({
"plan_type": "plus",
@@ -295,6 +298,467 @@ async fn gateway_refreshes_admin_provider_quota_locally_for_codex_with_trusted_a
upstream_handle.abort();
}
#[test]
fn gateway_codex_quota_refresh_persists_after_automatic_oauth_token_refresh() {
run_provider_quota_test(
"gateway_codex_quota_refresh_persists_after_automatic_oauth_token_refresh",
gateway_codex_quota_refresh_persists_after_automatic_oauth_token_refresh_impl,
);
}
async fn gateway_codex_quota_refresh_persists_after_automatic_oauth_token_refresh_impl() {
let token_hits = Arc::new(Mutex::new(0usize));
let token_server = Router::new().route(
"/oauth/token",
post({
let token_hits = Arc::clone(&token_hits);
move || {
let token_hits = Arc::clone(&token_hits);
async move {
*token_hits.lock().expect("mutex should lock") += 1;
Json(json!({
"access_token": "refreshed-codex-access-token",
"refresh_token": "rotated-codex-refresh-token",
"token_type": "Bearer",
"expires_in": 3_600,
"account_id": "acct-quota-refresh",
"plan_type": "plus"
}))
}
}
}),
);
let seen_requests = Arc::new(Mutex::new(Vec::<(String, String)>::new()));
let execution_runtime = Router::new().route(
"/v1/execute/sync",
any({
let seen_requests = Arc::clone(&seen_requests);
move |request: Request| {
let seen_requests = Arc::clone(&seen_requests);
async move {
let plan: aether_contracts::ExecutionPlan = serde_json::from_slice(
&to_bytes(request.into_body(), usize::MAX)
.await
.expect("body should read"),
)
.expect("plan should parse");
seen_requests.lock().expect("mutex should lock").push((
plan.url.clone(),
plan.headers
.get("authorization")
.cloned()
.unwrap_or_default(),
));
let body_json = match plan.url.as_str() {
"https://chatgpt.com/backend-api/wham/usage" => json!({
"plan_type": "plus",
"rate_limit": {
"primary_window": {
"used_percent": 23.0,
"reset_at": 1_900_000_000u64,
"window_minutes": 300
}
}
}),
"https://chatgpt.com/backend-api/wham/rate-limit-reset-credits" => {
json!({"available_count": 0, "credits": []})
}
url => panic!("unexpected execution runtime URL: {url}"),
};
let result = aether_contracts::ExecutionResult {
request_id: plan.request_id,
candidate_id: None,
status_code: 200,
headers: BTreeMap::new(),
response_observation: None,
body: Some(aether_contracts::ResponseBody {
json_body: Some(body_json),
body_bytes_b64: None,
}),
telemetry: None,
error: None,
};
(StatusCode::OK, Json(result))
}
}
}),
);
let mut key = sample_key(
"key-codex-expired-quota",
"provider-codex-expired-quota",
"openai:responses",
"expired-codex-access-token",
);
key.auth_type = "oauth".to_string();
key.expires_at_unix_secs = Some(1);
key.encrypted_auth_config = Some(
encrypt_python_fernet_plaintext(
DEVELOPMENT_ENCRYPTION_KEY,
&json!({
"provider_type": "codex",
"refresh_token": "expired-codex-refresh-token",
"expires_at": 1,
"account_id": "acct-quota-refresh",
"plan_type": "plus"
})
.to_string(),
)
.expect("auth config should encrypt"),
);
key.upstream_metadata = Some(json!({
"codex": {"credential_generation": "credential-quota-refresh"}
}));
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![StoredProviderCatalogProvider::new(
"provider-codex-expired-quota".to_string(),
"codex".to_string(),
Some("https://example.com".to_string()),
"codex".to_string(),
)
.expect("provider should build")],
vec![sample_endpoint(
"endpoint-codex-expired-quota",
"provider-codex-expired-quota",
"openai:responses",
"https://chatgpt.com/backend-api",
)],
vec![key],
));
let (token_url, token_handle) = start_server(token_server).await;
let (execution_runtime_url, execution_runtime_handle) = start_server(execution_runtime).await;
let oauth_refresh =
crate::provider_transport::LocalOAuthRefreshCoordinator::with_adapters_for_tests(vec![
Arc::new(
crate::provider_transport::oauth_refresh::GenericOAuthRefreshAdapter::default()
.with_token_url_for_tests("codex", format!("{token_url}/oauth/token")),
),
]);
let gateway = build_router_with_state(
build_state_with_execution_runtime_override(execution_runtime_url)
.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_oauth_refresh_coordinator_for_tests(oauth_refresh),
);
let (gateway_url, gateway_handle) = start_server(gateway).await;
let response = reqwest::Client::new()
.post(format!(
"{gateway_url}/api/admin/endpoints/providers/provider-codex-expired-quota/refresh-quota"
))
.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["success"], 1, "payload={payload}");
assert_eq!(payload["failed"], 0, "payload={payload}");
assert_eq!(payload["results"][0]["status"], "success");
assert_eq!(*token_hits.lock().expect("mutex should lock"), 1);
assert_eq!(
seen_requests.lock().expect("mutex should lock").as_slice(),
[
(
"https://chatgpt.com/backend-api/wham/usage".to_string(),
"Bearer refreshed-codex-access-token".to_string(),
),
(
"https://chatgpt.com/backend-api/wham/rate-limit-reset-credits".to_string(),
"Bearer refreshed-codex-access-token".to_string(),
),
]
);
let reloaded = provider_catalog_repository
.list_keys_by_ids(&["key-codex-expired-quota".to_string()])
.await
.expect("key should reload");
let persisted = reloaded.first().expect("key should remain installed");
let decrypted_api_key = decrypt_python_fernet_ciphertext(
DEVELOPMENT_ENCRYPTION_KEY,
persisted
.encrypted_api_key
.as_deref()
.expect("api key should persist"),
)
.expect("api key should decrypt");
assert_eq!(decrypted_api_key, "refreshed-codex-access-token");
let decrypted_auth_config = decrypt_python_fernet_ciphertext(
DEVELOPMENT_ENCRYPTION_KEY,
persisted
.encrypted_auth_config
.as_deref()
.expect("auth config should persist"),
)
.expect("auth config should decrypt");
let auth_config: serde_json::Value =
serde_json::from_str(&decrypted_auth_config).expect("auth config should parse");
assert_eq!(auth_config["refresh_token"], "rotated-codex-refresh-token");
assert_eq!(
persisted
.upstream_metadata
.as_ref()
.and_then(|metadata| metadata.pointer("/codex/credential_generation")),
Some(&json!("credential-quota-refresh"))
);
assert_eq!(
persisted
.upstream_metadata
.as_ref()
.and_then(|metadata| metadata.pointer("/codex/primary_used_percent")),
Some(&json!(23.0))
);
gateway_handle.abort();
execution_runtime_handle.abort();
token_handle.abort();
}
#[test]
fn gateway_codex_reset_credit_retries_until_same_window_usage_drop_is_authoritative() {
run_provider_quota_test(
"gateway_codex_reset_credit_retries_until_same_window_usage_drop_is_authoritative",
gateway_codex_reset_credit_retries_until_same_window_usage_drop_is_authoritative_impl,
);
}
async fn gateway_codex_reset_credit_retries_until_same_window_usage_drop_is_authoritative_impl() {
const RESET_FENCE_UNIX_MS: u64 = 1_800_000_000_000;
const RESET_AT_UNIX_SECS: u64 = 2_000_000_000;
let usage_hits = Arc::new(Mutex::new(0usize));
let detail_hits = Arc::new(Mutex::new(0usize));
let seen_urls = Arc::new(Mutex::new(Vec::<String>::new()));
let execution_runtime = Router::new().route(
"/v1/execute/sync",
any({
let usage_hits = Arc::clone(&usage_hits);
let detail_hits = Arc::clone(&detail_hits);
let seen_urls = Arc::clone(&seen_urls);
move |request: Request| {
let usage_hits = Arc::clone(&usage_hits);
let detail_hits = Arc::clone(&detail_hits);
let seen_urls = Arc::clone(&seen_urls);
async move {
let plan: aether_contracts::ExecutionPlan = serde_json::from_slice(
&to_bytes(request.into_body(), usize::MAX)
.await
.expect("body should read"),
)
.expect("plan should parse");
seen_urls
.lock()
.expect("mutex should lock")
.push(plan.url.clone());
assert_eq!(
plan.headers.get("authorization").map(String::as_str),
Some("Bearer codex-reset-access-token")
);
let (body_json, response_observation) = match plan.url.as_str() {
"https://chatgpt.com/backend-api/wham/rate-limit-reset-credits/consume" => {
assert_eq!(plan.method, "POST");
assert_eq!(
plan.body.json_body,
Some(json!({"redeem_request_id": "reset-e2e"}))
);
(
json!({"outcome": "reset"}),
Some(aether_contracts::ExecutionResponseObservation {
request_started_at_unix_ms: RESET_FENCE_UNIX_MS - 100,
response_headers_observed_at_unix_ms: RESET_FENCE_UNIX_MS,
request_order_id: "consume-reset-e2e".to_string(),
}),
)
}
"https://chatgpt.com/backend-api/wham/usage" => {
let hit = {
let mut hits = usage_hits.lock().expect("mutex should lock");
*hits += 1;
*hits
};
let used_percent = match hit {
1 => 100.0,
2 => 0.0,
_ => panic!("unexpected wham/usage request #{hit}"),
};
(
json!({
"plan_type": "plus",
"rate_limit": {
"primary_window": {
"used_percent": used_percent,
"reset_at": RESET_AT_UNIX_SECS,
"window_minutes": 300
}
}
}),
Some(aether_contracts::ExecutionResponseObservation {
request_started_at_unix_ms: RESET_FENCE_UNIX_MS
+ (hit as u64 * 1_000),
response_headers_observed_at_unix_ms: RESET_FENCE_UNIX_MS
+ (hit as u64 * 1_000)
+ 100,
request_order_id: format!("usage-reset-e2e-{hit}"),
}),
)
}
"https://chatgpt.com/backend-api/wham/rate-limit-reset-credits" => {
*detail_hits.lock().expect("mutex should lock") += 1;
(json!({"available_count": 0, "credits": []}), None)
}
url => panic!("unexpected execution runtime URL: {url}"),
};
let result = aether_contracts::ExecutionResult {
request_id: plan.request_id,
candidate_id: None,
status_code: 200,
headers: BTreeMap::new(),
response_observation,
body: Some(aether_contracts::ResponseBody {
json_body: Some(body_json),
body_bytes_b64: None,
}),
telemetry: None,
error: None,
};
(StatusCode::OK, Json(result))
}
}
}),
);
let mut key = sample_key(
"key-codex-reset",
"provider-codex-reset",
"openai:responses",
"codex-reset-access-token",
);
key.auth_type = "oauth".to_string();
key.expires_at_unix_secs = Some(4_102_444_800);
key.encrypted_auth_config = Some(
encrypt_python_fernet_plaintext(
DEVELOPMENT_ENCRYPTION_KEY,
&json!({
"provider_type": "codex",
"refresh_token": "codex-reset-refresh-token",
"expires_at": 4_102_444_800u64,
"account_id": "acct-reset-e2e",
"plan_type": "plus"
})
.to_string(),
)
.expect("auth config should encrypt"),
);
key.upstream_metadata = Some(json!({
"codex": {
"credential_generation": "credential-reset-e2e",
"plan_type": "plus",
"primary_used_percent": 100.0,
"primary_reset_at": RESET_AT_UNIX_SECS,
"primary_window_minutes": 300,
"updated_at": (RESET_FENCE_UNIX_MS / 1_000) - 1,
"account_quota_request_started_at_unix_ms": RESET_FENCE_UNIX_MS - 1_000,
"account_quota_request_id": "usage-before-reset"
}
}));
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![StoredProviderCatalogProvider::new(
"provider-codex-reset".to_string(),
"codex".to_string(),
Some("https://example.com".to_string()),
"codex".to_string(),
)
.expect("provider should build")],
vec![sample_endpoint(
"endpoint-codex-reset",
"provider-codex-reset",
"openai:responses",
"https://chatgpt.com/backend-api",
)],
vec![key],
));
let (execution_runtime_url, execution_runtime_handle) = start_server(execution_runtime).await;
let gateway = build_router_with_state(
build_state_with_execution_runtime_override(execution_runtime_url)
.with_data_state_for_tests(
GatewayDataState::with_provider_catalog_repository_for_tests(
provider_catalog_repository.clone(),
)
.with_encryption_key_for_tests(DEVELOPMENT_ENCRYPTION_KEY),
),
);
let (gateway_url, gateway_handle) = start_server(gateway).await;
let response = reqwest::Client::new()
.post(format!(
"{gateway_url}/api/admin/endpoints/keys/key-codex-reset/codex-reset-credit/consume"
))
.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!({
"idempotency_key": "reset-e2e",
"expected_credential_generation": "credential-reset-e2e"
}))
.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", "payload={payload}");
assert_eq!(payload["outcome"], "reset");
assert_eq!(payload["refresh_status"], "success");
assert_eq!(*usage_hits.lock().expect("mutex should lock"), 2);
assert_eq!(*detail_hits.lock().expect("mutex should lock"), 2);
assert_eq!(
seen_urls.lock().expect("mutex should lock").as_slice(),
[
"https://chatgpt.com/backend-api/wham/rate-limit-reset-credits/consume",
"https://chatgpt.com/backend-api/wham/usage",
"https://chatgpt.com/backend-api/wham/rate-limit-reset-credits",
"https://chatgpt.com/backend-api/wham/usage",
"https://chatgpt.com/backend-api/wham/rate-limit-reset-credits",
]
);
let reloaded = provider_catalog_repository
.list_keys_by_ids(&["key-codex-reset".to_string()])
.await
.expect("key should reload");
let codex = reloaded[0]
.upstream_metadata
.as_ref()
.and_then(|metadata| metadata.get("codex"))
.expect("codex metadata should persist");
assert_eq!(codex["primary_used_percent"], json!(0.0));
assert_eq!(codex["primary_reset_at"], json!(RESET_AT_UNIX_SECS));
assert_eq!(codex["account_quota_reset_pending"], json!(false));
assert_eq!(
codex["account_quota_reset_processed_ids"],
json!(["reset-e2e"])
);
gateway_handle.abort();
execution_runtime_handle.abort();
}
#[tokio::test]
async fn gateway_marks_codex_quota_exhausted_when_wham_usage_returns_payment_required() {
let upstream = Router::new().route(
@@ -318,6 +782,7 @@ async fn gateway_marks_codex_quota_exhausted_when_wham_usage_returns_payment_req
candidate_id: None,
status_code: 402,
headers: BTreeMap::new(),
response_observation: None,
body: Some(aether_contracts::ResponseBody {
json_body: Some(json!({
"error": {
@@ -439,6 +904,7 @@ async fn gateway_auto_removes_codex_key_when_quota_proves_oauth_invalid() {
candidate_id: None,
status_code: 401,
headers: BTreeMap::new(),
response_observation: None,
body: Some(aether_contracts::ResponseBody {
json_body: Some(json!({
"error": {
@@ -594,6 +1060,7 @@ async fn gateway_refreshes_admin_provider_quota_locally_for_requested_codex_keys
"1900000000".to_string(),
),
]),
response_observation: None,
body: Some(aether_contracts::ResponseBody {
json_body: Some(json!({
"plan_type": "plus",
@@ -740,6 +1207,7 @@ async fn gateway_refreshes_admin_provider_quota_for_codex_proxy_with_extended_ti
candidate_id: None,
status_code: 200,
headers: BTreeMap::new(),
response_observation: None,
body: Some(aether_contracts::ResponseBody {
json_body: Some(json!({
"plan_type": "plus",
@@ -897,6 +1365,7 @@ async fn gateway_refreshes_admin_provider_quota_locally_for_kiro_with_trusted_ad
candidate_id: None,
status_code: 200,
headers: BTreeMap::new(),
response_observation: None,
body: Some(aether_contracts::ResponseBody {
json_body: Some(json!({
"subscriptionInfo": {
@@ -1280,6 +1749,7 @@ async fn gateway_refresh_kiro_quota_reconciles_missing_fixed_endpoint_before_ref
candidate_id: None,
status_code: 200,
headers: BTreeMap::new(),
response_observation: None,
body: Some(aether_contracts::ResponseBody {
json_body: Some(json!({
"subscriptionInfo": {
@@ -1443,6 +1913,7 @@ async fn gateway_refreshes_admin_provider_quota_locally_for_gemini_cli_with_trus
candidate_id: None,
status_code: 200,
headers: BTreeMap::new(),
response_observation: None,
body: Some(aether_contracts::ResponseBody {
json_body: Some(json!({
"buckets": [
@@ -1840,6 +2311,7 @@ async fn gateway_refreshes_admin_provider_quota_locally_for_antigravity_with_tru
candidate_id: None,
status_code: 200,
headers: BTreeMap::new(),
response_observation: None,
body: Some(aether_contracts::ResponseBody {
json_body: Some(json!({
"models": {
apps/aether-gateway/src/tests/control/admin/oauth.rs
+85
View File
@@ -3440,11 +3440,16 @@ async fn gateway_completes_admin_provider_oauth_key_locally_with_trusted_admin_p
let token_hits_clone = Arc::clone(&token_hits);
let seen_token = Arc::new(Mutex::new(None::<SeenTokenRequest>));
let seen_token_clone = Arc::clone(&seen_token);
let namespace_race_repository = Arc::new(Mutex::new(
None::<Arc<InMemoryProviderCatalogReadRepository>>,
));
let namespace_race_repository_clone = Arc::clone(&namespace_race_repository);
let token_server = Router::new().route(
"/oauth/token",
any(move |request: Request| {
let token_hits_inner = Arc::clone(&token_hits_clone);
let seen_token_inner = Arc::clone(&seen_token_clone);
let namespace_race_repository_inner = Arc::clone(&namespace_race_repository_clone);
async move {
*token_hits_inner.lock().expect("mutex should lock") += 1;
let (parts, body) = request.into_parts();
@@ -3459,6 +3464,23 @@ async fn gateway_completes_admin_provider_oauth_key_locally_with_trusted_admin_p
body: String::from_utf8(raw_body.to_vec())
.expect("token request body should be utf8"),
});
let repository = namespace_race_repository_inner
.lock()
.expect("mutex should lock")
.clone()
.expect("provider catalog repository should be installed");
repository
.upsert_key_upstream_metadata_namespace(
"key-codex-oauth",
"codex",
&json!({
"primary_used_percent": 80.0,
"account_quota_request_id": "concurrent-refresh"
}),
Some(1_800_000_001),
)
.await
.expect("concurrent Codex namespace update should succeed");
Json(json!({
"access_token": "new-codex-access-token",
"refresh_token": "new-codex-refresh-token",
@@ -3492,6 +3514,33 @@ async fn gateway_completes_admin_provider_oauth_key_locally_with_trusted_admin_p
key.circuit_breaker_by_format = Some(json!({
"openai:chat": {"state": "open"}
}));
key.upstream_metadata = Some(json!({
"codex": {
"primary_used_percent": 100.0,
"primary_reset_at": 4_200_000_000u64,
"account_quota_reset_fence_unix_ms": 1_800_000_000_000u64,
"account_quota_reset_fence_id": "old-account-fence",
"account_quota_reset_processed_ids": ["old-account-reset"],
"account_quota_reset_pending": true
},
"unrelated_runtime": {
"preserved": true
}
}));
key.status_snapshot = Some(json!({
"oauth": {
"code": "invalid"
},
"quota": {
"provider_type": "codex",
"usage_ratio": 1.0,
"windows": [{
"kind": "primary",
"usage": 100.0,
"reset_at": 4_200_000_000u64
}]
}
}));
let score_identity = PoolMemberIdentity::provider_api_key("provider-codex", "key-codex-oauth");
let score_scope = provider_key_pool_score_scope();
@@ -3510,6 +3559,8 @@ async fn gateway_completes_admin_provider_oauth_key_locally_with_trusted_admin_p
vec![],
vec![key],
));
*namespace_race_repository.lock().expect("mutex should lock") =
Some(Arc::clone(&provider_catalog_repository));
let pool_score_repository =
Arc::new(InMemoryPoolMemberScoreRepository::seed(vec![invalid_score]));
@@ -3593,6 +3644,37 @@ async fn gateway_completes_admin_provider_oauth_key_locally_with_trusted_admin_p
assert_eq!(persisted.error_count, Some(0));
assert_eq!(persisted.health_by_format, Some(json!({})));
assert_eq!(persisted.circuit_breaker_by_format, Some(json!({})));
assert_eq!(
persisted
.upstream_metadata
.as_ref()
.and_then(Value::as_object)
.and_then(|metadata| metadata.get("codex"))
.and_then(Value::as_object)
.map(|codex| codex.keys().cloned().collect::<Vec<_>>()),
Some(vec!["credential_generation".to_string()]),
"explicit Codex reauthorization must not carry quota/reset state across accounts"
);
assert!(persisted
.upstream_metadata
.as_ref()
.and_then(|metadata| metadata.pointer("/codex/credential_generation"))
.and_then(Value::as_str)
.is_some_and(|generation| !generation.is_empty()));
assert_eq!(
persisted
.upstream_metadata
.as_ref()
.and_then(|metadata| metadata.get("unrelated_runtime")),
Some(&json!({"preserved": true}))
);
assert_eq!(
persisted
.status_snapshot
.as_ref()
.and_then(|snapshot| snapshot.get("quota")),
Some(&Value::Null)
);
let scores = pool_score_repository
.get_pool_member_scores_by_ids(&GetPoolMemberScoresByIdsQuery {
ids: vec![provider_key_pool_score_id(&score_identity, &score_scope)],
@@ -6164,6 +6246,7 @@ async fn gateway_refreshes_admin_provider_oauth_key_locally_with_trusted_admin_p
candidate_id: None,
status_code: 401,
headers: std::collections::BTreeMap::new(),
response_observation: None,
body: None,
telemetry: None,
error: None,
@@ -6523,6 +6606,7 @@ async fn gateway_manual_codex_oauth_refresh_reconciles_missing_fixed_endpoint_im
candidate_id: None,
status_code: 200,
headers: std::collections::BTreeMap::new(),
response_observation: None,
body: Some(aether_contracts::ResponseBody {
json_body: Some(json!({
"plan_type": "plus",
@@ -6757,6 +6841,7 @@ async fn run_gateway_manual_kiro_oauth_refresh_maintenance_endpoint_test(
candidate_id: None,
status_code: 200,
headers: std::collections::BTreeMap::new(),
response_observation: None,
body: Some(aether_contracts::ResponseBody {
json_body: Some(json!({
"subscriptionInfo": {
apps/aether-gateway/src/tests/control/admin/system_import.rs
+36
View File
@@ -2317,6 +2317,17 @@ async fn gateway_overwrites_oauth_provider_key_credentials_from_admin_system_imp
)
.expect("auth config should encrypt"),
);
existing_key.upstream_metadata = Some(json!({
"codex": {
"credential_generation": "generation-before-import",
"primary_used_percent": 90.0,
},
"unrelated": {"preserved": true},
}));
existing_key.status_snapshot = Some(json!({
"oauth": {"status": "invalid"},
"quota": {"used_ratio": 0.9},
}));
let score_identity =
PoolMemberIdentity::provider_api_key("provider-codex-existing", "key-codex-existing");
@@ -2417,6 +2428,31 @@ async fn gateway_overwrites_oauth_provider_key_credentials_from_admin_system_imp
assert_eq!(key.error_count, Some(0));
assert_eq!(key.health_by_format, Some(json!({})));
assert_eq!(key.circuit_breaker_by_format, Some(json!({})));
let codex = key
.upstream_metadata
.as_ref()
.and_then(|metadata| metadata.get("codex"))
.and_then(Value::as_object)
.expect("Codex metadata should exist");
assert_eq!(codex.len(), 1);
assert_ne!(
codex
.get(aether_admin::provider::quota::CODEX_CREDENTIAL_GENERATION_KEY)
.and_then(Value::as_str),
Some("generation-before-import")
);
assert_eq!(
key.upstream_metadata
.as_ref()
.and_then(|metadata| metadata.pointer("/unrelated/preserved")),
Some(&json!(true))
);
assert_eq!(
key.status_snapshot
.as_ref()
.and_then(|snapshot| snapshot.get("quota")),
Some(&Value::Null)
);
assert_eq!(
decrypt_python_fernet_ciphertext(
DEVELOPMENT_ENCRYPTION_KEY,
apps/aether-gateway/src/tests/control/proxy/embeddings.rs
+5
View File
@@ -657,6 +657,7 @@ fn embedding_execution_result(plan: &ExecutionPlan) -> ExecutionResult {
candidate_id: plan.candidate_id.clone(),
status_code: 200,
headers: BTreeMap::from([("content-type".to_string(), "application/json".to_string())]),
response_observation: None,
body: Some(ResponseBody {
json_body: Some(json!({
"object": "list",
@@ -679,6 +680,7 @@ fn gemini_embedding_execution_result(plan: &ExecutionPlan) -> ExecutionResult {
candidate_id: plan.candidate_id.clone(),
status_code: 200,
headers: BTreeMap::from([("content-type".to_string(), "application/json".to_string())]),
response_observation: None,
body: Some(ResponseBody {
json_body: Some(json!({
"model": "gemini-embedding-2-preview",
@@ -703,6 +705,7 @@ fn vertex_gemini_embedding_execution_result(plan: &ExecutionPlan) -> ExecutionRe
candidate_id: plan.candidate_id.clone(),
status_code: 200,
headers: BTreeMap::from([("content-type".to_string(), "application/json".to_string())]),
response_observation: None,
body: Some(ResponseBody {
json_body: Some(json!({
"predictions": [
@@ -727,6 +730,7 @@ fn gemini_batch_embedding_execution_result(plan: &ExecutionPlan) -> ExecutionRes
candidate_id: plan.candidate_id.clone(),
status_code: 200,
headers: BTreeMap::from([("content-type".to_string(), "application/json".to_string())]),
response_observation: None,
body: Some(ResponseBody {
json_body: Some(json!({
"model": "gemini-embedding-2-preview",
@@ -752,6 +756,7 @@ fn aliyun_embedding_execution_result(plan: &ExecutionPlan) -> ExecutionResult {
candidate_id: plan.candidate_id.clone(),
status_code: 200,
headers: BTreeMap::from([("content-type".to_string(), "application/json".to_string())]),
response_observation: None,
body: Some(ResponseBody {
json_body: Some(json!({
"output": {
apps/aether-gateway/src/tests/control/proxy/rerank.rs
+1
View File
@@ -151,6 +151,7 @@ fn rerank_execution_result(plan: &ExecutionPlan) -> ExecutionResult {
candidate_id: plan.candidate_id.clone(),
status_code: 200,
headers: BTreeMap::from([("content-type".to_string(), "application/json".to_string())]),
response_observation: None,
body: Some(ResponseBody {
json_body: Some(json!({
"model": "upstream-rerank",
apps/aether-gateway/src/tests/files/stream.rs
+1
View File
@@ -148,6 +148,7 @@ async fn gateway_executes_gemini_files_download_via_local_decision_gate_with_loc
"content-type".to_string(),
"application/octet-stream".to_string(),
)]),
response_observation: None,
},
},
StreamFrame {
apps/aether-gateway/src/tests/usage/local.rs
+2
View File
@@ -774,6 +774,7 @@ async fn gateway_records_failed_usage_when_all_local_openai_chat_candidates_exha
"content-type".to_string(),
"application/json".to_string(),
)]),
response_observation: None,
body: Some(aether_contracts::ResponseBody {
json_body: Some(json!({
"error": {
@@ -1056,6 +1057,7 @@ async fn sync_transport_error_policy_stops_or_retries_candidates_end_to_end_impl
"content-type".to_string(),
"application/json".to_string(),
)]),
response_observation: None,
body: Some(aether_contracts::ResponseBody {
json_body: Some(json!({
"id": "chatcmpl-transport-policy",
apps/aether-gateway/src/tests/video/stream.rs
+1
View File
@@ -212,6 +212,7 @@ async fn gateway_executes_openai_video_content_from_reconstructed_data_task_with
"content-type".to_string(),
"video/mp4".to_string(),
)]),
response_observation: None,
},
},
StreamFrame {