feat(providers): add credential-fenced OAuth key cleanup

This commit is contained in:
ZheFox
2026-07-28 11:32:11 +08:00
parent 531cf11025
commit f0f87b56a3
17 changed files with 702 additions and 142 deletions
@@ -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
}
@@ -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
}
@@ -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]
@@ -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,
};
@@ -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,
@@ -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");
@@ -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,