i/Aether
1
0
Fork 0
mirror of https://github.com/fawney19/Aether.git synced 2026-10-11 19:59:50 +08:00
Code Issues Packages Projects Releases Wiki Activity

Unify quota snapshots and oauth refresh handling

Browse Source
This commit is contained in:
fawney19
2026-04-17 18:22:41 +08:00
parent b8702ae124
commit 7eae1f90f6
38 changed files with 3435 additions and 458 deletions
Download Patch File Download Diff File Expand all files Collapse all files
apps/aether-gateway/src/handlers/admin/provider/oauth/dispatch/batch/execution.rs
+8 -4
View File
@@ -11,7 +11,8 @@ use crate::handlers::admin::provider::oauth::provisioning::{
update_existing_provider_oauth_catalog_key,
};
use crate::handlers::admin::provider::oauth::runtime::{
provider_oauth_runtime_endpoint_for_provider, refresh_provider_oauth_account_state_after_update,
provider_oauth_runtime_endpoint_for_provider,
spawn_provider_oauth_account_state_refresh_after_update,
};
use crate::handlers::admin::provider::oauth::state::{
admin_provider_oauth_template, exchange_admin_provider_oauth_refresh_token,
@@ -264,9 +265,12 @@ pub(super) async fn execute_admin_provider_oauth_batch_import(
}
};
let _ =
refresh_provider_oauth_account_state_after_update(state, &provider, &persisted_key.id)
.await;
spawn_provider_oauth_account_state_refresh_after_update(
state.cloned_app(),
provider.clone(),
persisted_key.id.clone(),
request_proxy.clone(),
);
success += 1;
results.push(json!({
apps/aether-gateway/src/handlers/admin/provider/oauth/dispatch/batch/kiro_import.rs
+8 -4
View File
@@ -5,7 +5,8 @@ use crate::handlers::admin::provider::oauth::provisioning::{
update_existing_provider_oauth_catalog_key,
};
use crate::handlers::admin::provider::oauth::runtime::{
provider_oauth_runtime_endpoint_for_provider, refresh_provider_oauth_account_state_after_update,
provider_oauth_runtime_endpoint_for_provider,
spawn_provider_oauth_account_state_refresh_after_update,
};
use crate::handlers::admin::provider::oauth::state::decode_jwt_claims;
use crate::handlers::admin::provider::shared::support::ADMIN_PROVIDER_OAUTH_DATA_UNAVAILABLE_DETAIL;
@@ -548,9 +549,12 @@ pub(super) async fn execute_admin_provider_oauth_kiro_batch_import(
.get("auth_method")
.cloned()
.unwrap_or(serde_json::Value::Null);
let _ =
refresh_provider_oauth_account_state_after_update(state, &provider, &persisted_key.id)
.await;
spawn_provider_oauth_account_state_refresh_after_update(
state.cloned_app(),
provider.clone(),
persisted_key.id.clone(),
request_proxy.clone(),
);
success += 1;
results.push(json!({
apps/aether-gateway/src/handlers/admin/provider/oauth/dispatch/complete/key.rs
+2 -1
View File
@@ -148,7 +148,7 @@ pub(super) async fn handle_admin_provider_oauth_complete_key(
&callback.code,
&callback.state_nonce,
state_data.pkce_verifier.as_deref(),
request_proxy,
request_proxy.clone(),
)
.await
{
@@ -243,6 +243,7 @@ pub(super) async fn handle_admin_provider_oauth_complete_key(
&provider,
&endpoint,
vec![refreshed_key],
request_proxy.clone(),
)
.await?
{
apps/aether-gateway/src/handlers/admin/provider/oauth/dispatch/complete/provider.rs
+9 -5
View File
@@ -5,7 +5,8 @@ use super::super::super::provisioning::{
provider_oauth_active_api_formats, update_existing_provider_oauth_catalog_key,
};
use super::super::super::runtime::{
provider_oauth_runtime_endpoint_for_provider, refresh_provider_oauth_account_state_after_update,
provider_oauth_runtime_endpoint_for_provider,
spawn_provider_oauth_account_state_refresh_after_update,
};
use super::super::super::state::{
admin_provider_oauth_template, build_admin_provider_oauth_backend_unavailable_response,
@@ -137,7 +138,7 @@ pub(super) async fn handle_admin_provider_oauth_complete_provider(
&callback.code,
&callback.state_nonce,
state_data.pkce_verifier.as_deref(),
request_proxy,
request_proxy.clone(),
)
.await
{
@@ -231,9 +232,12 @@ pub(super) async fn handle_admin_provider_oauth_complete_provider(
}
};
let _ = state
.refresh_provider_oauth_account_state_after_update(&provider, &persisted_key.id)
.await;
spawn_provider_oauth_account_state_refresh_after_update(
state.cloned_app(),
provider.clone(),
persisted_key.id.clone(),
request_proxy.clone(),
);
Ok(Json(json!({
"key_id": persisted_key.id,
apps/aether-gateway/src/handlers/admin/provider/oauth/dispatch/device/poll.rs
+9 -5
View File
@@ -8,7 +8,8 @@ use crate::handlers::admin::provider::oauth::provisioning::{
update_existing_provider_oauth_catalog_key,
};
use crate::handlers::admin::provider::oauth::runtime::{
provider_oauth_runtime_endpoint_for_provider, refresh_provider_oauth_account_state_after_update,
provider_oauth_runtime_endpoint_for_provider,
spawn_provider_oauth_account_state_refresh_after_update,
};
use crate::handlers::admin::provider::oauth::state::{
build_admin_provider_oauth_backend_unavailable_response, build_kiro_device_key_name,
@@ -149,7 +150,7 @@ pub(super) async fn handle_admin_provider_oauth_device_poll(
&session.client_id,
&session.client_secret,
&session.device_code,
request_proxy,
request_proxy.clone(),
)
.await
{
@@ -315,9 +316,12 @@ pub(super) async fn handle_admin_provider_oauth_device_poll(
}
};
let _ = state
.refresh_provider_oauth_account_state_after_update(&provider, &persisted_key.id)
.await;
spawn_provider_oauth_account_state_refresh_after_update(
state.cloned_app(),
provider.clone(),
persisted_key.id.clone(),
request_proxy.clone(),
);
session.status = "authorized".to_string();
session.key_id = Some(persisted_key.id.clone());
apps/aether-gateway/src/handlers/admin/provider/oauth/dispatch/import.rs
+13 -5
View File
@@ -5,7 +5,8 @@ use super::super::provisioning::{
provider_oauth_active_api_formats, update_existing_provider_oauth_catalog_key,
};
use super::super::runtime::{
provider_oauth_runtime_endpoint_for_provider, refresh_provider_oauth_account_state_after_update,
provider_oauth_runtime_endpoint_for_provider,
spawn_provider_oauth_account_state_refresh_after_update,
};
use super::super::state::{
admin_provider_oauth_template, build_admin_provider_oauth_backend_unavailable_response,
@@ -120,7 +121,11 @@ pub(super) async fn handle_admin_provider_oauth_import_refresh_token(
.await;
let token_payload = match state
.exchange_admin_provider_oauth_refresh_token(template, refresh_token_input, request_proxy)
.exchange_admin_provider_oauth_refresh_token(
template,
refresh_token_input,
request_proxy.clone(),
)
.await
{
Ok(payload) => payload,
@@ -218,9 +223,12 @@ pub(super) async fn handle_admin_provider_oauth_import_refresh_token(
}
};
let _ = state
.refresh_provider_oauth_account_state_after_update(&provider, &persisted_key.id)
.await;
spawn_provider_oauth_account_state_refresh_after_update(
state.cloned_app(),
provider.clone(),
persisted_key.id.clone(),
request_proxy.clone(),
);
Ok(Json(json!({
"key_id": persisted_key.id,
apps/aether-gateway/src/handlers/admin/provider/oauth/dispatch/refresh/execution.rs
+1 -1
View File
@@ -94,7 +94,7 @@ pub(super) async fn execute_admin_provider_oauth_refresh(
refreshed_key.encrypted_auth_config.as_deref(),
);
let (account_state_recheck_attempted, account_state_recheck_error) = state
.refresh_provider_oauth_account_state_after_update(&provider, &key_id)
.refresh_provider_oauth_account_state_after_update(&provider, &key_id, None)
.await?;
Ok(RefreshDispatch::Continue(RefreshSuccessContext {
apps/aether-gateway/src/handlers/admin/provider/oauth/quota/antigravity.rs
+23 -8
View File
@@ -1,14 +1,14 @@
use super::shared::{
coerce_json_f64, coerce_json_string, default_provider_quota_execution_timeouts,
execute_provider_quota_plan, extract_execution_error_message,
persist_provider_quota_refresh_state, quota_refresh_success_invalid_state,
ProviderQuotaExecutionOutcome,
build_quota_snapshot_payload, coerce_json_f64, coerce_json_string,
default_provider_quota_execution_timeouts, execute_provider_quota_plan,
extract_execution_error_message, persist_provider_quota_refresh_state,
quota_refresh_success_invalid_state, ProviderQuotaExecutionOutcome,
};
use crate::handlers::admin::provider::shared::payloads::ANTIGRAVITY_FETCH_AVAILABLE_MODELS_PATH;
use crate::handlers::admin::request::{AdminAppState, AdminGatewayProviderTransportSnapshot};
use crate::GatewayError;
use aether_admin::provider::quota::parse_antigravity_usage_response;
use aether_contracts::{ExecutionPlan, RequestBody};
use aether_contracts::{ExecutionPlan, ProxySnapshot, RequestBody};
use aether_data_contracts::repository::provider_catalog::{
StoredProviderCatalogEndpoint, StoredProviderCatalogKey, StoredProviderCatalogProvider,
};
@@ -22,6 +22,7 @@ async fn execute_antigravity_quota_plan(
authorization: (String, String),
project_id: &str,
mut identity_headers: BTreeMap<String, String>,
proxy_override: Option<&ProxySnapshot>,
) -> Result<ProviderQuotaExecutionOutcome, GatewayError> {
let mut headers = std::mem::take(&mut identity_headers);
headers.insert("authorization".to_string(), authorization.1);
@@ -32,9 +33,14 @@ async fn execute_antigravity_quota_plan(
.or_insert_with(|| "antigravity".to_string());
let body = json!({ "project": project_id });
let proxy = state
.resolve_transport_proxy_snapshot_with_tunnel_affinity(transport)
.await;
let proxy = match proxy_override {
Some(proxy) => Some(proxy.clone()),
None => {
state
.resolve_transport_proxy_snapshot_with_tunnel_affinity(transport)
.await
}
};
let timeouts = state
.resolve_transport_execution_timeouts(transport)
.or(Some(default_provider_quota_execution_timeouts(
@@ -78,6 +84,7 @@ pub(crate) async fn refresh_antigravity_provider_quota_locally(
provider: &StoredProviderCatalogProvider,
endpoint: &StoredProviderCatalogEndpoint,
keys: Vec<StoredProviderCatalogKey>,
proxy_override: Option<ProxySnapshot>,
) -> Result<Option<serde_json::Value>, GatewayError> {
let mut results = Vec::new();
let mut success_count = 0usize;
@@ -134,6 +141,7 @@ pub(crate) async fn refresh_antigravity_provider_quota_locally(
authorization,
&project_id,
identity_headers,
proxy_override.as_ref(),
)
.await?
{
@@ -250,6 +258,13 @@ pub(crate) async fn refresh_antigravity_provider_quota_locally(
{
payload.insert("metadata".to_string(), metadata);
}
if let Some(quota_snapshot) = build_quota_snapshot_payload(
"antigravity",
key.status_snapshot.as_ref(),
metadata_update.as_ref(),
) {
payload.insert("quota_snapshot".to_string(), quota_snapshot);
}
results.push(serde_json::Value::Object(payload));
}
apps/aether-gateway/src/handlers/admin/provider/oauth/quota/codex/mod.rs
+30 -17
View File
@@ -13,13 +13,15 @@ use self::parse::{
};
use self::plan::{build_codex_refresh_headers, execute_codex_quota_plan};
use super::shared::{
extract_execution_error_message, persist_provider_quota_refresh_state,
provider_auto_remove_banned_keys, quota_refresh_success_invalid_state,
should_auto_remove_structured_reason, ProviderQuotaExecutionOutcome,
build_quota_snapshot_payload, extract_execution_error_message,
persist_provider_quota_refresh_state, provider_auto_remove_banned_keys,
quota_refresh_success_invalid_state, should_auto_remove_structured_reason,
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::{
StoredProviderCatalogEndpoint, StoredProviderCatalogKey, StoredProviderCatalogProvider,
};
@@ -31,6 +33,7 @@ pub(crate) async fn refresh_codex_provider_quota_locally(
provider: &StoredProviderCatalogProvider,
endpoint: &StoredProviderCatalogEndpoint,
keys: Vec<StoredProviderCatalogKey>,
proxy_override: Option<ProxySnapshot>,
) -> Result<Option<serde_json::Value>, GatewayError> {
let auto_remove_abnormal_keys = provider_auto_remove_banned_keys(provider.config.as_ref());
let mut results = Vec::new();
@@ -77,20 +80,23 @@ pub(crate) async fn refresh_codex_provider_quota_locally(
}
};
let result = match execute_codex_quota_plan(state, &transport, headers).await? {
ProviderQuotaExecutionOutcome::Response(result) => result,
ProviderQuotaExecutionOutcome::Failure(detail) => {
failed_count += 1;
results.push(json!({
"key_id": key.id,
"key_name": key.name,
"status": "error",
"message": format!("wham/usage 请求执行失败: {detail}"),
"status_code": 502,
}));
continue;
}
};
let result =
match execute_codex_quota_plan(state, &transport, headers, proxy_override.as_ref())
.await?
{
ProviderQuotaExecutionOutcome::Response(result) => result,
ProviderQuotaExecutionOutcome::Failure(detail) => {
failed_count += 1;
results.push(json!({
"key_id": key.id,
"key_name": key.name,
"status": "error",
"message": format!("wham/usage 请求执行失败: {detail}"),
"status_code": 502,
}));
continue;
}
};
let now_unix_secs = SystemTime::now()
.duration_since(UNIX_EPOCH)
.ok()
@@ -278,6 +284,13 @@ pub(crate) async fn refresh_codex_provider_quota_locally(
{
payload.insert("metadata".to_string(), metadata_update);
}
if let Some(quota_snapshot) = build_quota_snapshot_payload(
"codex",
key.status_snapshot.as_ref(),
metadata_update.as_ref(),
) {
payload.insert("quota_snapshot".to_string(), quota_snapshot);
}
if auto_removed {
payload.insert("auto_removed".to_string(), json!(true));
}
apps/aether-gateway/src/handlers/admin/provider/oauth/quota/codex/plan.rs
+10 -4
View File
@@ -6,7 +6,7 @@ use super::parse::normalize_codex_plan_type;
use crate::handlers::admin::provider::shared::payloads::CODEX_WHAM_USAGE_URL;
use crate::handlers::admin::request::{AdminAppState, AdminGatewayProviderTransportSnapshot};
use crate::GatewayError;
use aether_contracts::{ExecutionPlan, RequestBody};
use aether_contracts::{ExecutionPlan, ProxySnapshot, RequestBody};
use std::collections::BTreeMap;
pub(super) fn build_codex_refresh_headers(
@@ -60,10 +60,16 @@ pub(super) async fn execute_codex_quota_plan(
state: &AdminAppState<'_>,
transport: &AdminGatewayProviderTransportSnapshot,
headers: BTreeMap<String, String>,
proxy_override: Option<&ProxySnapshot>,
) -> Result<ProviderQuotaExecutionOutcome, GatewayError> {
let proxy = state
.resolve_transport_proxy_snapshot_with_tunnel_affinity(transport)
.await;
let proxy = match proxy_override {
Some(proxy) => Some(proxy.clone()),
None => {
state
.resolve_transport_proxy_snapshot_with_tunnel_affinity(transport)
.await
}
};
let timeouts = state
.resolve_transport_execution_timeouts(transport)
.or(Some(default_provider_quota_execution_timeouts(
apps/aether-gateway/src/handlers/admin/provider/oauth/quota/kiro/mod.rs
+47 -17
View File
@@ -4,11 +4,13 @@ mod plan;
use self::parse::parse_kiro_usage_response;
use self::plan::execute_kiro_quota_plan;
use super::shared::{
extract_execution_error_message, persist_provider_quota_refresh_state,
quota_refresh_success_invalid_state, ProviderQuotaExecutionOutcome,
build_quota_snapshot_payload, extract_execution_error_message,
persist_provider_quota_refresh_state, quota_refresh_success_invalid_state,
ProviderQuotaExecutionOutcome,
};
use crate::handlers::admin::request::AdminAppState;
use crate::GatewayError;
use aether_contracts::ProxySnapshot;
use aether_data_contracts::repository::provider_catalog::{
StoredProviderCatalogEndpoint, StoredProviderCatalogKey, StoredProviderCatalogProvider,
};
@@ -20,6 +22,7 @@ pub(crate) async fn refresh_kiro_provider_quota_locally(
provider: &StoredProviderCatalogProvider,
endpoint: &StoredProviderCatalogEndpoint,
keys: Vec<StoredProviderCatalogKey>,
proxy_override: Option<ProxySnapshot>,
) -> Result<Option<serde_json::Value>, GatewayError> {
let mut results = Vec::new();
let mut success_count = 0usize;
@@ -57,20 +60,22 @@ pub(crate) async fn refresh_kiro_provider_quota_locally(
continue;
};
let result = match execute_kiro_quota_plan(state, &transport, &auth).await? {
ProviderQuotaExecutionOutcome::Response(result) => result,
ProviderQuotaExecutionOutcome::Failure(detail) => {
failed_count += 1;
results.push(json!({
"key_id": key.id,
"key_name": key.name,
"status": "error",
"message": format!("getUsageLimits 请求执行失败: {detail}"),
"status_code": 502,
}));
continue;
}
};
let result =
match execute_kiro_quota_plan(state, &transport, &auth, proxy_override.as_ref()).await?
{
ProviderQuotaExecutionOutcome::Response(result) => result,
ProviderQuotaExecutionOutcome::Failure(detail) => {
failed_count += 1;
results.push(json!({
"key_id": key.id,
"key_name": key.name,
"status": "error",
"message": format!("getUsageLimits 请求执行失败: {detail}"),
"status_code": 502,
}));
continue;
}
};
let now_unix_secs = SystemTime::now()
.duration_since(UNIX_EPOCH)
@@ -93,7 +98,25 @@ pub(crate) async fn refresh_kiro_provider_quota_locally(
metadata_update = parse_kiro_usage_response(body_json, now_unix_secs)
.map(|metadata| json!({ "kiro": metadata }));
if metadata_update.is_some() {
let auth_config_json = auth.auth_config.to_json_value().to_string();
let mut auth_config_object = transport
.key
.decrypted_auth_config
.as_deref()
.and_then(|raw| serde_json::from_str::<serde_json::Value>(raw).ok())
.and_then(|value| value.as_object().cloned())
.unwrap_or_default();
if let Some(refreshed_auth_config) =
auth.auth_config.to_json_value().as_object()
{
for (key, value) in refreshed_auth_config {
auth_config_object.insert(key.clone(), value.clone());
}
}
auth_config_object
.entry("provider_type".to_string())
.or_insert_with(|| json!("kiro"));
let auth_config_json =
serde_json::Value::Object(auth_config_object).to_string();
if let Some(auth_config_json) =
state.encrypt_catalog_secret_with_fallbacks(auth_config_json.as_str())
{
@@ -185,6 +208,13 @@ pub(crate) async fn refresh_kiro_provider_quota_locally(
{
payload.insert("metadata".to_string(), metadata);
}
if let Some(quota_snapshot) = build_quota_snapshot_payload(
"kiro",
key.status_snapshot.as_ref(),
metadata_update.as_ref(),
) {
payload.insert("quota_snapshot".to_string(), quota_snapshot);
}
results.push(serde_json::Value::Object(payload));
}
apps/aether-gateway/src/handlers/admin/provider/oauth/quota/kiro/plan.rs
+10 -4
View File
@@ -7,7 +7,7 @@ use crate::handlers::admin::request::{
AdminAppState, AdminGatewayProviderTransportSnapshot, AdminKiroRequestAuth,
};
use crate::GatewayError;
use aether_contracts::{ExecutionPlan, RequestBody};
use aether_contracts::{ExecutionPlan, ProxySnapshot, RequestBody};
use std::collections::BTreeMap;
use url::form_urlencoded;
use uuid::Uuid;
@@ -66,10 +66,16 @@ pub(super) async fn execute_kiro_quota_plan(
state: &AdminAppState<'_>,
transport: &AdminGatewayProviderTransportSnapshot,
auth: &AdminKiroRequestAuth,
proxy_override: Option<&ProxySnapshot>,
) -> Result<ProviderQuotaExecutionOutcome, GatewayError> {
let proxy = state
.resolve_transport_proxy_snapshot_with_tunnel_affinity(transport)
.await;
let proxy = match proxy_override {
Some(proxy) => Some(proxy.clone()),
None => {
state
.resolve_transport_proxy_snapshot_with_tunnel_affinity(transport)
.await
}
};
let timeouts = state
.resolve_transport_execution_timeouts(transport)
.or(Some(default_provider_quota_execution_timeouts(
apps/aether-gateway/src/handlers/admin/provider/oauth/quota/shared.rs
+29
View File
@@ -2,6 +2,7 @@ use crate::handlers::admin::provider::shared::payloads::{
OAUTH_ACCOUNT_BLOCK_PREFIX, OAUTH_REFRESH_FAILED_PREFIX,
};
use crate::handlers::admin::request::{AdminAppState, AdminGatewayProviderTransportSnapshot};
use crate::handlers::shared::sync_provider_key_quota_status_snapshot;
use crate::GatewayError;
use aether_admin::provider::quota as admin_provider_quota_pure;
use aether_contracts::{ExecutionPlan, ExecutionResult, ExecutionTimeouts, ProxySnapshot};
@@ -89,6 +90,20 @@ pub(super) fn coerce_json_string(value: Option<&serde_json::Value>) -> Option<St
admin_provider_quota_pure::coerce_json_string(value)
}
pub(super) fn build_quota_snapshot_payload(
provider_type: &str,
current_status_snapshot: Option<&serde_json::Value>,
metadata_update: Option<&serde_json::Value>,
) -> Option<serde_json::Value> {
let updated_snapshot = sync_provider_key_quota_status_snapshot(
current_status_snapshot,
provider_type,
metadata_update,
"refresh_api",
)?;
updated_snapshot.get("quota").cloned()
}
pub(crate) async fn persist_provider_quota_refresh_state(
state: &AdminAppState<'_>,
key_id: &str,
@@ -106,17 +121,31 @@ pub(crate) async fn persist_provider_quota_refresh_state(
return Ok(false);
};
let mut quota_snapshot_provider_type = None::<&str>;
if let Some(metadata_update) = metadata_update {
latest_key.upstream_metadata = Some(merge_upstream_metadata(
latest_key.upstream_metadata.as_ref(),
metadata_update,
));
quota_snapshot_provider_type = metadata_update.as_object().and_then(|object| {
["codex", "kiro", "antigravity", "gemini_cli"]
.into_iter()
.find(|provider_type| object.contains_key(*provider_type))
});
}
if let Some(encrypted_auth_config) = encrypted_auth_config {
latest_key.encrypted_auth_config = Some(encrypted_auth_config);
}
latest_key.oauth_invalid_at_unix_secs = oauth_invalid_at_unix_secs;
latest_key.oauth_invalid_reason = oauth_invalid_reason;
if let Some(provider_type) = quota_snapshot_provider_type {
latest_key.status_snapshot = sync_provider_key_quota_status_snapshot(
latest_key.status_snapshot.as_ref(),
provider_type,
latest_key.upstream_metadata.as_ref(),
"refresh_api",
);
}
latest_key.updated_at_unix_secs = SystemTime::now()
.duration_since(UNIX_EPOCH)
.ok()
apps/aether-gateway/src/handlers/admin/provider/oauth/runtime.rs
+45 -5
View File
@@ -3,7 +3,8 @@ use super::quota::codex::refresh_codex_provider_quota_locally;
use super::quota::kiro::refresh_kiro_provider_quota_locally;
use crate::handlers::admin::request::AdminAppState;
use crate::provider_key_auth::provider_key_is_oauth_managed;
use crate::GatewayError;
use crate::{AppState, GatewayError};
use aether_contracts::ProxySnapshot;
use aether_data_contracts::repository::provider_catalog::{
StoredProviderCatalogEndpoint, StoredProviderCatalogProvider,
};
@@ -65,6 +66,7 @@ pub(crate) async fn refresh_provider_oauth_account_state_after_update(
state: &AdminAppState<'_>,
provider: &StoredProviderCatalogProvider,
key_id: &str,
proxy_override: Option<&ProxySnapshot>,
) -> Result<(bool, Option<String>), GatewayError> {
let provider_type = provider.provider_type.trim().to_ascii_lowercase();
if !matches!(provider_type.as_str(), "codex" | "kiro" | "antigravity") {
@@ -90,16 +92,37 @@ pub(crate) async fn refresh_provider_oauth_account_state_after_update(
return Ok((false, None));
}
let proxy_override = proxy_override.cloned();
let payload = match provider_type.as_str() {
"codex" => {
refresh_codex_provider_quota_locally(state, provider, &endpoint, vec![key]).await?
refresh_codex_provider_quota_locally(
state,
provider,
&endpoint,
vec![key],
proxy_override.clone(),
)
.await?
}
"kiro" => {
refresh_kiro_provider_quota_locally(state, provider, &endpoint, vec![key]).await?
refresh_kiro_provider_quota_locally(
state,
provider,
&endpoint,
vec![key],
proxy_override.clone(),
)
.await?
}
"antigravity" => {
refresh_antigravity_provider_quota_locally(state, provider, &endpoint, vec![key])
.await?
refresh_antigravity_provider_quota_locally(
state,
provider,
&endpoint,
vec![key],
proxy_override,
)
.await?
}
_ => None,
};
@@ -123,3 +146,20 @@ pub(crate) async fn refresh_provider_oauth_account_state_after_update(
};
Ok((true, error))
}
pub(crate) fn spawn_provider_oauth_account_state_refresh_after_update(
app: AppState,
provider: StoredProviderCatalogProvider,
key_id: String,
proxy_override: Option<ProxySnapshot>,
) {
tokio::spawn(async move {
let _ = refresh_provider_oauth_account_state_after_update(
&AdminAppState::new(&app),
&provider,
&key_id,
proxy_override.as_ref(),
)
.await;
});
}