mirror of
https://github.com/fawney19/Aether.git
synced 2026-09-10 05:00:19 +08:00
feat(providers): add credential-fenced OAuth key cleanup
This commit is contained in:
@@ -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,
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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();
|
||||
}
|
||||
|
||||
@@ -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 {
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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());
|
||||
}
|
||||
}
|
||||
|
||||
@@ -3336,7 +3336,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();
|
||||
|
||||
@@ -3355,17 +3355,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]
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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()];
|
||||
|
||||
@@ -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();
|
||||
|
||||
Reference in New Issue
Block a user