mirror of
https://github.com/fawney19/Aether.git
synced 2026-10-11 03:39:49 +08:00
fix(provider): harden Agent Identity OAuth lifecycle
This commit is contained in:
@@ -69,11 +69,16 @@ impl<'a> AdminAppState<'a> {
|
||||
)
|
||||
}
|
||||
|
||||
pub(crate) fn masked_catalog_api_key(
|
||||
pub(crate) fn masked_catalog_api_key_for_provider(
|
||||
&self,
|
||||
key: &aether_data_contracts::repository::provider_catalog::StoredProviderCatalogKey,
|
||||
provider_type: &str,
|
||||
) -> String {
|
||||
crate::handlers::admin::shared::masked_catalog_api_key(self.app, key)
|
||||
crate::handlers::admin::shared::masked_catalog_api_key_for_provider(
|
||||
self.app,
|
||||
key,
|
||||
provider_type,
|
||||
)
|
||||
}
|
||||
|
||||
pub(crate) async fn build_admin_provider_keys_payload(
|
||||
|
||||
@@ -227,6 +227,7 @@ impl<'a> AdminAppState<'a> {
|
||||
.compare_and_update_provider_catalog_key_adaptive_state(
|
||||
&ProviderCatalogKeyAdaptiveStateUpdate {
|
||||
key_id: key_id.to_string(),
|
||||
expected_encrypted_auth_config: None,
|
||||
expected,
|
||||
next,
|
||||
status_snapshot_patch: serde_json::json!({
|
||||
@@ -257,6 +258,33 @@ impl<'a> AdminAppState<'a> {
|
||||
) -> Result<
|
||||
Option<aether_data_contracts::repository::provider_catalog::StoredProviderCatalogKey>,
|
||||
GatewayError,
|
||||
> {
|
||||
self.reset_provider_catalog_key_recovery_state_inner(key_id, None)
|
||||
.await
|
||||
}
|
||||
|
||||
pub(crate) async fn reset_provider_catalog_key_recovery_state_fenced(
|
||||
&self,
|
||||
key_id: &str,
|
||||
expected_encrypted_auth_config: &str,
|
||||
) -> Result<
|
||||
Option<aether_data_contracts::repository::provider_catalog::StoredProviderCatalogKey>,
|
||||
GatewayError,
|
||||
> {
|
||||
self.reset_provider_catalog_key_recovery_state_inner(
|
||||
key_id,
|
||||
Some(expected_encrypted_auth_config),
|
||||
)
|
||||
.await
|
||||
}
|
||||
|
||||
async fn reset_provider_catalog_key_recovery_state_inner(
|
||||
&self,
|
||||
key_id: &str,
|
||||
expected_auth_config: Option<&str>,
|
||||
) -> Result<
|
||||
Option<aether_data_contracts::repository::provider_catalog::StoredProviderCatalogKey>,
|
||||
GatewayError,
|
||||
> {
|
||||
use aether_data_contracts::repository::provider_catalog::ProviderCatalogKeyHealthStateUpdate;
|
||||
|
||||
@@ -271,6 +299,11 @@ impl<'a> AdminAppState<'a> {
|
||||
else {
|
||||
return Ok(None);
|
||||
};
|
||||
if expected_auth_config
|
||||
.is_some_and(|expected| current.encrypted_auth_config.as_deref() != Some(expected))
|
||||
{
|
||||
return Ok(None);
|
||||
}
|
||||
if current.health_by_format.as_ref() == Some(&empty)
|
||||
&& current.circuit_breaker_by_format.as_ref() == Some(&empty)
|
||||
{
|
||||
@@ -282,6 +315,7 @@ impl<'a> AdminAppState<'a> {
|
||||
.compare_and_update_provider_catalog_key_health_state(
|
||||
&ProviderCatalogKeyHealthStateUpdate {
|
||||
key_id: key_id.to_string(),
|
||||
expected_encrypted_auth_config: expected_auth_config.map(ToOwned::to_owned),
|
||||
expected_health_by_format: current.health_by_format,
|
||||
expected_circuit_breaker_by_format: current.circuit_breaker_by_format,
|
||||
health_by_format: Some(empty.clone()),
|
||||
@@ -299,15 +333,24 @@ impl<'a> AdminAppState<'a> {
|
||||
"provider key {key_id} health state changed repeatedly while resetting OAuth recovery state"
|
||||
)));
|
||||
}
|
||||
if !self.reset_provider_catalog_key_error_count(key_id).await? {
|
||||
if expected_auth_config.is_none()
|
||||
&& !self.reset_provider_catalog_key_error_count(key_id).await?
|
||||
{
|
||||
return Ok(None);
|
||||
}
|
||||
|
||||
Ok(self
|
||||
let current = self
|
||||
.read_provider_catalog_keys_by_ids(&[key_id.to_string()])
|
||||
.await?
|
||||
.into_iter()
|
||||
.next())
|
||||
.next();
|
||||
if current.as_ref().is_some_and(|key| {
|
||||
expected_auth_config
|
||||
.is_some_and(|expected| key.encrypted_auth_config.as_deref() != Some(expected))
|
||||
}) {
|
||||
return Ok(None);
|
||||
}
|
||||
Ok(current)
|
||||
}
|
||||
|
||||
pub(crate) async fn update_provider_catalog_key_status_snapshot(
|
||||
|
||||
@@ -90,6 +90,7 @@ impl<'a> AdminAppState<'a> {
|
||||
provider_id: &str,
|
||||
provider_type: &str,
|
||||
pkce_verifier: Option<&str>,
|
||||
expected_encrypted_auth_config: Option<&str>,
|
||||
) -> Result<String, GatewayError> {
|
||||
let nonce = aether_admin::provider::state::generate_provider_oauth_nonce();
|
||||
let payload = json!({
|
||||
@@ -98,6 +99,7 @@ impl<'a> AdminAppState<'a> {
|
||||
"provider_id": provider_id,
|
||||
"provider_type": provider_type,
|
||||
"pkce_verifier": pkce_verifier,
|
||||
"expected_encrypted_auth_config": expected_encrypted_auth_config,
|
||||
"created_at": aether_admin::provider::state::current_unix_secs(),
|
||||
});
|
||||
let key = provider_oauth_state_storage_key(&nonce);
|
||||
|
||||
@@ -16,6 +16,10 @@ 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 {
|
||||
provider.provider_type.trim().eq_ignore_ascii_case("codex")
|
||||
}
|
||||
|
||||
impl<'a> AdminAppState<'a> {
|
||||
pub(crate) async fn clear_admin_provider_pool_cooldown(&self, provider_id: &str, key_id: &str) {
|
||||
crate::handlers::admin::provider::pool::runtime::clear_admin_provider_pool_cooldown(
|
||||
@@ -368,6 +372,13 @@ 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?
|
||||
@@ -408,6 +419,10 @@ 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?;
|
||||
@@ -849,3 +864,26 @@ impl<'a> AdminAppState<'a> {
|
||||
.into_response())
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod automatic_cleanup_tests {
|
||||
use super::provider_skips_automatic_key_cleanup;
|
||||
use aether_data_contracts::repository::provider_catalog::StoredProviderCatalogProvider;
|
||||
|
||||
fn provider(provider_type: &str) -> StoredProviderCatalogProvider {
|
||||
StoredProviderCatalogProvider::new(
|
||||
format!("provider-{provider_type}"),
|
||||
provider_type.to_string(),
|
||||
None,
|
||||
provider_type.to_string(),
|
||||
)
|
||||
.expect("provider should build")
|
||||
}
|
||||
|
||||
#[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")));
|
||||
}
|
||||
}
|
||||
|
||||
@@ -153,6 +153,7 @@ impl<'a> AdminAppState<'a> {
|
||||
.compare_and_update_provider_catalog_key_adaptive_state(
|
||||
&ProviderCatalogKeyAdaptiveStateUpdate {
|
||||
key_id: key.id.clone(),
|
||||
expected_encrypted_auth_config: None,
|
||||
expected,
|
||||
next,
|
||||
status_snapshot_patch: json!({
|
||||
|
||||
Reference in New Issue
Block a user