i/Aether
1
0
Fork 0
mirror of https://github.com/fawney19/Aether.git synced 2026-09-12 22:20:19 +08:00
Code Issues Packages Projects Releases Wiki Activity

Merge pull request #695 from zhefox/main

Browse Source 
Enhance provider capabilities and clean up OAuth keys
This commit is contained in:
ZheFox
2026-07-28 13:58:33 +08:00
committed by GitHub
parent 4148ab1931 9d9316e434
commit 9a88f966d8
24 changed files with 857 additions and 173 deletions
Download Patch File Download Diff File Expand all files Collapse all files
apps/aether-gateway/src/data/state/catalog.rs
+23 -6
View File
@@ -2,12 +2,13 @@ use super::{
ApiKeyLastUsedDelta, DataLayerError, GatewayDataState, GeminiFileMappingListQuery,
GeminiFileMappingStats, ProviderCatalogKeyAdaptiveStateUpdate,
ProviderCatalogKeyHealthStateUpdate, ProviderCatalogKeyListQuery,
ProviderCatalogKeyOAuthRuntimeStateCasUpdate, ProviderCatalogKeyRuntimeMetadataUpdate,
ProviderCatalogKeyStatusSnapshotUpdate, PublicHealthStatusCount, PublicHealthTimelineBucket,
StoredGeminiFileMapping, StoredGeminiFileMappingListPage, StoredProviderCatalogEndpoint,
StoredProviderCatalogKey, StoredProviderCatalogKeyMaintenanceSummary,
StoredProviderCatalogKeyPage, StoredProviderCatalogKeyStats, StoredProviderCatalogProvider,
StoredRequestCandidate, UpsertGeminiFileMappingRecord, UpsertRequestCandidateRecord,
ProviderCatalogKeyOAuthCredentialCasDelete, ProviderCatalogKeyOAuthRuntimeStateCasUpdate,
ProviderCatalogKeyRuntimeMetadataUpdate, ProviderCatalogKeyStatusSnapshotUpdate,
PublicHealthStatusCount, PublicHealthTimelineBucket, StoredGeminiFileMapping,
StoredGeminiFileMappingListPage, StoredProviderCatalogEndpoint, StoredProviderCatalogKey,
StoredProviderCatalogKeyMaintenanceSummary, StoredProviderCatalogKeyPage,
StoredProviderCatalogKeyStats, StoredProviderCatalogProvider, StoredRequestCandidate,
UpsertGeminiFileMappingRecord, UpsertRequestCandidateRecord,
};
impl GatewayDataState {
@@ -690,6 +691,22 @@ impl GatewayDataState {
Ok(deleted)
}
pub(crate) async fn compare_and_delete_provider_catalog_key_oauth_credential(
&self,
delete: &ProviderCatalogKeyOAuthCredentialCasDelete,
) -> Result<bool, DataLayerError> {
let deleted = match &self.provider_catalog_writer {
Some(repository) => {
repository
.compare_and_delete_key_oauth_credential(delete)
.await
}
None => Ok(false),
}?;
self.clear_provider_catalog_cache();
Ok(deleted)
}
pub(crate) async fn clear_provider_catalog_key_oauth_invalid_marker(
&self,
key_id: &str,
apps/aether-gateway/src/data/state/mod.rs
+6 -5
View File
@@ -122,11 +122,12 @@ use aether_data_contracts::repository::pool_scores::{
};
use aether_data_contracts::repository::provider_catalog::{
ProviderCatalogKeyAdaptiveStateUpdate, ProviderCatalogKeyHealthStateUpdate,
ProviderCatalogKeyListQuery, ProviderCatalogKeyOAuthRuntimeStateCasUpdate,
ProviderCatalogKeyRuntimeMetadataUpdate, ProviderCatalogKeyStatusSnapshotUpdate,
ProviderCatalogReadRepository, ProviderCatalogWriteRepository, StoredProviderCatalogEndpoint,
StoredProviderCatalogKey, StoredProviderCatalogKeyMaintenanceSummary,
StoredProviderCatalogKeyPage, StoredProviderCatalogKeyStats, StoredProviderCatalogProvider,
ProviderCatalogKeyListQuery, ProviderCatalogKeyOAuthCredentialCasDelete,
ProviderCatalogKeyOAuthRuntimeStateCasUpdate, ProviderCatalogKeyRuntimeMetadataUpdate,
ProviderCatalogKeyStatusSnapshotUpdate, ProviderCatalogReadRepository,
ProviderCatalogWriteRepository, StoredProviderCatalogEndpoint, StoredProviderCatalogKey,
StoredProviderCatalogKeyMaintenanceSummary, StoredProviderCatalogKeyPage,
StoredProviderCatalogKeyStats, StoredProviderCatalogProvider,
};
use aether_data_contracts::repository::quota::{
ProviderQuotaReadRepository, ProviderQuotaWriteRepository, StoredProviderQuotaSnapshot,
apps/aether-gateway/src/execution_runtime/oauth_retry.rs
+2 -8
View File
@@ -312,7 +312,7 @@ mod tests {
}
#[tokio::test]
async fn retains_codex_key_after_request_proven_terminal_refresh_failure() {
async fn auto_removes_codex_key_after_request_proven_terminal_refresh_failure() {
let token_hits = Arc::new(Mutex::new(0usize));
let token_hits_clone = Arc::clone(&token_hits);
let token_server = Router::new().route(
@@ -467,13 +467,7 @@ mod tests {
.list_keys_by_ids(&["key-codex-oauth-retry".to_string()])
.await
.expect("keys should read");
assert_eq!(keys.len(), 1);
assert!(keys[0].oauth_invalid_at_unix_secs.is_some());
assert!(keys[0]
.oauth_invalid_reason
.as_deref()
.is_some_and(|reason| reason.contains("[REFRESH_FAILED]")
&& reason.contains("Token 续期失败 (401)")));
assert!(keys.is_empty());
token_handle.abort();
}
apps/aether-gateway/src/handlers/admin/provider/oauth/dispatch/import.rs
+8 -1
View File
@@ -344,10 +344,15 @@ async fn resolve_admin_provider_oauth_codex_access_token_agent_identity_import(
request_proxy: Option<ProxySnapshot>,
) -> Result<AdminProviderOAuthSingleImportTokens, Response<Body>> {
let executor = crate::oauth::GatewayOAuthHttpExecutor::new(*state);
let is_fedramp_account = identity_hints
.get("is_fedramp")
.and_then(serde_json::Value::as_bool)
.unwrap_or(false);
let mut auth_config =
aether_provider_transport::register_codex_agent_identity_from_access_token(
&executor,
access_token,
aether_provider_transport::CodexAgentIdentityRegistrationOptions { is_fedramp_account },
OAuthNetworkContext::provider_operation(request_proxy),
)
.await
@@ -1334,7 +1339,8 @@ mod tests {
"https://api.openai.com/auth": {
"chatgpt_account_id": "account-1",
"chatgpt_user_id": "user-1",
"chatgpt_plan_type": "plus"
"chatgpt_plan_type": "plus",
"chatgpt_account_is_fedramp": true
},
"https://api.openai.com/profile": {
"email": "agent@example.com"
@@ -1347,6 +1353,7 @@ mod tests {
assert_eq!(hints.get("account_id"), Some(&json!("account-1")));
assert_eq!(hints.get("user_id"), Some(&json!("user-1")));
assert_eq!(hints.get("plan_type"), Some(&json!("plus")));
assert_eq!(hints.get("is_fedramp"), Some(&json!(true)));
assert_eq!(hints.get("email"), Some(&json!("agent@example.com")));
assert!(!hints.contains_key("access_token"));
assert!(!hints.contains_key("id_token"));
apps/aether-gateway/src/handlers/admin/provider/oauth/quota/codex/mod.rs
+58 -7
View File
@@ -20,14 +20,17 @@ use self::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, quota_key_auto_removed,
quota_refresh_success_invalid_state, ProviderQuotaExecutionOutcome,
persist_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,
};
use crate::handlers::admin::request::AdminAppState;
use crate::provider_key_auth::provider_key_is_oauth_managed;
use crate::GatewayError;
use aether_contracts::ProxySnapshot;
use aether_data_contracts::repository::provider_catalog::{
ProviderCatalogKeyOAuthCredentialCasDelete, ProviderCatalogKeyOAuthCredentialFence,
StoredProviderCatalogEndpoint, StoredProviderCatalogKey, StoredProviderCatalogProvider,
};
use axum::http::StatusCode;
@@ -695,20 +698,68 @@ pub(crate) async fn refresh_codex_provider_quota_locally(
}));
continue;
}
// Codex quota responses never auto-delete keys. Without a repository
// conditional delete, any read-then-delete sequence could remove a
// replacement Agent Identity installed while the response was in flight.
let auto_removed_hard_banned = false;
let credential_cas_delete = quota_auth_config_fence.as_ref().map(|auth_config| {
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(),
},
}
});
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 auto_removed_hard_banned = if should_auto_remove_hard_banned {
match credential_cas_delete.as_ref() {
Some(delete) => {
state
.compare_and_delete_provider_catalog_key_oauth_credential(delete)
.await?
}
None => false,
}
} else {
false
};
if auto_removed_hard_banned {
auto_removed_count += 1;
auto_removed_hard_banned_count += 1;
}
let auto_removed_quota_exhausted = false;
let auto_removed_quota_exhausted = if !auto_removed_hard_banned
&& status == "quota_exhausted"
&& provider_auto_remove_quota_exhausted_keys(provider.config.as_ref())
{
match credential_cas_delete.as_ref() {
Some(delete) => {
state
.compare_and_delete_provider_catalog_key_oauth_credential(delete)
.await?
}
None => false,
}
} else {
false
};
if auto_removed_quota_exhausted {
auto_removed_count += 1;
status = "quota_exhausted".to_string();
}
let auto_removed = auto_removed_hard_banned || auto_removed_quota_exhausted;
if auto_removed {
let deleted_key_ids = [key.id.clone()];
state
.cleanup_deleted_provider_catalog_refs(&provider.id, false, &[], &deleted_key_ids)
.await?;
}
let refresh_fixed =
status == "success" && had_oauth_refresh_issue && oauth_invalid_reason.is_none();
if refresh_fixed {
apps/aether-gateway/src/handlers/admin/request/provider/catalog.rs
+9
View File
@@ -389,6 +389,15 @@ impl<'a> AdminAppState<'a> {
self.app.delete_provider_catalog_key(key_id).await
}
pub(crate) async fn compare_and_delete_provider_catalog_key_oauth_credential(
&self,
delete: &aether_data_contracts::repository::provider_catalog::ProviderCatalogKeyOAuthCredentialCasDelete,
) -> Result<bool, GatewayError> {
self.app
.compare_and_delete_provider_catalog_key_oauth_credential(delete)
.await
}
pub(crate) async fn create_provider_catalog_provider(
&self,
provider: &aether_data_contracts::repository::provider_catalog::StoredProviderCatalogProvider,
apps/aether-gateway/src/handlers/admin/request/provider/tasks.rs
+120 -52
View File
@@ -16,7 +16,7 @@ use serde_json::{json, Map, Value};
use std::collections::{BTreeMap, BTreeSet};
use std::time::{SystemTime, UNIX_EPOCH};
fn provider_skips_automatic_key_cleanup(provider: &StoredProviderCatalogProvider) -> bool {
fn provider_requires_credential_cas_cleanup(provider: &StoredProviderCatalogProvider) -> bool {
provider.provider_type.trim().eq_ignore_ascii_case("codex")
}
@@ -372,13 +372,6 @@ impl<'a> AdminAppState<'a> {
) -> Result<usize, GatewayError> {
use aether_admin::provider::pool as admin_provider_pool_pure;
// Codex OAuth credentials can be replaced by a long-lived Agent
// Identity under the same key id. Until deletes support an auth_config
// CAS, automatic cleanup must retain every Codex key.
if provider_skips_automatic_key_cleanup(provider) {
return Ok(0);
}
let banned_keys = self
.list_provider_catalog_keys_by_provider_ids(std::slice::from_ref(&provider.id))
.await?
@@ -389,27 +382,25 @@ impl<'a> AdminAppState<'a> {
return Ok(0);
}
let deleted_key_ids = banned_keys
.iter()
.map(|key| key.id.clone())
.collect::<Vec<_>>();
let mut deleted_key_ids = Vec::new();
for key in &banned_keys {
self.clear_admin_provider_pool_cooldown(&provider.id, &key.id)
.await;
self.reset_admin_provider_pool_cost(&provider.id, &key.id)
.await;
}
let mut affected = 0usize;
for key_id in &deleted_key_ids {
if self.delete_provider_catalog_key(key_id).await? {
affected += 1;
if self
.delete_provider_catalog_key_for_automatic_cleanup(provider, key)
.await?
{
self.clear_admin_provider_pool_cooldown(&provider.id, &key.id)
.await;
self.reset_admin_provider_pool_cost(&provider.id, &key.id)
.await;
deleted_key_ids.push(key.id.clone());
}
}
self.cleanup_deleted_provider_catalog_refs(&provider.id, false, &[], &deleted_key_ids)
.await?;
if !deleted_key_ids.is_empty() {
self.cleanup_deleted_provider_catalog_refs(&provider.id, false, &[], &deleted_key_ids)
.await?;
}
Ok(affected)
Ok(deleted_key_ids.len())
}
pub(crate) async fn cleanup_quota_exhausted_provider_catalog_keys(
@@ -419,10 +410,6 @@ impl<'a> AdminAppState<'a> {
) -> Result<usize, GatewayError> {
use aether_admin::provider::pool as admin_provider_pool_pure;
if provider_skips_automatic_key_cleanup(provider) {
return Ok(0);
}
let keys = self
.list_provider_catalog_keys_by_provider_ids(std::slice::from_ref(&provider.id))
.await?;
@@ -487,27 +474,48 @@ impl<'a> AdminAppState<'a> {
return Ok(0);
}
let deleted_key_ids = exhausted_keys
.iter()
.map(|key| key.id.clone())
.collect::<Vec<_>>();
let mut deleted_key_ids = Vec::new();
for key in exhausted_keys {
self.clear_admin_provider_pool_cooldown(&provider.id, &key.id)
.await;
self.reset_admin_provider_pool_cost(&provider.id, &key.id)
.await;
}
let mut affected = 0usize;
for key_id in &deleted_key_ids {
if self.delete_provider_catalog_key(key_id).await? {
affected += 1;
if self
.delete_provider_catalog_key_for_automatic_cleanup(provider, key)
.await?
{
self.clear_admin_provider_pool_cooldown(&provider.id, &key.id)
.await;
self.reset_admin_provider_pool_cost(&provider.id, &key.id)
.await;
deleted_key_ids.push(key.id.clone());
}
}
self.cleanup_deleted_provider_catalog_refs(&provider.id, false, &[], &deleted_key_ids)
.await?;
if !deleted_key_ids.is_empty() {
self.cleanup_deleted_provider_catalog_refs(&provider.id, false, &[], &deleted_key_ids)
.await?;
}
Ok(affected)
Ok(deleted_key_ids.len())
}
async fn delete_provider_catalog_key_for_automatic_cleanup(
&self,
provider: &StoredProviderCatalogProvider,
key: &StoredProviderCatalogKey,
) -> Result<bool, GatewayError> {
if !provider_requires_credential_cas_cleanup(provider) {
return self.delete_provider_catalog_key(&key.id).await;
}
self.compare_and_delete_provider_catalog_key_oauth_credential(
&aether_data_contracts::repository::provider_catalog::ProviderCatalogKeyOAuthCredentialCasDelete {
key_id: key.id.clone(),
expected_encrypted_auth_config: key.encrypted_auth_config.clone(),
expected_credential: aether_data_contracts::repository::provider_catalog::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(),
},
},
)
.await
}
pub(crate) async fn cleanup_provider_catalog_key_if_current<F>(
@@ -867,8 +875,12 @@ impl<'a> AdminAppState<'a> {
#[cfg(test)]
mod automatic_cleanup_tests {
use super::provider_skips_automatic_key_cleanup;
use aether_data_contracts::repository::provider_catalog::StoredProviderCatalogProvider;
use super::{provider_requires_credential_cas_cleanup, AdminAppState};
use aether_data::repository::provider_catalog::InMemoryProviderCatalogReadRepository;
use aether_data_contracts::repository::provider_catalog::{
ProviderCatalogReadRepository, StoredProviderCatalogKey, StoredProviderCatalogProvider,
};
use std::sync::Arc;
fn provider(provider_type: &str) -> StoredProviderCatalogProvider {
StoredProviderCatalogProvider::new(
@@ -881,9 +893,65 @@ mod automatic_cleanup_tests {
}
#[test]
fn codex_automatic_cleanup_is_disabled_for_replaceable_agent_credentials() {
assert!(provider_skips_automatic_key_cleanup(&provider("codex")));
assert!(provider_skips_automatic_key_cleanup(&provider("CoDeX")));
assert!(!provider_skips_automatic_key_cleanup(&provider("kiro")));
fn codex_automatic_cleanup_uses_credential_cas() {
assert!(provider_requires_credential_cas_cleanup(&provider("codex")));
assert!(provider_requires_credential_cas_cleanup(&provider("CoDeX")));
assert!(!provider_requires_credential_cas_cleanup(&provider("kiro")));
}
#[tokio::test]
async fn cleanup_removes_existing_hard_invalid_codex_key() {
let provider = provider("codex");
let mut key = StoredProviderCatalogKey::new(
"key-codex-invalid".to_string(),
provider.id.clone(),
"Invalid Codex OAuth".to_string(),
"oauth".to_string(),
None,
true,
)
.expect("key should build")
.with_transport_fields(
None,
"encrypted-api-key".to_string(),
Some("encrypted-auth-config".to_string()),
None,
None,
None,
None,
None,
None,
)
.expect("key transport should build");
key.oauth_invalid_at_unix_secs = Some(1);
key.oauth_invalid_reason = Some(
"[OAUTH_EXPIRED] Your authentication token has been invalidated. Please try signing in again."
.to_string(),
);
let repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![provider.clone()],
vec![],
vec![key],
));
let state = crate::AppState::new()
.expect("state should build")
.with_data_state_for_tests(
crate::data::GatewayDataState::with_provider_catalog_repository_for_tests(
repository.clone(),
),
);
let admin_state = AdminAppState::new(&state);
let affected = admin_state
.cleanup_known_banned_provider_catalog_keys(&provider)
.await
.expect("cleanup should succeed");
assert_eq!(affected, 1);
assert!(repository
.list_keys_by_ids(&["key-codex-invalid".to_string()])
.await
.expect("keys should read")
.is_empty());
}
}
apps/aether-gateway/src/orchestration/effects.rs
+4 -10
View File
@@ -3581,7 +3581,7 @@ mod tests {
}
#[tokio::test]
async fn oauth_invalidation_retains_inactive_pat_owner_when_auto_remove_is_enabled() {
async fn oauth_invalidation_auto_removes_inactive_pat_owner() {
let state = codex_state_with_auto_remove();
let plan = sample_codex_plan();
@@ -3600,17 +3600,11 @@ mod tests {
)
.await;
let stored_key = state
let stored_keys = state
.read_provider_catalog_keys_by_ids(std::slice::from_ref(&plan.key_id))
.await
.expect("provider catalog keys should load")
.into_iter()
.next()
.expect("runtime invalidation must not race-delete a replacement key");
assert_eq!(
stored_key.oauth_invalid_reason.as_deref(),
Some("[OAUTH_EXPIRED] Personal access token owner is inactive.")
);
.expect("provider catalog keys should load");
assert!(stored_keys.is_empty());
}
#[tokio::test]
apps/aether-gateway/src/state/catalog.rs
+32
View File
@@ -888,6 +888,38 @@ impl AppState {
Ok(deleted)
}
pub(crate) async fn compare_and_delete_provider_catalog_key_oauth_credential(
&self,
delete: &provider_catalog::ProviderCatalogKeyOAuthCredentialCasDelete,
) -> Result<bool, GatewayError> {
let deleted = self
.data
.compare_and_delete_provider_catalog_key_oauth_credential(delete)
.await
.map_err(|err| GatewayError::Internal(err.to_string()))?;
if deleted {
if let Err(err) = self
.data
.delete_pool_member_scores_for_member(
&pool_scores::PoolMemberIdentity::provider_api_key(
delete.expected_credential.provider_id.clone(),
delete.key_id.clone(),
),
)
.await
{
warn!(
provider_id = %delete.expected_credential.provider_id,
key_id = %delete.key_id,
error = ?err,
"gateway provider catalog OAuth credential CAS delete: failed to delete pool member scores"
);
}
self.invalidate_provider_routing_caches();
}
Ok(deleted)
}
pub(crate) async fn clear_provider_catalog_key_oauth_invalid_marker(
&self,
key_id: &str,
apps/aether-gateway/src/state/oauth.rs
+79 -17
View File
@@ -17,8 +17,9 @@ use aether_contracts::{
EXECUTION_REQUEST_FOLLOW_REDIRECTS_HEADER, EXECUTION_REQUEST_HTTP1_ONLY_HEADER,
};
use aether_data_contracts::repository::provider_catalog::{
ProviderCatalogKeyOAuthCredentialFence, ProviderCatalogKeyOAuthRuntimeStateCasUpdate,
ProviderCatalogKeyStatusSnapshotUpdate, StoredProviderCatalogKey,
ProviderCatalogKeyOAuthCredentialCasDelete, ProviderCatalogKeyOAuthCredentialFence,
ProviderCatalogKeyOAuthRuntimeStateCasUpdate, ProviderCatalogKeyStatusSnapshotUpdate,
StoredProviderCatalogKey,
};
use aether_runtime_state::RuntimeLockLease;
use base64::{engine::general_purpose::STANDARD, Engine as _};
@@ -1988,8 +1989,8 @@ impl AppState {
expected_encrypted_auth_config: Some(
expected_credential_fence.encrypted_auth_config.clone(),
),
expected_credential: Some(expected_credential_fence.credential),
encrypted_auth_config: expected_credential_fence.encrypted_auth_config,
expected_credential: Some(expected_credential_fence.credential.clone()),
encrypted_auth_config: expected_credential_fence.encrypted_auth_config.clone(),
encrypted_api_key_update: None,
expires_at_unix_secs_update: None,
oauth_invalid_at_unix_secs: latest_key.oauth_invalid_at_unix_secs,
@@ -2002,12 +2003,68 @@ impl AppState {
},
)
.await?;
let auto_removed = if updated
&& admin_provider_quota_pure::provider_auto_remove_banned_keys(
transport.provider.config.as_ref(),
)
&& admin_provider_quota_pure::should_auto_remove_oauth_invalid_key(
&latest_key,
None,
true,
now_unix_secs,
) {
self.delete_provider_transport_oauth_credential_fenced(
key_id,
&expected_credential_fence,
)
.await?
} else {
false
};
if updated {
let _ = self.invalidate_local_oauth_refresh_entry(key_id).await;
}
tracing::info!(
key_id,
provider_id = %transport.provider.id,
provider_type,
updated,
auto_removed,
"gateway fenced OAuth invalidation persisted"
);
Ok(updated)
}
async fn delete_provider_transport_oauth_credential_fenced(
&self,
key_id: &str,
expected: &ProviderTransportCredentialFence,
) -> Result<bool, GatewayError> {
let deleted = self
.compare_and_delete_provider_catalog_key_oauth_credential(
&ProviderCatalogKeyOAuthCredentialCasDelete {
key_id: key_id.to_string(),
expected_encrypted_auth_config: Some(expected.encrypted_auth_config.clone()),
expected_credential: expected.credential.clone(),
},
)
.await?;
if !deleted {
return Ok(false);
}
let deleted_key_ids = [key_id.to_string()];
self.cleanup_deleted_provider_catalog_refs(
&expected.credential.provider_id,
false,
&[],
&deleted_key_ids,
)
.await?;
self.clear_provider_transport_snapshot_cache();
let _ = self.invalidate_local_oauth_refresh_entry(key_id).await;
Ok(true)
}
async fn persist_local_oauth_refresh_entry(
&self,
transport: &provider_transport::GatewayProviderTransportSnapshot,
@@ -2482,24 +2539,29 @@ impl AppState {
}
}
// Codex credentials are replaceable under a stable key id. Without a
// conditional delete, refresh failure handling must retain them after
// writing the generation-fenced marker.
let auto_removed = if expected_credential_fence.is_none()
&& !transport
.provider
.provider_type
.trim()
.eq_ignore_ascii_case("codex")
&& admin_provider_quota_pure::provider_auto_remove_banned_keys(
transport.provider.config.as_ref(),
)
let should_auto_remove = admin_provider_quota_pure::provider_auto_remove_banned_keys(
transport.provider.config.as_ref(),
)
&& admin_provider_quota_pure::should_auto_remove_oauth_invalid_key(
&latest_key,
None,
access_token_invalid_proven,
now_unix_secs,
) {
);
let auto_removed = if !should_auto_remove {
false
} else if let Some(expected_credential_fence) = expected_credential_fence {
self.delete_provider_transport_oauth_credential_fenced(
key_id,
&expected_credential_fence,
)
.await?
} else if !transport
.provider
.provider_type
.trim()
.eq_ignore_ascii_case("codex")
{
self.clear_provider_transport_snapshot_cache();
if self.delete_provider_catalog_key(key_id).await? {
let deleted_key_ids = [key_id.to_string()];
apps/aether-gateway/src/tests/control/admin/endpoints/quota.rs
+4 -8
View File
@@ -417,7 +417,7 @@ async fn gateway_marks_codex_quota_exhausted_when_wham_usage_returns_payment_req
}
#[tokio::test]
async fn gateway_retains_codex_key_when_quota_only_reports_oauth_invalid() {
async fn gateway_auto_removes_codex_key_when_quota_proves_oauth_invalid() {
let upstream = Router::new().route(
"/api/admin/endpoints/providers/provider-codex/refresh-quota",
any(move |_request: Request| async move {
@@ -521,19 +521,15 @@ async fn gateway_retains_codex_key_when_quota_only_reports_oauth_invalid() {
let payload: serde_json::Value = response.json().await.expect("json body should parse");
assert_eq!(payload["success"], 0);
assert_eq!(payload["failed"], 1);
assert_eq!(payload["auto_removed"], 0);
assert_eq!(payload["auto_removed"], 1);
assert_eq!(payload["results"][0]["status"], "auth_invalid");
assert!(payload["results"][0].get("auto_removed").is_none());
assert_eq!(payload["results"][0]["auto_removed"], true);
let reloaded = provider_catalog_repository
.list_keys_by_ids(&["key-codex-expired".to_string()])
.await
.expect("keys should read");
assert_eq!(reloaded.len(), 1);
assert!(reloaded[0]
.oauth_invalid_reason
.as_deref()
.is_some_and(|reason| reason.starts_with("[OAUTH_EXPIRED]")));
assert!(reloaded.is_empty());
gateway_handle.abort();
execution_runtime_handle.abort();
crates/aether-data/adapters/mysql/src/provider_catalog.rs
+56 -4
View File
@@ -10,10 +10,10 @@ use sqlx::{
use aether_data_contracts::repository::provider_catalog::{
ProviderCatalogKeyAdaptiveStateUpdate, ProviderCatalogKeyHealthStateUpdate,
ProviderCatalogKeyListOrder, ProviderCatalogKeyListQuery,
ProviderCatalogKeyOAuthRuntimeStateCasUpdate, ProviderCatalogKeyRuntimeMetadataUpdate,
ProviderCatalogKeyStatusSnapshotUpdate, ProviderCatalogReadRepository,
ProviderCatalogUpstreamMetadataNamespaceUpdate, ProviderCatalogWriteRepository,
StoredProviderCatalogEndpoint, StoredProviderCatalogKey,
ProviderCatalogKeyOAuthCredentialCasDelete, ProviderCatalogKeyOAuthRuntimeStateCasUpdate,
ProviderCatalogKeyRuntimeMetadataUpdate, ProviderCatalogKeyStatusSnapshotUpdate,
ProviderCatalogReadRepository, ProviderCatalogUpstreamMetadataNamespaceUpdate,
ProviderCatalogWriteRepository, StoredProviderCatalogEndpoint, StoredProviderCatalogKey,
StoredProviderCatalogKeyMaintenanceSummary, StoredProviderCatalogKeyPage,
StoredProviderCatalogKeyStats, StoredProviderCatalogProvider,
};
@@ -994,6 +994,51 @@ WHERE id = ?
Ok(rows_affected > 0)
}
pub async fn compare_and_delete_key_oauth_credential(
&self,
delete: &ProviderCatalogKeyOAuthCredentialCasDelete,
) -> Result<bool, DataLayerError> {
validate_non_empty(&delete.key_id, "provider catalog key_id")?;
let expected = &delete.expected_credential;
if expected
.encrypted_api_key
.as_deref()
.is_some_and(|value| value.trim().is_empty())
|| expected.auth_type.trim().is_empty()
|| expected.provider_id.trim().is_empty()
|| expected.provider_type.trim().is_empty()
{
return Err(DataLayerError::InvalidInput(
"provider catalog OAuth credential CAS delete contains empty fields".to_string(),
));
}
let mut builder = QueryBuilder::<MySql>::new("DELETE FROM provider_api_keys WHERE id = ");
builder
.push_bind(&delete.key_id)
.push(" AND BINARY auth_config <=> BINARY ")
.push_bind(delete.expected_encrypted_auth_config.as_deref())
.push(" AND BINARY api_key <=> BINARY ")
.push_bind(expected.encrypted_api_key.as_deref())
.push(" AND BINARY auth_type = BINARY ")
.push_bind(&expected.auth_type)
.push(" AND BINARY provider_id = BINARY ")
.push_bind(&expected.provider_id)
.push(
" AND EXISTS (SELECT 1 FROM providers WHERE \
BINARY providers.id = BINARY provider_api_keys.provider_id \
AND BINARY providers.provider_type = BINARY ",
)
.push_bind(&expected.provider_type)
.push(")");
let rows_affected = builder
.build()
.execute(&self.pool)
.await
.map_sql_err()?
.rows_affected();
Ok(rows_affected > 0)
}
pub async fn update_key_upstream_metadata(
&self,
key_id: &str,
@@ -1928,6 +1973,13 @@ impl ProviderCatalogWriteRepository for MysqlProviderCatalogReadRepository {
Self::delete_key(self, key_id).await
}
async fn compare_and_delete_key_oauth_credential(
&self,
delete: &ProviderCatalogKeyOAuthCredentialCasDelete,
) -> Result<bool, DataLayerError> {
Self::compare_and_delete_key_oauth_credential(self, delete).await
}
async fn clear_key_oauth_invalid_marker(&self, key_id: &str) -> Result<bool, DataLayerError> {
Self::clear_key_oauth_invalid_marker(self, key_id).await
}
crates/aether-data/adapters/postgres/src/provider_catalog.rs
+58 -4
View File
@@ -11,10 +11,10 @@ use sqlx::{
use aether_data_contracts::repository::provider_catalog::{
ProviderCatalogKeyAdaptiveStateUpdate, ProviderCatalogKeyHealthStateUpdate,
ProviderCatalogKeyListOrder, ProviderCatalogKeyListQuery,
ProviderCatalogKeyOAuthRuntimeStateCasUpdate, ProviderCatalogKeyRuntimeMetadataUpdate,
ProviderCatalogKeyStatusSnapshotUpdate, ProviderCatalogReadRepository,
ProviderCatalogUpstreamMetadataNamespaceUpdate, ProviderCatalogWriteRepository,
StoredProviderCatalogEndpoint, StoredProviderCatalogKey,
ProviderCatalogKeyOAuthCredentialCasDelete, ProviderCatalogKeyOAuthRuntimeStateCasUpdate,
ProviderCatalogKeyRuntimeMetadataUpdate, ProviderCatalogKeyStatusSnapshotUpdate,
ProviderCatalogReadRepository, ProviderCatalogUpstreamMetadataNamespaceUpdate,
ProviderCatalogWriteRepository, StoredProviderCatalogEndpoint, StoredProviderCatalogKey,
StoredProviderCatalogKeyMaintenanceSummary, StoredProviderCatalogKeyPage,
StoredProviderCatalogKeyStats, StoredProviderCatalogProvider,
};
@@ -2075,6 +2075,53 @@ WHERE id = $1
Ok(rows_affected > 0)
}
pub async fn compare_and_delete_key_oauth_credential(
&self,
delete: &ProviderCatalogKeyOAuthCredentialCasDelete,
) -> Result<bool, DataLayerError> {
let expected = &delete.expected_credential;
if delete.key_id.trim().is_empty()
|| expected
.encrypted_api_key
.as_deref()
.is_some_and(|value| value.trim().is_empty())
|| expected.auth_type.trim().is_empty()
|| expected.provider_id.trim().is_empty()
|| expected.provider_type.trim().is_empty()
{
return Err(DataLayerError::InvalidInput(
"provider catalog OAuth credential CAS delete contains empty fields".to_string(),
));
}
let rows_affected = sqlx::query(
r#"
DELETE FROM provider_api_keys
WHERE id = $1
AND auth_config IS NOT DISTINCT FROM $2
AND api_key IS NOT DISTINCT FROM $3
AND auth_type = $4
AND provider_id = $5
AND EXISTS (
SELECT 1
FROM providers
WHERE providers.id = provider_api_keys.provider_id
AND providers.provider_type = $6
)
"#,
)
.bind(&delete.key_id)
.bind(delete.expected_encrypted_auth_config.as_deref())
.bind(expected.encrypted_api_key.as_deref())
.bind(&expected.auth_type)
.bind(&expected.provider_id)
.bind(&expected.provider_type)
.execute(&self.pool)
.await
.map_postgres_err()?
.rows_affected();
Ok(rows_affected > 0)
}
pub async fn update_key_upstream_metadata(
&self,
key_id: &str,
@@ -2692,6 +2739,13 @@ impl ProviderCatalogWriteRepository for SqlxProviderCatalogReadRepository {
Self::delete_key(self, key_id).await
}
async fn compare_and_delete_key_oauth_credential(
&self,
delete: &ProviderCatalogKeyOAuthCredentialCasDelete,
) -> Result<bool, DataLayerError> {
Self::compare_and_delete_key_oauth_credential(self, delete).await
}
async fn clear_key_oauth_invalid_marker(&self, key_id: &str) -> Result<bool, DataLayerError> {
Self::clear_key_oauth_invalid_marker(self, key_id).await
}
crates/aether-data/adapters/sqlite/src/provider_catalog.rs
+97 -8
View File
@@ -10,10 +10,10 @@ use sqlx::{
use aether_data_contracts::repository::provider_catalog::{
ProviderCatalogKeyAdaptiveStateUpdate, ProviderCatalogKeyHealthStateUpdate,
ProviderCatalogKeyListOrder, ProviderCatalogKeyListQuery,
ProviderCatalogKeyOAuthRuntimeStateCasUpdate, ProviderCatalogKeyRuntimeMetadataUpdate,
ProviderCatalogKeyStatusSnapshotUpdate, ProviderCatalogReadRepository,
ProviderCatalogUpstreamMetadataNamespaceUpdate, ProviderCatalogWriteRepository,
StoredProviderCatalogEndpoint, StoredProviderCatalogKey,
ProviderCatalogKeyOAuthCredentialCasDelete, ProviderCatalogKeyOAuthRuntimeStateCasUpdate,
ProviderCatalogKeyRuntimeMetadataUpdate, ProviderCatalogKeyStatusSnapshotUpdate,
ProviderCatalogReadRepository, ProviderCatalogUpstreamMetadataNamespaceUpdate,
ProviderCatalogWriteRepository, StoredProviderCatalogEndpoint, StoredProviderCatalogKey,
StoredProviderCatalogKeyMaintenanceSummary, StoredProviderCatalogKeyPage,
StoredProviderCatalogKeyStats, StoredProviderCatalogProvider,
};
@@ -1175,6 +1175,53 @@ WHERE id = ?
Ok(rows_affected > 0)
}
pub async fn compare_and_delete_key_oauth_credential(
&self,
delete: &ProviderCatalogKeyOAuthCredentialCasDelete,
) -> Result<bool, DataLayerError> {
validate_non_empty(&delete.key_id, "provider catalog key_id")?;
let expected = &delete.expected_credential;
if expected
.encrypted_api_key
.as_deref()
.is_some_and(|value| value.trim().is_empty())
|| expected.auth_type.trim().is_empty()
|| expected.provider_id.trim().is_empty()
|| expected.provider_type.trim().is_empty()
{
return Err(DataLayerError::InvalidInput(
"provider catalog OAuth credential CAS delete contains empty fields".to_string(),
));
}
let rows_affected = sqlx::query(
r#"
DELETE FROM provider_api_keys
WHERE id = ?
AND auth_config IS ?
AND api_key IS ?
AND auth_type = ?
AND provider_id = ?
AND EXISTS (
SELECT 1
FROM providers
WHERE providers.id = provider_api_keys.provider_id
AND providers.provider_type = ?
)
"#,
)
.bind(&delete.key_id)
.bind(delete.expected_encrypted_auth_config.as_deref())
.bind(expected.encrypted_api_key.as_deref())
.bind(&expected.auth_type)
.bind(&expected.provider_id)
.bind(&expected.provider_type)
.execute(&self.pool)
.await
.map_sql_err()?
.rows_affected();
Ok(rows_affected > 0)
}
pub async fn update_key_upstream_metadata(
&self,
key_id: &str,
@@ -2121,6 +2168,13 @@ impl ProviderCatalogWriteRepository for SqliteProviderCatalogReadRepository {
Self::delete_key(self, key_id).await
}
async fn compare_and_delete_key_oauth_credential(
&self,
delete: &ProviderCatalogKeyOAuthCredentialCasDelete,
) -> Result<bool, DataLayerError> {
Self::compare_and_delete_key_oauth_credential(self, delete).await
}
async fn clear_key_oauth_invalid_marker(&self, key_id: &str) -> Result<bool, DataLayerError> {
Self::clear_key_oauth_invalid_marker(self, key_id).await
}
@@ -2967,10 +3021,10 @@ mod tests {
use aether_data_contracts::repository::provider_catalog::{
ProviderCatalogKeyAdaptiveState, ProviderCatalogKeyAdaptiveStateUpdate,
ProviderCatalogKeyHealthStateUpdate, ProviderCatalogKeyListOrder,
ProviderCatalogKeyListQuery, ProviderCatalogKeyOAuthCredentialFence,
ProviderCatalogKeyOAuthRuntimeStateCasUpdate, ProviderCatalogKeyRuntimeMetadataUpdate,
ProviderCatalogUpstreamMetadataNamespaceUpdate, StoredProviderCatalogEndpoint,
StoredProviderCatalogKey, StoredProviderCatalogProvider,
ProviderCatalogKeyListQuery, ProviderCatalogKeyOAuthCredentialCasDelete,
ProviderCatalogKeyOAuthCredentialFence, ProviderCatalogKeyOAuthRuntimeStateCasUpdate,
ProviderCatalogKeyRuntimeMetadataUpdate, ProviderCatalogUpstreamMetadataNamespaceUpdate,
StoredProviderCatalogEndpoint, StoredProviderCatalogKey, StoredProviderCatalogProvider,
};
use serde_json::json;
@@ -3391,6 +3445,41 @@ mod tests {
stored_after_stale.upstream_metadata.as_ref().unwrap()["codex"]["remaining"],
3
);
let stale_delete = ProviderCatalogKeyOAuthCredentialCasDelete {
key_id: stored_after_stale.id.clone(),
expected_encrypted_auth_config: Some("encrypted-auth-v1".to_string()),
expected_credential: ProviderCatalogKeyOAuthCredentialFence {
encrypted_api_key: Some("encrypted-api-key".to_string()),
auth_type: "oauth".to_string(),
provider_id: "oauth-cas-provider".to_string(),
provider_type: "codex".to_string(),
},
};
assert!(!repository
.compare_and_delete_key_oauth_credential(&stale_delete)
.await
.expect("stale credential delete should conflict"));
let current_delete = ProviderCatalogKeyOAuthCredentialCasDelete {
key_id: stored_after_stale.id.clone(),
expected_encrypted_auth_config: stored_after_stale.encrypted_auth_config.clone(),
expected_credential: ProviderCatalogKeyOAuthCredentialFence {
encrypted_api_key: stored_after_stale.encrypted_api_key.clone(),
auth_type: stored_after_stale.auth_type.clone(),
provider_id: stored_after_stale.provider_id.clone(),
provider_type: "codex".to_string(),
},
};
assert!(repository
.compare_and_delete_key_oauth_credential(&current_delete)
.await
.expect("current credential generation should delete"));
assert!(repository
.list_keys_by_ids(&[stored_after_stale.id])
.await
.expect("deleted key lookup should succeed")
.is_empty());
}
#[tokio::test]
crates/aether-data/contracts/src/repository/provider_catalog/mod.rs
+5 -4
View File
@@ -5,10 +5,11 @@ pub use snapshot::ProviderCatalogSnapshot;
pub use types::{
ProviderCatalogKeyAdaptiveState, ProviderCatalogKeyAdaptiveStateUpdate,
ProviderCatalogKeyHealthStateUpdate, ProviderCatalogKeyListOrder, ProviderCatalogKeyListQuery,
ProviderCatalogKeyOAuthCredentialFence, ProviderCatalogKeyOAuthRuntimeStateCasUpdate,
ProviderCatalogKeyRuntimeMetadataUpdate, ProviderCatalogKeyStatusSnapshotUpdate,
ProviderCatalogReadRepository, ProviderCatalogUpstreamMetadataNamespaceUpdate,
ProviderCatalogWriteRepository, StoredProviderCatalogEndpoint, StoredProviderCatalogKey,
ProviderCatalogKeyOAuthCredentialCasDelete, ProviderCatalogKeyOAuthCredentialFence,
ProviderCatalogKeyOAuthRuntimeStateCasUpdate, ProviderCatalogKeyRuntimeMetadataUpdate,
ProviderCatalogKeyStatusSnapshotUpdate, ProviderCatalogReadRepository,
ProviderCatalogUpstreamMetadataNamespaceUpdate, ProviderCatalogWriteRepository,
StoredProviderCatalogEndpoint, StoredProviderCatalogKey,
StoredProviderCatalogKeyMaintenanceSummary, StoredProviderCatalogKeyPage,
StoredProviderCatalogKeyStats, StoredProviderCatalogProvider,
};
crates/aether-data/contracts/src/repository/provider_catalog/types.rs
+19
View File
@@ -79,6 +79,15 @@ pub struct ProviderCatalogKeyOAuthCredentialFence {
pub provider_type: String,
}
/// Atomic key deletion fenced by the exact OAuth credential generation that
/// produced the terminal failure.
#[derive(Debug, Clone, PartialEq, serde::Serialize, serde::Deserialize)]
pub struct ProviderCatalogKeyOAuthCredentialCasDelete {
pub key_id: String,
pub expected_encrypted_auth_config: Option<String>,
pub expected_credential: ProviderCatalogKeyOAuthCredentialFence,
}
/// Agent/runtime-owned OAuth state update fenced by the exact encrypted
/// auth_config and, when supplied, credential context observed before the
/// refresh started. Repositories must update only these fields and return
@@ -853,6 +862,16 @@ pub trait ProviderCatalogWriteRepository: Send + Sync {
async fn delete_key(&self, key_id: &str) -> Result<bool, crate::DataLayerError>;
async fn compare_and_delete_key_oauth_credential(
&self,
_delete: &ProviderCatalogKeyOAuthCredentialCasDelete,
) -> Result<bool, crate::DataLayerError> {
Err(crate::DataLayerError::InvalidConfiguration(
"provider catalog OAuth credential CAS deletes are not supported by this repository"
.to_string(),
))
}
async fn clear_key_oauth_invalid_marker(
&self,
key_id: &str,
crates/aether-data/runtime/src/repository/provider_catalog/memory.rs
+127 -6
View File
@@ -8,8 +8,9 @@ use serde_json::{json, Map, Value};
use super::{
ProviderCatalogKeyAdaptiveState, ProviderCatalogKeyAdaptiveStateUpdate,
ProviderCatalogKeyHealthStateUpdate, ProviderCatalogKeyListQuery,
ProviderCatalogKeyOAuthRuntimeStateCasUpdate, ProviderCatalogKeyRuntimeMetadataUpdate,
ProviderCatalogKeyStatusSnapshotUpdate, ProviderCatalogReadRepository, ProviderCatalogSnapshot,
ProviderCatalogKeyOAuthCredentialCasDelete, ProviderCatalogKeyOAuthRuntimeStateCasUpdate,
ProviderCatalogKeyRuntimeMetadataUpdate, ProviderCatalogKeyStatusSnapshotUpdate,
ProviderCatalogReadRepository, ProviderCatalogSnapshot,
ProviderCatalogUpstreamMetadataNamespaceUpdate, ProviderCatalogWriteRepository,
StoredProviderCatalogEndpoint, StoredProviderCatalogKey,
StoredProviderCatalogKeyMaintenanceSummary, StoredProviderCatalogKeyPage,
@@ -702,6 +703,46 @@ impl ProviderCatalogWriteRepository for InMemoryProviderCatalogReadRepository {
Ok(index.keys.remove(key_id).is_some())
}
async fn compare_and_delete_key_oauth_credential(
&self,
delete: &ProviderCatalogKeyOAuthCredentialCasDelete,
) -> Result<bool, DataLayerError> {
let expected = &delete.expected_credential;
if delete.key_id.trim().is_empty()
|| expected
.encrypted_api_key
.as_deref()
.is_some_and(|value| value.trim().is_empty())
|| expected.auth_type.trim().is_empty()
|| expected.provider_id.trim().is_empty()
|| expected.provider_type.trim().is_empty()
{
return Err(DataLayerError::InvalidInput(
"provider catalog OAuth credential CAS delete contains empty fields".to_string(),
));
}
let mut index = self
.index
.write()
.expect("provider catalog repository lock");
let Some(key) = index.keys.get(&delete.key_id) else {
return Ok(false);
};
let provider_type_matches = index
.providers
.get(&key.provider_id)
.is_some_and(|provider| provider.provider_type == expected.provider_type);
if key.encrypted_auth_config != delete.expected_encrypted_auth_config
|| key.encrypted_api_key != expected.encrypted_api_key
|| key.auth_type != expected.auth_type
|| key.provider_id != expected.provider_id
|| !provider_type_matches
{
return Ok(false);
}
Ok(index.keys.remove(&delete.key_id).is_some())
}
async fn clear_key_oauth_invalid_marker(&self, key_id: &str) -> Result<bool, DataLayerError> {
let mut index = self
.index
@@ -1196,10 +1237,11 @@ mod tests {
use crate::repository::provider_catalog::{
ProviderCatalogKeyAdaptiveState, ProviderCatalogKeyAdaptiveStateUpdate,
ProviderCatalogKeyHealthStateUpdate, ProviderCatalogKeyListOrder,
ProviderCatalogKeyListQuery, ProviderCatalogKeyOAuthCredentialFence,
ProviderCatalogKeyOAuthRuntimeStateCasUpdate, ProviderCatalogKeyRuntimeMetadataUpdate,
ProviderCatalogReadRepository, ProviderCatalogWriteRepository,
StoredProviderCatalogEndpoint, StoredProviderCatalogKey, StoredProviderCatalogProvider,
ProviderCatalogKeyListQuery, ProviderCatalogKeyOAuthCredentialCasDelete,
ProviderCatalogKeyOAuthCredentialFence, ProviderCatalogKeyOAuthRuntimeStateCasUpdate,
ProviderCatalogKeyRuntimeMetadataUpdate, ProviderCatalogReadRepository,
ProviderCatalogWriteRepository, StoredProviderCatalogEndpoint, StoredProviderCatalogKey,
StoredProviderCatalogProvider,
};
use crate::repository::usage::ProviderApiKeyUsageDelta;
use serde_json::{json, Value};
@@ -1516,6 +1558,85 @@ mod tests {
.expect("provider type mismatch should be a CAS miss"));
}
#[tokio::test]
async fn oauth_credential_cas_delete_rejects_replacement_generation() {
let repository = InMemoryProviderCatalogReadRepository::seed(
vec![sample_provider("provider-1")],
vec![],
vec![sample_key("key-1", "provider-1")
.with_transport_fields(
None,
"ciphertext-api-1".to_string(),
Some("ciphertext-auth-1".to_string()),
None,
None,
None,
None,
None,
None,
)
.expect("key transport should build")],
);
let stale_delete = ProviderCatalogKeyOAuthCredentialCasDelete {
key_id: "key-1".to_string(),
expected_encrypted_auth_config: Some("ciphertext-auth-1".to_string()),
expected_credential: ProviderCatalogKeyOAuthCredentialFence {
encrypted_api_key: Some("ciphertext-api-1".to_string()),
auth_type: "api_key".to_string(),
provider_id: "provider-1".to_string(),
provider_type: "custom".to_string(),
},
};
let mut replacement = repository
.list_keys_by_ids(&["key-1".to_string()])
.await
.expect("key should load")
.pop()
.expect("key should exist");
replacement.encrypted_api_key = Some("ciphertext-api-2".to_string());
replacement.encrypted_auth_config = Some("ciphertext-auth-2".to_string());
replacement.auth_type = "oauth".to_string();
repository
.update_key(&replacement)
.await
.expect("replacement should persist");
assert!(!repository
.compare_and_delete_key_oauth_credential(&stale_delete)
.await
.expect("stale delete should be a CAS miss"));
assert_eq!(
repository
.list_keys_by_ids(&["key-1".to_string()])
.await
.expect("replacement should load")[0]
.encrypted_auth_config
.as_deref(),
Some("ciphertext-auth-2")
);
let current_delete = ProviderCatalogKeyOAuthCredentialCasDelete {
key_id: "key-1".to_string(),
expected_encrypted_auth_config: Some("ciphertext-auth-2".to_string()),
expected_credential: ProviderCatalogKeyOAuthCredentialFence {
encrypted_api_key: Some("ciphertext-api-2".to_string()),
auth_type: "oauth".to_string(),
provider_id: "provider-1".to_string(),
provider_type: "custom".to_string(),
},
};
assert!(repository
.compare_and_delete_key_oauth_credential(&current_delete)
.await
.expect("current generation should delete"));
assert!(repository
.list_keys_by_ids(&["key-1".to_string()])
.await
.expect("deleted key lookup should succeed")
.is_empty());
}
#[tokio::test]
async fn materializes_codex_window_usage_stats_delta_in_memory() {
let mut key = sample_key("key-1", "provider-1");
crates/aether-data/runtime/src/repository/provider_catalog/mod.rs
+3 -3
View File
@@ -4,9 +4,9 @@ mod memory;
pub(crate) use aether_data_contracts::repository::provider_catalog::{
ProviderCatalogKeyAdaptiveState, ProviderCatalogKeyAdaptiveStateUpdate,
ProviderCatalogKeyHealthStateUpdate, ProviderCatalogKeyListOrder, ProviderCatalogKeyListQuery,
ProviderCatalogKeyOAuthCredentialFence, ProviderCatalogKeyOAuthRuntimeStateCasUpdate,
ProviderCatalogKeyRuntimeMetadataUpdate, ProviderCatalogKeyStatusSnapshotUpdate,
ProviderCatalogReadRepository, ProviderCatalogSnapshot,
ProviderCatalogKeyOAuthCredentialCasDelete, ProviderCatalogKeyOAuthCredentialFence,
ProviderCatalogKeyOAuthRuntimeStateCasUpdate, ProviderCatalogKeyRuntimeMetadataUpdate,
ProviderCatalogKeyStatusSnapshotUpdate, ProviderCatalogReadRepository, ProviderCatalogSnapshot,
ProviderCatalogUpstreamMetadataNamespaceUpdate, ProviderCatalogWriteRepository,
StoredProviderCatalogEndpoint, StoredProviderCatalogKey,
StoredProviderCatalogKeyMaintenanceSummary, StoredProviderCatalogKeyPage,
crates/aether-provider/transport/src/agent_identity.rs
+89 -23
View File
@@ -38,7 +38,7 @@ const CODEX_AGENT_IDENTITY_AUTH_API_BASE_URL: &str = "https://auth.openai.com/ap
const AUTHORIZATION_HEADER: &str = "authorization";
const ASSERTION_PREFIX: &str = "AgentAssertion ";
const CODEX_AGENT_IDENTITY_AGENT_HARNESS_ID: &str = "codex-cli";
const CODEX_AGENT_IDENTITY_RUNNING_LOCATION: &str = "local";
const CODEX_AGENT_IDENTITY_RESPONSES_API_CAPABILITY: &str = "responsesapi";
/// The AgentAssertion scheme is generated internally after an Agent Identity
/// task has been registered. Keep the scheme check separate from envelope
@@ -54,8 +54,11 @@ pub enum CodexAgentIdentityEnrollmentError {
MissingSessionToken,
#[error("Agent Identity 注册请求失败")]
RegistrationRequestFailed,
#[error("Agent Identity 注册被拒绝（HTTP {status_code}）")]
RegistrationRejected { status_code: u16 },
#[error("Agent Identity 注册被拒绝（HTTP {status_code}）{reason}")]
RegistrationRejected {
status_code: u16,
reason: &'static str,
},
#[error("Agent Identity 注册响应无效")]
InvalidRegistrationResponse,
#[error("Agent Identity 密钥生成失败")]
@@ -68,6 +71,11 @@ pub enum CodexAgentIdentityEnrollmentError {
InvalidTaskRegistrationResponse,
}
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
pub struct CodexAgentIdentityRegistrationOptions {
pub is_fedramp_account: bool,
}
#[derive(Clone)]
struct AgentIdentityCredentials {
runtime_id: String,
@@ -598,16 +606,29 @@ fn agent_runtime_id_from_registration_response(body: &str) -> Result<String, ()>
.ok_or(())
}
fn agent_identity_registration_rejection_reason(status_code: u16) -> &'static str {
match status_code {
401 => "：AccessToken 无效或已过期",
403 => {
"：账户、工作区或出口区域未获服务端 Agent Identity 资格；该资格无法由 Aether 本地开通"
}
429 => "：请求过于频繁，请稍后重试",
_ => "",
}
}
/// Uses a ChatGPT access token once to register a fresh Agent Identity. The returned config
/// contains only the generated signing credentials and is deliberately free of the access token.
pub async fn register_codex_agent_identity_from_access_token(
executor: &dyn OAuthHttpExecutor,
access_token: &str,
options: CodexAgentIdentityRegistrationOptions,
network: OAuthNetworkContext,
) -> Result<Map<String, Value>, CodexAgentIdentityEnrollmentError> {
register_codex_agent_identity_from_access_token_with_auth_api_base_url(
executor,
access_token,
options,
network,
CODEX_AGENT_IDENTITY_AUTH_API_BASE_URL,
)
@@ -647,6 +668,7 @@ async fn create_codex_agent_identity_from_session_token_with_auth_api_base_url(
let mut auth_config = register_codex_agent_identity_from_access_token_with_auth_api_base_url(
executor,
session_token,
CodexAgentIdentityRegistrationOptions::default(),
network.clone(),
auth_api_base_url,
)
@@ -696,6 +718,7 @@ async fn create_codex_agent_identity_from_session_token_with_auth_api_base_url(
async fn register_codex_agent_identity_from_access_token_with_auth_api_base_url(
executor: &dyn OAuthHttpExecutor,
access_token: &str,
options: CodexAgentIdentityRegistrationOptions,
network: OAuthNetworkContext,
auth_api_base_url: &str,
) -> Result<Map<String, Value>, CodexAgentIdentityEnrollmentError> {
@@ -712,35 +735,41 @@ async fn register_codex_agent_identity_from_access_token_with_auth_api_base_url(
let agent_public_key = agent_identity_ssh_public_key(&signing_key);
let registration_url = agent_registration_url(auth_api_base_url)
.map_err(|_| CodexAgentIdentityEnrollmentError::RegistrationRequestFailed)?;
let mut headers = BTreeMap::from([
("accept".to_string(), "application/json".to_string()),
("content-type".to_string(), "application/json".to_string()),
(
"authorization".to_string(),
format!("Bearer {access_token}"),
),
(
"user-agent".to_string(),
aether_ai_formats::CODEX_CLIENT_USER_AGENT.to_string(),
),
(
"originator".to_string(),
aether_ai_formats::CODEX_CLIENT_ORIGINATOR.to_string(),
),
]);
if options.is_fedramp_account {
headers.insert("x-openai-fedramp".to_string(), "true".to_string());
}
let response = executor
.execute(OAuthHttpRequest {
request_id: CODEX_AGENT_IDENTITY_AGENT_REGISTRATION_REQUEST_ID.to_string(),
method: reqwest::Method::POST,
url: registration_url,
headers: BTreeMap::from([
("accept".to_string(), "application/json".to_string()),
("content-type".to_string(), "application/json".to_string()),
(
"authorization".to_string(),
format!("Bearer {access_token}"),
),
(
"user-agent".to_string(),
aether_ai_formats::CODEX_CLIENT_USER_AGENT.to_string(),
),
(
"originator".to_string(),
aether_ai_formats::CODEX_CLIENT_ORIGINATOR.to_string(),
),
]),
headers,
content_type: Some("application/json".to_string()),
json_body: Some(json!({
"abom": {
"agent_version": aether_ai_formats::CODEX_CLIENT_VERSION,
"agent_harness_id": CODEX_AGENT_IDENTITY_AGENT_HARNESS_ID,
"running_location": CODEX_AGENT_IDENTITY_RUNNING_LOCATION,
"running_location": format!("cli-{}", std::env::consts::OS),
},
"agent_public_key": agent_public_key,
"capabilities": [CODEX_AGENT_IDENTITY_RESPONSES_API_CAPABILITY],
"ttl": null,
})),
body_bytes: None,
network,
@@ -751,6 +780,7 @@ async fn register_codex_agent_identity_from_access_token_with_auth_api_base_url(
if !(200..300).contains(&response.status_code) {
return Err(CodexAgentIdentityEnrollmentError::RegistrationRejected {
status_code: response.status_code,
reason: agent_identity_registration_rejection_reason(response.status_code),
});
}
let agent_runtime_id = agent_runtime_id_from_registration_response(response.body_text.as_str())
@@ -1034,7 +1064,8 @@ mod tests {
register_codex_agent_identity_from_access_token_with_auth_api_base_url,
task_id_from_registration_response, validate_codex_agent_identity_auth_config,
with_agent_identity_task_id, CodexAgentIdentityEnrollmentError,
CodexAgentIdentityRefreshAdapter, CODEX_AGENT_IDENTITY_CACHED_ENTRY_PROVIDER_TYPE,
CodexAgentIdentityRefreshAdapter, CodexAgentIdentityRegistrationOptions,
CODEX_AGENT_IDENTITY_CACHED_ENTRY_PROVIDER_TYPE,
};
use crate::oauth_refresh::{
LocalOAuthHttpExecutor, LocalOAuthHttpRequest, LocalOAuthHttpResponse,
@@ -1416,6 +1447,20 @@ mod tests {
.and_then(|abom| abom.get("agent_harness_id")),
Some(&json!("codex-cli"))
);
assert_eq!(
requests[0]
.json_body
.as_ref()
.and_then(|body| body.get("capabilities")),
Some(&json!(["responsesapi"]))
);
assert_eq!(
requests[0]
.json_body
.as_ref()
.and_then(|body| body.get("ttl")),
Some(&json!(null))
);
assert!(requests[0]
.json_body
.as_ref()
@@ -1443,6 +1488,9 @@ mod tests {
let config = register_codex_agent_identity_from_access_token_with_auth_api_base_url(
&executor,
"access-token-for-test-only",
CodexAgentIdentityRegistrationOptions {
is_fedramp_account: true,
},
OAuthNetworkContext::direct_identity(),
"https://auth.test/api/accounts",
)
@@ -1458,6 +1506,21 @@ mod tests {
let requests = requests.lock().expect("recording lock should hold");
assert_eq!(requests.len(), 1);
assert!(requests[0].url.ends_with("/v1/agent/register"));
assert_eq!(
requests[0]
.headers
.get("x-openai-fedramp")
.map(String::as_str),
Some("true")
);
assert_eq!(
requests[0]
.json_body
.as_ref()
.and_then(|body| body.get("abom"))
.and_then(|abom| abom.get("running_location")),
Some(&json!(format!("cli-{}", std::env::consts::OS)))
);
}
#[tokio::test]
@@ -1465,7 +1528,7 @@ mod tests {
let executor = RecordingEnrollmentExecutor {
requests: Arc::new(Mutex::new(Vec::new())),
responses: Arc::new(Mutex::new(vec![OAuthHttpResponse {
status_code: 401,
status_code: 403,
body_text: r#"{"detail":"session-token-for-test-only"}"#.to_string(),
json_body: None,
}])),
@@ -1482,7 +1545,10 @@ mod tests {
assert_eq!(
error,
CodexAgentIdentityEnrollmentError::RegistrationRejected { status_code: 401 }
CodexAgentIdentityEnrollmentError::RegistrationRejected {
status_code: 403,
reason: "：账户、工作区或出口区域未获服务端 Agent Identity 资格；该资格无法由 Aether 本地开通",
}
);
assert!(!error.to_string().contains("session-token-for-test-only"));
assert!(!error.to_string().contains("detail"));
crates/aether-provider/transport/src/lib.rs
+4 -3
View File
@@ -43,9 +43,10 @@ pub use agent_identity::{
is_codex_agent_identity_cached_entry, is_codex_agent_identity_invalid_task_response,
is_codex_agent_identity_transport, register_codex_agent_identity_from_access_token,
validate_codex_agent_identity_auth_config, CodexAgentIdentityEnrollmentError,
CodexAgentIdentityRefreshAdapter, CODEX_AGENT_IDENTITY_AGENT_REGISTRATION_REQUEST_ID,
CODEX_AGENT_IDENTITY_AUTH_MODE, CODEX_AGENT_IDENTITY_CACHED_ENTRY_PROVIDER_TYPE,
CODEX_AGENT_IDENTITY_PROVIDER_TYPE, CODEX_AGENT_IDENTITY_TASK_REGISTRATION_REQUEST_ID,
CodexAgentIdentityRefreshAdapter, CodexAgentIdentityRegistrationOptions,
CODEX_AGENT_IDENTITY_AGENT_REGISTRATION_REQUEST_ID, CODEX_AGENT_IDENTITY_AUTH_MODE,
CODEX_AGENT_IDENTITY_CACHED_ENTRY_PROVIDER_TYPE, CODEX_AGENT_IDENTITY_PROVIDER_TYPE,
CODEX_AGENT_IDENTITY_TASK_REGISTRATION_REQUEST_ID,
};
pub use anthropic_compat::{
resolve_anthropic_compatibility_profile, validate_anthropic_compatibility_profile_config,
frontend/src/features/providers/components/ProviderDetailDrawer.vue
+15 -3
View File
@@ -189,10 +189,10 @@
class="grid gap-3"
:class="isCodexTeamPlan(key) ? 'grid-cols-2' : 'grid-cols-1'"
>
<!-- 周限额 -->
<!-- 主限额 -->
<ProviderQuotaProgressRow
v-if="getCodexQuotaDisplay(key)?.primary_used_percent !== undefined"
:label="legacyT('周限额')"
:label="legacyT(getCodexPrimaryQuotaLabel(key))"
:used-percent="getCodexQuotaDisplay(key)?.primary_used_percent || 0"
:remaining-percent="toCodexRemainingPercent(getCodexQuotaDisplay(key)?.primary_used_percent)"
:meter-class="getQuotaRemainingClass(getCodexQuotaDisplay(key)?.primary_used_percent || 0)"
@@ -1021,6 +1021,10 @@ import { formatApiFormatShort } from '@/api/endpoints/types/api-format'
import { isOAuthAccountProviderType, isKeyManagedProviderType } from '../utils/providerTypeUtils'
import { getOAuthOrgBadge } from '@/utils/oauthIdentity'
import { getOAuthRefreshFeedback } from '@/utils/oauthRefreshFeedback'
import {
getCodexPrimaryQuotaWindow,
getCodexQuotaWindowLimitLabel,
} from '@/utils/codexQuotaWindow'
import { formatCompactNumber } from '@/utils/format'
import {
canEditOAuthCredential,
@@ -1945,7 +1949,7 @@ function getCodexQuotaDisplayFromSnapshot(quota: QuotaStatusSnapshot | null | un
if (updatedAt !== undefined) display.updated_at = updatedAt
if (quota.plan_type) display.plan_type = quota.plan_type
const primaryWindow = getQuotaWindow(quota, 'weekly')
const primaryWindow = getCodexPrimaryQuotaWindow(quota.windows)
const primaryUsedPercent = getQuotaWindowUsedPercent(primaryWindow)
if (primaryUsedPercent !== undefined) display.primary_used_percent = primaryUsedPercent
const primaryResetAt = getQuotaWindowResetAt(primaryWindow)
@@ -2004,6 +2008,14 @@ function getCodexQuotaDisplay(key: EndpointAPIKey): CodexUpstreamMetadata | null
return mergeCodexQuotaDisplays(snapshotDisplay, metadataDisplay)
}
function getCodexPrimaryQuotaLabel(key: EndpointAPIKey): string {
return getCodexQuotaWindowLimitLabel({
code: 'weekly',
label: '周',
window_minutes: getCodexQuotaDisplay(key)?.primary_window_minutes,
}) || '周限额'
}
function hasCodexQuotaDisplayData(key: EndpointAPIKey): boolean {
const codex = getCodexQuotaDisplay(key)
return !!codex && (
frontend/src/i18n/messages.ts
+1
View File
@@ -1831,6 +1831,7 @@ const legacyExactEnglishMessages: Record<string, string> = {
'账号配额': 'Account quota',
'模型配额': 'Model quota',
'周限额': 'Weekly limit',
'月限额': 'Monthly limit',
'5H限额': '5H limit',
'Spark 周': 'Spark weekly',
'Spark 5H': 'Spark 5H',
frontend/src/utils/__tests__/codexQuotaWindow.spec.ts
+21 -1
View File
@@ -1,6 +1,10 @@
import { describe, expect, it } from 'vitest'
import { getCodexQuotaWindowPresentation } from '../codexQuotaWindow'
import {
getCodexPrimaryQuotaWindow,
getCodexQuotaWindowLimitLabel,
getCodexQuotaWindowPresentation,
} from '../codexQuotaWindow'
describe('getCodexQuotaWindowPresentation', () => {
it.each([
@@ -25,6 +29,22 @@ describe('getCodexQuotaWindowPresentation', () => {
expect(windows.sort((a, b) => a.sortOrder - b.sortOrder).map(item => item.label)).toEqual(['5H', '周'])
})
it('builds the provider limit label from the actual window duration', () => {
expect(getCodexQuotaWindowLimitLabel({ code: 'weekly', window_minutes: 10_080 })).toBe('周限额')
expect(getCodexQuotaWindowLimitLabel({ code: 'weekly', window_minutes: 43_800 })).toBe('月限额')
})
it('selects a monthly primary window over a zero-minute weekly placeholder', () => {
const monthly = { code: 'monthly', label: '月', window_minutes: 43_800, used_ratio: 0.02 }
const selected = getCodexPrimaryQuotaWindow([
monthly,
{ code: 'weekly', label: '周', window_minutes: 0, used_ratio: 1 },
])
expect(selected).toEqual(monthly)
expect(getCodexQuotaWindowLimitLabel(selected!)).toBe('月限额')
})
it('drops zero-minute placeholder windows', () => {
expect(getCodexQuotaWindowPresentation({
code: 'weekly',
frontend/src/utils/codexQuotaWindow.ts
+17
View File
@@ -60,3 +60,20 @@ export function getCodexQuotaWindowPresentation(
sortOrder: (isSpark ? 10_000_000 : 0) + (hasExplicitWindowMinutes ? windowMinutes : fallbackOrder),
}
}
export function getCodexQuotaWindowLimitLabel(window: QuotaWindowSnapshot): string | null {
const presentation = getCodexQuotaWindowPresentation(window)
return presentation ? `${presentation.label}限额` : null
}
export function getCodexPrimaryQuotaWindow(
windows: QuotaWindowSnapshot[] | null | undefined,
): QuotaWindowSnapshot | null {
if (!Array.isArray(windows)) return null
for (const code of ['monthly', 'weekly']) {
const window = windows.find(candidate => String(candidate.code || '').trim().toLowerCase() === code)
if (window && getCodexQuotaWindowPresentation(window)) return window
}
return null
}