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

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/handlers/admin/provider/endpoint_keys/mutations/codex_reset_credit.rs
+36 -3
View File
@@ -55,6 +55,33 @@ pub(super) async fn maybe_handle(
if idempotency_key.is_empty() {
return Ok(Some(bad_request_response("idempotency_key 不能为空")));
}
if idempotency_key.len() > 256 {
return Ok(Some(bad_request_response(
"idempotency_key 不能超过 256 个字节",
)));
}
let expected_credential_generation = match payload.expected_credential_generation {
serde_json::Value::Null => None,
serde_json::Value::String(value) => {
let value = value.trim().to_string();
if value.is_empty() {
return Ok(Some(bad_request_response(
"expected_credential_generation 不能为空字符串",
)));
}
if value.len() > 256 {
return Ok(Some(bad_request_response(
"expected_credential_generation 不能超过 256 个字节",
)));
}
Some(value)
}
_ => {
return Ok(Some(bad_request_response(
"expected_credential_generation 必须是字符串或 null",
)));
}
};
let Some(key) = state
.read_provider_catalog_keys_by_ids(std::slice::from_ref(&key_id))
@@ -93,9 +120,15 @@ pub(super) async fn maybe_handle(
)));
};
let (status, payload) =
consume_codex_reset_credit_locally(state, &provider, &endpoint, key, &idempotency_key)
.await?;
let (status, payload) = consume_codex_reset_credit_locally(
state,
&provider,
&endpoint,
key,
&idempotency_key,
expected_credential_generation.as_deref(),
)
.await?;
Ok(Some((status, Json(payload)).into_response()))
}
apps/aether-gateway/src/handlers/admin/provider/endpoint_keys/mutations/update.rs
+31 -2
View File
@@ -1,7 +1,10 @@
use crate::handlers::admin::admin_provider_pool_config;
use crate::handlers::admin::provider::shared::paths::admin_update_key_id;
use crate::handlers::admin::provider::shared::payloads::AdminProviderKeyUpdatePatch;
use crate::handlers::admin::provider::write::keys::admin_provider_key_update_requires_immediate_model_fetch;
use crate::handlers::admin::provider::write::keys::{
admin_provider_key_update_requires_immediate_model_fetch,
build_provider_catalog_key_admin_cas_update,
};
use crate::handlers::admin::request::{AdminAppState, AdminRequestContext};
use crate::maintenance::ensure_provider_key_pool_scores_for_keys;
use crate::provider_key_auth::provider_key_effective_api_formats;
@@ -82,7 +85,25 @@ pub(super) async fn maybe_handle(
Ok(record) => record,
Err(detail) => return Ok(Some(bad_request_response(detail))),
};
let Some(mut updated) = state.update_provider_catalog_key(&updated_record).await? else {
let admin_update = build_provider_catalog_key_admin_cas_update(
&existing_key,
updated_record.clone(),
&provider.provider_type,
);
if !state
.compare_and_update_provider_catalog_key_admin_state(&admin_update)
.await?
{
return Ok(Some(conflict_response(
"Key 凭据或配置已被其他请求更新，请刷新后重试",
)));
}
let Some(mut updated) = state
.read_provider_catalog_keys_by_ids(std::slice::from_ref(&key_id))
.await?
.into_iter()
.next()
else {
return Ok(None);
};
if updated_record.learned_rpm_limit != existing_key.learned_rpm_limit {
@@ -183,3 +204,11 @@ fn not_found_response(detail: impl Into<String>) -> Response<Body> {
)
.into_response()
}
fn conflict_response(detail: impl Into<String>) -> Response<Body> {
(
http::StatusCode::CONFLICT,
Json(json!({ "detail": detail.into() })),
)
.into_response()
}
apps/aether-gateway/src/handlers/admin/provider/oauth/dispatch/complete/key.rs
+209 -21
View File
@@ -22,16 +22,60 @@ use crate::handlers::admin::request::{AdminAppState, AdminRequestContext};
use crate::handlers::shared::sync_provider_key_oauth_status_snapshot;
use crate::provider_key_auth::provider_key_is_oauth_managed;
use crate::GatewayError;
use aether_data_contracts::repository::provider_catalog::ProviderCatalogKeyOAuthRuntimeStateCasUpdate;
use aether_data_contracts::repository::provider_catalog::{
ProviderCatalogKeyOAuthRuntimeStateCasUpdate,
ProviderCatalogUpstreamMetadataNamespaceExpectation,
};
use axum::{
body::{Body, Bytes},
http,
response::{IntoResponse, Response},
Json,
};
use serde_json::json;
use serde_json::{json, Value};
use std::time::{SystemTime, UNIX_EPOCH};
const CODEX_OAUTH_COMPLETE_NAMESPACE_CAS_MAX_RETRIES: usize = 3;
const CODEX_CREDENTIAL_GENERATION_KEY: &str = "credential_generation";
#[derive(Debug, PartialEq)]
enum CodexOAuthCompleteCasMissAction {
AlreadyCompleted,
RetryNamespace(Option<Value>),
Conflict,
}
fn codex_oauth_complete_cas_miss_action(
latest_encrypted_auth_config: Option<&str>,
latest_upstream_metadata: Option<&Value>,
latest_status_snapshot: Option<&Value>,
expected_encrypted_auth_config: Option<&str>,
persisted_encrypted_auth_config: &str,
expected_codex_metadata_value: Option<&Value>,
replacement_codex_metadata_value: &Value,
) -> CodexOAuthCompleteCasMissAction {
let latest_codex_metadata_value = latest_upstream_metadata
.and_then(Value::as_object)
.and_then(|metadata| metadata.get("codex"))
.cloned();
let quota_is_cleared = latest_status_snapshot
.and_then(Value::as_object)
.and_then(|snapshot| snapshot.get("quota"))
== Some(&Value::Null);
if latest_encrypted_auth_config == Some(persisted_encrypted_auth_config)
&& latest_codex_metadata_value.as_ref() == Some(replacement_codex_metadata_value)
&& quota_is_cleared
{
return CodexOAuthCompleteCasMissAction::AlreadyCompleted;
}
if latest_encrypted_auth_config != expected_encrypted_auth_config
|| latest_codex_metadata_value.as_ref() == expected_codex_metadata_value
{
return CodexOAuthCompleteCasMissAction::Conflict;
}
CodexOAuthCompleteCasMissAction::RetryNamespace(latest_codex_metadata_value)
}
pub(super) async fn handle_admin_provider_oauth_complete_key(
state: &AdminAppState<'_>,
request_context: &AdminRequestContext<'_>,
@@ -277,29 +321,98 @@ pub(super) async fn handle_admin_provider_oauth_complete_key(
.and_then(|snapshot| snapshot.get("oauth"))
.cloned()
.unwrap_or(serde_json::Value::Null);
let mut expected_codex_metadata_value = key
.upstream_metadata
.as_ref()
.and_then(serde_json::Value::as_object)
.and_then(|metadata| metadata.get("codex"))
.cloned();
let mut status_snapshot_patch =
serde_json::Map::from_iter([("oauth".to_string(), oauth_status)]);
if provider_type == "codex" {
status_snapshot_patch.insert("quota".to_string(), serde_json::Value::Null);
}
let persisted_encrypted_auth_config = recovered_key
.encrypted_auth_config
.clone()
.expect("recovered auth config should be present");
let updated_result = state
.app()
.compare_and_update_provider_catalog_key_oauth_runtime_state(
&ProviderCatalogKeyOAuthRuntimeStateCasUpdate {
key_id: key_id.clone(),
expected_encrypted_auth_config: state_data.expected_encrypted_auth_config,
expected_credential: None,
encrypted_auth_config: persisted_encrypted_auth_config.clone(),
encrypted_api_key_update: Some(encrypted_api_key),
expires_at_unix_secs_update: Some(expires_at),
oauth_invalid_at_unix_secs: None,
oauth_invalid_reason: None,
reset_error_count: true,
upstream_metadata_patch: None,
status_snapshot_patch: json!({ "oauth": oauth_status }),
updated_at_unix_secs: Some(now_unix_secs),
},
)
.await;
let replacement_codex_metadata_value = json!({
CODEX_CREDENTIAL_GENERATION_KEY: uuid::Uuid::now_v7().to_string()
});
let expected_encrypted_auth_config = state_data.expected_encrypted_auth_config.clone();
let updated_result: Result<bool, GatewayError> = async {
let max_namespace_retries = if provider_type == "codex" {
CODEX_OAUTH_COMPLETE_NAMESPACE_CAS_MAX_RETRIES
} else {
0
};
for retry in 0..=max_namespace_retries {
let updated = state
.app()
.compare_and_update_provider_catalog_key_oauth_runtime_state(
&ProviderCatalogKeyOAuthRuntimeStateCasUpdate {
key_id: key_id.clone(),
expected_encrypted_auth_config: expected_encrypted_auth_config.clone(),
expected_credential: None,
expected_upstream_metadata_namespace: (provider_type == "codex").then(
|| ProviderCatalogUpstreamMetadataNamespaceExpectation {
namespace: "codex".to_string(),
expected_value: expected_codex_metadata_value.clone(),
},
),
encrypted_auth_config: persisted_encrypted_auth_config.clone(),
encrypted_api_key_update: Some(encrypted_api_key.clone()),
expires_at_unix_secs_update: Some(expires_at),
oauth_invalid_at_unix_secs: None,
oauth_invalid_reason: None,
reset_error_count: true,
upstream_metadata_patch: (provider_type == "codex")
.then(|| json!({"codex": replacement_codex_metadata_value.clone()})),
upstream_metadata_namespace_to_remove: None,
status_snapshot_patch: serde_json::Value::Object(
status_snapshot_patch.clone(),
),
updated_at_unix_secs: Some(now_unix_secs),
},
)
.await?;
if updated {
return Ok(true);
}
if provider_type != "codex" {
return Ok(false);
}
let Some(latest_key) = state
.read_provider_catalog_keys_by_ids(std::slice::from_ref(&key_id))
.await?
.into_iter()
.next()
else {
return Ok(false);
};
match codex_oauth_complete_cas_miss_action(
latest_key.encrypted_auth_config.as_deref(),
latest_key.upstream_metadata.as_ref(),
latest_key.status_snapshot.as_ref(),
expected_encrypted_auth_config.as_deref(),
&persisted_encrypted_auth_config,
expected_codex_metadata_value.as_ref(),
&replacement_codex_metadata_value,
) {
CodexOAuthCompleteCasMissAction::AlreadyCompleted => return Ok(true),
CodexOAuthCompleteCasMissAction::Conflict => return Ok(false),
CodexOAuthCompleteCasMissAction::RetryNamespace(latest_codex_metadata_value) => {
if retry == max_namespace_retries {
return Ok(false);
}
expected_codex_metadata_value = latest_codex_metadata_value;
}
}
}
Ok(false)
}
.await;
let _ = state
.app()
.invalidate_local_oauth_refresh_entry(&key_id)
@@ -397,3 +510,78 @@ pub(super) async fn handle_admin_provider_oauth_complete_key(
}))
.into_response())
}
#[cfg(test)]
mod tests {
use super::{codex_oauth_complete_cas_miss_action, CodexOAuthCompleteCasMissAction};
use serde_json::json;
#[test]
fn codex_oauth_complete_retries_only_when_namespace_changed() {
let expected_codex = json!({"request_id": "old"});
let replacement_codex = json!({"credential_generation": "generation-new"});
let latest_metadata = json!({
"codex": {"request_id": "new"},
"unrelated": {"preserved": true}
});
assert_eq!(
codex_oauth_complete_cas_miss_action(
Some("old-auth"),
Some(&latest_metadata),
None,
Some("old-auth"),
"new-auth",
Some(&expected_codex),
&replacement_codex,
),
CodexOAuthCompleteCasMissAction::RetryNamespace(Some(json!({
"request_id": "new"
})))
);
assert_eq!(
codex_oauth_complete_cas_miss_action(
Some("old-auth"),
Some(&json!({"codex": expected_codex.clone()})),
None,
Some("old-auth"),
"new-auth",
Some(&expected_codex),
&replacement_codex,
),
CodexOAuthCompleteCasMissAction::Conflict
);
}
#[test]
fn codex_oauth_complete_accepts_an_ambiguous_success_but_rejects_auth_rotation() {
let replacement_codex = json!({"credential_generation": "generation-new"});
assert_eq!(
codex_oauth_complete_cas_miss_action(
Some("new-auth"),
Some(&json!({
"codex": replacement_codex.clone(),
"unrelated": {"preserved": true}
})),
Some(&json!({"quota": null})),
Some("old-auth"),
"new-auth",
Some(&json!({"request_id": "old"})),
&replacement_codex,
),
CodexOAuthCompleteCasMissAction::AlreadyCompleted
);
assert_eq!(
codex_oauth_complete_cas_miss_action(
Some("other-auth"),
Some(&json!({"codex": {"request_id": "new"}})),
Some(&json!({"quota": null})),
Some("old-auth"),
"new-auth",
Some(&json!({"request_id": "old"})),
&replacement_codex,
),
CodexOAuthCompleteCasMissAction::Conflict
);
}
}
apps/aether-gateway/src/handlers/admin/provider/oauth/provisioning.rs
+169 -7
View File
@@ -12,6 +12,7 @@ use crate::ai_serving::{
build_provider_key_pool_score_upsert, provider_key_pool_score_id, provider_key_pool_score_scope,
};
use crate::handlers::admin::admin_provider_pool_config;
use crate::handlers::admin::provider::write::keys::build_provider_catalog_key_admin_cas_update;
use crate::handlers::admin::request::AdminAppState;
use crate::provider_key_auth::provider_active_api_formats;
use crate::GatewayError;
@@ -286,6 +287,61 @@ fn grok_oauth_catalog_key_fingerprint(
grok_browser_transport_fingerprint_from_auth_config(auth_config)
}
pub(crate) fn rotate_codex_credential_generation(
key: &mut StoredProviderCatalogKey,
provider_type: &str,
) {
if !provider_type.trim().eq_ignore_ascii_case("codex") {
return;
}
let mut upstream_metadata = key
.upstream_metadata
.as_ref()
.and_then(Value::as_object)
.cloned()
.unwrap_or_default();
upstream_metadata.insert(
"codex".to_string(),
json!({
aether_admin::provider::quota::CODEX_CREDENTIAL_GENERATION_KEY:
Uuid::now_v7().to_string(),
}),
);
key.upstream_metadata = Some(Value::Object(upstream_metadata));
if let Some(mut status_snapshot) = key
.status_snapshot
.as_ref()
.and_then(Value::as_object)
.cloned()
{
status_snapshot.insert("quota".to_string(), Value::Null);
key.status_snapshot = Some(Value::Object(status_snapshot));
}
}
pub(crate) fn ensure_codex_credential_generation_rotated(
key: &mut StoredProviderCatalogKey,
provider_type: &str,
previous_generation: Option<&str>,
) {
if !provider_type.trim().eq_ignore_ascii_case("codex") {
return;
}
let current_generation = key
.upstream_metadata
.as_ref()
.and_then(Value::as_object)
.and_then(|metadata| metadata.get("codex"))
.and_then(|codex| aether_admin::provider::quota::codex_credential_generation(Some(codex)));
let already_rotated = current_generation.is_some() && current_generation != previous_generation;
if !already_rotated {
rotate_codex_credential_generation(key, provider_type);
}
}
pub(crate) async fn create_provider_oauth_catalog_key(
state: &AdminAppState<'_>,
provider_id: &str,
@@ -344,6 +400,7 @@ pub(crate) async fn create_provider_oauth_catalog_key(
record.circuit_breaker_by_format = Some(json!({}));
record.created_at_unix_ms = Some(now_unix_secs);
record.updated_at_unix_secs = Some(now_unix_secs);
rotate_codex_credential_generation(&mut record, provider_type);
let created = state.create_provider_catalog_key(&record).await?;
if let Some(key) = created.as_ref() {
let _ = state
@@ -395,17 +452,23 @@ pub(crate) async fn update_existing_provider_oauth_catalog_key(
updated.proxy = Some(proxy);
}
updated.updated_at_unix_secs = Some(now_unix_secs);
if state.update_provider_catalog_key(&updated).await?.is_none() {
return Ok(None);
}
rotate_codex_credential_generation(&mut updated, provider_type);
let admin_update =
build_provider_catalog_key_admin_cas_update(existing_key, updated.clone(), provider_type);
if !state
.clear_provider_catalog_key_oauth_invalid_marker(&updated.id)
.compare_and_update_provider_catalog_key_admin_state(&admin_update)
.await?
{
return Ok(None);
}
let persisted = state
.reset_provider_catalog_key_recovery_state(&updated.id)
.reset_provider_catalog_key_recovery_state_fenced(
&updated.id,
updated
.encrypted_auth_config
.as_deref()
.expect("OAuth update always supplies encrypted auth_config"),
)
.await?;
if let Some(key) = persisted.as_ref() {
let _ = state
@@ -502,10 +565,12 @@ fn provider_oauth_catalog_key_api_formats(
#[cfg(test)]
mod tests {
use super::{
grok_oauth_catalog_key_fingerprint, provider_oauth_token_payload_expires_at_unix_secs,
ensure_codex_credential_generation_rotated, grok_oauth_catalog_key_fingerprint,
provider_oauth_token_payload_expires_at_unix_secs, rotate_codex_credential_generation,
};
use aether_data_contracts::repository::provider_catalog::StoredProviderCatalogKey;
use base64::{engine::general_purpose::URL_SAFE_NO_PAD, Engine as _};
use serde_json::json;
use serde_json::{json, Value};
fn sample_unsigned_jwt(payload: serde_json::Value) -> String {
let header = URL_SAFE_NO_PAD.encode(r#"{"alg":"none","typ":"JWT"}"#);
@@ -609,4 +674,101 @@ mod tests {
assert!(grok_oauth_catalog_key_fingerprint("openai", auth_config).is_none());
}
#[test]
fn codex_credential_rotation_replaces_quota_namespace_and_preserves_unrelated_state() {
let mut key = StoredProviderCatalogKey::new(
"key".to_string(),
"provider".to_string(),
"Codex".to_string(),
"oauth".to_string(),
None,
true,
)
.expect("key should build");
key.upstream_metadata = Some(json!({
"codex": {
"credential_generation": "old-generation",
"primary_used_percent": 75.0,
},
"unrelated": {"preserved": true},
}));
key.status_snapshot = Some(json!({
"oauth": {"status": "valid"},
"quota": {"used_ratio": 0.75},
}));
rotate_codex_credential_generation(&mut key, "codex");
let codex = key
.upstream_metadata
.as_ref()
.and_then(|metadata| metadata.get("codex"))
.and_then(Value::as_object)
.expect("codex namespace 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("old-generation")
);
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!(
key.status_snapshot
.as_ref()
.and_then(|snapshot| snapshot.pointer("/oauth/status")),
Some(&json!("valid"))
);
}
#[test]
fn codex_credential_rotation_ensure_does_not_rotate_twice_in_one_write() {
let mut key = StoredProviderCatalogKey::new(
"key".to_string(),
"provider".to_string(),
"Codex".to_string(),
"oauth".to_string(),
None,
true,
)
.expect("key should build");
key.upstream_metadata = Some(json!({
"codex": {"credential_generation": "generation-before-write"}
}));
rotate_codex_credential_generation(&mut key, "codex");
let builder_generation = key
.upstream_metadata
.as_ref()
.and_then(|metadata| metadata.pointer("/codex/credential_generation"))
.and_then(Value::as_str)
.expect("builder should rotate the generation")
.to_string();
ensure_codex_credential_generation_rotated(
&mut key,
"codex",
Some("generation-before-write"),
);
assert_eq!(
key.upstream_metadata
.as_ref()
.and_then(|metadata| metadata.pointer("/codex/credential_generation"))
.and_then(Value::as_str),
Some(builder_generation.as_str())
);
}
}
apps/aether-gateway/src/handlers/admin/provider/oauth/quota/chatgpt_web.rs
+1
View File
@@ -470,6 +470,7 @@ mod tests {
candidate_id: None,
status_code: 403,
headers: BTreeMap::new(),
response_observation: None,
body: Some(ResponseBody {
json_body: None,
body_bytes_b64: Some(base64::engine::general_purpose::STANDARD.encode(body)),
apps/aether-gateway/src/handlers/admin/provider/oauth/quota/codex/mod.rs
+694 -80
View File
@@ -18,24 +18,229 @@ use self::plan::{
execute_codex_reset_credit_plan,
};
use super::shared::{
build_quota_snapshot_payload, extract_execution_error_message,
oauth_refresh_auto_removed_result, persist_fenced_provider_quota_refresh_state,
persist_provider_quota_refresh_state, provider_auto_remove_banned_keys,
build_quota_snapshot_payload, complete_codex_account_reset, extract_execution_error_message,
oauth_refresh_auto_removed_result, persist_codex_provider_quota_refresh_state,
persist_fenced_provider_quota_refresh_state, provider_auto_remove_banned_keys,
provider_auto_remove_quota_exhausted_keys, quota_key_auto_removed,
quota_refresh_success_invalid_state, should_auto_remove_oauth_invalid_key,
ProviderQuotaExecutionOutcome,
quota_refresh_success_invalid_state, reserve_codex_account_reset,
should_auto_remove_oauth_invalid_key, CodexAccountResetCompleteResult,
CodexAccountResetReserveResult, CodexAccountResetTerminal, ProviderQuotaExecutionOutcome,
};
use crate::handlers::admin::request::AdminAppState;
use crate::handlers::admin::request::{AdminAppState, AdminGatewayProviderTransportSnapshot};
use crate::provider_key_auth::provider_key_is_oauth_managed;
use crate::state::ProviderTransportCredentialFence;
use crate::GatewayError;
use aether_contracts::ProxySnapshot;
use aether_data_contracts::repository::provider_catalog::{
ProviderCatalogKeyOAuthCredentialCasDelete, ProviderCatalogKeyOAuthCredentialFence,
StoredProviderCatalogEndpoint, StoredProviderCatalogKey, StoredProviderCatalogProvider,
ProviderCatalogKeyOAuthCredentialCasDelete,
ProviderCatalogUpstreamMetadataNamespaceExpectation, StoredProviderCatalogEndpoint,
StoredProviderCatalogKey, StoredProviderCatalogProvider,
};
use axum::http::StatusCode;
use serde_json::{json, Map, Value};
use std::time::{SystemTime, UNIX_EPOCH};
const CODEX_OAUTH_CREDENTIAL_STABILIZATION_ATTEMPTS: usize = 3;
const CODEX_RESET_QUOTA_RECONCILIATION_DELAYS_MS: [u64; 4] = [1_000, 2_000, 4_000, 8_000];
enum CodexOAuthRequestPreparation {
Ready {
transport: AdminGatewayProviderTransportSnapshot,
auth: (String, String),
credential_fence: ProviderTransportCredentialFence,
},
MissingAuth,
Conflict,
}
async fn prepare_codex_oauth_request(
state: &AdminAppState<'_>,
initial_transport: &AdminGatewayProviderTransportSnapshot,
) -> Result<CodexOAuthRequestPreparation, GatewayError> {
for _ in 0..CODEX_OAUTH_CREDENTIAL_STABILIZATION_ATTEMPTS {
let Some(transport) = state
.read_provider_transport_snapshot_uncached(
&initial_transport.provider.id,
&initial_transport.endpoint.id,
&initial_transport.key.id,
)
.await?
else {
return Ok(CodexOAuthRequestPreparation::Conflict);
};
if !crate::state::provider_transport_context_allows_credential_rotation(
initial_transport,
&transport,
) {
return Ok(CodexOAuthRequestPreparation::Conflict);
}
let Some(before_fence) = state
.app()
.capture_provider_transport_credential_fence(&transport)
.await?
else {
continue;
};
let resolved_auth = state.resolve_local_oauth_header_auth(&transport).await?;
let Some(current_transport) = state
.read_provider_transport_snapshot_uncached(
&initial_transport.provider.id,
&initial_transport.endpoint.id,
&initial_transport.key.id,
)
.await?
else {
return Ok(CodexOAuthRequestPreparation::Conflict);
};
if !crate::state::provider_transport_context_allows_credential_rotation(
initial_transport,
&current_transport,
) {
return Ok(CodexOAuthRequestPreparation::Conflict);
}
let Some(after_fence) = state
.app()
.capture_provider_transport_credential_fence(&current_transport)
.await?
else {
continue;
};
if before_fence != after_fence {
continue;
}
return Ok(match resolved_auth {
Some(auth) => CodexOAuthRequestPreparation::Ready {
transport: current_transport,
auth,
credential_fence: after_fence,
},
None => CodexOAuthRequestPreparation::MissingAuth,
});
}
Ok(CodexOAuthRequestPreparation::Conflict)
}
fn codex_reset_refresh_succeeded(payload: Option<&Value>, key_id: &str) -> bool {
payload
.and_then(|payload| payload.get("results"))
.and_then(Value::as_array)
.into_iter()
.flatten()
.filter_map(Value::as_object)
.find(|item| item.get("key_id").and_then(Value::as_str) == Some(key_id))
.and_then(|item| item.get("status"))
.and_then(Value::as_str)
.is_some_and(|status| status.eq_ignore_ascii_case("success"))
}
async fn codex_reset_fence_is_still_pending(
state: &AdminAppState<'_>,
key_id: &str,
expected_credential: &ProviderTransportCredentialFence,
reset_fence: &super::shared::CodexAccountResetFence,
) -> Result<bool, GatewayError> {
let Some(key) = state
.read_provider_catalog_keys_by_ids(&[key_id.to_string()])
.await?
.into_iter()
.next()
else {
return Ok(false);
};
if key.encrypted_auth_config.as_deref()
!= Some(expected_credential.encrypted_auth_config.as_str())
|| key.encrypted_api_key != expected_credential.credential.encrypted_api_key
|| key.auth_type != expected_credential.credential.auth_type
|| key.provider_id != expected_credential.credential.provider_id
{
return Ok(false);
}
let provider_type_matches = state
.read_provider_catalog_providers_by_ids(std::slice::from_ref(&key.provider_id))
.await?
.into_iter()
.next()
.is_some_and(|provider| {
provider.provider_type == expected_credential.credential.provider_type
});
if !provider_type_matches {
return Ok(false);
}
let codex = key
.upstream_metadata
.as_ref()
.and_then(Value::as_object)
.and_then(|metadata| metadata.get("codex"))
.and_then(Value::as_object);
Ok(codex.is_some_and(|codex| {
codex
.get(aether_admin::provider::quota::CODEX_QUOTA_ACCOUNT_RESET_FENCE_ID_KEY)
.and_then(Value::as_str)
== Some(reset_fence.id.as_str())
&& codex
.get(aether_admin::provider::quota::CODEX_QUOTA_ACCOUNT_RESET_GENERATION_KEY)
.and_then(aether_admin::provider::quota::coerce_json_u64)
== Some(reset_fence.generation)
&& codex
.get(
aether_admin::provider::quota::CODEX_QUOTA_ACCOUNT_RESET_PENDING_GENERATION_KEY,
)
.and_then(aether_admin::provider::quota::coerce_json_u64)
== Some(reset_fence.generation)
&& codex
.get(aether_admin::provider::quota::CODEX_QUOTA_ACCOUNT_RESET_PENDING_KEY)
.and_then(Value::as_bool)
== Some(true)
}))
}
async fn refresh_codex_quota_after_reset_until_settled(
state: &AdminAppState<'_>,
provider: &StoredProviderCatalogProvider,
endpoint: &StoredProviderCatalogEndpoint,
key: &StoredProviderCatalogKey,
reset_fence: &super::shared::CodexAccountResetFence,
expected_credential: &ProviderTransportCredentialFence,
) -> Result<Option<Value>, GatewayError> {
let mut latest_payload = None;
for attempt in 0..=CODEX_RESET_QUOTA_RECONCILIATION_DELAYS_MS.len() {
if attempt > 0 {
tokio::time::sleep(std::time::Duration::from_millis(
CODEX_RESET_QUOTA_RECONCILIATION_DELAYS_MS[attempt - 1],
))
.await;
}
if !codex_reset_fence_is_still_pending(state, &key.id, expected_credential, reset_fence)
.await?
{
break;
}
let payload = refresh_codex_provider_quota_locally_with_reset_fence(
state,
provider,
endpoint,
vec![key.clone()],
None,
Some(reset_fence.id.as_str()),
Some(reset_fence.generation),
Some(expected_credential),
)
.await?;
let refresh_succeeded = codex_reset_refresh_succeeded(payload.as_ref(), &key.id);
latest_payload = payload;
if !refresh_succeeded
|| !codex_reset_fence_is_still_pending(state, &key.id, expected_credential, reset_fence)
.await?
{
break;
}
}
Ok(latest_payload)
}
fn merge_codex_quota_metadata(
header_metadata: Option<&serde_json::Value>,
@@ -53,6 +258,26 @@ fn merge_codex_quota_metadata(
serde_json::Value::Object(merged)
}
fn codex_quota_window_coverage(
body_json: Option<&Value>,
) -> aether_admin::provider::quota::CodexQuotaWindowCoverage {
let body = body_json.and_then(Value::as_object);
let has_account_snapshot = body
.and_then(|body| body.get("rate_limit"))
.and_then(Value::as_object)
.is_some();
let has_spark_snapshot = body
.and_then(|body| body.get("additional_rate_limits"))
.and_then(Value::as_array)
.is_some();
match (has_account_snapshot, has_spark_snapshot) {
(true, true) => aether_admin::provider::quota::CodexQuotaWindowCoverage::FullSnapshot,
(true, false) => aether_admin::provider::quota::CodexQuotaWindowCoverage::AccountSnapshot,
_ => aether_admin::provider::quota::CodexQuotaWindowCoverage::Patch,
}
}
fn truncate_codex_reset_credit_detail_error(message: impl Into<String>) -> String {
let message = message.into();
let mut sanitized = message.replace('\n', " ");
@@ -198,6 +423,10 @@ fn codex_consume_success_status(outcome: &str) -> &'static str {
}
}
fn codex_reset_credit_outcome_allows_usage_drop(outcome: &str) -> bool {
matches!(outcome, "reset" | "already_redeemed")
}
fn codex_extract_refresh_result_fields(
refresh_payload: Option<&Value>,
key_id: &str,
@@ -247,12 +476,74 @@ fn codex_extract_refresh_result_fields(
)
}
async fn finish_codex_reset_replay(
state: &AdminAppState<'_>,
provider: &StoredProviderCatalogProvider,
endpoint: &StoredProviderCatalogEndpoint,
key: &StoredProviderCatalogKey,
credential: &ProviderTransportCredentialFence,
terminal: CodexAccountResetTerminal,
) -> Result<(StatusCode, Value), GatewayError> {
let mut refresh_status = "skipped".to_string();
let mut refresh_error = None;
let mut metadata = None;
let mut quota_snapshot = None;
if codex_reset_credit_outcome_allows_usage_drop(&terminal.outcome) {
let fence = super::shared::CodexAccountResetFence {
unix_ms: crate::clock::current_unix_ms(),
id: format!("reset:{}", terminal.idempotency_key),
generation: terminal.generation,
};
if codex_reset_fence_is_still_pending(state, &key.id, credential, &fence).await? {
match refresh_codex_quota_after_reset_until_settled(
state, provider, endpoint, key, &fence, credential,
)
.await
{
Ok(payload) => {
(refresh_status, refresh_error, metadata, quota_snapshot) =
codex_extract_refresh_result_fields(payload.as_ref(), &key.id);
}
Err(err) => {
refresh_status = "failed".to_string();
refresh_error =
Some(truncate_codex_reset_credit_detail_error(err.into_message()));
}
}
}
}
let mut payload = Map::new();
payload.insert("key_id".to_string(), json!(key.id));
payload.insert(
"status".to_string(),
json!(codex_consume_success_status(&terminal.outcome)),
);
payload.insert("outcome".to_string(), json!(terminal.outcome));
payload.insert(
"idempotency_key".to_string(),
json!(terminal.idempotency_key),
);
payload.insert("replay".to_string(), json!(true));
payload.insert("refresh_status".to_string(), json!(refresh_status));
if let Some(refresh_error) = refresh_error {
payload.insert("refresh_error".to_string(), json!(refresh_error));
}
if let Some(metadata) = metadata {
payload.insert("metadata".to_string(), metadata);
}
if let Some(quota_snapshot) = quota_snapshot {
payload.insert("quota_snapshot".to_string(), quota_snapshot);
}
Ok((StatusCode::OK, Value::Object(payload)))
}
pub(crate) async fn consume_codex_reset_credit_locally(
state: &AdminAppState<'_>,
provider: &StoredProviderCatalogProvider,
endpoint: &StoredProviderCatalogEndpoint,
key: StoredProviderCatalogKey,
idempotency_key: &str,
expected_credential_generation: Option<&str>,
) -> Result<(StatusCode, Value), GatewayError> {
let transport = match state
.read_provider_transport_snapshot(&provider.id, &endpoint.id, &key.id)
@@ -273,22 +564,48 @@ pub(crate) async fn consume_codex_reset_credit_locally(
};
let is_oauth_managed = provider_key_is_oauth_managed(&key, provider.provider_type.as_str());
let resolved_oauth_auth = if is_oauth_managed {
state.resolve_local_oauth_header_auth(&transport).await?
} else {
None
};
if is_oauth_managed && resolved_oauth_auth.is_none() {
if !is_oauth_managed {
return Ok((
StatusCode::BAD_REQUEST,
json!({
"key_id": key.id,
"status": "error",
"outcome": "error",
"message": "缺少 Codex OAuth 认证信息，请先重新授权/刷新 Token",
"message": "Codex reset credit 仅支持 OAuth 托管账号",
}),
));
}
let (transport, resolved_oauth_auth, reset_credential_fence) =
match prepare_codex_oauth_request(state, &transport).await? {
CodexOAuthRequestPreparation::Ready {
transport,
auth,
credential_fence,
} => (transport, Some(auth), credential_fence),
CodexOAuthRequestPreparation::MissingAuth => {
return Ok((
StatusCode::BAD_REQUEST,
json!({
"key_id": key.id,
"status": "error",
"outcome": "error",
"message": "缺少 Codex OAuth 认证信息，请先重新授权/刷新 Token",
}),
));
}
CodexOAuthRequestPreparation::Conflict => {
return Ok((
StatusCode::CONFLICT,
json!({
"key_id": key.id,
"status": "error",
"outcome": "error",
"idempotency_key": idempotency_key,
"message": "Codex credential changed before reset credit could be consumed",
}),
));
}
};
let request_spec = match build_codex_reset_credit_consume_request_spec(
&transport,
@@ -309,6 +626,79 @@ pub(crate) async fn consume_codex_reset_credit_locally(
}
};
let reservation = match reserve_codex_account_reset(
state,
&key.id,
reset_credential_fence.encrypted_auth_config.as_str(),
&reset_credential_fence.credential,
expected_credential_generation,
idempotency_key,
)
.await?
{
Some(CodexAccountResetReserveResult::Reserved(reservation)) => reservation,
Some(CodexAccountResetReserveResult::Replay(terminal)) => {
return finish_codex_reset_replay(
state,
provider,
endpoint,
&key,
&reset_credential_fence,
terminal,
)
.await;
}
Some(CodexAccountResetReserveResult::LegacyReplay) => {
return Ok((
StatusCode::OK,
json!({
"key_id": key.id,
"status": "success",
"outcome": "historical_replay",
"idempotency_key": idempotency_key,
"refresh_status": "skipped",
}),
));
}
Some(CodexAccountResetReserveResult::Busy(active)) => {
return Ok((
StatusCode::CONFLICT,
json!({
"key_id": key.id,
"status": "error",
"outcome": "busy",
"idempotency_key": idempotency_key,
"active_idempotency_key": active.idempotency_key,
"message": "Another Codex reset credit operation is unresolved",
}),
));
}
Some(CodexAccountResetReserveResult::CredentialGenerationMismatch) => {
return Ok((
StatusCode::CONFLICT,
json!({
"key_id": key.id,
"status": "error",
"outcome": "credential_changed",
"idempotency_key": idempotency_key,
"message": "Codex credential changed since this reset request was prepared",
}),
));
}
None => {
return Ok((
StatusCode::CONFLICT,
json!({
"key_id": key.id,
"status": "error",
"outcome": "error",
"idempotency_key": idempotency_key,
"message": "Codex reset reservation could not be persisted",
}),
));
}
};
let result =
match execute_codex_reset_credit_plan(state, &transport, request_spec, None).await? {
ProviderQuotaExecutionOutcome::Response(result) => result,
@@ -331,11 +721,11 @@ pub(crate) async fn consume_codex_reset_credit_locally(
.and_then(|body| body.json_body.as_ref());
let outcome = normalize_codex_reset_credit_consume_outcome(body_json)
.unwrap_or_else(|| "unknown".to_string());
let known_non_error_outcome = matches!(
let known_terminal_outcome = matches!(
outcome.as_str(),
"reset" | "already_redeemed" | "nothing_to_reset" | "no_credit"
);
if result.status_code >= 400 && !known_non_error_outcome {
if !known_terminal_outcome {
let detail = extract_execution_error_message(&result)
.unwrap_or_else(|| format!("HTTP {}", result.status_code));
return Ok((
@@ -345,19 +735,88 @@ pub(crate) async fn consume_codex_reset_credit_locally(
"status": "error",
"outcome": "error",
"idempotency_key": idempotency_key,
"message": format!("reset credit consume 返回状态码 {}: {detail}", result.status_code),
"message": format!("reset credit consume outcome is ambiguous: {detail}"),
"status_code": result.status_code,
}),
));
}
let (refresh_status, refresh_error, metadata, quota_snapshot) =
match refresh_codex_provider_quota_locally(
let fence_unix_ms = result
.response_observation
.as_ref()
.map(|observation| observation.response_headers_observed_at_unix_ms)
.unwrap_or_else(crate::clock::current_unix_ms);
let Some(completed) = complete_codex_account_reset(
state,
&key.id,
reset_credential_fence.encrypted_auth_config.as_str(),
&reset_credential_fence.credential,
&reservation,
&outcome,
fence_unix_ms,
)
.await?
else {
return Ok((
StatusCode::CONFLICT,
json!({
"key_id": key.id,
"status": "error",
"outcome": "error",
"idempotency_key": idempotency_key,
"message": "Codex reset completion could not be persisted",
}),
));
};
let (effective_outcome, reset_fence) = match completed {
CodexAccountResetCompleteResult::Activated(fence) => (outcome.clone(), Some(fence)),
CodexAccountResetCompleteResult::Noop(terminal)
| CodexAccountResetCompleteResult::Replay(terminal) => {
let fence =
codex_reset_credit_outcome_allows_usage_drop(&terminal.outcome).then(|| {
super::shared::CodexAccountResetFence {
unix_ms: fence_unix_ms,
id: format!("reset:{}", terminal.idempotency_key),
generation: terminal.generation,
}
});
(terminal.outcome, fence)
}
};
let (refresh_status, refresh_error, metadata, quota_snapshot) = match reset_fence.as_ref() {
Some(reset_fence) => {
match refresh_codex_quota_after_reset_until_settled(
state,
provider,
endpoint,
&key,
reset_fence,
&reset_credential_fence,
)
.await
{
Ok(refresh_payload) => {
codex_extract_refresh_result_fields(refresh_payload.as_ref(), &key.id)
}
Err(err) => (
"failed".to_string(),
Some(truncate_codex_reset_credit_detail_error(err.into_message())),
None,
None,
),
}
}
None => match refresh_codex_provider_quota_locally_with_reset_fence(
state,
provider,
endpoint,
vec![key.clone()],
None,
None,
None,
Some(&reset_credential_fence),
)
.await
{
@@ -370,15 +829,16 @@ pub(crate) async fn consume_codex_reset_credit_locally(
None,
None,
),
};
},
};
let mut payload = Map::new();
payload.insert("key_id".to_string(), json!(key.id));
payload.insert(
"status".to_string(),
json!(codex_consume_success_status(&outcome)),
json!(codex_consume_success_status(&effective_outcome)),
);
payload.insert("outcome".to_string(), json!(outcome));
payload.insert("outcome".to_string(), json!(effective_outcome));
payload.insert("idempotency_key".to_string(), json!(idempotency_key));
payload.insert("refresh_status".to_string(), json!(refresh_status));
if let Some(refresh_error) = refresh_error {
@@ -400,6 +860,29 @@ pub(crate) async fn refresh_codex_provider_quota_locally(
endpoint: &StoredProviderCatalogEndpoint,
keys: Vec<StoredProviderCatalogKey>,
proxy_override: Option<ProxySnapshot>,
) -> Result<Option<serde_json::Value>, GatewayError> {
refresh_codex_provider_quota_locally_with_reset_fence(
state,
provider,
endpoint,
keys,
proxy_override,
None,
None,
None,
)
.await
}
async fn refresh_codex_provider_quota_locally_with_reset_fence(
state: &AdminAppState<'_>,
provider: &StoredProviderCatalogProvider,
endpoint: &StoredProviderCatalogEndpoint,
keys: Vec<StoredProviderCatalogKey>,
proxy_override: Option<ProxySnapshot>,
account_reset_fence_id: Option<&str>,
authoritative_reset_generation: Option<u64>,
expected_reset_credential: Option<&crate::state::ProviderTransportCredentialFence>,
) -> Result<Option<serde_json::Value>, GatewayError> {
let mut results = Vec::new();
let mut success_count = 0usize;
@@ -412,7 +895,7 @@ pub(crate) async fn refresh_codex_provider_quota_locally(
for key in keys {
let had_oauth_refresh_issue =
codex_oauth_refresh_issue_reason(key.oauth_invalid_reason.as_deref());
let transport = match state
let initial_transport = match state
.read_provider_transport_snapshot(&provider.id, &endpoint.id, &key.id)
.await?
{
@@ -429,48 +912,70 @@ pub(crate) async fn refresh_codex_provider_quota_locally(
}
};
let is_oauth_managed = provider_key_is_oauth_managed(&key, provider.provider_type.as_str());
let quota_auth_config_fence = if is_oauth_managed {
match state
.app()
.capture_provider_transport_auth_config_fence(&transport)
.await?
{
Some(ciphertext) => Some(ciphertext),
None => {
let (transport, resolved_oauth_auth, quota_credential_fence) = if is_oauth_managed {
match prepare_codex_oauth_request(state, &initial_transport).await? {
CodexOAuthRequestPreparation::Ready {
transport,
auth,
credential_fence,
} => (transport, Some(auth), Some(credential_fence)),
CodexOAuthRequestPreparation::MissingAuth => {
failed_count += 1;
results.push(json!({
"key_id": key.id,
"key_name": key.name,
"status": "error",
"message": "OAuth credential changed before quota refresh",
"message": "缺少 Codex OAuth 认证信息，请先重新授权/刷新 Token",
}));
continue;
}
CodexOAuthRequestPreparation::Conflict => {
if quota_key_auto_removed(state, &key.id).await? {
auto_removed_count += 1;
results.push(oauth_refresh_auto_removed_result(&key));
} else {
failed_count += 1;
results.push(json!({
"key_id": key.id,
"key_name": key.name,
"status": "error",
"message": "OAuth credential changed before quota refresh",
}));
}
continue;
}
}
} else {
None
(initial_transport, None, None)
};
let resolved_oauth_auth = if is_oauth_managed {
state.resolve_local_oauth_header_auth(&transport).await?
} else {
None
};
if is_oauth_managed && quota_key_auto_removed(state, &key.id).await? {
auto_removed_count += 1;
results.push(oauth_refresh_auto_removed_result(&key));
continue;
}
if is_oauth_managed && resolved_oauth_auth.is_none() {
failed_count += 1;
results.push(json!({
"key_id": key.id,
"key_name": key.name,
"status": "error",
"message": "缺少 Codex OAuth 认证信息，请先重新授权/刷新 Token",
}));
continue;
if let Some(expected_reset_credential) = expected_reset_credential {
if quota_credential_fence.as_ref() != Some(expected_reset_credential) {
failed_count += 1;
results.push(json!({
"key_id": key.id,
"key_name": key.name,
"status": "error",
"message": "Codex credential changed after reset credit was consumed",
}));
continue;
}
}
let transport_codex_metadata = transport
.key
.upstream_metadata
.as_ref()
.and_then(Value::as_object)
.and_then(|metadata| metadata.get("codex"));
let observed_reset_generation = authoritative_reset_generation.or_else(|| {
Some(
aether_admin::provider::quota::codex_quota_account_reset_generation(
transport_codex_metadata,
),
)
});
let observed_credential_generation =
aether_admin::provider::quota::codex_credential_generation(transport_codex_metadata)
.map(ToOwned::to_owned);
let request_spec =
match build_codex_quota_request_spec(&transport, resolved_oauth_auth.clone()) {
@@ -487,6 +992,8 @@ pub(crate) async fn refresh_codex_provider_quota_locally(
}
};
let quota_request_fallback_started_at_unix_ms = crate::clock::current_unix_ms();
let quota_request_fallback_order_id = uuid::Uuid::now_v7().to_string();
let result = match execute_codex_quota_plan(
state,
&transport,
@@ -508,16 +1015,25 @@ pub(crate) async fn refresh_codex_provider_quota_locally(
continue;
}
};
let now_unix_secs = SystemTime::now()
.duration_since(UNIX_EPOCH)
.ok()
.map(|duration| duration.as_secs())
.unwrap_or(0);
let quota_response_fallback_observed_at_unix_ms = crate::clock::current_unix_ms();
let quota_response_observation = result.response_observation.as_ref();
let quota_request_started_at_unix_ms = quota_response_observation
.map(|observation| observation.request_started_at_unix_ms)
.unwrap_or(quota_request_fallback_started_at_unix_ms);
let quota_response_observed_at_unix_ms = quota_response_observation
.map(|observation| observation.response_headers_observed_at_unix_ms)
.unwrap_or(quota_response_fallback_observed_at_unix_ms);
let quota_request_order_id = quota_response_observation
.map(|observation| observation.request_order_id.as_str())
.unwrap_or(quota_request_fallback_order_id.as_str());
let now_unix_secs = quota_response_observed_at_unix_ms / 1_000;
let header_metadata = parse_codex_usage_headers(&result.headers, now_unix_secs);
let mut metadata_update = header_metadata
.as_ref()
.map(|metadata| json!({ "codex": metadata }));
let mut quota_window_coverage =
aether_admin::provider::quota::CodexQuotaWindowCoverage::Patch;
let (mut oauth_invalid_at_unix_secs, mut oauth_invalid_reason) = (None, None);
let mut status = "error".to_string();
let mut message = None::<String>;
@@ -544,6 +1060,7 @@ pub(crate) async fn refresh_codex_provider_quota_locally(
now_unix_secs,
)
.await?;
quota_window_coverage = codex_quota_window_coverage(Some(body_json));
metadata_update = Some(json!({
"codex": codex_metadata
}));
@@ -586,6 +1103,8 @@ pub(crate) async fn refresh_codex_provider_quota_locally(
}
402 => {
if codex_looks_like_workspace_deactivated(err_msg.as_deref()) {
quota_window_coverage =
aether_admin::provider::quota::CodexQuotaWindowCoverage::Patch;
let mut codex_meta = metadata_update
.as_ref()
.and_then(|value| value.get("codex"))
@@ -627,6 +1146,8 @@ pub(crate) async fn refresh_codex_provider_quota_locally(
oauth_invalid_reason = reason;
status = "workspace_deactivated".to_string();
} else {
quota_window_coverage =
aether_admin::provider::quota::CodexQuotaWindowCoverage::Patch;
let plan_type = transport
.key
.decrypted_auth_config
@@ -667,24 +1188,45 @@ pub(crate) async fn refresh_codex_provider_quota_locally(
}
}
let persisted = if let Some(expected_auth_config) = quota_auth_config_fence.as_deref() {
let persisted = if let Some(expected_credential) = quota_credential_fence.as_ref() {
persist_fenced_provider_quota_refresh_state(
state,
&key.id,
expected_auth_config,
expected_credential.encrypted_auth_config.as_str(),
metadata_update.as_ref(),
oauth_invalid_at_unix_secs,
oauth_invalid_reason.clone(),
aether_admin::provider::quota::CodexQuotaMergeContext {
observed_at_unix_secs: now_unix_secs,
request_started_at_unix_ms: Some(quota_request_started_at_unix_ms),
request_order_id: Some(quota_request_order_id),
observed_reset_generation,
authoritative_reset_generation,
observed_credential_generation: observed_credential_generation.as_deref(),
account_reset_fence_id,
coverage: quota_window_coverage,
},
Some(&expected_credential.credential),
)
.await?
} else {
persist_provider_quota_refresh_state(
persist_codex_provider_quota_refresh_state(
state,
&key.id,
metadata_update.as_ref(),
oauth_invalid_at_unix_secs,
oauth_invalid_reason.clone(),
None,
aether_admin::provider::quota::CodexQuotaMergeContext {
observed_at_unix_secs: now_unix_secs,
request_started_at_unix_ms: Some(quota_request_started_at_unix_ms),
request_order_id: Some(quota_request_order_id),
observed_reset_generation,
authoritative_reset_generation,
observed_credential_generation: observed_credential_generation.as_deref(),
account_reset_fence_id,
coverage: quota_window_coverage,
},
)
.await?
};
@@ -698,26 +1240,57 @@ pub(crate) async fn refresh_codex_provider_quota_locally(
}));
continue;
}
let credential_cas_delete = quota_auth_config_fence.as_ref().map(|auth_config| {
let persisted_key = state
.read_provider_catalog_keys_by_ids(&[key.id.clone()])
.await?
.into_iter()
.next();
let persisted_codex_metadata = persisted_key
.as_ref()
.and_then(|key| key.upstream_metadata.as_ref())
.and_then(serde_json::Value::as_object)
.and_then(|metadata| metadata.get("codex"))
.cloned();
if let Some(codex_metadata) = persisted_codex_metadata.as_ref() {
metadata_update = Some(json!({"codex": codex_metadata}));
}
let persisted_codex_object = persisted_codex_metadata
.as_ref()
.and_then(serde_json::Value::as_object);
let request_owns_persisted_oauth_state = quota_credential_fence.is_none()
|| (persisted_codex_object
.and_then(|codex| codex.get("oauth_state_request_started_at_unix_ms"))
.and_then(aether_admin::provider::quota::coerce_json_u64)
== Some(quota_request_started_at_unix_ms)
&& persisted_codex_object
.and_then(|codex| codex.get("oauth_state_request_id"))
.and_then(serde_json::Value::as_str)
== Some(quota_request_order_id));
let credential_cas_delete = quota_credential_fence.as_ref().map(|credential_fence| {
ProviderCatalogKeyOAuthCredentialCasDelete {
key_id: key.id.clone(),
expected_encrypted_auth_config: Some(auth_config.clone()),
expected_credential: ProviderCatalogKeyOAuthCredentialFence {
encrypted_api_key: key.encrypted_api_key.clone(),
auth_type: key.auth_type.clone(),
provider_id: key.provider_id.clone(),
provider_type: provider.provider_type.clone(),
},
expected_encrypted_auth_config: Some(
credential_fence.encrypted_auth_config.clone(),
),
expected_credential: credential_fence.credential.clone(),
expected_upstream_metadata_namespace: Some(
ProviderCatalogUpstreamMetadataNamespaceExpectation {
namespace: "codex".to_string(),
expected_value: persisted_codex_metadata.clone(),
},
),
}
});
let should_auto_remove_hard_banned =
provider_auto_remove_banned_keys(provider.config.as_ref())
&& should_auto_remove_oauth_invalid_key(
&key,
oauth_invalid_reason.as_deref(),
matches!(status_code, Some(401 | 403)),
now_unix_secs,
);
let should_auto_remove_hard_banned = request_owns_persisted_oauth_state
&& provider_auto_remove_banned_keys(provider.config.as_ref())
&& should_auto_remove_oauth_invalid_key(
persisted_key.as_ref().unwrap_or(&key),
persisted_key
.as_ref()
.and_then(|key| key.oauth_invalid_reason.as_deref()),
matches!(status_code, Some(401 | 403)),
now_unix_secs,
);
let auto_removed_hard_banned = if should_auto_remove_hard_banned {
match credential_cas_delete.as_ref() {
Some(delete) => {
@@ -735,6 +1308,7 @@ pub(crate) async fn refresh_codex_provider_quota_locally(
auto_removed_hard_banned_count += 1;
}
let auto_removed_quota_exhausted = if !auto_removed_hard_banned
&& request_owns_persisted_oauth_state
&& status == "quota_exhausted"
&& provider_auto_remove_quota_exhausted_keys(provider.config.as_ref())
{
@@ -798,7 +1372,9 @@ pub(crate) async fn refresh_codex_provider_quota_locally(
}
if let Some(quota_snapshot) = build_quota_snapshot_payload(
"codex",
key.status_snapshot.as_ref(),
persisted_key
.as_ref()
.and_then(|key| key.status_snapshot.as_ref()),
metadata_update.as_ref(),
) {
payload.insert("quota_snapshot".to_string(), quota_snapshot);
@@ -866,6 +1442,19 @@ mod tests {
);
}
#[test]
fn codex_reset_credit_only_allows_usage_drop_after_confirmed_redemption() {
assert!(codex_reset_credit_outcome_allows_usage_drop("reset"));
assert!(codex_reset_credit_outcome_allows_usage_drop(
"already_redeemed"
));
assert!(!codex_reset_credit_outcome_allows_usage_drop(
"nothing_to_reset"
));
assert!(!codex_reset_credit_outcome_allows_usage_drop("no_credit"));
assert!(!codex_reset_credit_outcome_allows_usage_drop("unknown"));
}
#[test]
fn codex_reset_credit_detail_failure_records_attempt_time() {
let mut metadata = Map::new();
@@ -880,4 +1469,29 @@ mod tests {
Some(&json!(1_777_000_000u64))
);
}
#[test]
fn codex_quota_coverage_only_replaces_observed_window_families() {
assert_eq!(
codex_quota_window_coverage(Some(&json!({"credits":{"balance":5}}))),
aether_admin::provider::quota::CodexQuotaWindowCoverage::Patch
);
assert_eq!(
codex_quota_window_coverage(None),
aether_admin::provider::quota::CodexQuotaWindowCoverage::Patch
);
assert_eq!(
codex_quota_window_coverage(Some(&json!({
"rate_limit":{"primary_window":{}}
}))),
aether_admin::provider::quota::CodexQuotaWindowCoverage::AccountSnapshot
);
assert_eq!(
codex_quota_window_coverage(Some(&json!({
"rate_limit":{"primary_window":{}},
"additional_rate_limits":[]
}))),
aether_admin::provider::quota::CodexQuotaWindowCoverage::FullSnapshot
);
}
}
apps/aether-gateway/src/handlers/admin/provider/oauth/quota/grok.rs
+1
View File
@@ -796,6 +796,7 @@ mod tests {
candidate_id: None,
status_code: 403,
headers: BTreeMap::new(),
response_observation: None,
body: Some(ResponseBody {
json_body: None,
body_bytes_b64: Some(base64::engine::general_purpose::STANDARD.encode(body)),
apps/aether-gateway/src/handlers/admin/provider/oauth/quota/shared.rs
+2139 -53
View File
File diff suppressed because it is too large Load Diff
apps/aether-gateway/src/handlers/admin/provider/query/models/model_test/tests.rs
+5
View File
@@ -178,6 +178,7 @@ fn provider_query_execution_json_body_decodes_stream_encoded_json_response() {
"content-type".to_string(),
"application/json".to_string(),
)]),
response_observation: None,
body: Some(aether_contracts::ResponseBody {
json_body: None,
body_bytes_b64: Some(encoded_body),
@@ -436,6 +437,7 @@ fn provider_query_standard_test_aggregates_responses_stream_body() {
candidate_id: Some("candidate-0".to_string()),
status_code: 200,
headers: BTreeMap::new(),
response_observation: None,
body: Some(aether_contracts::ResponseBody {
json_body: None,
body_bytes_b64: Some(
@@ -468,6 +470,7 @@ fn provider_query_standard_test_aggregates_responses_image_generation_call() {
candidate_id: Some("candidate-0".to_string()),
status_code: 200,
headers: BTreeMap::new(),
response_observation: None,
body: Some(aether_contracts::ResponseBody {
json_body: None,
body_bytes_b64: Some(
@@ -601,6 +604,7 @@ fn provider_query_search_success_requires_non_empty_output() {
candidate_id: Some("candidate-0".to_string()),
status_code: 200,
headers: BTreeMap::new(),
response_observation: None,
body: Some(aether_contracts::ResponseBody {
json_body: Some(body),
body_bytes_b64: None,
@@ -751,6 +755,7 @@ fn provider_query_standard_test_rejects_gemini_success_without_visible_output()
candidate_id: Some("candidate-0".to_string()),
status_code: 200,
headers: BTreeMap::new(),
response_observation: None,
body: Some(aether_contracts::ResponseBody {
json_body: Some(json!({
"candidates": [{
apps/aether-gateway/src/handlers/admin/provider/shared/payloads.rs
+35
View File
@@ -122,6 +122,7 @@ pub(crate) struct AdminProviderQuotaRefreshRequest {
#[derive(Debug, Deserialize)]
pub(crate) struct AdminCodexResetCreditConsumeRequest {
pub(crate) idempotency_key: String,
pub(crate) expected_credential_generation: serde_json::Value,
}
#[derive(Debug, Deserialize)]
@@ -339,3 +340,37 @@ pub(crate) struct AdminImportProviderModelsRequest {
)]
pub(crate) price_per_request: Option<f64>,
}
#[cfg(test)]
mod tests {
use super::AdminCodexResetCreditConsumeRequest;
#[test]
fn codex_reset_credit_consume_requires_an_explicit_credential_generation() {
assert!(
serde_json::from_value::<AdminCodexResetCreditConsumeRequest>(
serde_json::json!({"idempotency_key":"reset-old-client"}),
)
.is_err()
);
let legacy_account =
serde_json::from_value::<AdminCodexResetCreditConsumeRequest>(serde_json::json!({
"idempotency_key":"reset-legacy-account",
"expected_credential_generation":null,
}))
.expect("explicit null should fence an account without a generation");
assert!(legacy_account.expected_credential_generation.is_null());
let generated_account =
serde_json::from_value::<AdminCodexResetCreditConsumeRequest>(serde_json::json!({
"idempotency_key":"reset-generated-account",
"expected_credential_generation":"credential-v2",
}))
.expect("string generation should deserialize");
assert_eq!(
generated_account.expected_credential_generation,
serde_json::json!("credential-v2")
);
}
}
apps/aether-gateway/src/handlers/admin/provider/write/keys/create.rs
+2
View File
@@ -1,3 +1,4 @@
use crate::handlers::admin::provider::oauth::provisioning::rotate_codex_credential_generation;
use crate::handlers::admin::provider::shared::payloads::AdminProviderKeyCreateRequest;
use crate::handlers::admin::provider::write::normalize::{
normalize_allow_auth_channel_mismatch_formats, normalize_api_format_json_object_keys,
@@ -216,6 +217,7 @@ pub(crate) async fn build_admin_create_provider_key_record(
)?;
key.created_at_unix_ms = Some(now_unix_secs);
key.updated_at_unix_secs = Some(now_unix_secs);
rotate_codex_credential_generation(&mut key, &provider.provider_type);
Ok(key)
}
apps/aether-gateway/src/handlers/admin/provider/write/keys/mod.rs
+1
View File
@@ -6,6 +6,7 @@ pub(crate) use self::update::build_admin_update_provider_key_record;
pub(crate) use self::update::{
admin_provider_key_update_requires_immediate_model_fetch,
build_admin_update_provider_key_record_with_existing_keys,
build_provider_catalog_key_admin_cas_update,
};
mod batch;
apps/aether-gateway/src/handlers/admin/provider/write/keys/update.rs
+51
View File
@@ -1,3 +1,4 @@
use crate::handlers::admin::provider::oauth::provisioning::rotate_codex_credential_generation;
use crate::handlers::admin::provider::shared::payloads::AdminProviderKeyUpdatePatch;
use crate::handlers::admin::provider::write::normalize::{
normalize_allow_auth_channel_mismatch_formats, normalize_api_format_json_object_keys,
@@ -13,6 +14,7 @@ use crate::handlers::admin::shared::{
use crate::handlers::shared::normalize_optional_api_key_concurrent_limit;
use crate::provider_key_auth::provider_key_is_oauth_managed;
use aether_data_contracts::repository::provider_catalog::{
ProviderCatalogKeyAdminCasUpdate, ProviderCatalogKeyOAuthCredentialFence,
StoredProviderCatalogKey, StoredProviderCatalogProvider,
};
use aether_provider_transport::provider_types::provider_type_is_fixed;
@@ -368,6 +370,12 @@ pub(crate) fn build_admin_update_provider_key_record_with_existing_keys(
.duration_since(UNIX_EPOCH)
.ok()
.map(|duration| duration.as_secs());
let credential_identity_changed = !updated.auth_type.eq_ignore_ascii_case(&existing.auth_type)
|| updated.encrypted_api_key != existing.encrypted_api_key
|| updated.encrypted_auth_config != existing.encrypted_auth_config;
if credential_identity_changed {
rotate_codex_credential_generation(&mut updated, &provider.provider_type);
}
Ok(updated)
}
@@ -382,6 +390,49 @@ pub(crate) fn admin_provider_key_update_requires_immediate_model_fetch(
&& (!existing.auto_fetch_models || filters_changed || locked_models_changed)
}
pub(crate) fn build_provider_catalog_key_admin_cas_update(
existing: &StoredProviderCatalogKey,
updated: StoredProviderCatalogKey,
provider_type: &str,
) -> ProviderCatalogKeyAdminCasUpdate {
let previous_generation = existing
.upstream_metadata
.as_ref()
.and_then(|metadata| metadata.pointer("/codex/credential_generation"))
.and_then(serde_json::Value::as_str);
let next_generation = updated
.upstream_metadata
.as_ref()
.and_then(|metadata| metadata.pointer("/codex/credential_generation"))
.and_then(serde_json::Value::as_str);
let credential_changed = existing.auth_type != updated.auth_type
|| existing.encrypted_api_key != updated.encrypted_api_key
|| existing.encrypted_auth_config != updated.encrypted_auth_config;
let codex_rotation = provider_type
.trim()
.eq_ignore_ascii_case("codex")
.then(|| next_generation.filter(|next| Some(*next) != previous_generation))
.flatten()
.map(|generation| {
json!({
aether_admin::provider::quota::CODEX_CREDENTIAL_GENERATION_KEY: generation,
})
});
ProviderCatalogKeyAdminCasUpdate {
expected_encrypted_auth_config: existing.encrypted_auth_config.clone(),
expected_credential: ProviderCatalogKeyOAuthCredentialFence {
encrypted_api_key: existing.encrypted_api_key.clone(),
auth_type: existing.auth_type.clone(),
provider_id: existing.provider_id.clone(),
provider_type: provider_type.to_string(),
},
key: updated,
codex_rotation,
reset_oauth_runtime: credential_changed,
}
}
fn raw_secret_auth_type(value: &str) -> bool {
matches!(
value.trim().to_ascii_lowercase().as_str(),
apps/aether-gateway/src/handlers/admin/request/provider/catalog.rs
+9
View File
@@ -186,6 +186,15 @@ impl<'a> AdminAppState<'a> {
self.app.update_provider_catalog_key(key).await
}
pub(crate) async fn compare_and_update_provider_catalog_key_admin_state(
&self,
update: &aether_data_contracts::repository::provider_catalog::ProviderCatalogKeyAdminCasUpdate,
) -> Result<bool, GatewayError> {
self.app
.compare_and_update_provider_catalog_key_admin_state(update)
.await
}
pub(crate) async fn compare_and_update_provider_catalog_key_adaptive_state(
&self,
update: &aether_data_contracts::repository::provider_catalog::ProviderCatalogKeyAdaptiveStateUpdate,
apps/aether-gateway/src/handlers/admin/request/provider/tasks.rs
+1
View File
@@ -513,6 +513,7 @@ impl<'a> AdminAppState<'a> {
provider_id: key.provider_id.clone(),
provider_type: provider.provider_type.clone(),
},
expected_upstream_metadata_namespace: None,
},
)
.await
apps/aether-gateway/src/handlers/admin/request/system/import.rs
+61 -11
View File
@@ -4,10 +4,12 @@ use crate::api::ai::admin_endpoint_signature_parts;
use crate::handlers::admin::admin_provider_pool_config;
use crate::handlers::admin::model::ADMIN_EXTERNAL_MODELS_PROXY_NODE_CONFIG_KEY;
use crate::handlers::admin::provider::endpoints_admin::payloads::AdminProviderEndpointUpdatePatch;
use crate::handlers::admin::provider::oauth::provisioning::ensure_codex_credential_generation_rotated;
use crate::handlers::admin::provider::shared::payloads::{
AdminProviderCreateRequest, AdminProviderKeyCreateRequest, AdminProviderKeyUpdatePatch,
AdminProviderUpdatePatch,
};
use crate::handlers::admin::provider::write::keys::build_provider_catalog_key_admin_cas_update;
use crate::handlers::admin::shared::{
normalize_json_array, normalize_json_object, normalize_string_list,
};
@@ -377,10 +379,14 @@ fn normalize_import_key_raw_payload(
fn apply_imported_oauth_key_credentials(
state: &AdminAppState<'_>,
provider_type: &str,
previous_codex_credential_generation: Option<&str>,
raw_key: &Map<String, Value>,
normalized_auth_config: Option<&Value>,
record: &mut aether_data_contracts::repository::provider_catalog::StoredProviderCatalogKey,
) -> Result<bool, String> {
let previous_encrypted_api_key = record.encrypted_api_key.clone();
let previous_encrypted_auth_config = record.encrypted_auth_config.clone();
let mut credentials_supplied = false;
let mut api_key_supplied = false;
if let Some(api_key_value) = raw_key.get("api_key") {
@@ -424,10 +430,19 @@ fn apply_imported_oauth_key_credentials(
api_key_supplied,
);
let credential_material_changed = record.encrypted_api_key != previous_encrypted_api_key
|| record.encrypted_auth_config != previous_encrypted_auth_config;
if credentials_supplied {
record.oauth_invalid_at_unix_secs = None;
record.oauth_invalid_reason = None;
}
if credential_material_changed {
ensure_codex_credential_generation_rotated(
record,
provider_type,
previous_codex_credential_generation,
);
}
Ok(credentials_supplied)
}
@@ -1861,6 +1876,15 @@ impl<'a> AdminAppState<'a> {
if let Some(existing_index) = existing_key_index {
let existing_key = existing_keys[existing_index].clone();
let previous_codex_credential_generation = existing_key
.upstream_metadata
.as_ref()
.and_then(Value::as_object)
.and_then(|metadata| metadata.get("codex"))
.and_then(|codex| {
aether_admin::provider::quota::codex_credential_generation(Some(codex))
})
.map(ToOwned::to_owned);
match merge_mode {
AdminImportMergeMode::Skip => {
stats.keys.skipped += 1;
@@ -1890,6 +1914,8 @@ impl<'a> AdminAppState<'a> {
let oauth_credentials_supplied = if auth_type == "oauth" {
invalid!(apply_imported_oauth_key_credentials(
self,
&provider.provider_type,
previous_codex_credential_generation.as_deref(),
&raw_key,
normalized_auth_config.as_ref(),
&mut updated,
@@ -1903,8 +1929,31 @@ impl<'a> AdminAppState<'a> {
imported_key.fingerprint.clone(),
"fingerprint",
));
let Some(mut persisted) =
self.update_provider_catalog_key(&updated).await?
let admin_update = build_provider_catalog_key_admin_cas_update(
&existing_key,
updated.clone(),
&provider.provider_type,
);
if !self
.compare_and_update_provider_catalog_key_admin_state(&admin_update)
.await?
{
return Ok(Err((
http::StatusCode::CONFLICT,
json!({
"detail": format!(
"Provider '{provider_name}' 的 Key 已被其他请求更新，请重试"
)
}),
)));
}
let Some(mut persisted) = self
.read_provider_catalog_keys_by_ids(std::slice::from_ref(
&updated.id,
))
.await?
.into_iter()
.next()
else {
return Ok(Err(invalid_request(format!(
"更新 Provider '{provider_name}' 的 Key 失败"
@@ -1926,16 +1975,15 @@ impl<'a> AdminAppState<'a> {
persisted = reloaded;
}
if oauth_credentials_supplied {
if !self
.clear_provider_catalog_key_oauth_invalid_marker(&updated.id)
.await?
{
return Ok(Err(invalid_request(format!(
"更新 Provider '{provider_name}' 的 Key 失败"
))));
}
let Some(reloaded) = self
.reset_provider_catalog_key_recovery_state(&updated.id)
.reset_provider_catalog_key_recovery_state_fenced(
&updated.id,
updated.encrypted_auth_config.as_deref().ok_or_else(|| {
GatewayError::Internal(format!(
"OAuth Provider '{provider_name}' imported without auth_config"
))
})?,
)
.await?
else {
return Ok(Err(invalid_request(format!(
@@ -1975,6 +2023,8 @@ impl<'a> AdminAppState<'a> {
let oauth_credentials_supplied = if auth_type == "oauth" {
invalid!(apply_imported_oauth_key_credentials(
self,
&provider.provider_type,
None,
&raw_key,
normalized_auth_config.as_ref(),
&mut record,