fix(gateway): spawn oauth account refresh asynchronously

This commit is contained in:
zhefox
2026-05-25 15:09:41 +08:00
parent db4f3fd210
commit b72abec2fc
2 changed files with 16 additions and 49 deletions
@@ -1,7 +1,9 @@
use super::super::super::errors::build_internal_control_error_response;
use super::super::super::provisioning::provider_oauth_token_payload_expires_at_unix_secs;
use super::super::super::quota::codex::refresh_codex_provider_quota_locally;
use super::super::super::runtime::resolve_provider_oauth_runtime_endpoints;
use super::super::super::runtime::{
resolve_provider_oauth_runtime_endpoints,
spawn_provider_oauth_account_state_refresh_after_update,
};
use super::super::super::state::{
admin_provider_oauth_template, enrich_admin_provider_oauth_auth_config,
is_fixed_provider_type_for_provider_oauth, json_non_empty_string,
@@ -219,50 +221,20 @@ pub(super) async fn handle_admin_provider_oauth_complete_key(
));
}
let mut account_state_recheck_attempted = false;
let mut account_state_recheck_error = None::<String>;
if provider_type == "codex" {
if let Some(endpoint) = runtime_endpoint {
let refreshed_key = state
.read_provider_catalog_keys_by_ids(std::slice::from_ref(&key_id))
.await?
.into_iter()
.next()
.unwrap_or_else(|| key.clone());
if let Some(result) = refresh_codex_provider_quota_locally(
state,
&provider,
&endpoint,
vec![refreshed_key],
request_proxy.clone(),
)
.await?
{
account_state_recheck_attempted = true;
let success = result
.get("success")
.and_then(serde_json::Value::as_u64)
.unwrap_or(0);
if success == 0 {
account_state_recheck_error = result
.get("results")
.and_then(serde_json::Value::as_array)
.and_then(|results| results.first())
.and_then(|value| value.get("message"))
.and_then(serde_json::Value::as_str)
.map(ToOwned::to_owned);
}
}
}
}
spawn_provider_oauth_account_state_refresh_after_update(
state.cloned_app(),
provider.clone(),
key_id.clone(),
request_proxy.clone(),
);
Ok(Json(json!({
"provider_type": provider_type,
"expires_at": expires_at,
"has_refresh_token": refresh_token.is_some(),
"email": auth_config.get("email").cloned().unwrap_or(serde_json::Value::Null),
"account_state_recheck_attempted": account_state_recheck_attempted,
"account_state_recheck_error": account_state_recheck_error,
"account_state_recheck_attempted": false,
"account_state_recheck_error": serde_json::Value::Null,
}))
.into_response())
}
@@ -2874,15 +2874,10 @@ async fn gateway_completes_admin_provider_oauth_key_locally_with_trusted_admin_p
assert_eq!(payload["has_refresh_token"], true);
assert_eq!(payload["expires_at"], 4_102_444_800u64);
assert_eq!(payload["email"], "alice@example.com");
assert_eq!(payload["account_state_recheck_attempted"], true);
let account_state_recheck_error = payload["account_state_recheck_error"]
.as_str()
.expect("account_state_recheck_error should be string when recheck is attempted");
assert!(
account_state_recheck_error == "wham/usage API 返回状态码 401"
|| account_state_recheck_error == "wham/usage API 返回状态码 403"
|| account_state_recheck_error.starts_with("wham/usage 请求执行失败:"),
"unexpected account_state_recheck_error: {account_state_recheck_error}"
assert_eq!(payload["account_state_recheck_attempted"], false);
assert_eq!(
payload["account_state_recheck_error"],
serde_json::Value::Null
);
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
assert_eq!(*token_hits.lock().expect("mutex should lock"), 1);