use super::quota::dispatch::refresh_provider_pool_quota_locally; use super::quota::shared::provider_quota_refresh_endpoint_for_provider; use super::quota::shared::provider_type_supports_quota_refresh; use crate::handlers::admin::provider::write::provider::reconcile_admin_fixed_provider_template_endpoints; use crate::handlers::admin::request::AdminAppState; use crate::provider_key_auth::provider_key_is_oauth_managed; use crate::task_runtime::{spawn_fire_and_forget, TASK_KEY_PROVIDER_OAUTH_ACCOUNT_REFRESH}; use crate::{AppState, GatewayError}; use aether_contracts::ProxySnapshot; use aether_data_contracts::repository::provider_catalog::{ StoredProviderCatalogEndpoint, StoredProviderCatalogProvider, }; pub(crate) fn provider_oauth_runtime_endpoint_for_provider( provider_type: &str, endpoints: &[StoredProviderCatalogEndpoint], ) -> Option { select_provider_oauth_runtime_endpoint(provider_type, endpoints, false) } pub(crate) fn provider_oauth_maintenance_endpoint_for_provider( provider_type: &str, endpoints: &[StoredProviderCatalogEndpoint], ) -> Option { select_provider_oauth_runtime_endpoint(provider_type, endpoints, true) } fn matching_endpoint( endpoints: &[StoredProviderCatalogEndpoint], include_inactive: bool, predicate: F, ) -> Option where F: Fn(&StoredProviderCatalogEndpoint) -> bool, { endpoints .iter() .find(|endpoint| endpoint.is_active && predicate(endpoint)) .cloned() .or_else(|| { include_inactive.then(|| { endpoints .iter() .find(|endpoint| !endpoint.is_active && predicate(endpoint)) .cloned() })? }) } fn select_provider_oauth_runtime_endpoint( provider_type: &str, endpoints: &[StoredProviderCatalogEndpoint], include_inactive: bool, ) -> Option { let provider_type = provider_type.trim().to_ascii_lowercase(); match provider_type.as_str() { "codex" => matching_endpoint(endpoints, include_inactive, |endpoint| { crate::ai_serving::is_openai_responses_format(&endpoint.api_format) }), "chatgpt_web" => matching_endpoint(endpoints, include_inactive, |endpoint| { endpoint .api_format .trim() .eq_ignore_ascii_case("openai:image") }), "antigravity" => matching_endpoint(endpoints, include_inactive, |endpoint| { endpoint .api_format .trim() .eq_ignore_ascii_case("gemini:generate_content") }), "kiro" => matching_endpoint(endpoints, include_inactive, |endpoint| { endpoint .api_format .trim() .eq_ignore_ascii_case("claude:messages") }) .or_else(|| matching_endpoint(endpoints, include_inactive, |_| true)), "claude_code" => matching_endpoint(endpoints, include_inactive, |endpoint| { endpoint .api_format .trim() .eq_ignore_ascii_case("claude:messages") }), "gemini_cli" => matching_endpoint(endpoints, include_inactive, |endpoint| { endpoint .api_format .trim() .eq_ignore_ascii_case("gemini:generate_content") }), "vertex_ai" => matching_endpoint(endpoints, include_inactive, |endpoint| { endpoint .api_format .trim() .eq_ignore_ascii_case("gemini:generate_content") }) .or_else(|| { matching_endpoint(endpoints, include_inactive, |endpoint| { endpoint .api_format .trim() .eq_ignore_ascii_case("claude:messages") }) }), _ => matching_endpoint(endpoints, include_inactive, |_| true), } } #[derive(Debug, Clone)] pub(crate) struct ProviderOAuthRuntimeEndpoints { pub(crate) endpoints: Vec, pub(crate) runtime_endpoint: Option, } async fn resolve_provider_runtime_endpoints_with_selector( state: &AdminAppState<'_>, provider: &StoredProviderCatalogProvider, provider_type: &str, endpoint_selector: fn( &str, &[StoredProviderCatalogEndpoint], bool, ) -> Option, ) -> Result { let mut endpoints = state .list_provider_catalog_endpoints_by_provider_ids(std::slice::from_ref(&provider.id)) .await?; let mut runtime_endpoint = endpoint_selector(provider_type, &endpoints, true); if runtime_endpoint.is_none() && state .fixed_provider_template(&provider.provider_type) .is_some() && state.has_provider_catalog_data_writer() { reconcile_admin_fixed_provider_template_endpoints(state, provider).await?; endpoints = state .list_provider_catalog_endpoints_by_provider_ids(std::slice::from_ref(&provider.id)) .await?; runtime_endpoint = endpoint_selector(provider_type, &endpoints, true); } Ok(ProviderOAuthRuntimeEndpoints { endpoints, runtime_endpoint, }) } pub(crate) async fn resolve_provider_oauth_runtime_endpoints( state: &AdminAppState<'_>, provider: &StoredProviderCatalogProvider, provider_type: &str, ) -> Result { resolve_provider_runtime_endpoints_with_selector( state, provider, provider_type, select_provider_oauth_runtime_endpoint, ) .await } async fn resolve_provider_quota_runtime_endpoints( state: &AdminAppState<'_>, provider: &StoredProviderCatalogProvider, provider_type: &str, ) -> Result { resolve_provider_runtime_endpoints_with_selector( state, provider, provider_type, provider_quota_refresh_endpoint_for_provider, ) .await } 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), GatewayError> { let provider_type = provider.provider_type.trim().to_ascii_lowercase(); if !provider_type_supports_quota_refresh(&provider_type) { return Ok((false, None)); } let ProviderOAuthRuntimeEndpoints { runtime_endpoint, .. } = resolve_provider_quota_runtime_endpoints(state, provider, &provider_type).await?; let Some(endpoint) = runtime_endpoint else { return Ok((false, None)); }; let Some(key) = state .read_provider_catalog_keys_by_ids(&[key_id.to_string()]) .await? .into_iter() .next() else { return Ok((false, None)); }; if provider_type == "kiro" && !provider_key_is_oauth_managed(&key, provider_type.as_str()) { return Ok((false, None)); } let payload = refresh_provider_pool_quota_locally( state, provider, &endpoint, &provider_type, vec![key], proxy_override.cloned(), ) .await?; let Some(payload) = payload else { return Ok((false, None)); }; let success = payload .get("success") .and_then(serde_json::Value::as_u64) .unwrap_or(0); let error = if success == 0 { payload .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) } else { None }; Ok((true, error)) } pub(crate) fn spawn_provider_oauth_account_state_refresh_after_update( app: AppState, provider: StoredProviderCatalogProvider, key_id: String, proxy_override: Option, ) { spawn_fire_and_forget(TASK_KEY_PROVIDER_OAUTH_ACCOUNT_REFRESH, async move { let _ = refresh_provider_oauth_account_state_after_update( &AdminAppState::new(&app), &provider, &key_id, proxy_override.as_ref(), ) .await; }); }