i/Aether
1
0
Fork 0
mirror of https://github.com/fawney19/Aether.git synced 2026-09-02 17:30:23 +08:00
Code Issues Packages Projects Releases Wiki Activity

refactor: extract provider pool abstractions

Browse Source
This commit is contained in:
fawney19
2026-05-13 18:19:15 +08:00
parent 3c2497f019
commit 5d1460e051
55 changed files with 3469 additions and 2184 deletions
Download Patch File Download Diff File Expand all files Collapse all files

6
apps/aether-gateway/src/handlers/admin/mod.rs
View File

@@ -25,10 +25,8 @@ pub(crate) use self::provider::oauth::errors::build_internal_control_error_respo
pub(crate) use self::provider::oauth::provisioning::{
create_provider_oauth_catalog_key, update_existing_provider_oauth_catalog_key,
};
pub(crate) use self::provider::oauth::quota::antigravity::refresh_antigravity_provider_quota_locally;
pub(crate) use self::provider::oauth::quota::chatgpt_web::refresh_chatgpt_web_provider_quota_locally;
pub(crate) use self::provider::oauth::quota::codex::refresh_codex_provider_quota_locally;
pub(crate) use self::provider::oauth::quota::kiro::refresh_kiro_provider_quota_locally;
pub(crate) use self::provider::oauth::quota::dispatch::refresh_provider_pool_quota_locally;
pub(crate) use self::provider::oauth::quota::shared::provider_quota_refresh_endpoint_for_provider;
pub(crate) use self::provider::oauth::quota::shared::provider_type_supports_quota_refresh;
pub(crate) use self::provider::oauth::runtime::{
provider_oauth_maintenance_endpoint_for_provider, provider_oauth_runtime_endpoint_for_provider,

53
apps/aether-gateway/src/handlers/admin/provider/endpoint_keys/quota.rs
View File

@@ -13,15 +13,12 @@ use axum::{
use serde_json::json;
use std::collections::{BTreeMap, BTreeSet};
use super::super::oauth::quota::antigravity::refresh_antigravity_provider_quota_locally;
use super::super::oauth::quota::chatgpt_web::refresh_chatgpt_web_provider_quota_locally;
use super::super::oauth::quota::codex::refresh_codex_provider_quota_locally;
use super::super::oauth::quota::kiro::refresh_kiro_provider_quota_locally;
use super::super::oauth::quota::dispatch::refresh_provider_pool_quota_locally;
use super::super::oauth::quota::shared::normalize_string_id_list;
use super::super::oauth::quota::shared::{
provider_quota_refresh_endpoint_for_provider, provider_quota_refresh_missing_endpoint_message,
provider_type_supports_quota_refresh, unsupported_provider_quota_refresh_message,
};
use super::super::oauth::runtime::provider_oauth_maintenance_endpoint_for_provider;
use super::super::write::provider::reconcile_admin_fixed_provider_template_endpoints;
fn unsupported_provider_quota_refresh_response(provider_type: &str) -> Response<Body> {
@@ -115,7 +112,7 @@ pub(super) async fn maybe_handle(
.list_provider_catalog_endpoints_by_provider_ids(std::slice::from_ref(&provider_id))
.await?;
let mut endpoint =
provider_oauth_maintenance_endpoint_for_provider(&normalized_provider_type, &endpoints);
provider_quota_refresh_endpoint_for_provider(&normalized_provider_type, &endpoints, true);
if endpoint.is_none() && is_fixed_provider {
if !state.has_provider_catalog_data_writer() {
@@ -136,8 +133,11 @@ pub(super) async fn maybe_handle(
endpoints = state
.list_provider_catalog_endpoints_by_provider_ids(std::slice::from_ref(&provider_id))
.await?;
endpoint =
provider_oauth_maintenance_endpoint_for_provider(&normalized_provider_type, &endpoints);
endpoint = provider_quota_refresh_endpoint_for_provider(
&normalized_provider_type,
&endpoints,
true,
);
}
if !provider_type_supports_quota_refresh(&normalized_provider_type) {
@@ -147,15 +147,7 @@ pub(super) async fn maybe_handle(
}
let Some(endpoint) = endpoint else {
let detail = match normalized_provider_type.as_str() {
"codex" => "找不到有效的 openai:responses 端点",
"antigravity" => "找不到有效的 gemini:generate_content 端点",
"kiro" => "找不到有效的 Kiro 端点",
"chatgpt_web" => "找不到有效的 openai:image 端点",
"claude_code" => "找不到有效的 claude:messages 端点",
"gemini_cli" | "vertex_ai" => "找不到有效的 gemini:generate_content 端点",
_ => "找不到有效端点",
};
let detail = provider_quota_refresh_missing_endpoint_message(&normalized_provider_type);
return Ok(Some(
(
http::StatusCode::BAD_REQUEST,
@@ -231,23 +223,16 @@ pub(super) async fn maybe_handle(
));
}
let Some(payload) = (match normalized_provider_type.as_str() {
"codex" => {
refresh_codex_provider_quota_locally(state, &provider, &endpoint, keys, None).await?
}
"kiro" => {
refresh_kiro_provider_quota_locally(state, &provider, &endpoint, keys, None).await?
}
"antigravity" => {
refresh_antigravity_provider_quota_locally(state, &provider, &endpoint, keys, None)
.await?
}
"chatgpt_web" => {
refresh_chatgpt_web_provider_quota_locally(state, &provider, &endpoint, keys, None)
.await?
}
_ => None,
}) else {
let Some(payload) = refresh_provider_pool_quota_locally(
state,
&provider,
&endpoint,
&normalized_provider_type,
keys,
None,
)
.await?
else {
return Ok(None);
};
Ok(Some(Json(payload).into_response()))

325
apps/aether-gateway/src/handlers/admin/provider/oauth/dispatch/kiro.rs
View File

@@ -5,13 +5,12 @@ use crate::handlers::admin::provider::shared::payloads::{
use crate::handlers::admin::request::{AdminAppState, AdminKiroAuthConfig};
use crate::provider_transport::kiro::{build_kiro_request_auth_from_config, KiroRequestAuth};
use aether_contracts::ProxySnapshot;
use serde_json::{json, Value};
use std::time::{SystemTime, UNIX_EPOCH};
use aether_oauth::core::OAuthError;
use aether_oauth::provider::providers::KiroProviderOAuthAdapter;
use aether_oauth::provider::ProviderOAuthTransportContext;
use serde_json::Value;
use url::form_urlencoded;
const KIRO_IDC_AMZ_USER_AGENT: &str =
"aws-sdk-js/3.738.0 ua/2.1 os/other lang/js md/browser#unknown_unknown api/sso-oidc#3.738.0 m/E KiroIDE";
pub(super) fn admin_provider_oauth_kiro_refresh_base_url_override(
state: &AdminAppState<'_>,
override_key: &str,
@@ -21,29 +20,6 @@ pub(super) fn admin_provider_oauth_kiro_refresh_base_url_override(
(!normalized.is_empty()).then(|| normalized.to_string())
}
fn admin_provider_oauth_kiro_build_refresh_url(
auth_config: &AdminKiroAuthConfig,
override_base_url: Option<&str>,
path: &str,
default_host: impl FnOnce(&str) -> String,
) -> String {
if let Some(base_url) = override_base_url
.map(str::trim)
.filter(|value| !value.is_empty())
{
return format!("{}/{}", base_url.trim_end_matches('/'), path);
}
let region = auth_config.effective_auth_region();
default_host(region)
}
fn admin_provider_oauth_kiro_effective_host(url: &str, fallback_host: String) -> String {
reqwest::Url::parse(url)
.ok()
.and_then(|value| value.host_str().map(ToOwned::to_owned))
.unwrap_or(fallback_host)
}
fn admin_provider_oauth_kiro_ide_tag(kiro_version: &str, machine_id: &str) -> String {
if machine_id.trim().is_empty() {
format!("KiroIDE-{kiro_version}")
@@ -52,41 +28,49 @@ fn admin_provider_oauth_kiro_ide_tag(kiro_version: &str, machine_id: &str) -> St
}
}
fn admin_provider_oauth_kiro_refresh_expires_at(payload: &Value) -> u64 {
let expires_in = payload
.get("expiresIn")
.and_then(|value| {
value
.as_u64()
.or_else(|| value.as_str()?.parse::<u64>().ok())
})
.unwrap_or(3600);
SystemTime::now()
.duration_since(UNIX_EPOCH)
.ok()
.map(|value| value.as_secs())
.unwrap_or_default()
.saturating_add(expires_in)
fn admin_provider_oauth_kiro_refresh_context(
proxy: Option<ProxySnapshot>,
) -> ProviderOAuthTransportContext {
ProviderOAuthTransportContext {
provider_id: String::new(),
provider_type: "kiro".to_string(),
endpoint_id: None,
key_id: None,
auth_type: Some("oauth".to_string()),
decrypted_api_key: None,
decrypted_auth_config: None,
provider_config: None,
endpoint_config: None,
key_config: None,
network: aether_oauth::network::OAuthNetworkContext::provider_operation(proxy),
}
}
fn admin_provider_oauth_kiro_refresh_response_json(
body_text: &str,
json_body: Option<Value>,
) -> Result<Value, String> {
json_body
.or_else(|| serde_json::from_str::<Value>(body_text).ok())
.ok_or_else(|| "refresh 接口返回了非 JSON 响应".to_string())
}
fn admin_provider_oauth_kiro_refresh_error_detail(
status: http::StatusCode,
body_text: &str,
fn admin_provider_oauth_kiro_refresh_error(
auth_config: &AdminKiroAuthConfig,
error: OAuthError,
) -> String {
let detail = body_text.trim();
if detail.is_empty() {
format!("HTTP {}", status.as_u16())
let prefix = if auth_config.is_idc_auth() {
"IDC refresh"
} else {
detail.to_string()
"social refresh"
};
match error {
OAuthError::HttpStatus {
status_code,
body_excerpt,
} => {
let detail = body_excerpt.trim();
if detail.is_empty() {
format!("{prefix} 失败: HTTP {status_code}")
} else {
format!("{prefix} 失败: {detail}")
}
}
OAuthError::Transport(message) => format!("{prefix} 请求失败: {message}"),
OAuthError::InvalidRequest(message) => format!("{prefix} 参数无效: {message}"),
OAuthError::InvalidResponse(message) => format!("{prefix} 返回无效响应: {message}"),
error => format!("{prefix} 失败: {error}"),
}
}
@@ -97,216 +81,25 @@ pub(super) async fn refresh_admin_provider_oauth_kiro_auth_config(
social_refresh_base_url: Option<&str>,
idc_refresh_base_url: Option<&str>,
) -> Result<AdminKiroAuthConfig, String> {
if auth_config.is_idc_auth() {
let fallback_host = format!("oidc.{}.amazonaws.com", auth_config.effective_auth_region());
let url = admin_provider_oauth_kiro_build_refresh_url(
let adapter = KiroProviderOAuthAdapter::default().with_refresh_base_urls(
social_refresh_base_url
.map(str::trim)
.filter(|value| !value.is_empty())
.map(ToOwned::to_owned),
idc_refresh_base_url
.map(str::trim)
.filter(|value| !value.is_empty())
.map(ToOwned::to_owned),
);
let ctx = admin_provider_oauth_kiro_refresh_context(proxy);
adapter
.refresh_auth_config(
&crate::oauth::GatewayOAuthHttpExecutor::new(*state),
&ctx,
auth_config,
idc_refresh_base_url,
"token",
|region| format!("https://oidc.{region}.amazonaws.com/token"),
);
let host = admin_provider_oauth_kiro_effective_host(&url, fallback_host);
let headers = reqwest::header::HeaderMap::from_iter([
(
reqwest::header::CONTENT_TYPE,
reqwest::header::HeaderValue::from_static("application/json"),
),
(
reqwest::header::HOST,
reqwest::header::HeaderValue::from_str(&host)
.map_err(|_| "IDC host 无效".to_string())?,
),
(
reqwest::header::HeaderName::from_static("x-amz-user-agent"),
reqwest::header::HeaderValue::from_static(KIRO_IDC_AMZ_USER_AGENT),
),
(
reqwest::header::USER_AGENT,
reqwest::header::HeaderValue::from_static("node"),
),
(
reqwest::header::ACCEPT,
reqwest::header::HeaderValue::from_static("*/*"),
),
]);
let response = state
.execute_admin_provider_oauth_http_request(
"kiro_batch_refresh:idc",
reqwest::Method::POST,
&url,
&headers,
Some("application/json"),
Some(json!({
"clientId": auth_config
.client_id
.as_deref()
.map(str::trim)
.unwrap_or_default(),
"clientSecret": auth_config
.client_secret
.as_deref()
.map(str::trim)
.unwrap_or_default(),
"refreshToken": auth_config
.refresh_token
.as_deref()
.map(str::trim)
.unwrap_or_default(),
"grantType": "refresh_token",
})),
None,
proxy.clone(),
)
.await
.map_err(|err| format!("IDC refresh 请求失败: {err}"))?;
if !response.status.is_success() {
return Err(format!(
"IDC refresh 失败: {}",
admin_provider_oauth_kiro_refresh_error_detail(
response.status,
&response.body_text
)
));
}
let payload = admin_provider_oauth_kiro_refresh_response_json(
&response.body_text,
response.json_body,
)?;
let access_token = payload
.get("accessToken")
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
.ok_or_else(|| "IDC refresh 返回了空 accessToken".to_string())?;
let mut refreshed = auth_config.clone();
refreshed.access_token = Some(access_token.to_string());
refreshed.expires_at = Some(admin_provider_oauth_kiro_refresh_expires_at(&payload));
if refreshed
.machine_id
.as_deref()
.map(str::trim)
.is_none_or(|value| value.is_empty())
{
refreshed.machine_id =
crate::provider_transport::kiro::generate_machine_id(auth_config, None);
}
if let Some(refresh_token) = payload
.get("refreshToken")
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
{
refreshed.refresh_token = Some(refresh_token.to_string());
}
return Ok(refreshed);
}
let machine_id = crate::provider_transport::kiro::generate_machine_id(auth_config, None)
.ok_or_else(|| "缺少 machine_id 种子，无法刷新 social token".to_string())?;
let fallback_host = format!(
"prod.{}.auth.desktop.kiro.dev",
auth_config.effective_auth_region()
);
let url = admin_provider_oauth_kiro_build_refresh_url(
auth_config,
social_refresh_base_url,
"refreshToken",
|region| format!("https://prod.{region}.auth.desktop.kiro.dev/refreshToken"),
);
let host = admin_provider_oauth_kiro_effective_host(&url, fallback_host);
let user_agent =
admin_provider_oauth_kiro_ide_tag(auth_config.effective_kiro_version(), &machine_id);
let headers = reqwest::header::HeaderMap::from_iter([
(
reqwest::header::USER_AGENT,
reqwest::header::HeaderValue::from_str(&user_agent)
.map_err(|_| "Kiro User-Agent 无效".to_string())?,
),
(
reqwest::header::HOST,
reqwest::header::HeaderValue::from_str(&host)
.map_err(|_| "Kiro host 无效".to_string())?,
),
(
reqwest::header::ACCEPT,
reqwest::header::HeaderValue::from_static("application/json, text/plain, */*"),
),
(
reqwest::header::CONTENT_TYPE,
reqwest::header::HeaderValue::from_static("application/json"),
),
(
reqwest::header::CONNECTION,
reqwest::header::HeaderValue::from_static("close"),
),
(
reqwest::header::ACCEPT_ENCODING,
reqwest::header::HeaderValue::from_static("gzip, compress, deflate, br"),
),
]);
let response = state
.execute_admin_provider_oauth_http_request(
"kiro_batch_refresh:social",
reqwest::Method::POST,
&url,
&headers,
Some("application/json"),
Some(json!({
"refreshToken": auth_config
.refresh_token
.as_deref()
.map(str::trim)
.unwrap_or_default(),
})),
None,
proxy,
)
.await
.map_err(|err| format!("social refresh 请求失败: {err}"))?;
if !response.status.is_success() {
return Err(format!(
"social refresh 失败: {}",
admin_provider_oauth_kiro_refresh_error_detail(response.status, &response.body_text)
));
}
let payload =
admin_provider_oauth_kiro_refresh_response_json(&response.body_text, response.json_body)?;
let access_token = payload
.get("accessToken")
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
.ok_or_else(|| "social refresh 返回了空 accessToken".to_string())?;
let mut refreshed = auth_config.clone();
refreshed.access_token = Some(access_token.to_string());
refreshed.expires_at = Some(admin_provider_oauth_kiro_refresh_expires_at(&payload));
if refreshed
.machine_id
.as_deref()
.map(str::trim)
.is_none_or(|value| value.is_empty())
{
refreshed.machine_id = Some(machine_id);
}
if let Some(refresh_token) = payload
.get("refreshToken")
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
{
refreshed.refresh_token = Some(refresh_token.to_string());
}
if let Some(profile_arn) = payload
.get("profileArn")
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
{
refreshed.profile_arn = Some(profile_arn.to_string());
}
Ok(refreshed)
.map_err(|error| admin_provider_oauth_kiro_refresh_error(auth_config, error))
}
fn build_kiro_usage_url(auth: &KiroRequestAuth) -> String {

58
apps/aether-gateway/src/handlers/admin/provider/oauth/quota/antigravity.rs
View File

@@ -1,17 +1,17 @@
use super::shared::{
build_quota_snapshot_payload, coerce_json_f64, coerce_json_string,
default_provider_quota_execution_timeouts, execute_provider_quota_plan,
build_provider_quota_execution_plan, 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, ProxySnapshot, RequestBody};
use aether_contracts::ProxySnapshot;
use aether_data_contracts::repository::provider_catalog::{
StoredProviderCatalogEndpoint, StoredProviderCatalogKey, StoredProviderCatalogProvider,
};
use aether_provider_pool::build_antigravity_pool_quota_request;
use serde_json::json;
use std::collections::BTreeMap;
use std::time::{SystemTime, UNIX_EPOCH};
@@ -21,18 +21,9 @@ async fn execute_antigravity_quota_plan(
transport: &AdminGatewayProviderTransportSnapshot,
authorization: (String, String),
project_id: &str,
mut identity_headers: BTreeMap<String, String>,
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);
headers.insert("content-type".to_string(), "application/json".to_string());
headers.insert("accept".to_string(), "application/json".to_string());
headers
.entry("user-agent".to_string())
.or_insert_with(|| "antigravity".to_string());
let body = json!({ "project": project_id });
let proxy = match proxy_override {
Some(proxy) => Some(proxy.clone()),
None => {
@@ -46,35 +37,20 @@ async fn execute_antigravity_quota_plan(
.or(Some(default_provider_quota_execution_timeouts(
proxy.as_ref(),
)));
let plan = ExecutionPlan {
request_id: format!("antigravity-quota:{}", transport.key.id),
candidate_id: None,
provider_name: Some("antigravity".to_string()),
provider_id: transport.provider.id.clone(),
endpoint_id: transport.endpoint.id.clone(),
key_id: transport.key.id.clone(),
method: "POST".to_string(),
url: format!(
"{}{}",
transport.endpoint.base_url.trim_end_matches('/'),
ANTIGRAVITY_FETCH_AVAILABLE_MODELS_PATH
),
headers,
content_type: Some("application/json".to_string()),
content_encoding: None,
body: RequestBody {
json_body: Some(body),
body_bytes_b64: None,
body_ref: None,
},
stream: false,
client_api_format: "gemini:generate_content".to_string(),
provider_api_format: "antigravity:fetch_available_models".to_string(),
model_name: Some("fetchAvailableModels".to_string()),
let spec = build_antigravity_pool_quota_request(
&transport.key.id,
&transport.endpoint.base_url,
authorization,
project_id,
identity_headers,
);
let plan = build_provider_quota_execution_plan(
transport,
spec,
proxy,
transport_profile: state.resolve_transport_profile(transport),
state.resolve_transport_profile(transport),
timeouts,
};
);
execute_provider_quota_plan(state, transport, plan, "antigravity").await
}

246
apps/aether-gateway/src/handlers/admin/provider/oauth/quota/chatgpt_web.rs
View File

@@ -10,94 +10,18 @@ use crate::handlers::admin::provider::shared::payloads::{
use crate::handlers::admin::request::{AdminAppState, AdminGatewayProviderTransportSnapshot};
use crate::GatewayError;
use aether_admin::provider::quota::parse_chatgpt_web_conversation_init_response;
use aether_contracts::{
ExecutionPlan, ProxySnapshot, RequestBody, EXECUTION_REQUEST_ACCEPT_INVALID_CERTS_HEADER,
};
use aether_contracts::ProxySnapshot;
use aether_data_contracts::repository::provider_catalog::{
StoredProviderCatalogEndpoint, StoredProviderCatalogKey, StoredProviderCatalogProvider,
};
use aether_provider_pool::{
build_chatgpt_web_pool_quota_request, enrich_chatgpt_web_quota_metadata,
normalize_chatgpt_web_image_quota_limit,
};
use serde_json::json;
use std::collections::BTreeMap;
use std::time::{SystemTime, UNIX_EPOCH};
const CHATGPT_WEB_DEFAULT_BASE_URL: &str = "https://chatgpt.com";
const CHATGPT_WEB_CONVERSATION_INIT_PATH: &str = "/backend-api/conversation/init";
const CHATGPT_WEB_USER_AGENT: &str = "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/143.0.0.0 Safari/537.36 Edg/143.0.0.0";
const CHATGPT_WEB_CLIENT_VERSION: &str = "prod-be885abbfcfe7b1f511e88b3003d9ee44757fbad";
const CHATGPT_WEB_BUILD_NUMBER: &str = "5955942";
const CHATGPT_WEB_SEC_CH_UA: &str =
r#""Microsoft Edge";v="143", "Chromium";v="143", "Not A(Brand";v="24""#;
const PLACEHOLDER_API_KEY: &str = "__placeholder__";
const CHATGPT_WEB_FREE_IMAGE_QUOTA_LIMIT: f64 = 25.0;
fn chatgpt_web_base_url(endpoint: &StoredProviderCatalogEndpoint) -> String {
let base_url = endpoint.base_url.trim().trim_end_matches('/');
if base_url.is_empty() {
CHATGPT_WEB_DEFAULT_BASE_URL.to_string()
} else {
base_url.to_string()
}
}
fn build_chatgpt_web_quota_headers(
authorization: (String, String),
base_url: &str,
) -> BTreeMap<String, String> {
let device_id = uuid::Uuid::new_v4().to_string();
let session_id = uuid::Uuid::new_v4().to_string();
let mut headers = BTreeMap::from([
("accept".to_string(), "application/json".to_string()),
("content-type".to_string(), "application/json".to_string()),
("user-agent".to_string(), CHATGPT_WEB_USER_AGENT.to_string()),
("origin".to_string(), base_url.to_string()),
("referer".to_string(), format!("{base_url}/")),
(
"accept-language".to_string(),
"zh-CN,zh;q=0.9,en;q=0.8,en-US;q=0.7".to_string(),
),
("cache-control".to_string(), "no-cache".to_string()),
("pragma".to_string(), "no-cache".to_string()),
("priority".to_string(), "u=1, i".to_string()),
("sec-ch-ua".to_string(), CHATGPT_WEB_SEC_CH_UA.to_string()),
("sec-ch-ua-arch".to_string(), r#""x86""#.to_string()),
("sec-ch-ua-bitness".to_string(), r#""64""#.to_string()),
("sec-ch-ua-mobile".to_string(), "?0".to_string()),
("sec-ch-ua-model".to_string(), r#""""#.to_string()),
("sec-ch-ua-platform".to_string(), r#""Windows""#.to_string()),
(
"sec-ch-ua-platform-version".to_string(),
r#""19.0.0""#.to_string(),
),
("sec-fetch-dest".to_string(), "empty".to_string()),
("sec-fetch-mode".to_string(), "cors".to_string()),
("sec-fetch-site".to_string(), "same-origin".to_string()),
("oai-device-id".to_string(), device_id),
("oai-session-id".to_string(), session_id),
("oai-language".to_string(), "zh-CN".to_string()),
(
"oai-client-version".to_string(),
CHATGPT_WEB_CLIENT_VERSION.to_string(),
),
(
"oai-client-build-number".to_string(),
CHATGPT_WEB_BUILD_NUMBER.to_string(),
),
(
"x-openai-target-path".to_string(),
CHATGPT_WEB_CONVERSATION_INIT_PATH.to_string(),
),
(
"x-openai-target-route".to_string(),
CHATGPT_WEB_CONVERSATION_INIT_PATH.to_string(),
),
(
EXECUTION_REQUEST_ACCEPT_INVALID_CERTS_HEADER.to_string(),
"true".to_string(),
),
]);
headers.insert(authorization.0.to_ascii_lowercase(), authorization.1);
headers
}
fn chatgpt_web_auth_config(
transport: &AdminGatewayProviderTransportSnapshot,
@@ -111,133 +35,6 @@ fn chatgpt_web_auth_config(
.and_then(|value| serde_json::from_str::<serde_json::Value>(value).ok())
}
fn chatgpt_web_auth_config_string(
auth_config: Option<&serde_json::Value>,
fields: &[&str],
) -> Option<String> {
let object = auth_config.and_then(serde_json::Value::as_object)?;
fields.iter().find_map(|field| {
object
.get(*field)
.and_then(serde_json::Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
.map(ToOwned::to_owned)
})
}
fn enrich_chatgpt_web_quota_metadata(
metadata: &mut serde_json::Value,
auth_config: Option<&serde_json::Value>,
) {
let Some(object) = metadata.as_object_mut() else {
return;
};
for (target, fields) in [
("plan_type", &["plan_type", "tier", "plan"][..]),
("email", &["email"][..]),
("account_id", &["account_id", "accountId"][..]),
("account_user_id", &["account_user_id", "accountUserId"][..]),
("user_id", &["user_id", "userId"][..]),
] {
if object.contains_key(target) {
continue;
}
if let Some(value) = chatgpt_web_auth_config_string(auth_config, fields) {
object.insert(target.to_string(), json!(value));
}
}
}
fn chatgpt_web_json_number(value: Option<&serde_json::Value>) -> Option<f64> {
let value = value?;
if let Some(number) = value.as_f64() {
return number.is_finite().then_some(number);
}
value
.as_str()
.map(str::trim)
.filter(|value| !value.is_empty())
.and_then(|value| value.parse::<f64>().ok())
.filter(|value| value.is_finite())
}
fn chatgpt_web_json_string(value: Option<&serde_json::Value>) -> Option<&str> {
value
.and_then(serde_json::Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
}
fn existing_chatgpt_web_image_quota_limit(
upstream_metadata: Option<&serde_json::Value>,
) -> Option<f64> {
upstream_metadata
.and_then(serde_json::Value::as_object)
.and_then(|metadata| metadata.get("chatgpt_web"))
.and_then(serde_json::Value::as_object)
.and_then(|bucket| chatgpt_web_json_number(bucket.get("image_quota_total")))
.filter(|value| *value > 0.0)
}
fn infer_chatgpt_web_image_quota_limit(
plan_type: Option<&str>,
remaining: Option<f64>,
existing_limit: Option<f64>,
) -> Option<f64> {
let normalized_plan = plan_type.unwrap_or_default().trim().to_ascii_lowercase();
if normalized_plan == "free" {
return Some(CHATGPT_WEB_FREE_IMAGE_QUOTA_LIMIT);
}
if let Some(existing_limit) = existing_limit.filter(|value| *value > 0.0) {
return Some(existing_limit);
}
remaining.filter(|value| *value > 0.0)
}
fn normalize_chatgpt_web_image_quota_limit(
metadata: &mut serde_json::Value,
upstream_metadata: Option<&serde_json::Value>,
) {
let existing_limit = existing_chatgpt_web_image_quota_limit(upstream_metadata);
let Some(object) = metadata.as_object_mut() else {
return;
};
let remaining = chatgpt_web_json_number(object.get("image_quota_remaining"));
let explicit_limit =
chatgpt_web_json_number(object.get("image_quota_total")).filter(|value| *value > 0.0);
let plan_type = chatgpt_web_json_string(object.get("plan_type"));
let is_free_plan = plan_type.is_some_and(|value| value.trim().eq_ignore_ascii_case("free"));
let limit = if is_free_plan {
Some(CHATGPT_WEB_FREE_IMAGE_QUOTA_LIMIT)
} else {
explicit_limit
.or_else(|| infer_chatgpt_web_image_quota_limit(plan_type, remaining, existing_limit))
};
if let Some(limit) = limit {
object.insert("image_quota_total".to_string(), json!(limit));
if !object.contains_key("image_quota_used") {
if let Some(remaining) = remaining {
object.insert(
"image_quota_used".to_string(),
json!((limit - remaining).max(0.0)),
);
} else if object
.get("image_quota_blocked")
.and_then(serde_json::Value::as_bool)
== Some(true)
{
object.insert("image_quota_used".to_string(), json!(limit));
}
}
}
}
async fn resolve_chatgpt_web_quota_auth(
state: &AdminAppState<'_>,
transport: &AdminGatewayProviderTransportSnapshot,
@@ -262,7 +59,6 @@ async fn execute_chatgpt_web_quota_plan(
authorization: (String, String),
proxy_override: Option<&ProxySnapshot>,
) -> Result<ProviderQuotaExecutionOutcome, GatewayError> {
let base_url = chatgpt_web_base_url(endpoint);
let proxy = match proxy_override {
Some(proxy) => Some(proxy.clone()),
None => {
@@ -276,33 +72,15 @@ async fn execute_chatgpt_web_quota_plan(
.or(Some(default_provider_quota_execution_timeouts(
proxy.as_ref(),
)));
let plan = ExecutionPlan {
request_id: format!("chatgpt-web-quota:{}", transport.key.id),
candidate_id: None,
provider_name: Some("chatgpt_web".to_string()),
provider_id: transport.provider.id.clone(),
endpoint_id: transport.endpoint.id.clone(),
key_id: transport.key.id.clone(),
method: "POST".to_string(),
url: format!("{base_url}{CHATGPT_WEB_CONVERSATION_INIT_PATH}"),
headers: build_chatgpt_web_quota_headers(authorization, base_url.as_str()),
content_type: Some("application/json".to_string()),
content_encoding: None,
body: RequestBody::from_json(json!({
"gizmo_id": serde_json::Value::Null,
"requested_default_model": serde_json::Value::Null,
"conversation_id": serde_json::Value::Null,
"timezone_offset_min": -480,
"system_hints": ["picture_v2"],
})),
stream: false,
client_api_format: "openai:image".to_string(),
provider_api_format: "chatgpt_web:conversation_init".to_string(),
model_name: Some("chatgpt-web-conversation-init".to_string()),
let spec =
build_chatgpt_web_pool_quota_request(&transport.key.id, &endpoint.base_url, authorization);
let plan = super::shared::build_provider_quota_execution_plan(
transport,
spec,
proxy,
transport_profile: state.resolve_transport_profile(transport),
state.resolve_transport_profile(transport),
timeouts,
};
);
execute_provider_quota_plan(state, transport, plan, "chatgpt_web").await
}

44
apps/aether-gateway/src/handlers/admin/provider/oauth/quota/codex/mod.rs
View File

@@ -11,7 +11,7 @@ use self::parse::{
build_codex_quota_exhausted_fallback_metadata, parse_codex_usage_headers,
parse_codex_wham_usage_response,
};
use self::plan::{build_codex_refresh_headers, execute_codex_quota_plan};
use self::plan::{build_codex_quota_request_spec, execute_codex_quota_plan};
use super::shared::{
build_quota_snapshot_payload, extract_execution_error_message,
persist_provider_quota_refresh_state, provider_auto_remove_banned_keys,
@@ -82,8 +82,8 @@ pub(crate) async fn refresh_codex_provider_quota_locally(
None
};
let headers = match build_codex_refresh_headers(&transport, resolved_oauth_auth) {
Ok(headers) => headers,
let request_spec = match build_codex_quota_request_spec(&transport, resolved_oauth_auth) {
Ok(request_spec) => request_spec,
Err(message) => {
failed_count += 1;
results.push(json!({
@@ -96,23 +96,27 @@ pub(crate) async fn refresh_codex_provider_quota_locally(
}
};
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 result = match execute_codex_quota_plan(
state,
&transport,
request_spec,
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()

86
apps/aether-gateway/src/handlers/admin/provider/oauth/quota/codex/plan.rs
View File

@@ -1,65 +1,33 @@
use super::super::shared::{
default_provider_quota_execution_timeouts, execute_provider_quota_plan,
ProviderQuotaExecutionOutcome,
build_provider_quota_execution_plan, default_provider_quota_execution_timeouts,
execute_provider_quota_plan, ProviderQuotaExecutionOutcome,
};
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, ProxySnapshot, RequestBody};
use std::collections::BTreeMap;
use aether_contracts::ProxySnapshot;
use aether_provider_pool::{build_codex_pool_quota_request, ProviderPoolQuotaRequestSpec};
pub(super) fn build_codex_refresh_headers(
pub(super) fn build_codex_quota_request_spec(
transport: &AdminGatewayProviderTransportSnapshot,
resolved_oauth_auth: Option<(String, String)>,
) -> Result<BTreeMap<String, String>, String> {
let mut headers = BTreeMap::new();
headers.insert("accept".to_string(), "application/json".to_string());
if let Some((name, value)) = resolved_oauth_auth {
headers.insert(name.to_ascii_lowercase(), value);
} else {
let decrypted_key = transport.key.decrypted_api_key.trim();
if decrypted_key.is_empty() || decrypted_key == "__placeholder__" {
return Err("缺少 OAuth 认证信息，请先授权/刷新 Token".to_string());
}
headers.insert(
"authorization".to_string(),
format!("Bearer {decrypted_key}"),
);
}
) -> Result<ProviderPoolQuotaRequestSpec, String> {
let auth_config = transport
.key
.decrypted_auth_config
.as_deref()
.and_then(|raw| serde_json::from_str::<serde_json::Value>(raw).ok());
let oauth_plan_type = normalize_codex_plan_type(
auth_config
.as_ref()
.and_then(|value| value.get("plan_type"))
.and_then(serde_json::Value::as_str),
);
let oauth_account_id = auth_config
.as_ref()
.and_then(|value| value.get("account_id"))
.and_then(serde_json::Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty());
if oauth_account_id.is_some() && oauth_plan_type.as_deref() != Some("free") {
headers.insert(
"chatgpt-account-id".to_string(),
oauth_account_id.unwrap_or_default().to_string(),
);
}
Ok(headers)
build_codex_pool_quota_request(
&transport.key.id,
resolved_oauth_auth,
Some(transport.key.decrypted_api_key.as_str()),
auth_config.as_ref(),
)
}
pub(super) async fn execute_codex_quota_plan(
state: &AdminAppState<'_>,
transport: &AdminGatewayProviderTransportSnapshot,
headers: BTreeMap<String, String>,
spec: ProviderPoolQuotaRequestSpec,
proxy_override: Option<&ProxySnapshot>,
) -> Result<ProviderQuotaExecutionOutcome, GatewayError> {
let proxy = match proxy_override {
@@ -75,30 +43,12 @@ pub(super) async fn execute_codex_quota_plan(
.or(Some(default_provider_quota_execution_timeouts(
proxy.as_ref(),
)));
let plan = ExecutionPlan {
request_id: format!("codex-quota:{}", transport.key.id),
candidate_id: None,
provider_name: Some("codex".to_string()),
provider_id: transport.provider.id.clone(),
endpoint_id: transport.endpoint.id.clone(),
key_id: transport.key.id.clone(),
method: "GET".to_string(),
url: CODEX_WHAM_USAGE_URL.to_string(),
headers,
content_type: None,
content_encoding: None,
body: RequestBody {
json_body: None,
body_bytes_b64: None,
body_ref: None,
},
stream: false,
client_api_format: "openai:responses".to_string(),
provider_api_format: "openai:responses".to_string(),
model_name: Some("codex-wham-usage".to_string()),
let plan = build_provider_quota_execution_plan(
transport,
spec,
proxy,
transport_profile: state.resolve_transport_profile(transport),
state.resolve_transport_profile(transport),
timeouts,
};
);
execute_provider_quota_plan(state, transport, plan, "codex").await
}

119
apps/aether-gateway/src/handlers/admin/provider/oauth/quota/dispatch.rs Normal file
View File

@@ -0,0 +1,119 @@
use std::future::Future;
use std::pin::Pin;
use super::antigravity::refresh_antigravity_provider_quota_locally;
use super::chatgpt_web::refresh_chatgpt_web_provider_quota_locally;
use super::codex::refresh_codex_provider_quota_locally;
use super::kiro::refresh_kiro_provider_quota_locally;
use crate::handlers::admin::request::AdminAppState;
use crate::GatewayError;
use aether_contracts::ProxySnapshot;
use aether_data_contracts::repository::provider_catalog::{
StoredProviderCatalogEndpoint, StoredProviderCatalogKey, StoredProviderCatalogProvider,
};
type ProviderQuotaRefreshFuture<'a> =
Pin<Box<dyn Future<Output = Result<Option<serde_json::Value>, GatewayError>> + Send + 'a>>;
type ProviderQuotaRefreshHandler = for<'a> fn(
&'a AdminAppState<'a>,
&'a StoredProviderCatalogProvider,
&'a StoredProviderCatalogEndpoint,
Vec<StoredProviderCatalogKey>,
Option<ProxySnapshot>,
) -> ProviderQuotaRefreshFuture<'a>;
const PROVIDER_QUOTA_REFRESH_HANDLERS: &[(&str, ProviderQuotaRefreshHandler)] = &[
(
"antigravity",
refresh_antigravity_provider_quota_locally_boxed,
),
(
"chatgpt_web",
refresh_chatgpt_web_provider_quota_locally_boxed,
),
("codex", refresh_codex_provider_quota_locally_boxed),
("kiro", refresh_kiro_provider_quota_locally_boxed),
];
pub(crate) async fn refresh_provider_pool_quota_locally(
state: &AdminAppState<'_>,
provider: &StoredProviderCatalogProvider,
endpoint: &StoredProviderCatalogEndpoint,
provider_type: &str,
keys: Vec<StoredProviderCatalogKey>,
proxy_override: Option<ProxySnapshot>,
) -> Result<Option<serde_json::Value>, GatewayError> {
let normalized_provider_type = provider_type.trim().to_ascii_lowercase();
let Some((_, handler)) = PROVIDER_QUOTA_REFRESH_HANDLERS
.iter()
.find(|(supported_provider_type, _)| *supported_provider_type == normalized_provider_type)
else {
return Ok(None);
};
handler(state, provider, endpoint, keys, proxy_override).await
}
fn refresh_antigravity_provider_quota_locally_boxed<'a>(
state: &'a AdminAppState<'a>,
provider: &'a StoredProviderCatalogProvider,
endpoint: &'a StoredProviderCatalogEndpoint,
keys: Vec<StoredProviderCatalogKey>,
proxy_override: Option<ProxySnapshot>,
) -> ProviderQuotaRefreshFuture<'a> {
Box::pin(refresh_antigravity_provider_quota_locally(
state,
provider,
endpoint,
keys,
proxy_override,
))
}
fn refresh_chatgpt_web_provider_quota_locally_boxed<'a>(
state: &'a AdminAppState<'a>,
provider: &'a StoredProviderCatalogProvider,
endpoint: &'a StoredProviderCatalogEndpoint,
keys: Vec<StoredProviderCatalogKey>,
proxy_override: Option<ProxySnapshot>,
) -> ProviderQuotaRefreshFuture<'a> {
Box::pin(refresh_chatgpt_web_provider_quota_locally(
state,
provider,
endpoint,
keys,
proxy_override,
))
}
fn refresh_codex_provider_quota_locally_boxed<'a>(
state: &'a AdminAppState<'a>,
provider: &'a StoredProviderCatalogProvider,
endpoint: &'a StoredProviderCatalogEndpoint,
keys: Vec<StoredProviderCatalogKey>,
proxy_override: Option<ProxySnapshot>,
) -> ProviderQuotaRefreshFuture<'a> {
Box::pin(refresh_codex_provider_quota_locally(
state,
provider,
endpoint,
keys,
proxy_override,
))
}
fn refresh_kiro_provider_quota_locally_boxed<'a>(
state: &'a AdminAppState<'a>,
provider: &'a StoredProviderCatalogProvider,
endpoint: &'a StoredProviderCatalogEndpoint,
keys: Vec<StoredProviderCatalogKey>,
proxy_override: Option<ProxySnapshot>,
) -> ProviderQuotaRefreshFuture<'a> {
Box::pin(refresh_kiro_provider_quota_locally(
state,
provider,
endpoint,
keys,
proxy_override,
))
}

102
apps/aether-gateway/src/handlers/admin/provider/oauth/quota/kiro/plan.rs
View File

@@ -1,66 +1,13 @@
use super::super::shared::default_provider_quota_execution_timeouts;
use super::super::shared::{execute_provider_quota_plan, ProviderQuotaExecutionOutcome};
use crate::handlers::admin::provider::shared::payloads::{
KIRO_USAGE_LIMITS_PATH, KIRO_USAGE_SDK_VERSION,
use super::super::shared::{
build_provider_quota_execution_plan, default_provider_quota_execution_timeouts,
execute_provider_quota_plan, ProviderQuotaExecutionOutcome,
};
use crate::handlers::admin::request::{
AdminAppState, AdminGatewayProviderTransportSnapshot, AdminKiroRequestAuth,
};
use crate::GatewayError;
use aether_contracts::{ExecutionPlan, ProxySnapshot, RequestBody};
use std::collections::BTreeMap;
use url::form_urlencoded;
use uuid::Uuid;
fn build_kiro_usage_headers(auth: &AdminKiroRequestAuth) -> BTreeMap<String, String> {
let kiro_version = auth.auth_config.effective_kiro_version();
let machine_id = auth.machine_id.trim();
let ide_tag = if machine_id.is_empty() {
format!("KiroIDE-{kiro_version}")
} else {
format!("KiroIDE-{kiro_version}-{machine_id}")
};
let host = format!(
"q.{}.amazonaws.com",
auth.auth_config.effective_api_region()
);
BTreeMap::from([
(
"x-amz-user-agent".to_string(),
format!("aws-sdk-js/{KIRO_USAGE_SDK_VERSION} {ide_tag}"),
),
(
"user-agent".to_string(),
format!(
"aws-sdk-js/{KIRO_USAGE_SDK_VERSION} ua/2.1 os/other#unknown lang/js md/nodejs#22.21.1 api/codewhispererruntime#1.0.0 m/N,E {ide_tag}"
),
),
("host".to_string(), host),
("amz-sdk-invocation-id".to_string(), Uuid::new_v4().to_string()),
("amz-sdk-request".to_string(), "attempt=1; max=1".to_string()),
("authorization".to_string(), auth.value.clone()),
("connection".to_string(), "close".to_string()),
])
}
fn build_kiro_usage_url(auth: &AdminKiroRequestAuth) -> String {
let host = format!(
"q.{}.amazonaws.com",
auth.auth_config.effective_api_region()
);
let mut serializer = form_urlencoded::Serializer::new(String::new());
serializer.append_pair("origin", "AI_EDITOR");
serializer.append_pair("resourceType", "AGENTIC_REQUEST");
serializer.append_pair("isEmailRequired", "true");
if let Some(profile_arn) = auth.auth_config.profile_arn_for_payload() {
serializer.append_pair("profileArn", profile_arn);
}
format!(
"https://{host}{KIRO_USAGE_LIMITS_PATH}?{}",
serializer.finish()
)
}
use aether_contracts::ProxySnapshot;
use aether_provider_pool::{build_kiro_pool_quota_request, KiroPoolQuotaAuthInput};
pub(super) async fn execute_kiro_quota_plan(
state: &AdminAppState<'_>,
@@ -81,31 +28,26 @@ pub(super) async fn execute_kiro_quota_plan(
.or(Some(default_provider_quota_execution_timeouts(
proxy.as_ref(),
)));
let plan = ExecutionPlan {
request_id: format!("kiro-quota:{}", transport.key.id),
candidate_id: None,
provider_name: Some("kiro".to_string()),
provider_id: transport.provider.id.clone(),
endpoint_id: transport.endpoint.id.clone(),
key_id: transport.key.id.clone(),
method: "GET".to_string(),
url: build_kiro_usage_url(auth),
headers: build_kiro_usage_headers(auth),
content_type: None,
content_encoding: None,
body: RequestBody {
json_body: None,
body_bytes_b64: None,
body_ref: None,
let spec = build_kiro_pool_quota_request(
&transport.key.id,
&KiroPoolQuotaAuthInput {
authorization_value: auth.value.clone(),
api_region: auth.auth_config.effective_api_region().to_string(),
kiro_version: auth.auth_config.effective_kiro_version().to_string(),
machine_id: auth.machine_id.clone(),
profile_arn: auth
.auth_config
.profile_arn_for_payload()
.map(str::to_string),
},
stream: false,
client_api_format: "claude:messages".to_string(),
provider_api_format: "kiro:usage".to_string(),
model_name: Some("kiro-usage-limits".to_string()),
);
let plan = build_provider_quota_execution_plan(
transport,
spec,
proxy,
transport_profile: state.resolve_transport_profile(transport),
state.resolve_transport_profile(transport),
timeouts,
};
);
execute_provider_quota_plan(state, transport, plan, "kiro").await
}

1
apps/aether-gateway/src/handlers/admin/provider/oauth/quota/mod.rs
View File

@@ -1,5 +1,6 @@
pub(crate) mod antigravity;
pub(crate) mod chatgpt_web;
pub(crate) mod codex;
pub(crate) mod dispatch;
pub(crate) mod kiro;
pub(crate) mod shared;

110
apps/aether-gateway/src/handlers/admin/provider/oauth/quota/shared.rs
View File

@@ -7,8 +7,14 @@ use crate::handlers::shared::{
};
use crate::GatewayError;
use aether_admin::provider::quota as admin_provider_quota_pure;
use aether_contracts::{ExecutionPlan, ExecutionResult, ExecutionTimeouts, ProxySnapshot};
use aether_data_contracts::repository::provider_catalog::StoredProviderCatalogKey;
use aether_contracts::{
ExecutionPlan, ExecutionResult, ExecutionTimeouts, ProxySnapshot, RequestBody,
ResolvedTransportProfile, EXECUTION_REQUEST_ACCEPT_INVALID_CERTS_HEADER,
};
use aether_data_contracts::repository::provider_catalog::{
StoredProviderCatalogEndpoint, StoredProviderCatalogKey,
};
use aether_provider_pool::{ProviderPoolQuotaRequestSpec, ProviderPoolService};
use std::time::{SystemTime, UNIX_EPOCH};
use tracing::warn;
@@ -51,22 +57,28 @@ pub(crate) fn normalize_string_id_list(values: Option<Vec<String>>) -> Option<Ve
}
pub(crate) fn provider_type_supports_quota_refresh(provider_type: &str) -> bool {
matches!(
provider_type.trim().to_ascii_lowercase().as_str(),
"codex" | "kiro" | "antigravity" | "chatgpt_web"
)
ProviderPoolService::with_builtin_adapters().supports_quota_refresh(provider_type)
}
pub(crate) fn unsupported_provider_quota_refresh_message(provider_type: &str) -> String {
match provider_type.trim().to_ascii_lowercase().as_str() {
"claude_code" => "Claude Code 暂不支持自动刷新额度：上游没有稳定可用的账号额度查询接口",
"gemini_cli" => {
"Gemini CLI 暂不支持自动刷新额度：当前只能通过模型同步/缓存快照展示已知配额信息"
}
"vertex_ai" => "Vertex AI 暂不支持自动刷新额度：额度属于 Google Cloud 项目/区域配额",
_ => "该 Provider 暂不支持自动刷新额度",
}
.to_string()
ProviderPoolService::with_builtin_adapters().quota_refresh_unsupported_message(provider_type)
}
pub(crate) fn provider_quota_refresh_endpoint_for_provider(
provider_type: &str,
endpoints: &[StoredProviderCatalogEndpoint],
include_inactive: bool,
) -> Option<StoredProviderCatalogEndpoint> {
ProviderPoolService::with_builtin_adapters().quota_refresh_endpoint_for_provider(
provider_type,
endpoints,
include_inactive,
)
}
pub(crate) fn provider_quota_refresh_missing_endpoint_message(provider_type: &str) -> String {
ProviderPoolService::with_builtin_adapters()
.quota_refresh_missing_endpoint_message(provider_type)
}
pub(super) fn coerce_json_u64(value: &serde_json::Value) -> Option<u64> {
@@ -125,6 +137,63 @@ pub(super) fn build_quota_snapshot_payload(
updated_snapshot.get("quota").cloned()
}
pub(super) fn build_provider_quota_execution_plan(
transport: &AdminGatewayProviderTransportSnapshot,
spec: ProviderPoolQuotaRequestSpec,
proxy: Option<ProxySnapshot>,
transport_profile: Option<ResolvedTransportProfile>,
timeouts: Option<ExecutionTimeouts>,
) -> ExecutionPlan {
let ProviderPoolQuotaRequestSpec {
request_id,
provider_name,
quota_kind: _,
method,
url,
mut headers,
content_type,
json_body,
client_api_format,
provider_api_format,
model_name,
accept_invalid_certs,
} = spec;
if accept_invalid_certs {
headers.insert(
EXECUTION_REQUEST_ACCEPT_INVALID_CERTS_HEADER.to_string(),
"true".to_string(),
);
}
let body = json_body
.map(RequestBody::from_json)
.unwrap_or(RequestBody {
json_body: None,
body_bytes_b64: None,
body_ref: None,
});
ExecutionPlan {
request_id,
candidate_id: None,
provider_name: Some(provider_name),
provider_id: transport.provider.id.clone(),
endpoint_id: transport.endpoint.id.clone(),
key_id: transport.key.id.clone(),
method,
url,
headers,
content_type,
content_encoding: None,
body,
stream: false,
client_api_format,
provider_api_format,
model_name,
proxy,
transport_profile,
timeouts,
}
}
pub(crate) async fn persist_provider_quota_refresh_state(
state: &AdminAppState<'_>,
key_id: &str,
@@ -142,24 +211,21 @@ pub(crate) async fn persist_provider_quota_refresh_state(
return Ok(false);
};
let mut quota_snapshot_provider_type = None::<&str>;
let mut quota_snapshot_provider_type = None::<String>;
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", "chatgpt_web"]
.into_iter()
.find(|provider_type| object.contains_key(*provider_type))
});
quota_snapshot_provider_type =
aether_provider_pool::provider_pool_quota_metadata_provider_type(metadata_update);
}
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 {
if let Some(provider_type) = quota_snapshot_provider_type.as_deref() {
latest_key.status_snapshot = sync_provider_key_quota_status_snapshot(
latest_key.status_snapshot.as_ref(),
provider_type,

108
apps/aether-gateway/src/handlers/admin/provider/oauth/runtime.rs
View File

@@ -1,7 +1,6 @@
use super::quota::antigravity::refresh_antigravity_provider_quota_locally;
use super::quota::chatgpt_web::refresh_chatgpt_web_provider_quota_locally;
use super::quota::codex::refresh_codex_provider_quota_locally;
use super::quota::kiro::refresh_kiro_provider_quota_locally;
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;
@@ -113,16 +112,20 @@ pub(crate) struct ProviderOAuthRuntimeEndpoints {
pub(crate) runtime_endpoint: Option<StoredProviderCatalogEndpoint>,
}
pub(crate) async fn resolve_provider_oauth_runtime_endpoints(
async fn resolve_provider_runtime_endpoints_with_selector(
state: &AdminAppState<'_>,
provider: &StoredProviderCatalogProvider,
provider_type: &str,
endpoint_selector: fn(
&str,
&[StoredProviderCatalogEndpoint],
bool,
) -> Option<StoredProviderCatalogEndpoint>,
) -> Result<ProviderOAuthRuntimeEndpoints, GatewayError> {
let mut endpoints = state
.list_provider_catalog_endpoints_by_provider_ids(std::slice::from_ref(&provider.id))
.await?;
let mut runtime_endpoint =
provider_oauth_maintenance_endpoint_for_provider(provider_type, &endpoints);
let mut runtime_endpoint = endpoint_selector(provider_type, &endpoints, true);
if runtime_endpoint.is_none()
&& state
.fixed_provider_template(&provider.provider_type)
@@ -133,8 +136,7 @@ pub(crate) async fn resolve_provider_oauth_runtime_endpoints(
endpoints = state
.list_provider_catalog_endpoints_by_provider_ids(std::slice::from_ref(&provider.id))
.await?;
runtime_endpoint =
provider_oauth_maintenance_endpoint_for_provider(provider_type, &endpoints);
runtime_endpoint = endpoint_selector(provider_type, &endpoints, true);
}
Ok(ProviderOAuthRuntimeEndpoints {
@@ -143,6 +145,34 @@ pub(crate) async fn resolve_provider_oauth_runtime_endpoints(
})
}
pub(crate) async fn resolve_provider_oauth_runtime_endpoints(
state: &AdminAppState<'_>,
provider: &StoredProviderCatalogProvider,
provider_type: &str,
) -> Result<ProviderOAuthRuntimeEndpoints, GatewayError> {
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<ProviderOAuthRuntimeEndpoints, GatewayError> {
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,
@@ -150,16 +180,13 @@ pub(crate) async fn refresh_provider_oauth_account_state_after_update(
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" | "chatgpt_web"
) {
if !provider_type_supports_quota_refresh(&provider_type) {
return Ok((false, None));
}
let ProviderOAuthRuntimeEndpoints {
runtime_endpoint, ..
} = resolve_provider_oauth_runtime_endpoints(state, provider, &provider_type).await?;
} = resolve_provider_quota_runtime_endpoints(state, provider, &provider_type).await?;
let Some(endpoint) = runtime_endpoint else {
return Ok((false, None));
};
@@ -175,50 +202,15 @@ 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],
proxy_override.clone(),
)
.await?
}
"kiro" => {
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],
proxy_override,
)
.await?
}
"chatgpt_web" => {
refresh_chatgpt_web_provider_quota_locally(
state,
provider,
&endpoint,
vec![key],
proxy_override,
)
.await?
}
_ => 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));
};

2
apps/aether-gateway/src/handlers/admin/provider/pool/config.rs
View File

@@ -1,7 +1,7 @@
use crate::handlers::admin::provider::shared::support::{
AdminProviderPoolConfig, AdminProviderPoolSchedulingPreset, AdminProviderPoolUnschedulableRule,
};
use aether_ai_serving::{PoolMemberScoreRules, PoolMemberScoreWeights};
use aether_pool_core::{PoolMemberScoreRules, PoolMemberScoreWeights};
use serde_json::{Map, Value};
const POOL_ALLOWED_SCHEDULING_PRESETS: &[&str] = &[

2
apps/aether-gateway/src/handlers/admin/provider/pool/runtime/writes.rs
View File

@@ -618,7 +618,7 @@ mod tests {
probe_concurrency: 4,
score_top_n: 128,
score_fallback_scan_limit: 1024,
score_rules: aether_ai_serving::PoolMemberScoreRules::default(),
score_rules: aether_pool_core::PoolMemberScoreRules::default(),
stream_timeout_threshold: 3,
stream_timeout_window_seconds: 1800,
stream_timeout_cooldown_seconds: 300,

79
apps/aether-gateway/src/handlers/admin/provider/pool_admin/payloads.rs
View File

@@ -95,27 +95,6 @@ fn admin_pool_oauth_organizations(
.unwrap_or_default()
}
fn admin_pool_normalize_oauth_plan_type(value: &str, provider_type: &str) -> Option<String> {
let mut normalized = value.trim().to_string();
if normalized.is_empty() {
return None;
}
let provider_type = provider_type.trim().to_ascii_lowercase();
if !provider_type.is_empty() && normalized.to_ascii_lowercase().starts_with(&provider_type) {
normalized = normalized[provider_type.len()..]
.trim_matches(|ch: char| [' ', ':', '-', '_'].contains(&ch))
.to_string();
}
let normalized = normalized.trim().to_ascii_lowercase();
if normalized.is_empty() {
None
} else {
Some(normalized)
}
}
fn admin_pool_derive_oauth_expires_at(
provider_type: &str,
key: &StoredProviderCatalogKey,
@@ -139,57 +118,6 @@ fn admin_pool_derive_oauth_expires_at(
None
}
fn admin_pool_derive_oauth_plan_type(
key: &StoredProviderCatalogKey,
provider_type: &str,
auth_config: Option<&serde_json::Map<String, serde_json::Value>>,
) -> Option<String> {
if !provider_key_auth_semantics(key, provider_type).oauth_managed() {
return None;
}
if let Some(upstream_metadata) = key
.upstream_metadata
.as_ref()
.and_then(serde_json::Value::as_object)
{
let provider_bucket = upstream_metadata
.get(&provider_type.trim().to_ascii_lowercase())
.and_then(serde_json::Value::as_object);
for source in provider_bucket
.into_iter()
.chain(std::iter::once(upstream_metadata))
{
for field in [
"plan_type",
"tier",
"subscription_title",
"subscription_plan",
] {
if let Some(value) = source.get(field).and_then(serde_json::Value::as_str) {
let normalized = admin_pool_normalize_oauth_plan_type(value, provider_type);
if normalized.is_some() {
return normalized;
}
}
}
}
}
if let Some(config) = auth_config {
for field in ["plan_type", "tier", "plan", "subscription_plan"] {
if let Some(value) = config.get(field).and_then(serde_json::Value::as_str) {
let normalized = admin_pool_normalize_oauth_plan_type(value, provider_type);
if normalized.is_some() {
return normalized;
}
}
}
}
None
}
fn admin_pool_format_percent(value: f64) -> String {
format!("{:.1}%", value.clamp(0.0, 100.0))
}
@@ -925,8 +853,11 @@ pub(super) fn build_admin_pool_key_payload(
let auth_config = state.parse_catalog_auth_config_json(key);
let oauth_expires_at =
admin_pool_derive_oauth_expires_at(provider_type, key, auth_config.as_ref());
let oauth_plan_type =
admin_pool_derive_oauth_plan_type(key, provider_type, auth_config.as_ref());
let oauth_plan_type = if auth_semantics.oauth_managed() {
aether_provider_pool::derive_plan_tier(provider_type, key, auth_config.as_ref())
} else {
None
};
let mut status_snapshot = provider_key_status_snapshot_payload(key, provider_type);
if provider_type.trim().eq_ignore_ascii_case("codex") {
admin_pool_apply_codex_window_usage_summaries(

67
apps/aether-gateway/src/handlers/admin/provider/pool_admin/selection.rs
View File

@@ -22,74 +22,17 @@ fn admin_pool_parse_auth_config_json(
.cloned()
}
fn admin_pool_derive_oauth_plan_type(
fn admin_pool_derive_plan_tier(
state: &AdminAppState<'_>,
key: &StoredProviderCatalogKey,
provider_type: &str,
) -> Option<String> {
let normalize = |value: &str| {
let mut text = value.trim().to_string();
if text.is_empty() {
return None;
}
let provider_type = provider_type.trim().to_ascii_lowercase();
if !provider_type.is_empty() && text.to_ascii_lowercase().starts_with(&provider_type) {
text = text[provider_type.len()..]
.trim_matches(|ch: char| [' ', ':', '-', '_'].contains(&ch))
.to_string();
}
if text.is_empty() {
None
} else {
Some(text.to_ascii_lowercase())
}
};
if !provider_key_is_oauth_managed(key, provider_type) {
return None;
}
if let Some(upstream_metadata) = key
.upstream_metadata
.as_ref()
.and_then(serde_json::Value::as_object)
{
let provider_bucket = upstream_metadata
.get(&provider_type.trim().to_ascii_lowercase())
.and_then(serde_json::Value::as_object);
for source in provider_bucket
.into_iter()
.chain(std::iter::once(upstream_metadata))
{
for plan_key in [
"plan_type",
"tier",
"subscription_title",
"subscription_plan",
] {
if let Some(value) = source.get(plan_key).and_then(serde_json::Value::as_str) {
if let Some(normalized) = normalize(value) {
return Some(normalized);
}
}
}
}
}
if let Some(auth_config) = admin_pool_parse_auth_config_json(state, key) {
for plan_key in ["plan_type", "tier", "plan", "subscription_plan"] {
if let Some(value) = auth_config
.get(plan_key)
.and_then(serde_json::Value::as_str)
{
if let Some(normalized) = normalize(value) {
return Some(normalized);
}
}
}
}
None
let auth_config = admin_pool_parse_auth_config_json(state, key);
aether_provider_pool::derive_plan_tier(provider_type, key, auth_config.as_ref())
}
pub(super) fn admin_pool_matches_quick_selector(
@@ -98,7 +41,7 @@ pub(super) fn admin_pool_matches_quick_selector(
provider_type: &str,
selector: &str,
) -> bool {
let oauth_plan_type = admin_pool_derive_oauth_plan_type(state, key, provider_type);
let oauth_plan_type = admin_pool_derive_plan_tier(state, key, provider_type);
admin_provider_pool_pure::admin_pool_matches_quick_selector(
key,
selector,
@@ -113,7 +56,7 @@ pub(super) fn admin_pool_matches_search(
provider_type: &str,
search: Option<&str>,
) -> bool {
let oauth_plan_type = admin_pool_derive_oauth_plan_type(state, key, provider_type);
let oauth_plan_type = admin_pool_derive_plan_tier(state, key, provider_type);
admin_provider_pool_pure::admin_pool_matches_search(key, search, oauth_plan_type.as_deref())
}

2
apps/aether-gateway/src/handlers/admin/provider/shared/support.rs
View File

@@ -1,6 +1,6 @@
use crate::handlers::admin::request::AdminAppState;
use crate::LocalProviderDeleteTaskState;
use aether_ai_serving::PoolMemberScoreRules;
use aether_pool_core::PoolMemberScoreRules;
use serde_json::json;
use std::collections::BTreeMap;