mirror of
https://github.com/fawney19/Aether.git
synced 2026-09-11 05:30:19 +08:00
fix: refine oauth auto-removal behavior
This commit is contained in:
+59
-2
@@ -1,7 +1,10 @@
|
||||
use super::super::super::errors::{
|
||||
merge_provider_oauth_refresh_failure_reason, normalize_provider_oauth_refresh_error_message,
|
||||
};
|
||||
use super::super::super::quota::shared::persist_provider_quota_refresh_state;
|
||||
use super::super::super::quota::shared::{
|
||||
persist_provider_quota_refresh_state, provider_auto_remove_banned_keys,
|
||||
should_auto_remove_oauth_refresh_failed_key,
|
||||
};
|
||||
use super::super::super::runtime::refresh_provider_oauth_account_state_after_update;
|
||||
use super::helpers::{self, RefreshDispatch, RefreshRequestContext, RefreshSuccessContext};
|
||||
use super::response;
|
||||
@@ -76,6 +79,42 @@ pub(super) async fn execute_admin_provider_oauth_refresh(
|
||||
None,
|
||||
)
|
||||
.await?;
|
||||
if provider_auto_remove_banned_keys(provider.config.as_ref()) {
|
||||
let now_unix_secs = helpers::unix_now_secs();
|
||||
let auto_removed = state
|
||||
.cleanup_provider_catalog_key_if_current(
|
||||
&provider,
|
||||
&key_id,
|
||||
|latest_key| {
|
||||
should_auto_remove_oauth_refresh_failed_key(
|
||||
latest_key,
|
||||
now_unix_secs,
|
||||
)
|
||||
},
|
||||
)
|
||||
.await?;
|
||||
if auto_removed {
|
||||
tracing::info!(
|
||||
trace_id = %trace_id,
|
||||
key_id = %key_id,
|
||||
provider_id = %provider.id,
|
||||
provider_type = %provider_type,
|
||||
event_name = "auto_removed_oauth_refresh_failed",
|
||||
"gateway manual provider oauth refresh auto-removed expired key"
|
||||
);
|
||||
return Ok(RefreshDispatch::Respond(
|
||||
response::oauth_refresh_auto_removed_response(&error_reason),
|
||||
));
|
||||
}
|
||||
}
|
||||
tracing::info!(
|
||||
trace_id = %trace_id,
|
||||
key_id = %key_id,
|
||||
provider_id = %provider.id,
|
||||
provider_type = %provider_type,
|
||||
event_name = "refresh_failed_retained",
|
||||
"gateway manual provider oauth refresh failure retained key"
|
||||
);
|
||||
}
|
||||
}
|
||||
return Ok(RefreshDispatch::Respond(
|
||||
@@ -124,9 +163,27 @@ pub(super) async fn execute_admin_provider_oauth_refresh(
|
||||
};
|
||||
|
||||
if !helpers::key_is_account_blocked(&key, OAUTH_ACCOUNT_BLOCK_PREFIX) {
|
||||
let _ = state
|
||||
let previous_oauth_refresh_issue = key
|
||||
.oauth_invalid_reason
|
||||
.as_deref()
|
||||
.is_some_and(|reason| {
|
||||
reason.lines().map(str::trim).any(|line| {
|
||||
line.starts_with("[OAUTH_EXPIRED]") || line.starts_with("[REFRESH_FAILED]")
|
||||
})
|
||||
});
|
||||
let cleared = state
|
||||
.clear_provider_catalog_key_oauth_invalid_marker(&key_id)
|
||||
.await?;
|
||||
if cleared && previous_oauth_refresh_issue {
|
||||
tracing::info!(
|
||||
trace_id = %trace_id,
|
||||
key_id = %key_id,
|
||||
provider_id = %provider.id,
|
||||
provider_type = %provider_type,
|
||||
event_name = "refresh_fixed",
|
||||
"gateway manual provider oauth refresh cleared oauth invalid marker"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
let refreshed_key = state
|
||||
|
||||
@@ -24,6 +24,15 @@ pub(super) fn oauth_refresh_failed_bad_request_response(
|
||||
)
|
||||
}
|
||||
|
||||
pub(super) fn oauth_refresh_auto_removed_response(error_reason: impl AsRef<str>) -> Response<Body> {
|
||||
Json(json!({
|
||||
"status": "auto_removed",
|
||||
"message": "已自动删除",
|
||||
"detail": format!("Token 刷新失败且 Access Token 已过期,已自动删除:{}", error_reason.as_ref()),
|
||||
}))
|
||||
.into_response()
|
||||
}
|
||||
|
||||
pub(super) fn oauth_refresh_failed_service_unavailable_response(
|
||||
error_reason: impl Into<String>,
|
||||
) -> Response<Body> {
|
||||
|
||||
@@ -1,6 +1,4 @@
|
||||
use crate::handlers::admin::provider::shared::payloads::{
|
||||
OAUTH_ACCOUNT_BLOCK_PREFIX, OAUTH_EXPIRED_PREFIX,
|
||||
};
|
||||
use crate::handlers::admin::provider::shared::payloads::OAUTH_ACCOUNT_BLOCK_PREFIX;
|
||||
use axum::{
|
||||
body::Body,
|
||||
http,
|
||||
@@ -146,9 +144,6 @@ pub(crate) fn merge_provider_oauth_refresh_failure_reason(
|
||||
if current_reason.is_empty() {
|
||||
return Some(refresh_reason.to_string());
|
||||
}
|
||||
if current_reason.starts_with(OAUTH_EXPIRED_PREFIX) {
|
||||
return Some(current_reason.to_string());
|
||||
}
|
||||
if oauth_invalid_reason_is_account_level_block(Some(current_reason)) {
|
||||
return None;
|
||||
}
|
||||
@@ -190,13 +185,13 @@ mod tests {
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn refresh_failure_does_not_replace_access_token_expired_marker() {
|
||||
fn refresh_failure_replaces_access_token_expired_marker() {
|
||||
assert_eq!(
|
||||
merge_provider_oauth_refresh_failure_reason(
|
||||
Some("[OAUTH_EXPIRED] access token invalid"),
|
||||
"[REFRESH_FAILED] Token 续期失败 (401): refresh_token 无效",
|
||||
),
|
||||
Some("[OAUTH_EXPIRED] access token invalid".to_string()),
|
||||
Some("[REFRESH_FAILED] Token 续期失败 (401): refresh_token 无效".to_string()),
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -189,7 +189,6 @@ pub(crate) async fn update_existing_provider_oauth_catalog_key(
|
||||
updated.encrypted_api_key = Some(encrypted_api_key);
|
||||
updated.encrypted_auth_config = Some(encrypted_auth_config);
|
||||
updated.api_formats = provider_oauth_catalog_key_api_formats(provider_type, api_formats);
|
||||
updated.is_active = true;
|
||||
updated.expires_at_unix_secs = expires_at_unix_secs;
|
||||
updated.oauth_invalid_at_unix_secs = None;
|
||||
updated.oauth_invalid_reason = None;
|
||||
|
||||
@@ -44,6 +44,15 @@ fn merge_codex_quota_metadata(
|
||||
serde_json::Value::Object(merged)
|
||||
}
|
||||
|
||||
fn codex_oauth_refresh_issue_reason(reason: Option<&str>) -> bool {
|
||||
reason.is_some_and(|reason| {
|
||||
reason
|
||||
.lines()
|
||||
.map(str::trim)
|
||||
.any(|line| line.starts_with("[OAUTH_EXPIRED]") || line.starts_with("[REFRESH_FAILED]"))
|
||||
})
|
||||
}
|
||||
|
||||
pub(crate) async fn refresh_codex_provider_quota_locally(
|
||||
state: &AdminAppState<'_>,
|
||||
provider: &StoredProviderCatalogProvider,
|
||||
@@ -56,8 +65,13 @@ pub(crate) async fn refresh_codex_provider_quota_locally(
|
||||
let mut success_count = 0usize;
|
||||
let mut failed_count = 0usize;
|
||||
let mut auto_removed_count = 0usize;
|
||||
let mut refresh_fixed_count = 0usize;
|
||||
let mut refresh_failed_retained_count = 0usize;
|
||||
let mut auto_removed_hard_banned_count = 0usize;
|
||||
|
||||
for key in keys {
|
||||
let had_oauth_refresh_issue =
|
||||
codex_oauth_refresh_issue_reason(key.oauth_invalid_reason.as_deref());
|
||||
let transport = match state
|
||||
.read_provider_transport_snapshot(&provider.id, &endpoint.id, &key.id)
|
||||
.await?
|
||||
@@ -261,13 +275,9 @@ pub(crate) async fn refresh_codex_provider_quota_locally(
|
||||
}
|
||||
}
|
||||
|
||||
let auto_removed = auto_remove_abnormal_keys
|
||||
let auto_remove_candidate = auto_remove_abnormal_keys
|
||||
&& should_auto_remove_structured_reason(oauth_invalid_reason.as_deref());
|
||||
if auto_removed {
|
||||
if state.delete_provider_catalog_key(&key.id).await? {
|
||||
auto_removed_count += 1;
|
||||
}
|
||||
} else if !persist_provider_quota_refresh_state(
|
||||
let persisted = persist_provider_quota_refresh_state(
|
||||
state,
|
||||
&key.id,
|
||||
metadata_update.as_ref(),
|
||||
@@ -275,8 +285,8 @@ pub(crate) async fn refresh_codex_provider_quota_locally(
|
||||
oauth_invalid_reason.clone(),
|
||||
None,
|
||||
)
|
||||
.await?
|
||||
{
|
||||
.await?;
|
||||
if !persisted {
|
||||
failed_count += 1;
|
||||
results.push(json!({
|
||||
"key_id": key.id,
|
||||
@@ -286,6 +296,29 @@ pub(crate) async fn refresh_codex_provider_quota_locally(
|
||||
}));
|
||||
continue;
|
||||
}
|
||||
let auto_removed = if auto_remove_candidate {
|
||||
state
|
||||
.cleanup_provider_catalog_key_if_current(provider, &key.id, |latest_key| {
|
||||
should_auto_remove_structured_reason(latest_key.oauth_invalid_reason.as_deref())
|
||||
})
|
||||
.await?
|
||||
} else {
|
||||
false
|
||||
};
|
||||
if auto_removed {
|
||||
auto_removed_count += 1;
|
||||
auto_removed_hard_banned_count += 1;
|
||||
}
|
||||
let refresh_fixed =
|
||||
status == "success" && had_oauth_refresh_issue && oauth_invalid_reason.is_none();
|
||||
if refresh_fixed {
|
||||
refresh_fixed_count += 1;
|
||||
}
|
||||
let refresh_failed_retained =
|
||||
status != "success" && oauth_invalid_reason.is_some() && !auto_removed;
|
||||
if refresh_failed_retained {
|
||||
refresh_failed_retained_count += 1;
|
||||
}
|
||||
|
||||
if status == "success" {
|
||||
success_count += 1;
|
||||
@@ -321,6 +354,13 @@ pub(crate) async fn refresh_codex_provider_quota_locally(
|
||||
}
|
||||
if auto_removed {
|
||||
payload.insert("auto_removed".to_string(), json!(true));
|
||||
payload.insert("auto_removed_hard_banned".to_string(), json!(true));
|
||||
}
|
||||
if refresh_fixed {
|
||||
payload.insert("refresh_fixed".to_string(), json!(true));
|
||||
}
|
||||
if refresh_failed_retained {
|
||||
payload.insert("refresh_failed_retained".to_string(), json!(true));
|
||||
}
|
||||
results.push(serde_json::Value::Object(payload));
|
||||
}
|
||||
@@ -331,5 +371,8 @@ pub(crate) async fn refresh_codex_provider_quota_locally(
|
||||
"total": results.len(),
|
||||
"results": results,
|
||||
"auto_removed": auto_removed_count,
|
||||
"refresh_fixed": refresh_fixed_count,
|
||||
"refresh_failed_retained": refresh_failed_retained_count,
|
||||
"auto_removed_hard_banned": auto_removed_hard_banned_count,
|
||||
})))
|
||||
}
|
||||
|
||||
@@ -44,7 +44,7 @@ pub(super) fn default_provider_quota_execution_timeouts(
|
||||
}
|
||||
}
|
||||
|
||||
pub(super) fn provider_auto_remove_banned_keys(config: Option<&serde_json::Value>) -> bool {
|
||||
pub(crate) fn provider_auto_remove_banned_keys(config: Option<&serde_json::Value>) -> bool {
|
||||
admin_provider_quota_pure::provider_auto_remove_banned_keys(config)
|
||||
}
|
||||
|
||||
@@ -52,6 +52,13 @@ pub(super) fn should_auto_remove_structured_reason(reason: Option<&str>) -> bool
|
||||
admin_provider_quota_pure::should_auto_remove_structured_reason(reason)
|
||||
}
|
||||
|
||||
pub(crate) fn should_auto_remove_oauth_refresh_failed_key(
|
||||
key: &StoredProviderCatalogKey,
|
||||
now_unix_secs: u64,
|
||||
) -> bool {
|
||||
admin_provider_quota_pure::should_auto_remove_oauth_refresh_failed_key(key, now_unix_secs)
|
||||
}
|
||||
|
||||
pub(crate) fn normalize_string_id_list(values: Option<Vec<String>>) -> Option<Vec<String>> {
|
||||
admin_provider_quota_pure::normalize_string_id_list(values)
|
||||
}
|
||||
|
||||
@@ -1,5 +1,7 @@
|
||||
use super::*;
|
||||
use aether_data_contracts::repository::provider_catalog::StoredProviderCatalogProvider;
|
||||
use aether_data_contracts::repository::provider_catalog::{
|
||||
StoredProviderCatalogKey, StoredProviderCatalogProvider,
|
||||
};
|
||||
use axum::{
|
||||
body::Body,
|
||||
http,
|
||||
@@ -264,6 +266,41 @@ impl<'a> AdminAppState<'a> {
|
||||
Ok(affected)
|
||||
}
|
||||
|
||||
pub(crate) async fn cleanup_provider_catalog_key_if_current<F>(
|
||||
&self,
|
||||
provider: &StoredProviderCatalogProvider,
|
||||
key_id: &str,
|
||||
should_delete: F,
|
||||
) -> Result<bool, GatewayError>
|
||||
where
|
||||
F: FnOnce(&StoredProviderCatalogKey) -> bool,
|
||||
{
|
||||
let key_ids = [key_id.to_string()];
|
||||
let Some(key) = self
|
||||
.read_provider_catalog_keys_by_ids(&key_ids)
|
||||
.await?
|
||||
.into_iter()
|
||||
.next()
|
||||
else {
|
||||
return Ok(false);
|
||||
};
|
||||
if key.provider_id != provider.id || !should_delete(&key) {
|
||||
return Ok(false);
|
||||
}
|
||||
|
||||
self.clear_admin_provider_pool_cooldown(&provider.id, &key.id)
|
||||
.await;
|
||||
self.reset_admin_provider_pool_cost(&provider.id, &key.id)
|
||||
.await;
|
||||
let deleted = self.delete_provider_catalog_key(&key.id).await?;
|
||||
if deleted {
|
||||
let deleted_key_ids = [key.id.clone()];
|
||||
self.cleanup_deleted_provider_catalog_refs(&provider.id, &[], &deleted_key_ids)
|
||||
.await?;
|
||||
}
|
||||
Ok(deleted)
|
||||
}
|
||||
|
||||
pub(crate) async fn build_admin_pool_batch_action_response(
|
||||
&self,
|
||||
provider_id: &str,
|
||||
|
||||
@@ -98,6 +98,10 @@ enum AccountSelfCheckOutcome {
|
||||
status_code: Option<u16>,
|
||||
message: String,
|
||||
},
|
||||
AutoRemoved {
|
||||
status_code: Option<u16>,
|
||||
message: String,
|
||||
},
|
||||
Failed {
|
||||
status_code: Option<u16>,
|
||||
message: String,
|
||||
@@ -112,6 +116,7 @@ impl AccountSelfCheckOutcome {
|
||||
match self {
|
||||
Self::Success { .. } => "success",
|
||||
Self::Blocked { .. } => "blocked",
|
||||
Self::AutoRemoved { .. } => "auto_removed",
|
||||
Self::Failed { .. } => "failed",
|
||||
Self::Skipped { .. } => "skipped",
|
||||
}
|
||||
@@ -121,6 +126,7 @@ impl AccountSelfCheckOutcome {
|
||||
match self {
|
||||
Self::Success { status_code, .. }
|
||||
| Self::Blocked { status_code, .. }
|
||||
| Self::AutoRemoved { status_code, .. }
|
||||
| Self::Failed { status_code, .. } => *status_code,
|
||||
Self::Skipped { .. } => None,
|
||||
}
|
||||
@@ -130,6 +136,7 @@ impl AccountSelfCheckOutcome {
|
||||
match self {
|
||||
Self::Success { message, .. } => message.as_deref(),
|
||||
Self::Blocked { message, .. }
|
||||
| Self::AutoRemoved { message, .. }
|
||||
| Self::Failed { message, .. }
|
||||
| Self::Skipped { message, .. } => Some(message.as_str()),
|
||||
}
|
||||
@@ -392,6 +399,10 @@ fn quota_payload_result_for_key(key_id: &str, payload: Option<Value>) -> Account
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
.map(ToOwned::to_owned);
|
||||
let auto_removed = item
|
||||
.get("auto_removed")
|
||||
.and_then(Value::as_bool)
|
||||
.unwrap_or(false);
|
||||
|
||||
if status == "success" {
|
||||
return AccountSelfCheckOutcome::Success {
|
||||
@@ -399,6 +410,12 @@ fn quota_payload_result_for_key(key_id: &str, payload: Option<Value>) -> Account
|
||||
message,
|
||||
};
|
||||
}
|
||||
if auto_removed {
|
||||
return AccountSelfCheckOutcome::AutoRemoved {
|
||||
status_code,
|
||||
message: message.unwrap_or_else(|| "已自动删除".to_string()),
|
||||
};
|
||||
}
|
||||
if quota_result_status_is_blocked(&status, status_code, message.as_deref()) {
|
||||
return AccountSelfCheckOutcome::Blocked {
|
||||
status_code,
|
||||
@@ -499,6 +516,11 @@ async fn record_score_probe_result_for_key(
|
||||
Some(PoolMemberHardState::Banned),
|
||||
PoolMemberProbeStatus::Failed,
|
||||
),
|
||||
AccountSelfCheckOutcome::AutoRemoved { .. } => (
|
||||
false,
|
||||
Some(PoolMemberHardState::Banned),
|
||||
PoolMemberProbeStatus::Failed,
|
||||
),
|
||||
AccountSelfCheckOutcome::Failed { .. } => (
|
||||
false,
|
||||
Some(PoolMemberHardState::Cooldown),
|
||||
@@ -570,6 +592,9 @@ fn update_summary_from_outcome(
|
||||
AccountSelfCheckOutcome::Blocked { .. } => {
|
||||
summary.blocked = summary.blocked.saturating_add(1);
|
||||
}
|
||||
AccountSelfCheckOutcome::AutoRemoved { .. } => {
|
||||
summary.auto_removed = summary.auto_removed.saturating_add(1);
|
||||
}
|
||||
AccountSelfCheckOutcome::Failed { .. } => {
|
||||
summary.failed = summary.failed.saturating_add(1);
|
||||
}
|
||||
|
||||
@@ -1211,6 +1211,7 @@ async fn perform_pool_quota_probe_for_provider(
|
||||
if auto_removed > 0 {
|
||||
summary.auto_removed += auto_removed;
|
||||
info!(
|
||||
event_name = "auto_removed_hard_banned",
|
||||
provider_id = %provider_short_id,
|
||||
provider_type,
|
||||
auto_removed,
|
||||
|
||||
@@ -289,9 +289,6 @@ fn merge_local_oauth_refresh_failure_reason(
|
||||
if current_reason.is_empty() {
|
||||
return Some(refresh_reason.to_string());
|
||||
}
|
||||
if current_reason.starts_with(OAUTH_EXPIRED_PREFIX) {
|
||||
return Some(current_reason.to_string());
|
||||
}
|
||||
if oauth_invalid_reason_is_account_block(Some(current_reason)) {
|
||||
return None;
|
||||
}
|
||||
@@ -1157,7 +1154,6 @@ impl AppState {
|
||||
|
||||
latest_key.encrypted_api_key = Some(encrypted_api_key);
|
||||
latest_key.encrypted_auth_config = encrypted_auth_config;
|
||||
latest_key.is_active = true;
|
||||
latest_key.expires_at_unix_secs = entry.expires_at_unix_secs;
|
||||
let (oauth_invalid_at_unix_secs, oauth_invalid_reason) =
|
||||
local_oauth_refresh_success_invalid_state(&latest_key);
|
||||
@@ -1717,13 +1713,13 @@ mod tests {
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn local_refresh_failure_does_not_replace_access_token_expired_marker() {
|
||||
fn local_refresh_failure_replaces_access_token_expired_marker() {
|
||||
assert_eq!(
|
||||
super::merge_local_oauth_refresh_failure_reason(
|
||||
Some("[OAUTH_EXPIRED] access token invalid"),
|
||||
"[REFRESH_FAILED] Token 续期失败 (401): refresh_token 无效",
|
||||
),
|
||||
Some("[OAUTH_EXPIRED] access token invalid".to_string()),
|
||||
Some("[REFRESH_FAILED] Token 续期失败 (401): refresh_token 无效".to_string()),
|
||||
);
|
||||
}
|
||||
|
||||
|
||||
@@ -5694,6 +5694,117 @@ async fn gateway_marks_manual_oauth_refresh_failures_as_invalid_in_pool_payload(
|
||||
token_handle.abort();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_auto_removes_manual_oauth_refresh_failure_after_access_token_expired() {
|
||||
let token_server = Router::new().route(
|
||||
"/oauth/token",
|
||||
post(move |_request: Request| async move {
|
||||
(
|
||||
StatusCode::UNAUTHORIZED,
|
||||
Json(json!({
|
||||
"error": {
|
||||
"message": "Could not validate your refresh token. Please try signing in again.",
|
||||
"type": "invalid_request_error",
|
||||
"param": serde_json::Value::Null,
|
||||
"code": "refresh_token_expired"
|
||||
}
|
||||
})),
|
||||
)
|
||||
}),
|
||||
);
|
||||
|
||||
let mut provider = sample_provider("provider-codex", "codex", 10).with_transport_fields(
|
||||
true,
|
||||
false,
|
||||
true,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
Some(json!({
|
||||
"pool_advanced": {
|
||||
"enabled": true,
|
||||
"auto_remove_banned_keys": true
|
||||
}
|
||||
})),
|
||||
);
|
||||
provider.provider_type = "codex".to_string();
|
||||
let endpoint = sample_endpoint(
|
||||
"endpoint-codex-cli",
|
||||
"provider-codex",
|
||||
"openai:responses",
|
||||
"https://chatgpt.com/backend-api/codex",
|
||||
);
|
||||
let mut key = sample_key(
|
||||
"key-codex-oauth-refresh-expired",
|
||||
"provider-codex",
|
||||
"openai:responses",
|
||||
"expired-codex-access-token",
|
||||
);
|
||||
key.auth_type = "oauth".to_string();
|
||||
key.expires_at_unix_secs = Some(1);
|
||||
key.encrypted_auth_config = Some(
|
||||
encrypt_python_fernet_plaintext(
|
||||
DEVELOPMENT_ENCRYPTION_KEY,
|
||||
r#"{"provider_type":"codex","refresh_token":"expired-refresh-token","email":"alice@example.com","account_id":"acct-codex-123","plan_type":"plus","expires_at":1}"#,
|
||||
)
|
||||
.expect("auth config ciphertext should build"),
|
||||
);
|
||||
|
||||
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
|
||||
vec![provider],
|
||||
vec![endpoint],
|
||||
vec![key],
|
||||
));
|
||||
let (token_url, token_handle) = start_server(token_server).await;
|
||||
let oauth_refresh =
|
||||
crate::provider_transport::LocalOAuthRefreshCoordinator::with_adapters_for_tests(vec![
|
||||
Arc::new(
|
||||
crate::provider_transport::oauth_refresh::GenericOAuthRefreshAdapter::default()
|
||||
.with_token_url_for_tests("codex", format!("{token_url}/oauth/token")),
|
||||
),
|
||||
]);
|
||||
let gateway = build_router_with_state(
|
||||
AppState::new()
|
||||
.expect("gateway should build")
|
||||
.with_data_state_for_tests(
|
||||
GatewayDataState::with_provider_catalog_repository_for_tests(
|
||||
provider_catalog_repository.clone(),
|
||||
)
|
||||
.with_encryption_key_for_tests(DEVELOPMENT_ENCRYPTION_KEY),
|
||||
)
|
||||
.with_oauth_refresh_coordinator_for_tests(oauth_refresh),
|
||||
);
|
||||
let (gateway_url, gateway_handle) = start_server(gateway).await;
|
||||
|
||||
let response = reqwest::Client::new()
|
||||
.post(format!(
|
||||
"{gateway_url}/api/admin/provider-oauth/keys/key-codex-oauth-refresh-expired/refresh"
|
||||
))
|
||||
.header(crate::constants::GATEWAY_HEADER, "rust-phase3b")
|
||||
.header(TRUSTED_ADMIN_USER_ID_HEADER, "admin-user-123")
|
||||
.header(TRUSTED_ADMIN_USER_ROLE_HEADER, "admin")
|
||||
.header(TRUSTED_ADMIN_SESSION_ID_HEADER, "session-123")
|
||||
.send()
|
||||
.await
|
||||
.expect("refresh request should succeed");
|
||||
|
||||
assert_eq!(response.status(), StatusCode::OK);
|
||||
let payload: serde_json::Value = response.json().await.expect("json body should parse");
|
||||
assert_eq!(payload["status"], json!("auto_removed"));
|
||||
assert_eq!(payload["message"], json!("已自动删除"));
|
||||
|
||||
let stored_keys = provider_catalog_repository
|
||||
.list_keys_by_ids(&["key-codex-oauth-refresh-expired".to_string()])
|
||||
.await
|
||||
.expect("keys should list");
|
||||
assert!(stored_keys.is_empty());
|
||||
|
||||
gateway_handle.abort();
|
||||
token_handle.abort();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn gateway_refreshes_admin_provider_oauth_key_locally_via_execution_runtime_provider_proxy_before_system_proxy(
|
||||
) {
|
||||
|
||||
@@ -26,6 +26,25 @@ pub fn should_auto_remove_structured_reason(reason: Option<&str>) -> bool {
|
||||
))
|
||||
}
|
||||
|
||||
pub fn should_auto_remove_oauth_refresh_failed_key(
|
||||
key: &StoredProviderCatalogKey,
|
||||
now_unix_secs: u64,
|
||||
) -> bool {
|
||||
if !key.auth_type.trim().eq_ignore_ascii_case("oauth") {
|
||||
return false;
|
||||
}
|
||||
let Some(reason) = key.oauth_invalid_reason.as_deref().map(str::trim) else {
|
||||
return false;
|
||||
};
|
||||
reason
|
||||
.lines()
|
||||
.map(str::trim)
|
||||
.any(|line| line.starts_with(OAUTH_REFRESH_FAILED_PREFIX))
|
||||
&& key
|
||||
.expires_at_unix_secs
|
||||
.is_none_or(|value| value == 0 || value <= now_unix_secs)
|
||||
}
|
||||
|
||||
pub fn normalize_string_id_list(values: Option<Vec<String>>) -> Option<Vec<String>> {
|
||||
let mut out = Vec::new();
|
||||
let mut seen = std::collections::BTreeSet::new();
|
||||
@@ -116,7 +135,7 @@ pub fn quota_refresh_success_invalid_state(
|
||||
.as_deref()
|
||||
.map(str::trim)
|
||||
.unwrap_or_default();
|
||||
if current_reason.starts_with(OAUTH_REFRESH_FAILED_PREFIX) {
|
||||
if current_reason.starts_with(OAUTH_ACCOUNT_BLOCK_PREFIX) {
|
||||
return (
|
||||
key.oauth_invalid_at_unix_secs,
|
||||
(!current_reason.is_empty()).then_some(current_reason.to_string()),
|
||||
@@ -435,8 +454,7 @@ fn codex_merge_invalid_reason(current: &str, candidate_reason: &str) -> String {
|
||||
return current.to_string();
|
||||
}
|
||||
if current.starts_with(OAUTH_EXPIRED_PREFIX)
|
||||
&& (candidate_reason.starts_with(OAUTH_REQUEST_FAILED_PREFIX)
|
||||
|| candidate_reason.starts_with(OAUTH_REFRESH_FAILED_PREFIX))
|
||||
&& candidate_reason.starts_with(OAUTH_REQUEST_FAILED_PREFIX)
|
||||
{
|
||||
return current.to_string();
|
||||
}
|
||||
@@ -914,8 +932,9 @@ mod tests {
|
||||
use super::{
|
||||
codex_build_invalid_state, codex_runtime_invalid_reason,
|
||||
parse_chatgpt_web_conversation_init_response, parse_codex_wham_usage_response,
|
||||
should_auto_remove_structured_reason, OAUTH_ACCOUNT_BLOCK_PREFIX, OAUTH_EXPIRED_PREFIX,
|
||||
OAUTH_REFRESH_FAILED_PREFIX, OAUTH_REQUEST_FAILED_PREFIX,
|
||||
quota_refresh_success_invalid_state, should_auto_remove_structured_reason,
|
||||
should_auto_remove_oauth_refresh_failed_key, OAUTH_ACCOUNT_BLOCK_PREFIX,
|
||||
OAUTH_EXPIRED_PREFIX, OAUTH_REFRESH_FAILED_PREFIX, OAUTH_REQUEST_FAILED_PREFIX,
|
||||
};
|
||||
use aether_data_contracts::repository::provider_catalog::StoredProviderCatalogKey;
|
||||
use serde_json::json;
|
||||
@@ -944,7 +963,7 @@ mod tests {
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn codex_invalid_state_keeps_oauth_expired_over_refresh_failure() {
|
||||
fn codex_invalid_state_replaces_oauth_expired_with_refresh_failure() {
|
||||
let mut key = StoredProviderCatalogKey::new(
|
||||
"key-1".to_string(),
|
||||
"provider-1".to_string(),
|
||||
@@ -964,10 +983,26 @@ mod tests {
|
||||
200,
|
||||
),
|
||||
(
|
||||
Some(100),
|
||||
Some(format!("{OAUTH_EXPIRED_PREFIX}session expired"))
|
||||
Some(200),
|
||||
Some(format!("{OAUTH_REFRESH_FAILED_PREFIX}Token 续期失败"))
|
||||
)
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn codex_invalid_state_keeps_oauth_expired_over_request_failure() {
|
||||
let mut key = StoredProviderCatalogKey::new(
|
||||
"key-1".to_string(),
|
||||
"provider-1".to_string(),
|
||||
"key-1".to_string(),
|
||||
"oauth".to_string(),
|
||||
None,
|
||||
true,
|
||||
)
|
||||
.expect("key should build");
|
||||
key.oauth_invalid_at_unix_secs = Some(100);
|
||||
key.oauth_invalid_reason = Some(format!("{OAUTH_EXPIRED_PREFIX}session expired"));
|
||||
|
||||
assert_eq!(
|
||||
codex_build_invalid_state(
|
||||
&key,
|
||||
@@ -1011,12 +1046,46 @@ mod tests {
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn auto_remove_structured_reason_removes_oauth_expired_token_invalid() {
|
||||
assert!(should_auto_remove_structured_reason(Some(
|
||||
fn auto_remove_structured_reason_keeps_oauth_expired_token_invalid() {
|
||||
assert!(!should_auto_remove_structured_reason(Some(
|
||||
"[OAUTH_EXPIRED] token invalidated"
|
||||
)));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn auto_remove_oauth_refresh_failed_requires_expired_access_token() {
|
||||
let mut key = StoredProviderCatalogKey::new(
|
||||
"key-1".to_string(),
|
||||
"provider-1".to_string(),
|
||||
"key-1".to_string(),
|
||||
"oauth".to_string(),
|
||||
None,
|
||||
true,
|
||||
)
|
||||
.expect("key should build");
|
||||
key.oauth_invalid_reason = Some(format!("{OAUTH_REFRESH_FAILED_PREFIX}Token 续期失败"));
|
||||
key.expires_at_unix_secs = Some(2_000);
|
||||
|
||||
assert!(!should_auto_remove_oauth_refresh_failed_key(&key, 1_000));
|
||||
assert!(should_auto_remove_oauth_refresh_failed_key(&key, 2_000));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn quota_refresh_success_clears_refresh_failed_marker() {
|
||||
let mut key = StoredProviderCatalogKey::new(
|
||||
"key-1".to_string(),
|
||||
"provider-1".to_string(),
|
||||
"key-1".to_string(),
|
||||
"oauth".to_string(),
|
||||
None,
|
||||
true,
|
||||
)
|
||||
.expect("key should build");
|
||||
key.oauth_invalid_reason = Some("[REFRESH_FAILED] Token 续期失败".to_string());
|
||||
|
||||
assert_eq!(quota_refresh_success_invalid_state(&key), (None, None));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn auto_remove_structured_reason_keeps_request_and_refresh_failures() {
|
||||
assert!(!should_auto_remove_structured_reason(Some(
|
||||
|
||||
@@ -42,7 +42,6 @@ const AUTO_REMOVABLE_ACCOUNT_STATE_CODES: &[&str] = &[
|
||||
"account_disabled",
|
||||
"workspace_deactivated",
|
||||
"account_forbidden",
|
||||
"oauth_token_invalid",
|
||||
];
|
||||
|
||||
#[derive(Debug, Clone, Default, PartialEq, Eq)]
|
||||
@@ -604,7 +603,7 @@ mod tests {
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn oauth_expired_state_is_auto_removed() {
|
||||
fn oauth_expired_state_is_not_auto_removed() {
|
||||
let state = resolve_pool_account_state(
|
||||
Some("codex"),
|
||||
None,
|
||||
@@ -613,7 +612,7 @@ mod tests {
|
||||
|
||||
assert!(state.blocked);
|
||||
assert_eq!(state.code.as_deref(), Some("oauth_token_invalid"));
|
||||
assert!(should_auto_remove_account_state(&state));
|
||||
assert!(!should_auto_remove_account_state(&state));
|
||||
}
|
||||
|
||||
#[test]
|
||||
|
||||
Reference in New Issue
Block a user