mirror of
https://github.com/fawney19/Aether.git
synced 2026-09-02 09:20:22 +08:00
feat(grok): add admin oauth and quota support
This commit is contained in:
@@ -64,6 +64,8 @@ tracing.workspace = true
|
|||||||
url.workspace = true
|
url.workspace = true
|
||||||
uuid.workspace = true
|
uuid.workspace = true
|
||||||
webpki-roots.workspace = true
|
webpki-roots.workspace = true
|
||||||
|
wreq.workspace = true
|
||||||
|
wreq-util.workspace = true
|
||||||
|
|
||||||
[target.'cfg(not(target_env = "msvc"))'.dependencies]
|
[target.'cfg(not(target_env = "msvc"))'.dependencies]
|
||||||
tikv-jemallocator = "0.6"
|
tikv-jemallocator = "0.6"
|
||||||
|
|||||||
@@ -38,7 +38,7 @@ pub(super) async fn build_admin_create_api_key_install_session_response(
|
|||||||
Err(_) => {
|
Err(_) => {
|
||||||
return Ok(build_admin_api_keys_bad_request_response(
|
return Ok(build_admin_api_keys_bad_request_response(
|
||||||
"请求数据验证失败",
|
"请求数据验证失败",
|
||||||
))
|
));
|
||||||
}
|
}
|
||||||
};
|
};
|
||||||
|
|
||||||
|
|||||||
@@ -249,9 +249,10 @@ async fn admin_monitoring_resilience_status_returns_local_payload() {
|
|||||||
let recommendations = payload["recommendations"]
|
let recommendations = payload["recommendations"]
|
||||||
.as_array()
|
.as_array()
|
||||||
.expect("recommendations should be array");
|
.expect("recommendations should be array");
|
||||||
assert!(recommendations.iter().any(|item| item
|
assert!(recommendations.iter().any(|item| {
|
||||||
.as_str()
|
item.as_str()
|
||||||
.is_some_and(|value| value.contains("prod-key"))));
|
.is_some_and(|value| value.contains("prod-key"))
|
||||||
|
}));
|
||||||
assert!(payload["timestamp"].as_str().is_some());
|
assert!(payload["timestamp"].as_str().is_some());
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|||||||
@@ -32,7 +32,7 @@ pub(super) async fn maybe_handle(
|
|||||||
.into_response(),
|
.into_response(),
|
||||||
));
|
));
|
||||||
};
|
};
|
||||||
let Some(_provider) = state
|
let Some(provider) = state
|
||||||
.read_provider_catalog_providers_by_ids(std::slice::from_ref(&provider_id))
|
.read_provider_catalog_providers_by_ids(std::slice::from_ref(&provider_id))
|
||||||
.await?
|
.await?
|
||||||
.into_iter()
|
.into_iter()
|
||||||
@@ -129,7 +129,9 @@ pub(super) async fn maybe_handle(
|
|||||||
Json(serde_json::Value::Array(
|
Json(serde_json::Value::Array(
|
||||||
created
|
created
|
||||||
.iter()
|
.iter()
|
||||||
.map(|model| build_admin_provider_model_response(model, now_unix_secs))
|
.map(|model| {
|
||||||
|
build_admin_provider_model_response(&provider, model, now_unix_secs)
|
||||||
|
})
|
||||||
.collect(),
|
.collect(),
|
||||||
))
|
))
|
||||||
.into_response(),
|
.into_response(),
|
||||||
|
|||||||
@@ -31,7 +31,7 @@ pub(super) async fn maybe_handle(
|
|||||||
.into_response(),
|
.into_response(),
|
||||||
));
|
));
|
||||||
};
|
};
|
||||||
let Some(_provider) = state
|
let Some(provider) = state
|
||||||
.read_provider_catalog_providers_by_ids(std::slice::from_ref(&provider_id))
|
.read_provider_catalog_providers_by_ids(std::slice::from_ref(&provider_id))
|
||||||
.await?
|
.await?
|
||||||
.into_iter()
|
.into_iter()
|
||||||
@@ -90,8 +90,12 @@ pub(super) async fn maybe_handle(
|
|||||||
.ok()
|
.ok()
|
||||||
.map(|duration| duration.as_secs())
|
.map(|duration| duration.as_secs())
|
||||||
.unwrap_or(0);
|
.unwrap_or(0);
|
||||||
Json(build_admin_provider_model_response(&created, now_unix_secs))
|
Json(build_admin_provider_model_response(
|
||||||
.into_response()
|
&provider,
|
||||||
|
&created,
|
||||||
|
now_unix_secs,
|
||||||
|
))
|
||||||
|
.into_response()
|
||||||
}
|
}
|
||||||
None => (
|
None => (
|
||||||
http::StatusCode::INTERNAL_SERVER_ERROR,
|
http::StatusCode::INTERNAL_SERVER_ERROR,
|
||||||
|
|||||||
@@ -1,9 +1,13 @@
|
|||||||
|
use crate::handlers::admin::provider::shared::model_test_capabilities::{
|
||||||
|
admin_provider_model_supports_image_generation, admin_provider_model_test_capabilities_payload,
|
||||||
|
};
|
||||||
use crate::handlers::admin::request::AdminAppState;
|
use crate::handlers::admin::request::AdminAppState;
|
||||||
use crate::GatewayError;
|
use crate::GatewayError;
|
||||||
use aether_admin::provider::models as admin_provider_models_pure;
|
use aether_admin::provider::models as admin_provider_models_pure;
|
||||||
use aether_data_contracts::repository::global_models::{
|
use aether_data_contracts::repository::global_models::{
|
||||||
AdminProviderModelListQuery, StoredAdminProviderModel,
|
AdminProviderModelListQuery, StoredAdminProviderModel,
|
||||||
};
|
};
|
||||||
|
use aether_data_contracts::repository::provider_catalog::StoredProviderCatalogProvider;
|
||||||
use std::time::{SystemTime, UNIX_EPOCH};
|
use std::time::{SystemTime, UNIX_EPOCH};
|
||||||
|
|
||||||
pub(super) fn admin_provider_model_effective_input_price(
|
pub(super) fn admin_provider_model_effective_input_price(
|
||||||
@@ -26,10 +30,32 @@ pub(super) fn admin_provider_model_effective_capability(
|
|||||||
}
|
}
|
||||||
|
|
||||||
pub(super) fn build_admin_provider_model_response(
|
pub(super) fn build_admin_provider_model_response(
|
||||||
|
provider: &StoredProviderCatalogProvider,
|
||||||
model: &StoredAdminProviderModel,
|
model: &StoredAdminProviderModel,
|
||||||
now_unix_secs: u64,
|
now_unix_secs: u64,
|
||||||
) -> serde_json::Value {
|
) -> serde_json::Value {
|
||||||
admin_provider_models_pure::build_admin_provider_model_response(model, now_unix_secs)
|
let mut payload =
|
||||||
|
admin_provider_models_pure::build_admin_provider_model_response(model, now_unix_secs);
|
||||||
|
let fallback_supports_image_generation = payload
|
||||||
|
.get("effective_supports_image_generation")
|
||||||
|
.and_then(serde_json::Value::as_bool)
|
||||||
|
.unwrap_or(false);
|
||||||
|
let supports_image_generation = admin_provider_model_supports_image_generation(
|
||||||
|
&provider.provider_type,
|
||||||
|
&model.provider_model_name,
|
||||||
|
fallback_supports_image_generation,
|
||||||
|
);
|
||||||
|
if let Some(object) = payload.as_object_mut() {
|
||||||
|
object.insert(
|
||||||
|
"model_test_capabilities".to_string(),
|
||||||
|
admin_provider_model_test_capabilities_payload(
|
||||||
|
&provider.provider_type,
|
||||||
|
&model.provider_model_name,
|
||||||
|
supports_image_generation,
|
||||||
|
),
|
||||||
|
);
|
||||||
|
}
|
||||||
|
payload
|
||||||
}
|
}
|
||||||
|
|
||||||
pub(super) async fn build_admin_provider_models_payload(
|
pub(super) async fn build_admin_provider_models_payload(
|
||||||
@@ -48,9 +74,10 @@ pub(super) async fn build_admin_provider_models_payload(
|
|||||||
.ok()?
|
.ok()?
|
||||||
.into_iter()
|
.into_iter()
|
||||||
.next()?;
|
.next()?;
|
||||||
|
let provider_id = provider.id.clone();
|
||||||
let mut models = state
|
let mut models = state
|
||||||
.list_admin_provider_models(&AdminProviderModelListQuery {
|
.list_admin_provider_models(&AdminProviderModelListQuery {
|
||||||
provider_id: provider.id,
|
provider_id,
|
||||||
is_active,
|
is_active,
|
||||||
offset: skip,
|
offset: skip,
|
||||||
limit,
|
limit,
|
||||||
@@ -70,7 +97,7 @@ pub(super) async fn build_admin_provider_models_payload(
|
|||||||
Some(serde_json::Value::Array(
|
Some(serde_json::Value::Array(
|
||||||
models
|
models
|
||||||
.iter()
|
.iter()
|
||||||
.map(|model| build_admin_provider_model_response(model, now_unix_secs))
|
.map(|model| build_admin_provider_model_response(&provider, model, now_unix_secs))
|
||||||
.collect(),
|
.collect(),
|
||||||
))
|
))
|
||||||
}
|
}
|
||||||
@@ -80,9 +107,15 @@ pub(super) async fn build_admin_provider_model_payload(
|
|||||||
provider_id: &str,
|
provider_id: &str,
|
||||||
model_id: &str,
|
model_id: &str,
|
||||||
) -> Option<serde_json::Value> {
|
) -> Option<serde_json::Value> {
|
||||||
if !state.has_global_model_data_reader() {
|
if !state.has_provider_catalog_data_reader() || !state.has_global_model_data_reader() {
|
||||||
return None;
|
return None;
|
||||||
}
|
}
|
||||||
|
let provider = state
|
||||||
|
.read_provider_catalog_providers_by_ids(&[provider_id.to_string()])
|
||||||
|
.await
|
||||||
|
.ok()?
|
||||||
|
.into_iter()
|
||||||
|
.next()?;
|
||||||
let model = state
|
let model = state
|
||||||
.get_admin_provider_model(provider_id, model_id)
|
.get_admin_provider_model(provider_id, model_id)
|
||||||
.await
|
.await
|
||||||
@@ -92,7 +125,11 @@ pub(super) async fn build_admin_provider_model_payload(
|
|||||||
.ok()
|
.ok()
|
||||||
.map(|duration| duration.as_secs())
|
.map(|duration| duration.as_secs())
|
||||||
.unwrap_or(0);
|
.unwrap_or(0);
|
||||||
Some(build_admin_provider_model_response(&model, now_unix_secs))
|
Some(build_admin_provider_model_response(
|
||||||
|
&provider,
|
||||||
|
&model,
|
||||||
|
now_unix_secs,
|
||||||
|
))
|
||||||
}
|
}
|
||||||
|
|
||||||
pub(super) async fn admin_provider_model_name_exists(
|
pub(super) async fn admin_provider_model_name_exists(
|
||||||
|
|||||||
@@ -33,6 +33,20 @@ pub(super) async fn maybe_handle(
|
|||||||
.into_response(),
|
.into_response(),
|
||||||
));
|
));
|
||||||
};
|
};
|
||||||
|
let Some(provider) = state
|
||||||
|
.read_provider_catalog_providers_by_ids(std::slice::from_ref(&provider_id))
|
||||||
|
.await?
|
||||||
|
.into_iter()
|
||||||
|
.next()
|
||||||
|
else {
|
||||||
|
return Ok(Some(
|
||||||
|
(
|
||||||
|
http::StatusCode::NOT_FOUND,
|
||||||
|
Json(json!({ "detail": format!("Provider {provider_id} 不存在") })),
|
||||||
|
)
|
||||||
|
.into_response(),
|
||||||
|
));
|
||||||
|
};
|
||||||
let Some(existing) = state
|
let Some(existing) = state
|
||||||
.get_admin_provider_model(&provider_id, &model_id)
|
.get_admin_provider_model(&provider_id, &model_id)
|
||||||
.await?
|
.await?
|
||||||
@@ -110,8 +124,12 @@ pub(super) async fn maybe_handle(
|
|||||||
.ok()
|
.ok()
|
||||||
.map(|duration| duration.as_secs())
|
.map(|duration| duration.as_secs())
|
||||||
.unwrap_or(0);
|
.unwrap_or(0);
|
||||||
Json(build_admin_provider_model_response(&updated, now_unix_secs))
|
Json(build_admin_provider_model_response(
|
||||||
.into_response()
|
&provider,
|
||||||
|
&updated,
|
||||||
|
now_unix_secs,
|
||||||
|
))
|
||||||
|
.into_response()
|
||||||
}
|
}
|
||||||
None => (
|
None => (
|
||||||
http::StatusCode::NOT_FOUND,
|
http::StatusCode::NOT_FOUND,
|
||||||
|
|||||||
@@ -1,3 +1,4 @@
|
|||||||
|
use super::super::helpers::admin_provider_oauth_key_name_from_auth_config;
|
||||||
use super::super::token_import::{
|
use super::super::token_import::{
|
||||||
build_provider_access_token_import_auth_config, provider_type_supports_access_token_import,
|
build_provider_access_token_import_auth_config, provider_type_supports_access_token_import,
|
||||||
};
|
};
|
||||||
@@ -24,13 +25,11 @@ use crate::handlers::admin::provider::oauth::runtime::{
|
|||||||
use crate::handlers::admin::provider::oauth::state::{
|
use crate::handlers::admin::provider::oauth::state::{
|
||||||
admin_provider_oauth_template, exchange_admin_provider_oauth_refresh_token,
|
admin_provider_oauth_template, exchange_admin_provider_oauth_refresh_token,
|
||||||
};
|
};
|
||||||
use crate::handlers::admin::provider::shared::support::ADMIN_PROVIDER_OAUTH_DATA_UNAVAILABLE_DETAIL;
|
|
||||||
use crate::handlers::admin::request::{AdminAppState, AdminProviderOAuthTemplate};
|
use crate::handlers::admin::request::{AdminAppState, AdminProviderOAuthTemplate};
|
||||||
use crate::GatewayError;
|
use crate::GatewayError;
|
||||||
use aether_admin::provider::oauth::parse_admin_provider_oauth_kiro_batch_import_entries;
|
use aether_admin::provider::oauth::parse_admin_provider_oauth_kiro_batch_import_entries;
|
||||||
use aether_contracts::ProxySnapshot;
|
use aether_contracts::ProxySnapshot;
|
||||||
use serde_json::{json, Map, Value};
|
use serde_json::{json, Map, Value};
|
||||||
use std::time::{SystemTime, UNIX_EPOCH};
|
|
||||||
|
|
||||||
struct AdminProviderOAuthResolvedBatchImport {
|
struct AdminProviderOAuthResolvedBatchImport {
|
||||||
access_token: String,
|
access_token: String,
|
||||||
@@ -45,7 +44,7 @@ pub(super) fn estimate_admin_provider_oauth_batch_import_total(
|
|||||||
if provider_type.eq_ignore_ascii_case("kiro") {
|
if provider_type.eq_ignore_ascii_case("kiro") {
|
||||||
parse_admin_provider_oauth_kiro_batch_import_entries(raw_credentials).len()
|
parse_admin_provider_oauth_kiro_batch_import_entries(raw_credentials).len()
|
||||||
} else {
|
} else {
|
||||||
parse_admin_provider_oauth_batch_import_entries(raw_credentials).len()
|
parse_admin_provider_oauth_batch_import_entries(provider_type, raw_credentials).len()
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -67,7 +66,8 @@ pub(super) async fn execute_admin_provider_oauth_batch_import_for_provider_type(
|
|||||||
)
|
)
|
||||||
.await
|
.await
|
||||||
} else {
|
} else {
|
||||||
let entries = parse_admin_provider_oauth_batch_import_entries(raw_credentials);
|
let entries =
|
||||||
|
parse_admin_provider_oauth_batch_import_entries(provider_type, raw_credentials);
|
||||||
execute_admin_provider_oauth_batch_import(
|
execute_admin_provider_oauth_batch_import(
|
||||||
state,
|
state,
|
||||||
provider_id,
|
provider_id,
|
||||||
@@ -82,7 +82,7 @@ pub(super) async fn execute_admin_provider_oauth_batch_import_for_provider_type(
|
|||||||
|
|
||||||
async fn resolve_admin_provider_oauth_batch_import_tokens(
|
async fn resolve_admin_provider_oauth_batch_import_tokens(
|
||||||
state: &AdminAppState<'_>,
|
state: &AdminAppState<'_>,
|
||||||
template: AdminProviderOAuthTemplate,
|
template: Option<AdminProviderOAuthTemplate>,
|
||||||
provider_type: &str,
|
provider_type: &str,
|
||||||
entry: &AdminProviderOAuthBatchImportEntry,
|
entry: &AdminProviderOAuthBatchImportEntry,
|
||||||
request_proxy: Option<ProxySnapshot>,
|
request_proxy: Option<ProxySnapshot>,
|
||||||
@@ -99,6 +99,29 @@ async fn resolve_admin_provider_oauth_batch_import_tokens(
|
|||||||
.filter(|value| !value.is_empty());
|
.filter(|value| !value.is_empty());
|
||||||
|
|
||||||
if let Some(refresh_token) = refresh_token {
|
if let Some(refresh_token) = refresh_token {
|
||||||
|
let Some(template) = template else {
|
||||||
|
if provider_type_supports_access_token_import(provider_type) {
|
||||||
|
if let Some(access_token) = access_token {
|
||||||
|
let (auth_config, expires_at) = build_provider_access_token_import_auth_config(
|
||||||
|
provider_type,
|
||||||
|
access_token,
|
||||||
|
Some(refresh_token),
|
||||||
|
entry.expires_at,
|
||||||
|
Some("Provider 不支持 Refresh Token 交换,已回退为 Session Token 导入"),
|
||||||
|
);
|
||||||
|
return Ok(AdminProviderOAuthResolvedBatchImport {
|
||||||
|
access_token: access_token.to_string(),
|
||||||
|
auth_config,
|
||||||
|
expires_at,
|
||||||
|
});
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return Err(
|
||||||
|
"该 Provider 不支持 Refresh Token 导入,请提供 sso_token 或 access_token"
|
||||||
|
.to_string(),
|
||||||
|
);
|
||||||
|
};
|
||||||
|
|
||||||
let token_payload = match exchange_admin_provider_oauth_refresh_token(
|
let token_payload = match exchange_admin_provider_oauth_refresh_token(
|
||||||
state,
|
state,
|
||||||
template,
|
template,
|
||||||
@@ -152,7 +175,7 @@ async fn resolve_admin_provider_oauth_batch_import_tokens(
|
|||||||
|
|
||||||
if let Some(access_token) = access_token {
|
if let Some(access_token) = access_token {
|
||||||
if !provider_type_supports_access_token_import(provider_type) {
|
if !provider_type_supports_access_token_import(provider_type) {
|
||||||
return Err("Access Token 导入仅支持 Codex / ChatGPT Web Provider".to_string());
|
return Err("Access Token 导入仅支持 Codex / ChatGPT Web / Grok Provider".to_string());
|
||||||
}
|
}
|
||||||
let (auth_config, expires_at) = build_provider_access_token_import_auth_config(
|
let (auth_config, expires_at) = build_provider_access_token_import_auth_config(
|
||||||
provider_type,
|
provider_type,
|
||||||
@@ -204,25 +227,7 @@ pub(super) async fn execute_admin_provider_oauth_batch_import(
|
|||||||
});
|
});
|
||||||
};
|
};
|
||||||
|
|
||||||
let Some(template) = admin_provider_oauth_template(provider_type) else {
|
let template = admin_provider_oauth_template(provider_type);
|
||||||
return Ok(AdminProviderOAuthBatchImportOutcome {
|
|
||||||
total: entries.len(),
|
|
||||||
success: 0,
|
|
||||||
failed: entries.len(),
|
|
||||||
results: entries
|
|
||||||
.iter()
|
|
||||||
.enumerate()
|
|
||||||
.map(|(index, _)| {
|
|
||||||
json!({
|
|
||||||
"index": index,
|
|
||||||
"status": "error",
|
|
||||||
"error": ADMIN_PROVIDER_OAUTH_DATA_UNAVAILABLE_DETAIL,
|
|
||||||
"replaced": false,
|
|
||||||
})
|
|
||||||
})
|
|
||||||
.collect(),
|
|
||||||
});
|
|
||||||
};
|
|
||||||
|
|
||||||
let endpoint_resolution =
|
let endpoint_resolution =
|
||||||
resolve_provider_oauth_runtime_endpoints(state, &provider, provider_type).await?;
|
resolve_provider_oauth_runtime_endpoints(state, &provider, provider_type).await?;
|
||||||
@@ -340,24 +345,11 @@ pub(super) async fn execute_admin_provider_oauth_batch_import(
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
} else {
|
} else {
|
||||||
let key_name = auth_config
|
let key_name = admin_provider_oauth_key_name_from_auth_config(
|
||||||
.get("email")
|
provider_type,
|
||||||
.and_then(serde_json::Value::as_str)
|
&auth_config,
|
||||||
.map(str::trim)
|
Some(index),
|
||||||
.filter(|value| !value.is_empty())
|
);
|
||||||
.map(|email| format!("{provider_type}_{email}"))
|
|
||||||
.unwrap_or_else(|| {
|
|
||||||
format!(
|
|
||||||
"{}_{}_{}",
|
|
||||||
provider_type,
|
|
||||||
SystemTime::now()
|
|
||||||
.duration_since(UNIX_EPOCH)
|
|
||||||
.ok()
|
|
||||||
.map(|duration| duration.as_secs())
|
|
||||||
.unwrap_or(0),
|
|
||||||
index
|
|
||||||
)
|
|
||||||
});
|
|
||||||
match create_provider_oauth_catalog_key(
|
match create_provider_oauth_catalog_key(
|
||||||
state,
|
state,
|
||||||
provider_id,
|
provider_id,
|
||||||
|
|||||||
@@ -8,7 +8,7 @@ use super::parse::{
|
|||||||
};
|
};
|
||||||
use crate::handlers::admin::provider::oauth::errors::build_internal_control_error_response;
|
use crate::handlers::admin::provider::oauth::errors::build_internal_control_error_response;
|
||||||
use crate::handlers::admin::provider::oauth::state::{
|
use crate::handlers::admin::provider::oauth::state::{
|
||||||
admin_provider_oauth_template, build_admin_provider_oauth_backend_unavailable_response,
|
build_admin_provider_oauth_backend_unavailable_response,
|
||||||
is_fixed_provider_type_for_provider_oauth,
|
is_fixed_provider_type_for_provider_oauth,
|
||||||
};
|
};
|
||||||
use crate::handlers::admin::provider::shared::paths::admin_provider_oauth_batch_import_provider_id;
|
use crate::handlers::admin::provider::shared::paths::admin_provider_oauth_batch_import_provider_id;
|
||||||
@@ -60,10 +60,6 @@ pub(in super::super) async fn handle_admin_provider_oauth_batch_import(
|
|||||||
"该 Provider 不是固定类型,无法使用 provider-oauth",
|
"该 Provider 不是固定类型,无法使用 provider-oauth",
|
||||||
));
|
));
|
||||||
}
|
}
|
||||||
if provider_type != "kiro" && admin_provider_oauth_template(&provider_type).is_none() {
|
|
||||||
return Ok(build_admin_provider_oauth_backend_unavailable_response());
|
|
||||||
}
|
|
||||||
|
|
||||||
let total = estimate_admin_provider_oauth_batch_import_total(
|
let total = estimate_admin_provider_oauth_batch_import_total(
|
||||||
&provider_type,
|
&provider_type,
|
||||||
payload.credentials.as_str(),
|
payload.credentials.as_str(),
|
||||||
|
|||||||
@@ -1,4 +1,4 @@
|
|||||||
use super::super::token_import::{import_tokens_from_raw_token, normalize_single_import_tokens};
|
use super::super::token_import::{import_tokens_from_raw_token, normalize_provider_import_tokens};
|
||||||
use crate::handlers::admin::provider::oauth::errors::build_internal_control_error_response;
|
use crate::handlers::admin::provider::oauth::errors::build_internal_control_error_response;
|
||||||
use crate::handlers::admin::provider::oauth::state::{current_unix_secs, json_u64_value};
|
use crate::handlers::admin::provider::oauth::state::{current_unix_secs, json_u64_value};
|
||||||
use axum::{
|
use axum::{
|
||||||
@@ -25,8 +25,15 @@ pub(super) struct AdminProviderOAuthBatchImportEntry {
|
|||||||
pub account_id: Option<String>,
|
pub account_id: Option<String>,
|
||||||
pub account_user_id: Option<String>,
|
pub account_user_id: Option<String>,
|
||||||
pub plan_type: Option<String>,
|
pub plan_type: Option<String>,
|
||||||
|
pub pool_tier: Option<String>,
|
||||||
pub user_id: Option<String>,
|
pub user_id: Option<String>,
|
||||||
pub email: Option<String>,
|
pub email: Option<String>,
|
||||||
|
pub account_name: Option<String>,
|
||||||
|
pub sso_rw_token: Option<String>,
|
||||||
|
pub cf_cookies: Option<String>,
|
||||||
|
pub cf_clearance: Option<String>,
|
||||||
|
pub user_agent: Option<String>,
|
||||||
|
pub browser_profile: Option<String>,
|
||||||
}
|
}
|
||||||
|
|
||||||
#[derive(Debug, Clone)]
|
#[derive(Debug, Clone)]
|
||||||
@@ -67,16 +74,72 @@ fn coerce_admin_provider_oauth_import_str(value: Option<&serde_json::Value>) ->
|
|||||||
.map(ToOwned::to_owned)
|
.map(ToOwned::to_owned)
|
||||||
}
|
}
|
||||||
|
|
||||||
|
fn grok_cookie_value(raw: &str, name: &str) -> Option<String> {
|
||||||
|
raw.trim()
|
||||||
|
.strip_prefix("Cookie:")
|
||||||
|
.unwrap_or_else(|| raw.trim())
|
||||||
|
.split(';')
|
||||||
|
.filter_map(|segment| segment.trim().split_once('='))
|
||||||
|
.find_map(|(cookie_name, cookie_value)| {
|
||||||
|
cookie_name
|
||||||
|
.trim()
|
||||||
|
.eq_ignore_ascii_case(name)
|
||||||
|
.then(|| cookie_value.trim())
|
||||||
|
.filter(|value| !value.is_empty())
|
||||||
|
.map(ToOwned::to_owned)
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
fn grok_cookie_profile(raw: &str) -> Option<String> {
|
||||||
|
let raw = raw
|
||||||
|
.trim()
|
||||||
|
.strip_prefix("Cookie:")
|
||||||
|
.unwrap_or_else(|| raw.trim());
|
||||||
|
let parts = raw
|
||||||
|
.split(';')
|
||||||
|
.filter_map(|segment| {
|
||||||
|
let (cookie_name, cookie_value) = segment.trim().split_once('=')?;
|
||||||
|
let cookie_name = cookie_name.trim();
|
||||||
|
let cookie_value = cookie_value.trim();
|
||||||
|
if cookie_name.is_empty()
|
||||||
|
|| cookie_value.is_empty()
|
||||||
|
|| cookie_name.eq_ignore_ascii_case("sso")
|
||||||
|
|| cookie_name.eq_ignore_ascii_case("sso-rw")
|
||||||
|
{
|
||||||
|
return None;
|
||||||
|
}
|
||||||
|
Some(format!("{cookie_name}={cookie_value}"))
|
||||||
|
})
|
||||||
|
.collect::<Vec<_>>();
|
||||||
|
(!parts.is_empty()).then(|| parts.join("; "))
|
||||||
|
}
|
||||||
|
|
||||||
|
fn grok_cookie_session_token(provider_type: &str, raw: &str) -> Option<String> {
|
||||||
|
provider_type
|
||||||
|
.trim()
|
||||||
|
.eq_ignore_ascii_case("grok")
|
||||||
|
.then(|| grok_cookie_value(raw, "sso"))
|
||||||
|
.flatten()
|
||||||
|
}
|
||||||
|
|
||||||
fn extract_admin_provider_oauth_batch_import_entry(
|
fn extract_admin_provider_oauth_batch_import_entry(
|
||||||
|
provider_type: &str,
|
||||||
item: &serde_json::Value,
|
item: &serde_json::Value,
|
||||||
) -> Option<AdminProviderOAuthBatchImportEntry> {
|
) -> Option<AdminProviderOAuthBatchImportEntry> {
|
||||||
match item {
|
match item {
|
||||||
serde_json::Value::String(value) => {
|
serde_json::Value::String(value) => {
|
||||||
let refresh_token = value.trim();
|
let raw_token = value.trim();
|
||||||
if refresh_token.is_empty() {
|
if raw_token.is_empty() {
|
||||||
None
|
None
|
||||||
} else {
|
} else {
|
||||||
let (refresh_token, access_token) = import_tokens_from_raw_token(refresh_token);
|
let sso_from_cookie = grok_cookie_session_token(provider_type, raw_token);
|
||||||
|
let token_input = sso_from_cookie.as_deref().unwrap_or(raw_token);
|
||||||
|
let (refresh_token, access_token) = import_tokens_from_raw_token(token_input);
|
||||||
|
let (refresh_token, access_token) = normalize_provider_import_tokens(
|
||||||
|
provider_type,
|
||||||
|
refresh_token.as_deref(),
|
||||||
|
access_token.as_deref(),
|
||||||
|
);
|
||||||
Some(AdminProviderOAuthBatchImportEntry {
|
Some(AdminProviderOAuthBatchImportEntry {
|
||||||
refresh_token,
|
refresh_token,
|
||||||
access_token,
|
access_token,
|
||||||
@@ -84,8 +147,15 @@ fn extract_admin_provider_oauth_batch_import_entry(
|
|||||||
account_id: None,
|
account_id: None,
|
||||||
account_user_id: None,
|
account_user_id: None,
|
||||||
plan_type: None,
|
plan_type: None,
|
||||||
user_id: None,
|
pool_tier: None,
|
||||||
|
user_id: grok_cookie_value(raw_token, "x-userid"),
|
||||||
email: None,
|
email: None,
|
||||||
|
account_name: None,
|
||||||
|
sso_rw_token: grok_cookie_value(raw_token, "sso-rw"),
|
||||||
|
cf_cookies: grok_cookie_profile(raw_token),
|
||||||
|
cf_clearance: grok_cookie_value(raw_token, "cf_clearance"),
|
||||||
|
user_agent: None,
|
||||||
|
browser_profile: None,
|
||||||
})
|
})
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
@@ -100,8 +170,34 @@ fn extract_admin_provider_oauth_batch_import_entry(
|
|||||||
.get("access_token")
|
.get("access_token")
|
||||||
.or_else(|| object.get("accessToken")),
|
.or_else(|| object.get("accessToken")),
|
||||||
);
|
);
|
||||||
let (refresh_token, access_token) =
|
let grok_token_alias = if provider_type.trim().eq_ignore_ascii_case("grok") {
|
||||||
normalize_single_import_tokens(refresh_token.as_deref(), access_token.as_deref());
|
object.get("token")
|
||||||
|
} else {
|
||||||
|
None
|
||||||
|
};
|
||||||
|
let grok_cookie = if provider_type.trim().eq_ignore_ascii_case("grok") {
|
||||||
|
coerce_admin_provider_oauth_import_str(
|
||||||
|
object.get("cookie").or_else(|| object.get("cookieHeader")),
|
||||||
|
)
|
||||||
|
} else {
|
||||||
|
None
|
||||||
|
};
|
||||||
|
let session_token = coerce_admin_provider_oauth_import_str(
|
||||||
|
object
|
||||||
|
.get("sso_token")
|
||||||
|
.or_else(|| object.get("ssoToken"))
|
||||||
|
.or(grok_token_alias),
|
||||||
|
)
|
||||||
|
.or_else(|| {
|
||||||
|
grok_cookie
|
||||||
|
.as_deref()
|
||||||
|
.and_then(|cookie| grok_cookie_value(cookie, "sso"))
|
||||||
|
});
|
||||||
|
let (refresh_token, access_token) = normalize_provider_import_tokens(
|
||||||
|
provider_type,
|
||||||
|
refresh_token.as_deref(),
|
||||||
|
access_token.as_deref().or(session_token.as_deref()),
|
||||||
|
);
|
||||||
if refresh_token.is_none() && access_token.is_none() {
|
if refresh_token.is_none() && access_token.is_none() {
|
||||||
return None;
|
return None;
|
||||||
}
|
}
|
||||||
@@ -129,14 +225,65 @@ fn extract_admin_provider_oauth_batch_import_entry(
|
|||||||
.or_else(|| object.get("chatgptPlanType")),
|
.or_else(|| object.get("chatgptPlanType")),
|
||||||
)
|
)
|
||||||
.map(|value| value.to_ascii_lowercase());
|
.map(|value| value.to_ascii_lowercase());
|
||||||
|
let pool_tier = coerce_admin_provider_oauth_import_str(
|
||||||
|
object
|
||||||
|
.get("pool_tier")
|
||||||
|
.or_else(|| object.get("poolTier"))
|
||||||
|
.or_else(|| object.get("tier")),
|
||||||
|
)
|
||||||
|
.map(|value| value.to_ascii_lowercase());
|
||||||
let user_id = coerce_admin_provider_oauth_import_str(
|
let user_id = coerce_admin_provider_oauth_import_str(
|
||||||
object
|
object
|
||||||
.get("user_id")
|
.get("user_id")
|
||||||
.or_else(|| object.get("userId"))
|
.or_else(|| object.get("userId"))
|
||||||
.or_else(|| object.get("chatgpt_user_id"))
|
.or_else(|| object.get("chatgpt_user_id"))
|
||||||
.or_else(|| object.get("chatgptUserId")),
|
.or_else(|| object.get("chatgptUserId")),
|
||||||
);
|
)
|
||||||
|
.or_else(|| {
|
||||||
|
grok_cookie
|
||||||
|
.as_deref()
|
||||||
|
.and_then(|cookie| grok_cookie_value(cookie, "x-userid"))
|
||||||
|
});
|
||||||
let email = coerce_admin_provider_oauth_import_str(object.get("email"));
|
let email = coerce_admin_provider_oauth_import_str(object.get("email"));
|
||||||
|
let account_name = coerce_admin_provider_oauth_import_str(
|
||||||
|
object
|
||||||
|
.get("account_name")
|
||||||
|
.or_else(|| object.get("accountName")),
|
||||||
|
);
|
||||||
|
let sso_rw_token = coerce_admin_provider_oauth_import_str(
|
||||||
|
object
|
||||||
|
.get("sso_rw_token")
|
||||||
|
.or_else(|| object.get("ssoRwToken")),
|
||||||
|
)
|
||||||
|
.or_else(|| {
|
||||||
|
grok_cookie
|
||||||
|
.as_deref()
|
||||||
|
.and_then(|cookie| grok_cookie_value(cookie, "sso-rw"))
|
||||||
|
});
|
||||||
|
let cf_clearance = coerce_admin_provider_oauth_import_str(
|
||||||
|
object
|
||||||
|
.get("cf_clearance")
|
||||||
|
.or_else(|| object.get("cfClearance")),
|
||||||
|
)
|
||||||
|
.or_else(|| {
|
||||||
|
grok_cookie
|
||||||
|
.as_deref()
|
||||||
|
.and_then(|cookie| grok_cookie_value(cookie, "cf_clearance"))
|
||||||
|
});
|
||||||
|
let cf_cookies = coerce_admin_provider_oauth_import_str(
|
||||||
|
object.get("cf_cookies").or_else(|| object.get("cfCookies")),
|
||||||
|
)
|
||||||
|
.or_else(|| grok_cookie.as_deref().and_then(grok_cookie_profile));
|
||||||
|
let user_agent = coerce_admin_provider_oauth_import_str(
|
||||||
|
object.get("user_agent").or_else(|| object.get("userAgent")),
|
||||||
|
);
|
||||||
|
let browser_profile = coerce_admin_provider_oauth_import_str(
|
||||||
|
object
|
||||||
|
.get("browser_profile")
|
||||||
|
.or_else(|| object.get("browserProfile"))
|
||||||
|
.or_else(|| object.get("browser"))
|
||||||
|
.or_else(|| object.get("impersonate")),
|
||||||
|
);
|
||||||
Some(AdminProviderOAuthBatchImportEntry {
|
Some(AdminProviderOAuthBatchImportEntry {
|
||||||
refresh_token,
|
refresh_token,
|
||||||
access_token,
|
access_token,
|
||||||
@@ -144,8 +291,15 @@ fn extract_admin_provider_oauth_batch_import_entry(
|
|||||||
account_id,
|
account_id,
|
||||||
account_user_id,
|
account_user_id,
|
||||||
plan_type,
|
plan_type,
|
||||||
|
pool_tier,
|
||||||
user_id,
|
user_id,
|
||||||
email,
|
email,
|
||||||
|
account_name,
|
||||||
|
sso_rw_token,
|
||||||
|
cf_cookies,
|
||||||
|
cf_clearance,
|
||||||
|
user_agent,
|
||||||
|
browser_profile,
|
||||||
})
|
})
|
||||||
}
|
}
|
||||||
_ => None,
|
_ => None,
|
||||||
@@ -153,6 +307,7 @@ fn extract_admin_provider_oauth_batch_import_entry(
|
|||||||
}
|
}
|
||||||
|
|
||||||
pub(super) fn parse_admin_provider_oauth_batch_import_entries(
|
pub(super) fn parse_admin_provider_oauth_batch_import_entries(
|
||||||
|
provider_type: &str,
|
||||||
raw_credentials: &str,
|
raw_credentials: &str,
|
||||||
) -> Vec<AdminProviderOAuthBatchImportEntry> {
|
) -> Vec<AdminProviderOAuthBatchImportEntry> {
|
||||||
let raw = raw_credentials.trim();
|
let raw = raw_credentials.trim();
|
||||||
@@ -165,7 +320,9 @@ pub(super) fn parse_admin_provider_oauth_batch_import_entries(
|
|||||||
{
|
{
|
||||||
return items
|
return items
|
||||||
.iter()
|
.iter()
|
||||||
.filter_map(extract_admin_provider_oauth_batch_import_entry)
|
.filter_map(|item| {
|
||||||
|
extract_admin_provider_oauth_batch_import_entry(provider_type, item)
|
||||||
|
})
|
||||||
.collect();
|
.collect();
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
@@ -174,7 +331,7 @@ pub(super) fn parse_admin_provider_oauth_batch_import_entries(
|
|||||||
if let Ok(value @ serde_json::Value::Object(_)) =
|
if let Ok(value @ serde_json::Value::Object(_)) =
|
||||||
serde_json::from_str::<serde_json::Value>(raw)
|
serde_json::from_str::<serde_json::Value>(raw)
|
||||||
{
|
{
|
||||||
return extract_admin_provider_oauth_batch_import_entry(&value)
|
return extract_admin_provider_oauth_batch_import_entry(provider_type, &value)
|
||||||
.into_iter()
|
.into_iter()
|
||||||
.collect();
|
.collect();
|
||||||
}
|
}
|
||||||
@@ -183,18 +340,19 @@ pub(super) fn parse_admin_provider_oauth_batch_import_entries(
|
|||||||
raw.lines()
|
raw.lines()
|
||||||
.map(str::trim)
|
.map(str::trim)
|
||||||
.filter(|line| !line.is_empty() && !line.starts_with('#'))
|
.filter(|line| !line.is_empty() && !line.starts_with('#'))
|
||||||
.map(|token| {
|
.filter_map(|line| {
|
||||||
let (refresh_token, access_token) = import_tokens_from_raw_token(token);
|
if line.starts_with('{') {
|
||||||
AdminProviderOAuthBatchImportEntry {
|
return serde_json::from_str::<serde_json::Value>(line)
|
||||||
refresh_token,
|
.ok()
|
||||||
access_token,
|
.and_then(|value| {
|
||||||
expires_at: None,
|
extract_admin_provider_oauth_batch_import_entry(provider_type, &value)
|
||||||
account_id: None,
|
});
|
||||||
account_user_id: None,
|
|
||||||
plan_type: None,
|
|
||||||
user_id: None,
|
|
||||||
email: None,
|
|
||||||
}
|
}
|
||||||
|
|
||||||
|
extract_admin_provider_oauth_batch_import_entry(
|
||||||
|
provider_type,
|
||||||
|
&serde_json::Value::String(line.to_string()),
|
||||||
|
)
|
||||||
})
|
})
|
||||||
.collect()
|
.collect()
|
||||||
}
|
}
|
||||||
@@ -204,10 +362,8 @@ pub(super) fn apply_admin_provider_oauth_batch_import_hints(
|
|||||||
entry: &AdminProviderOAuthBatchImportEntry,
|
entry: &AdminProviderOAuthBatchImportEntry,
|
||||||
auth_config: &mut serde_json::Map<String, serde_json::Value>,
|
auth_config: &mut serde_json::Map<String, serde_json::Value>,
|
||||||
) {
|
) {
|
||||||
if !matches!(
|
let provider_type = provider_type.trim().to_ascii_lowercase();
|
||||||
provider_type.trim().to_ascii_lowercase().as_str(),
|
if !matches!(provider_type.as_str(), "codex" | "chatgpt_web" | "grok") {
|
||||||
"codex" | "chatgpt_web"
|
|
||||||
) {
|
|
||||||
return;
|
return;
|
||||||
}
|
}
|
||||||
if let Some(account_id) = entry.account_id.as_ref() {
|
if let Some(account_id) = entry.account_id.as_ref() {
|
||||||
@@ -225,6 +381,11 @@ pub(super) fn apply_admin_provider_oauth_batch_import_hints(
|
|||||||
.entry("plan_type".to_string())
|
.entry("plan_type".to_string())
|
||||||
.or_insert_with(|| json!(plan_type));
|
.or_insert_with(|| json!(plan_type));
|
||||||
}
|
}
|
||||||
|
if let Some(pool_tier) = entry.pool_tier.as_ref() {
|
||||||
|
auth_config
|
||||||
|
.entry("pool_tier".to_string())
|
||||||
|
.or_insert_with(|| json!(pool_tier));
|
||||||
|
}
|
||||||
if let Some(user_id) = entry.user_id.as_ref() {
|
if let Some(user_id) = entry.user_id.as_ref() {
|
||||||
auth_config
|
auth_config
|
||||||
.entry("user_id".to_string())
|
.entry("user_id".to_string())
|
||||||
@@ -235,6 +396,36 @@ pub(super) fn apply_admin_provider_oauth_batch_import_hints(
|
|||||||
.entry("email".to_string())
|
.entry("email".to_string())
|
||||||
.or_insert_with(|| json!(email));
|
.or_insert_with(|| json!(email));
|
||||||
}
|
}
|
||||||
|
if let Some(account_name) = entry.account_name.as_ref() {
|
||||||
|
auth_config
|
||||||
|
.entry("account_name".to_string())
|
||||||
|
.or_insert_with(|| json!(account_name));
|
||||||
|
}
|
||||||
|
if let Some(sso_rw_token) = entry.sso_rw_token.as_ref() {
|
||||||
|
auth_config
|
||||||
|
.entry("sso_rw_token".to_string())
|
||||||
|
.or_insert_with(|| json!(sso_rw_token));
|
||||||
|
}
|
||||||
|
if let Some(cf_cookies) = entry.cf_cookies.as_ref() {
|
||||||
|
auth_config
|
||||||
|
.entry("cf_cookies".to_string())
|
||||||
|
.or_insert_with(|| json!(cf_cookies));
|
||||||
|
}
|
||||||
|
if let Some(cf_clearance) = entry.cf_clearance.as_ref() {
|
||||||
|
auth_config
|
||||||
|
.entry("cf_clearance".to_string())
|
||||||
|
.or_insert_with(|| json!(cf_clearance));
|
||||||
|
}
|
||||||
|
if let Some(user_agent) = entry.user_agent.as_ref() {
|
||||||
|
auth_config
|
||||||
|
.entry("user_agent".to_string())
|
||||||
|
.or_insert_with(|| json!(user_agent));
|
||||||
|
}
|
||||||
|
if let Some(browser_profile) = entry.browser_profile.as_ref() {
|
||||||
|
auth_config
|
||||||
|
.entry("browser_profile".to_string())
|
||||||
|
.or_insert_with(|| json!(browser_profile));
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
pub(super) async fn extract_admin_provider_oauth_batch_error_detail(
|
pub(super) async fn extract_admin_provider_oauth_batch_error_detail(
|
||||||
@@ -337,6 +528,7 @@ mod tests {
|
|||||||
#[test]
|
#[test]
|
||||||
fn parses_access_token_only_entry() {
|
fn parses_access_token_only_entry() {
|
||||||
let entries = parse_admin_provider_oauth_batch_import_entries(
|
let entries = parse_admin_provider_oauth_batch_import_entries(
|
||||||
|
"codex",
|
||||||
r#"[{"accessToken":"at_1","expiresAt":2100000000,"accountId":"acc-1","email":"u@example.com"}]"#,
|
r#"[{"accessToken":"at_1","expiresAt":2100000000,"accountId":"acc-1","email":"u@example.com"}]"#,
|
||||||
);
|
);
|
||||||
|
|
||||||
@@ -356,10 +548,89 @@ mod tests {
|
|||||||
"exp": 2_000_000_000u64,
|
"exp": 2_000_000_000u64,
|
||||||
}));
|
}));
|
||||||
|
|
||||||
let entries = parse_admin_provider_oauth_batch_import_entries(&token);
|
let entries = parse_admin_provider_oauth_batch_import_entries("codex", &token);
|
||||||
|
|
||||||
assert_eq!(entries.len(), 1);
|
assert_eq!(entries.len(), 1);
|
||||||
assert_eq!(entries[0].refresh_token, None);
|
assert_eq!(entries[0].refresh_token, None);
|
||||||
assert_eq!(entries[0].access_token.as_deref(), Some(token.as_str()));
|
assert_eq!(entries[0].access_token.as_deref(), Some(token.as_str()));
|
||||||
}
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn parses_grok_jsonl_session_entries() {
|
||||||
|
let entries = parse_admin_provider_oauth_batch_import_entries(
|
||||||
|
"grok",
|
||||||
|
r#"{"sso_token":"sso-1","cf_clearance":"cf-1","pool_tier":"heavy","email":"grok@example.com","browser_profile":"chrome136"}"#,
|
||||||
|
);
|
||||||
|
|
||||||
|
assert_eq!(entries.len(), 1);
|
||||||
|
assert_eq!(entries[0].refresh_token, None);
|
||||||
|
assert_eq!(entries[0].access_token.as_deref(), Some("sso-1"));
|
||||||
|
assert_eq!(entries[0].cf_clearance.as_deref(), Some("cf-1"));
|
||||||
|
assert_eq!(entries[0].pool_tier.as_deref(), Some("heavy"));
|
||||||
|
assert_eq!(entries[0].email.as_deref(), Some("grok@example.com"));
|
||||||
|
assert_eq!(entries[0].browser_profile.as_deref(), Some("chrome136"));
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn parses_grok_token_alias_with_account_traits() {
|
||||||
|
let entries = parse_admin_provider_oauth_batch_import_entries(
|
||||||
|
"grok",
|
||||||
|
r#"[{"token":"sso-1","planType":"super","tier":"heavy","accountName":"Grok Heavy"}]"#,
|
||||||
|
);
|
||||||
|
|
||||||
|
assert_eq!(entries.len(), 1);
|
||||||
|
assert_eq!(entries[0].refresh_token, None);
|
||||||
|
assert_eq!(entries[0].access_token.as_deref(), Some("sso-1"));
|
||||||
|
assert_eq!(entries[0].plan_type.as_deref(), Some("super"));
|
||||||
|
assert_eq!(entries[0].pool_tier.as_deref(), Some("heavy"));
|
||||||
|
assert_eq!(entries[0].account_name.as_deref(), Some("Grok Heavy"));
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn parses_grok_plain_line_as_session_token() {
|
||||||
|
let entries = parse_admin_provider_oauth_batch_import_entries("grok", "opaque-sso-token");
|
||||||
|
|
||||||
|
assert_eq!(entries.len(), 1);
|
||||||
|
assert_eq!(entries[0].refresh_token, None);
|
||||||
|
assert_eq!(entries[0].access_token.as_deref(), Some("opaque-sso-token"));
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn parses_grok_cookie_line_as_session_metadata() {
|
||||||
|
let entries = parse_admin_provider_oauth_batch_import_entries(
|
||||||
|
"grok",
|
||||||
|
"i18nextLng=zh; cf_clearance=cf-1; sso-rw=rw-1; sso=sso-1; x-userid=user-1",
|
||||||
|
);
|
||||||
|
|
||||||
|
assert_eq!(entries.len(), 1);
|
||||||
|
assert_eq!(entries[0].refresh_token, None);
|
||||||
|
assert_eq!(entries[0].access_token.as_deref(), Some("sso-1"));
|
||||||
|
assert_eq!(entries[0].sso_rw_token.as_deref(), Some("rw-1"));
|
||||||
|
assert_eq!(
|
||||||
|
entries[0].cf_cookies.as_deref(),
|
||||||
|
Some("i18nextLng=zh; cf_clearance=cf-1; x-userid=user-1")
|
||||||
|
);
|
||||||
|
assert_eq!(entries[0].cf_clearance.as_deref(), Some("cf-1"));
|
||||||
|
assert_eq!(entries[0].user_id.as_deref(), Some("user-1"));
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn parses_grok_cookie_object_as_session_metadata() {
|
||||||
|
let entries = parse_admin_provider_oauth_batch_import_entries(
|
||||||
|
"grok",
|
||||||
|
r#"[{"cookie":"cf_clearance=cf-1; sso-rw=rw-1; sso=sso-1; x-userid=user-1","tier":"heavy"}]"#,
|
||||||
|
);
|
||||||
|
|
||||||
|
assert_eq!(entries.len(), 1);
|
||||||
|
assert_eq!(entries[0].refresh_token, None);
|
||||||
|
assert_eq!(entries[0].access_token.as_deref(), Some("sso-1"));
|
||||||
|
assert_eq!(entries[0].sso_rw_token.as_deref(), Some("rw-1"));
|
||||||
|
assert_eq!(
|
||||||
|
entries[0].cf_cookies.as_deref(),
|
||||||
|
Some("cf_clearance=cf-1; x-userid=user-1")
|
||||||
|
);
|
||||||
|
assert_eq!(entries[0].cf_clearance.as_deref(), Some("cf-1"));
|
||||||
|
assert_eq!(entries[0].user_id.as_deref(), Some("user-1"));
|
||||||
|
assert_eq!(entries[0].pool_tier.as_deref(), Some("heavy"));
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -10,7 +10,7 @@ use super::progress::{
|
|||||||
};
|
};
|
||||||
use crate::handlers::admin::provider::oauth::errors::build_internal_control_error_response;
|
use crate::handlers::admin::provider::oauth::errors::build_internal_control_error_response;
|
||||||
use crate::handlers::admin::provider::oauth::state::{
|
use crate::handlers::admin::provider::oauth::state::{
|
||||||
admin_provider_oauth_template, build_admin_provider_oauth_backend_unavailable_response,
|
build_admin_provider_oauth_backend_unavailable_response,
|
||||||
is_fixed_provider_type_for_provider_oauth,
|
is_fixed_provider_type_for_provider_oauth,
|
||||||
};
|
};
|
||||||
use crate::handlers::admin::provider::shared::paths::admin_provider_oauth_batch_import_task_provider_id;
|
use crate::handlers::admin::provider::shared::paths::admin_provider_oauth_batch_import_task_provider_id;
|
||||||
@@ -124,10 +124,6 @@ pub(in super::super) async fn handle_admin_provider_oauth_start_batch_import_tas
|
|||||||
"该 Provider 不是固定类型,无法使用 provider-oauth",
|
"该 Provider 不是固定类型,无法使用 provider-oauth",
|
||||||
));
|
));
|
||||||
}
|
}
|
||||||
if provider_type != "kiro" && admin_provider_oauth_template(&provider_type).is_none() {
|
|
||||||
return Ok(build_admin_provider_oauth_backend_unavailable_response());
|
|
||||||
}
|
|
||||||
|
|
||||||
let total = estimate_admin_provider_oauth_batch_import_total(
|
let total = estimate_admin_provider_oauth_batch_import_total(
|
||||||
&provider_type,
|
&provider_type,
|
||||||
payload.credentials.as_str(),
|
payload.credentials.as_str(),
|
||||||
|
|||||||
@@ -3,6 +3,8 @@ use axum::{
|
|||||||
body::Body,
|
body::Body,
|
||||||
response::{IntoResponse, Response},
|
response::{IntoResponse, Response},
|
||||||
};
|
};
|
||||||
|
use serde_json::{Map, Value};
|
||||||
|
use std::time::{SystemTime, UNIX_EPOCH};
|
||||||
|
|
||||||
pub(super) fn attach_admin_provider_oauth_audit_response(
|
pub(super) fn attach_admin_provider_oauth_audit_response(
|
||||||
response: Response<Body>,
|
response: Response<Body>,
|
||||||
@@ -19,3 +21,79 @@ pub(super) fn attach_admin_provider_oauth_audit_response(
|
|||||||
};
|
};
|
||||||
attach_admin_audit_response(response, event_name, action, target_type, &target_id)
|
attach_admin_audit_response(response, event_name, action, target_type, &target_id)
|
||||||
}
|
}
|
||||||
|
|
||||||
|
pub(super) fn admin_provider_oauth_key_name_from_auth_config(
|
||||||
|
provider_type: &str,
|
||||||
|
auth_config: &Map<String, Value>,
|
||||||
|
batch_index: Option<usize>,
|
||||||
|
) -> String {
|
||||||
|
let provider_type = provider_type.trim();
|
||||||
|
if let Some(email) = trimmed_auth_config_string(auth_config, "email") {
|
||||||
|
return format!("{provider_type}_{email}");
|
||||||
|
}
|
||||||
|
if provider_type.eq_ignore_ascii_case("grok") {
|
||||||
|
if let Some(user_id) = trimmed_auth_config_string(auth_config, "user_id") {
|
||||||
|
return format!("grok_{user_id}");
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
let timestamp = SystemTime::now()
|
||||||
|
.duration_since(UNIX_EPOCH)
|
||||||
|
.ok()
|
||||||
|
.map(|duration| duration.as_secs())
|
||||||
|
.unwrap_or(0);
|
||||||
|
match batch_index {
|
||||||
|
Some(index) => format!("{provider_type}_{timestamp}_{index}"),
|
||||||
|
None => format!("账号_{timestamp}"),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fn trimmed_auth_config_string(auth_config: &Map<String, Value>, key: &str) -> Option<String> {
|
||||||
|
auth_config
|
||||||
|
.get(key)
|
||||||
|
.and_then(Value::as_str)
|
||||||
|
.map(str::trim)
|
||||||
|
.filter(|value| !value.is_empty())
|
||||||
|
.map(ToOwned::to_owned)
|
||||||
|
}
|
||||||
|
|
||||||
|
#[cfg(test)]
|
||||||
|
mod tests {
|
||||||
|
use super::*;
|
||||||
|
use serde_json::{json, Map};
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn grok_default_key_name_uses_full_user_id() {
|
||||||
|
let mut auth_config = Map::new();
|
||||||
|
auth_config.insert(
|
||||||
|
"user_id".to_string(),
|
||||||
|
json!("1619039a-0191-4e0a-a490-8f4ad21262c9"),
|
||||||
|
);
|
||||||
|
|
||||||
|
assert_eq!(
|
||||||
|
admin_provider_oauth_key_name_from_auth_config("grok", &auth_config, None),
|
||||||
|
"grok_1619039a-0191-4e0a-a490-8f4ad21262c9"
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn default_key_name_prefers_email_over_grok_user_id() {
|
||||||
|
let mut auth_config = Map::new();
|
||||||
|
auth_config.insert("email".to_string(), json!("grok@example.com"));
|
||||||
|
auth_config.insert("user_id".to_string(), json!("user-1"));
|
||||||
|
|
||||||
|
assert_eq!(
|
||||||
|
admin_provider_oauth_key_name_from_auth_config("grok", &auth_config, None),
|
||||||
|
"grok_grok@example.com"
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn batch_default_key_name_keeps_existing_timestamp_shape() {
|
||||||
|
let auth_config = Map::new();
|
||||||
|
let name = admin_provider_oauth_key_name_from_auth_config("codex", &auth_config, Some(3));
|
||||||
|
|
||||||
|
assert!(name.starts_with("codex_"));
|
||||||
|
assert!(name.ends_with("_3"));
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|||||||
@@ -14,8 +14,9 @@ use super::super::state::{
|
|||||||
exchange_admin_provider_oauth_refresh_token, is_fixed_provider_type_for_provider_oauth,
|
exchange_admin_provider_oauth_refresh_token, is_fixed_provider_type_for_provider_oauth,
|
||||||
json_u64_value,
|
json_u64_value,
|
||||||
};
|
};
|
||||||
|
use super::helpers::admin_provider_oauth_key_name_from_auth_config;
|
||||||
use super::token_import::{
|
use super::token_import::{
|
||||||
build_provider_access_token_import_auth_config, normalize_single_import_tokens,
|
build_provider_access_token_import_auth_config, normalize_provider_import_tokens,
|
||||||
provider_type_supports_access_token_import,
|
provider_type_supports_access_token_import,
|
||||||
};
|
};
|
||||||
use crate::handlers::admin::provider::shared::paths::admin_provider_oauth_import_provider_id;
|
use crate::handlers::admin::provider::shared::paths::admin_provider_oauth_import_provider_id;
|
||||||
@@ -31,7 +32,6 @@ use axum::{
|
|||||||
Json,
|
Json,
|
||||||
};
|
};
|
||||||
use serde_json::json;
|
use serde_json::json;
|
||||||
use std::time::{SystemTime, UNIX_EPOCH};
|
|
||||||
|
|
||||||
struct AdminProviderOAuthSingleImportTokens {
|
struct AdminProviderOAuthSingleImportTokens {
|
||||||
access_token: String,
|
access_token: String,
|
||||||
@@ -72,7 +72,8 @@ fn apply_single_import_hints(
|
|||||||
payload: &serde_json::Map<String, serde_json::Value>,
|
payload: &serde_json::Map<String, serde_json::Value>,
|
||||||
auth_config: &mut serde_json::Map<String, serde_json::Value>,
|
auth_config: &mut serde_json::Map<String, serde_json::Value>,
|
||||||
) {
|
) {
|
||||||
if !provider_type_supports_access_token_import(provider_type) {
|
let provider_type = provider_type.trim().to_ascii_lowercase();
|
||||||
|
if !matches!(provider_type.as_str(), "codex" | "chatgpt_web" | "grok") {
|
||||||
return;
|
return;
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -110,19 +111,40 @@ fn apply_single_import_hints(
|
|||||||
&["user_id", "userId", "chatgpt_user_id", "chatgptUserId"][..],
|
&["user_id", "userId", "chatgpt_user_id", "chatgptUserId"][..],
|
||||||
),
|
),
|
||||||
("account_name", &["account_name", "accountName"][..]),
|
("account_name", &["account_name", "accountName"][..]),
|
||||||
|
("sso_rw_token", &["sso_rw_token", "ssoRwToken"][..]),
|
||||||
|
(
|
||||||
|
"cf_cookies",
|
||||||
|
&["cf_cookies", "cfCookies", "cookie", "cookieHeader"][..],
|
||||||
|
),
|
||||||
|
("cf_clearance", &["cf_clearance", "cfClearance"][..]),
|
||||||
|
("user_agent", &["user_agent", "userAgent"][..]),
|
||||||
|
(
|
||||||
|
"browser_profile",
|
||||||
|
&[
|
||||||
|
"browser_profile",
|
||||||
|
"browserProfile",
|
||||||
|
"browser",
|
||||||
|
"impersonate",
|
||||||
|
][..],
|
||||||
|
),
|
||||||
|
("pool_tier", &["pool_tier", "poolTier", "tier"][..]),
|
||||||
] {
|
] {
|
||||||
let Some(value) = import_payload_string_any(payload, keys) else {
|
let Some(value) = import_payload_string_any(payload, keys) else {
|
||||||
continue;
|
continue;
|
||||||
};
|
};
|
||||||
auth_config
|
auth_config.entry(target.to_string()).or_insert_with(|| {
|
||||||
.entry(target.to_string())
|
if target == "plan_type" || target == "pool_tier" {
|
||||||
.or_insert_with(|| json!(value));
|
json!(value.to_ascii_lowercase())
|
||||||
|
} else {
|
||||||
|
json!(value)
|
||||||
|
}
|
||||||
|
});
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
async fn resolve_admin_provider_oauth_single_import_tokens(
|
async fn resolve_admin_provider_oauth_single_import_tokens(
|
||||||
state: &AdminAppState<'_>,
|
state: &AdminAppState<'_>,
|
||||||
template: AdminProviderOAuthTemplate,
|
template: Option<AdminProviderOAuthTemplate>,
|
||||||
provider_type: &str,
|
provider_type: &str,
|
||||||
refresh_token: Option<&str>,
|
refresh_token: Option<&str>,
|
||||||
access_token: Option<&str>,
|
access_token: Option<&str>,
|
||||||
@@ -133,6 +155,32 @@ async fn resolve_admin_provider_oauth_single_import_tokens(
|
|||||||
.map(str::trim)
|
.map(str::trim)
|
||||||
.filter(|value| !value.is_empty())
|
.filter(|value| !value.is_empty())
|
||||||
{
|
{
|
||||||
|
let Some(template) = template else {
|
||||||
|
if provider_type_supports_access_token_import(provider_type) {
|
||||||
|
if let Some(access_token) = access_token
|
||||||
|
.map(str::trim)
|
||||||
|
.filter(|value| !value.is_empty())
|
||||||
|
{
|
||||||
|
let (auth_config, expires_at) = build_provider_access_token_import_auth_config(
|
||||||
|
provider_type,
|
||||||
|
access_token,
|
||||||
|
Some(refresh_token),
|
||||||
|
imported_expires_at,
|
||||||
|
Some("Provider 不支持 Refresh Token 交换,已回退为 Session Token 导入"),
|
||||||
|
);
|
||||||
|
return Ok(AdminProviderOAuthSingleImportTokens {
|
||||||
|
access_token: access_token.to_string(),
|
||||||
|
auth_config,
|
||||||
|
expires_at,
|
||||||
|
});
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return Err(build_internal_control_error_response(
|
||||||
|
http::StatusCode::BAD_REQUEST,
|
||||||
|
"该 Provider 不支持 Refresh Token 导入,请提供 sso_token 或 access_token",
|
||||||
|
));
|
||||||
|
};
|
||||||
|
|
||||||
let token_payload = match state
|
let token_payload = match state
|
||||||
.exchange_admin_provider_oauth_refresh_token(
|
.exchange_admin_provider_oauth_refresh_token(
|
||||||
template,
|
template,
|
||||||
@@ -200,7 +248,7 @@ async fn resolve_admin_provider_oauth_single_import_tokens(
|
|||||||
if !provider_type_supports_access_token_import(provider_type) {
|
if !provider_type_supports_access_token_import(provider_type) {
|
||||||
return Err(build_internal_control_error_response(
|
return Err(build_internal_control_error_response(
|
||||||
http::StatusCode::BAD_REQUEST,
|
http::StatusCode::BAD_REQUEST,
|
||||||
"Access Token 导入仅支持 Codex / ChatGPT Web Provider",
|
"Access Token 导入仅支持 Codex / ChatGPT Web / Grok Provider",
|
||||||
));
|
));
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -248,18 +296,11 @@ pub(super) async fn handle_admin_provider_oauth_import_refresh_token(
|
|||||||
}
|
}
|
||||||
};
|
};
|
||||||
let refresh_token_input = import_payload_string(&raw_payload, "refresh_token", "refreshToken");
|
let refresh_token_input = import_payload_string(&raw_payload, "refresh_token", "refreshToken");
|
||||||
let access_token_input = import_payload_string(&raw_payload, "access_token", "accessToken");
|
let access_token_input = import_payload_string_any(
|
||||||
let imported_expires_at = import_payload_u64(&raw_payload, "expires_at", "expiresAt");
|
&raw_payload,
|
||||||
let (refresh_token_input, access_token_input) = normalize_single_import_tokens(
|
&["access_token", "accessToken", "sso_token", "ssoToken"],
|
||||||
refresh_token_input.as_deref(),
|
|
||||||
access_token_input.as_deref(),
|
|
||||||
);
|
);
|
||||||
if refresh_token_input.is_none() && access_token_input.is_none() {
|
let imported_expires_at = import_payload_u64(&raw_payload, "expires_at", "expiresAt");
|
||||||
return Ok(build_internal_control_error_response(
|
|
||||||
http::StatusCode::BAD_REQUEST,
|
|
||||||
"Refresh Token 或 Access Token 不能为空",
|
|
||||||
));
|
|
||||||
}
|
|
||||||
let name = raw_payload
|
let name = raw_payload
|
||||||
.get("name")
|
.get("name")
|
||||||
.and_then(serde_json::Value::as_str)
|
.and_then(serde_json::Value::as_str)
|
||||||
@@ -285,6 +326,17 @@ pub(super) async fn handle_admin_provider_oauth_import_refresh_token(
|
|||||||
));
|
));
|
||||||
};
|
};
|
||||||
let provider_type = provider.provider_type.trim().to_ascii_lowercase();
|
let provider_type = provider.provider_type.trim().to_ascii_lowercase();
|
||||||
|
let (refresh_token_input, access_token_input) = normalize_provider_import_tokens(
|
||||||
|
&provider_type,
|
||||||
|
refresh_token_input.as_deref(),
|
||||||
|
access_token_input.as_deref(),
|
||||||
|
);
|
||||||
|
if refresh_token_input.is_none() && access_token_input.is_none() {
|
||||||
|
return Ok(build_internal_control_error_response(
|
||||||
|
http::StatusCode::BAD_REQUEST,
|
||||||
|
"Refresh Token、Access Token 或 sso_token 不能为空",
|
||||||
|
));
|
||||||
|
}
|
||||||
if !is_fixed_provider_type_for_provider_oauth(&provider_type) {
|
if !is_fixed_provider_type_for_provider_oauth(&provider_type) {
|
||||||
return Ok(build_internal_control_error_response(
|
return Ok(build_internal_control_error_response(
|
||||||
http::StatusCode::BAD_REQUEST,
|
http::StatusCode::BAD_REQUEST,
|
||||||
@@ -297,9 +349,10 @@ pub(super) async fn handle_admin_provider_oauth_import_refresh_token(
|
|||||||
"Kiro 不支持单条 Refresh Token 导入,请使用批量导入或设备授权。",
|
"Kiro 不支持单条 Refresh Token 导入,请使用批量导入或设备授权。",
|
||||||
));
|
));
|
||||||
}
|
}
|
||||||
let Some(template) = admin_provider_oauth_template(&provider_type) else {
|
let template = admin_provider_oauth_template(&provider_type);
|
||||||
|
if template.is_none() && !provider_type_supports_access_token_import(&provider_type) {
|
||||||
return Ok(build_admin_provider_oauth_backend_unavailable_response());
|
return Ok(build_admin_provider_oauth_backend_unavailable_response());
|
||||||
};
|
}
|
||||||
let endpoint_resolution =
|
let endpoint_resolution =
|
||||||
resolve_provider_oauth_runtime_endpoints(state, &provider, &provider_type).await?;
|
resolve_provider_oauth_runtime_endpoints(state, &provider, &provider_type).await?;
|
||||||
let endpoints = endpoint_resolution.endpoints;
|
let endpoints = endpoint_resolution.endpoints;
|
||||||
@@ -380,25 +433,9 @@ pub(super) async fn handle_admin_provider_oauth_import_refresh_token(
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
} else {
|
} else {
|
||||||
let name = name
|
let name = name.unwrap_or_else(|| {
|
||||||
.or_else(|| {
|
admin_provider_oauth_key_name_from_auth_config(&provider_type, &auth_config, None)
|
||||||
auth_config
|
});
|
||||||
.get("email")
|
|
||||||
.and_then(serde_json::Value::as_str)
|
|
||||||
.map(str::trim)
|
|
||||||
.filter(|value| !value.is_empty())
|
|
||||||
.map(ToOwned::to_owned)
|
|
||||||
})
|
|
||||||
.unwrap_or_else(|| {
|
|
||||||
format!(
|
|
||||||
"账号_{}",
|
|
||||||
SystemTime::now()
|
|
||||||
.duration_since(UNIX_EPOCH)
|
|
||||||
.ok()
|
|
||||||
.map(|duration| duration.as_secs())
|
|
||||||
.unwrap_or(0)
|
|
||||||
)
|
|
||||||
});
|
|
||||||
match state
|
match state
|
||||||
.create_provider_oauth_catalog_key(
|
.create_provider_oauth_catalog_key(
|
||||||
&provider_id,
|
&provider_id,
|
||||||
|
|||||||
@@ -81,6 +81,28 @@ pub(super) fn normalize_single_import_tokens(
|
|||||||
(refresh_token, access_token)
|
(refresh_token, access_token)
|
||||||
}
|
}
|
||||||
|
|
||||||
|
pub(super) fn normalize_provider_import_tokens(
|
||||||
|
provider_type: &str,
|
||||||
|
refresh_token: Option<&str>,
|
||||||
|
access_token: Option<&str>,
|
||||||
|
) -> (Option<String>, Option<String>) {
|
||||||
|
let provider_type = provider_type.trim().to_ascii_lowercase();
|
||||||
|
let refresh_token = refresh_token
|
||||||
|
.map(str::trim)
|
||||||
|
.filter(|value| !value.is_empty())
|
||||||
|
.map(ToOwned::to_owned);
|
||||||
|
let access_token = access_token
|
||||||
|
.map(str::trim)
|
||||||
|
.filter(|value| !value.is_empty())
|
||||||
|
.map(ToOwned::to_owned);
|
||||||
|
|
||||||
|
if provider_type == "grok" {
|
||||||
|
return (None, access_token.or(refresh_token));
|
||||||
|
}
|
||||||
|
|
||||||
|
normalize_single_import_tokens(refresh_token.as_deref(), access_token.as_deref())
|
||||||
|
}
|
||||||
|
|
||||||
pub(super) fn import_tokens_from_raw_token(token: &str) -> (Option<String>, Option<String>) {
|
pub(super) fn import_tokens_from_raw_token(token: &str) -> (Option<String>, Option<String>) {
|
||||||
if looks_like_access_token(token) {
|
if looks_like_access_token(token) {
|
||||||
(None, Some(token.trim().to_string()))
|
(None, Some(token.trim().to_string()))
|
||||||
@@ -98,7 +120,7 @@ pub(super) fn decode_access_token_expires_at(access_token: &str) -> Option<u64>
|
|||||||
pub(super) fn provider_type_supports_access_token_import(provider_type: &str) -> bool {
|
pub(super) fn provider_type_supports_access_token_import(provider_type: &str) -> bool {
|
||||||
matches!(
|
matches!(
|
||||||
provider_type.trim().to_ascii_lowercase().as_str(),
|
provider_type.trim().to_ascii_lowercase().as_str(),
|
||||||
"codex" | "chatgpt_web"
|
"codex" | "chatgpt_web" | "grok"
|
||||||
)
|
)
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -123,6 +145,11 @@ pub(super) fn build_provider_access_token_import_auth_config(
|
|||||||
auth_config.insert("refresh_token".to_string(), json!(refresh_token));
|
auth_config.insert("refresh_token".to_string(), json!(refresh_token));
|
||||||
}
|
}
|
||||||
|
|
||||||
|
if provider_type.trim().eq_ignore_ascii_case("grok") {
|
||||||
|
auth_config.insert("sso_token".to_string(), json!(access_token));
|
||||||
|
auth_config.insert("auth_method".to_string(), json!("sso_token"));
|
||||||
|
}
|
||||||
|
|
||||||
auth_config.insert(
|
auth_config.insert(
|
||||||
"access_token_import_temporary".to_string(),
|
"access_token_import_temporary".to_string(),
|
||||||
json!(refresh_token.is_none()),
|
json!(refresh_token.is_none()),
|
||||||
@@ -149,7 +176,7 @@ pub(super) fn build_provider_access_token_import_auth_config(
|
|||||||
mod tests {
|
mod tests {
|
||||||
use super::{
|
use super::{
|
||||||
build_provider_access_token_import_auth_config, decode_access_token_expires_at,
|
build_provider_access_token_import_auth_config, decode_access_token_expires_at,
|
||||||
looks_like_access_token, normalize_single_import_tokens,
|
looks_like_access_token, normalize_provider_import_tokens, normalize_single_import_tokens,
|
||||||
};
|
};
|
||||||
use base64::{engine::general_purpose::URL_SAFE_NO_PAD, Engine as _};
|
use base64::{engine::general_purpose::URL_SAFE_NO_PAD, Engine as _};
|
||||||
use serde_json::json;
|
use serde_json::json;
|
||||||
@@ -250,4 +277,34 @@ mod tests {
|
|||||||
Some(&json!(true))
|
Some(&json!(true))
|
||||||
);
|
);
|
||||||
}
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn normalize_grok_import_treats_opaque_session_as_access_token() {
|
||||||
|
let (refresh_token, access_token) =
|
||||||
|
normalize_provider_import_tokens("grok", Some("sso_session_token"), None);
|
||||||
|
assert!(refresh_token.is_none());
|
||||||
|
assert_eq!(access_token.as_deref(), Some("sso_session_token"));
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn builds_grok_auth_config_from_session_token() {
|
||||||
|
let (auth_config, expires_at) = build_provider_access_token_import_auth_config(
|
||||||
|
"grok",
|
||||||
|
"sso_session_token",
|
||||||
|
None,
|
||||||
|
Some(2_200_000_000),
|
||||||
|
None,
|
||||||
|
);
|
||||||
|
|
||||||
|
assert_eq!(expires_at, Some(2_200_000_000));
|
||||||
|
assert_eq!(
|
||||||
|
auth_config.get("sso_token"),
|
||||||
|
Some(&json!("sso_session_token"))
|
||||||
|
);
|
||||||
|
assert_eq!(auth_config.get("auth_method"), Some(&json!("sso_token")));
|
||||||
|
assert_eq!(
|
||||||
|
auth_config.get("expires_at"),
|
||||||
|
Some(&json!(2_200_000_000u64))
|
||||||
|
);
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -8,8 +8,10 @@ use crate::GatewayError;
|
|||||||
use aether_data_contracts::repository::provider_catalog::{
|
use aether_data_contracts::repository::provider_catalog::{
|
||||||
StoredProviderCatalogEndpoint, StoredProviderCatalogKey,
|
StoredProviderCatalogEndpoint, StoredProviderCatalogKey,
|
||||||
};
|
};
|
||||||
use aether_provider_transport::provider_types::provider_type_is_fixed;
|
use aether_provider_transport::{
|
||||||
use serde_json::json;
|
grok_browser_transport_fingerprint_from_auth_config, provider_types::provider_type_is_fixed,
|
||||||
|
};
|
||||||
|
use serde_json::{json, Map, Value};
|
||||||
use std::time::{SystemTime, UNIX_EPOCH};
|
use std::time::{SystemTime, UNIX_EPOCH};
|
||||||
use uuid::Uuid;
|
use uuid::Uuid;
|
||||||
|
|
||||||
@@ -92,6 +94,16 @@ pub(crate) fn build_provider_oauth_auth_config_from_token_payload(
|
|||||||
(auth_config, access_token, refresh_token, expires_at)
|
(auth_config, access_token, refresh_token, expires_at)
|
||||||
}
|
}
|
||||||
|
|
||||||
|
fn grok_oauth_catalog_key_fingerprint(
|
||||||
|
provider_type: &str,
|
||||||
|
auth_config: &Map<String, Value>,
|
||||||
|
) -> Option<Value> {
|
||||||
|
if !provider_type.trim().eq_ignore_ascii_case("grok") {
|
||||||
|
return None;
|
||||||
|
}
|
||||||
|
grok_browser_transport_fingerprint_from_auth_config(auth_config)
|
||||||
|
}
|
||||||
|
|
||||||
pub(crate) async fn create_provider_oauth_catalog_key(
|
pub(crate) async fn create_provider_oauth_catalog_key(
|
||||||
state: &AdminAppState<'_>,
|
state: &AdminAppState<'_>,
|
||||||
provider_id: &str,
|
provider_id: &str,
|
||||||
@@ -136,7 +148,7 @@ pub(crate) async fn create_provider_oauth_catalog_key(
|
|||||||
None,
|
None,
|
||||||
expires_at_unix_secs,
|
expires_at_unix_secs,
|
||||||
proxy,
|
proxy,
|
||||||
None,
|
grok_oauth_catalog_key_fingerprint(provider_type, auth_config),
|
||||||
)
|
)
|
||||||
.map_err(|err| GatewayError::Internal(err.to_string()))?;
|
.map_err(|err| GatewayError::Internal(err.to_string()))?;
|
||||||
record.internal_priority = 50;
|
record.internal_priority = 50;
|
||||||
@@ -193,6 +205,9 @@ pub(crate) async fn update_existing_provider_oauth_catalog_key(
|
|||||||
updated.expires_at_unix_secs = expires_at_unix_secs;
|
updated.expires_at_unix_secs = expires_at_unix_secs;
|
||||||
updated.oauth_invalid_at_unix_secs = None;
|
updated.oauth_invalid_at_unix_secs = None;
|
||||||
updated.oauth_invalid_reason = None;
|
updated.oauth_invalid_reason = None;
|
||||||
|
if updated.fingerprint.is_none() {
|
||||||
|
updated.fingerprint = grok_oauth_catalog_key_fingerprint(provider_type, auth_config);
|
||||||
|
}
|
||||||
updated.health_by_format = Some(json!({}));
|
updated.health_by_format = Some(json!({}));
|
||||||
updated.circuit_breaker_by_format = Some(json!({}));
|
updated.circuit_breaker_by_format = Some(json!({}));
|
||||||
updated.error_count = Some(0);
|
updated.error_count = Some(0);
|
||||||
@@ -223,7 +238,9 @@ fn provider_oauth_catalog_key_api_formats(
|
|||||||
|
|
||||||
#[cfg(test)]
|
#[cfg(test)]
|
||||||
mod tests {
|
mod tests {
|
||||||
use super::provider_oauth_token_payload_expires_at_unix_secs;
|
use super::{
|
||||||
|
grok_oauth_catalog_key_fingerprint, provider_oauth_token_payload_expires_at_unix_secs,
|
||||||
|
};
|
||||||
use base64::{engine::general_purpose::URL_SAFE_NO_PAD, Engine as _};
|
use base64::{engine::general_purpose::URL_SAFE_NO_PAD, Engine as _};
|
||||||
use serde_json::json;
|
use serde_json::json;
|
||||||
|
|
||||||
@@ -273,4 +290,60 @@ mod tests {
|
|||||||
Some(2_000_000_000)
|
Some(2_000_000_000)
|
||||||
);
|
);
|
||||||
}
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn grok_oauth_catalog_key_fingerprint_uses_browser_wreq_profile() {
|
||||||
|
let auth_config = json!({
|
||||||
|
"sso_token": "abc",
|
||||||
|
"browser_profile": "chrome-137",
|
||||||
|
});
|
||||||
|
let auth_config = auth_config.as_object().expect("object");
|
||||||
|
|
||||||
|
let fingerprint = grok_oauth_catalog_key_fingerprint("grok", auth_config)
|
||||||
|
.expect("fingerprint should resolve");
|
||||||
|
|
||||||
|
assert_eq!(
|
||||||
|
fingerprint["transport_profile"]["profile_id"],
|
||||||
|
json!("chrome137")
|
||||||
|
);
|
||||||
|
assert_eq!(
|
||||||
|
fingerprint["transport_profile"]["backend"],
|
||||||
|
json!("browser_wreq")
|
||||||
|
);
|
||||||
|
assert_eq!(
|
||||||
|
fingerprint["transport_profile"]["extra"]["browser_profile"],
|
||||||
|
json!("chrome137")
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn grok_oauth_catalog_key_fingerprint_infers_profile_from_user_agent() {
|
||||||
|
let auth_config = json!({
|
||||||
|
"sso_token": "abc",
|
||||||
|
"user_agent": "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/137.0.0.0 Safari/537.36",
|
||||||
|
});
|
||||||
|
let auth_config = auth_config.as_object().expect("object");
|
||||||
|
|
||||||
|
let fingerprint = grok_oauth_catalog_key_fingerprint("grok", auth_config)
|
||||||
|
.expect("fingerprint should resolve");
|
||||||
|
|
||||||
|
assert_eq!(
|
||||||
|
fingerprint["transport_profile"]["profile_id"],
|
||||||
|
json!("chrome137")
|
||||||
|
);
|
||||||
|
assert_eq!(
|
||||||
|
fingerprint["transport_profile"]["extra"]["browser_profile"],
|
||||||
|
json!("chrome137")
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn grok_oauth_catalog_key_fingerprint_ignores_non_grok_providers() {
|
||||||
|
let auth_config = json!({
|
||||||
|
"browser_profile": "chrome136",
|
||||||
|
});
|
||||||
|
let auth_config = auth_config.as_object().expect("object");
|
||||||
|
|
||||||
|
assert!(grok_oauth_catalog_key_fingerprint("openai", auth_config).is_none());
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -4,6 +4,7 @@ use std::pin::Pin;
|
|||||||
use super::antigravity::refresh_antigravity_provider_quota_locally;
|
use super::antigravity::refresh_antigravity_provider_quota_locally;
|
||||||
use super::chatgpt_web::refresh_chatgpt_web_provider_quota_locally;
|
use super::chatgpt_web::refresh_chatgpt_web_provider_quota_locally;
|
||||||
use super::codex::refresh_codex_provider_quota_locally;
|
use super::codex::refresh_codex_provider_quota_locally;
|
||||||
|
use super::grok::refresh_grok_provider_quota_locally;
|
||||||
use super::kiro::refresh_kiro_provider_quota_locally;
|
use super::kiro::refresh_kiro_provider_quota_locally;
|
||||||
use crate::handlers::admin::request::AdminAppState;
|
use crate::handlers::admin::request::AdminAppState;
|
||||||
use crate::GatewayError;
|
use crate::GatewayError;
|
||||||
@@ -33,6 +34,7 @@ const PROVIDER_QUOTA_REFRESH_HANDLERS: &[(&str, ProviderQuotaRefreshHandler)] =
|
|||||||
refresh_chatgpt_web_provider_quota_locally_boxed,
|
refresh_chatgpt_web_provider_quota_locally_boxed,
|
||||||
),
|
),
|
||||||
("codex", refresh_codex_provider_quota_locally_boxed),
|
("codex", refresh_codex_provider_quota_locally_boxed),
|
||||||
|
("grok", refresh_grok_provider_quota_locally_boxed),
|
||||||
("kiro", refresh_kiro_provider_quota_locally_boxed),
|
("kiro", refresh_kiro_provider_quota_locally_boxed),
|
||||||
];
|
];
|
||||||
|
|
||||||
@@ -117,3 +119,19 @@ fn refresh_kiro_provider_quota_locally_boxed<'a>(
|
|||||||
proxy_override,
|
proxy_override,
|
||||||
))
|
))
|
||||||
}
|
}
|
||||||
|
|
||||||
|
fn refresh_grok_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_grok_provider_quota_locally(
|
||||||
|
state,
|
||||||
|
provider,
|
||||||
|
endpoint,
|
||||||
|
keys,
|
||||||
|
proxy_override,
|
||||||
|
))
|
||||||
|
}
|
||||||
|
|||||||
@@ -0,0 +1,826 @@
|
|||||||
|
use super::shared::{
|
||||||
|
build_quota_snapshot_payload, 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::{
|
||||||
|
OAUTH_ACCOUNT_BLOCK_PREFIX, OAUTH_EXPIRED_PREFIX, OAUTH_REFRESH_FAILED_PREFIX,
|
||||||
|
};
|
||||||
|
use crate::handlers::admin::request::{AdminAppState, AdminGatewayProviderTransportSnapshot};
|
||||||
|
use crate::GatewayError;
|
||||||
|
use aether_contracts::{
|
||||||
|
ExecutionPlan, ExecutionResult, ProxySnapshot, RequestBody, ResolvedTransportProfile,
|
||||||
|
};
|
||||||
|
use aether_data_contracts::repository::provider_catalog::{
|
||||||
|
StoredProviderCatalogEndpoint, StoredProviderCatalogKey, StoredProviderCatalogProvider,
|
||||||
|
};
|
||||||
|
use aether_provider_pool::{
|
||||||
|
grok_pool_tier_from_quota_bucket, grok_supported_quota_windows_for_tier,
|
||||||
|
};
|
||||||
|
use aether_provider_transport::grok_browser_profile_metadata_from_resolved_transport_profile;
|
||||||
|
use base64::Engine as _;
|
||||||
|
use serde_json::json;
|
||||||
|
use std::collections::BTreeMap;
|
||||||
|
use std::time::{SystemTime, UNIX_EPOCH};
|
||||||
|
use uuid::Uuid;
|
||||||
|
|
||||||
|
const GROK_DEFAULT_BASE_URL: &str = "https://grok.com";
|
||||||
|
const GROK_RATE_LIMITS_PATH: &str = "/rest/rate-limits";
|
||||||
|
const GROK_STATSIG_ID: &str = "ZTpUeXBlRXJyb3I6IENhbm5vdCByZWFkIHByb3BlcnRpZXMgb2YgdW5kZWZpbmVkIChyZWFkaW5nICdjaGlsZE5vZGVzJyk=";
|
||||||
|
|
||||||
|
fn grok_base_url(endpoint: &StoredProviderCatalogEndpoint) -> String {
|
||||||
|
let base_url = endpoint.base_url.trim().trim_end_matches('/');
|
||||||
|
if base_url.is_empty() {
|
||||||
|
GROK_DEFAULT_BASE_URL.to_string()
|
||||||
|
} else {
|
||||||
|
base_url.to_string()
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fn grok_auth_config(
|
||||||
|
transport: &AdminGatewayProviderTransportSnapshot,
|
||||||
|
) -> Option<serde_json::Value> {
|
||||||
|
transport
|
||||||
|
.key
|
||||||
|
.decrypted_auth_config
|
||||||
|
.as_deref()
|
||||||
|
.map(str::trim)
|
||||||
|
.filter(|value| !value.is_empty())
|
||||||
|
.and_then(|value| serde_json::from_str::<serde_json::Value>(value).ok())
|
||||||
|
}
|
||||||
|
|
||||||
|
fn grok_auth_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 build_grok_quota_headers(
|
||||||
|
auth_config: Option<&serde_json::Value>,
|
||||||
|
transport_profile: Option<&ResolvedTransportProfile>,
|
||||||
|
base_url: &str,
|
||||||
|
) -> Option<BTreeMap<String, String>> {
|
||||||
|
let cookie = build_grok_quota_cookie(auth_config).unwrap_or_default();
|
||||||
|
let browser_profile =
|
||||||
|
grok_browser_profile_metadata_from_resolved_transport_profile(transport_profile?)?;
|
||||||
|
Some(BTreeMap::from([
|
||||||
|
("accept".to_string(), "*/*".to_string()),
|
||||||
|
(
|
||||||
|
"accept-language".to_string(),
|
||||||
|
"zh-CN,zh;q=0.9,en;q=0.8".to_string(),
|
||||||
|
),
|
||||||
|
(
|
||||||
|
"baggage".to_string(),
|
||||||
|
"sentry-environment=production,sentry-release=d6add6fb0460641fd482d767a335ef72b9b6abb8,sentry-public_key=b311e0f2690c81f25e2c4cf6d4f7ce1c".to_string(),
|
||||||
|
),
|
||||||
|
("content-type".to_string(), "application/json".to_string()),
|
||||||
|
("origin".to_string(), base_url.to_string()),
|
||||||
|
("priority".to_string(), "u=1, i".to_string()),
|
||||||
|
("referer".to_string(), format!("{base_url}/")),
|
||||||
|
("sec-ch-ua".to_string(), browser_profile.sec_ch_ua),
|
||||||
|
("sec-ch-ua-mobile".to_string(), "?0".to_string()),
|
||||||
|
("sec-ch-ua-model".to_string(), String::new()),
|
||||||
|
(
|
||||||
|
"sec-ch-ua-platform".to_string(),
|
||||||
|
browser_profile.sec_ch_ua_platform,
|
||||||
|
),
|
||||||
|
("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()),
|
||||||
|
("user-agent".to_string(), browser_profile.user_agent),
|
||||||
|
("cookie".to_string(), cookie),
|
||||||
|
("x-statsig-id".to_string(), GROK_STATSIG_ID.to_string()),
|
||||||
|
("x-xai-request-id".to_string(), Uuid::new_v4().to_string()),
|
||||||
|
]))
|
||||||
|
}
|
||||||
|
|
||||||
|
fn build_grok_quota_cookie(auth_config: Option<&serde_json::Value>) -> Option<String> {
|
||||||
|
let token = grok_auth_string(auth_config, &["sso_token", "access_token", "token"])?;
|
||||||
|
let token = strip_cookie_prefix(token.trim(), "sso=");
|
||||||
|
if token.is_empty() {
|
||||||
|
return None;
|
||||||
|
}
|
||||||
|
let sso_rw = grok_auth_string(auth_config, &["sso_rw_token", "ssoRwToken"])
|
||||||
|
.map(|value| strip_cookie_prefix(value.trim(), "sso-rw="))
|
||||||
|
.filter(|value| !value.is_empty())
|
||||||
|
.unwrap_or_else(|| token.clone());
|
||||||
|
|
||||||
|
let mut parts = vec![format!("sso={token}"), format!("sso-rw={sso_rw}")];
|
||||||
|
if let Some(extra_cookies) =
|
||||||
|
grok_auth_string(auth_config, &["cf_cookies", "cfCookies", "cookie"])
|
||||||
|
.and_then(|value| normalize_grok_extra_cookies(value.as_str()))
|
||||||
|
{
|
||||||
|
parts.push(extra_cookies);
|
||||||
|
}
|
||||||
|
let cf_clearance = grok_auth_string(auth_config, &["cf_clearance", "cfClearance"])
|
||||||
|
.map(|value| strip_cookie_prefix(value.trim(), "cf_clearance="))
|
||||||
|
.filter(|value| !value.is_empty());
|
||||||
|
if let Some(cf_clearance) = cf_clearance {
|
||||||
|
if !parts.iter().any(|part| part.contains("cf_clearance=")) {
|
||||||
|
parts.push(format!("cf_clearance={cf_clearance}"));
|
||||||
|
}
|
||||||
|
}
|
||||||
|
Some(parts.join("; "))
|
||||||
|
}
|
||||||
|
|
||||||
|
fn strip_cookie_prefix(value: &str, prefix: &str) -> String {
|
||||||
|
value
|
||||||
|
.strip_prefix(prefix)
|
||||||
|
.map(str::trim)
|
||||||
|
.unwrap_or(value)
|
||||||
|
.to_string()
|
||||||
|
}
|
||||||
|
|
||||||
|
fn normalize_grok_extra_cookies(value: &str) -> Option<String> {
|
||||||
|
let parts = value
|
||||||
|
.trim()
|
||||||
|
.trim_matches(';')
|
||||||
|
.split(';')
|
||||||
|
.filter_map(|segment| {
|
||||||
|
let (name, value) = segment.trim().split_once('=')?;
|
||||||
|
let name = name.trim();
|
||||||
|
let value = value.trim();
|
||||||
|
if name.is_empty()
|
||||||
|
|| value.is_empty()
|
||||||
|
|| name.eq_ignore_ascii_case("sso")
|
||||||
|
|| name.eq_ignore_ascii_case("sso-rw")
|
||||||
|
{
|
||||||
|
return None;
|
||||||
|
}
|
||||||
|
Some(format!("{name}={value}"))
|
||||||
|
})
|
||||||
|
.collect::<Vec<_>>();
|
||||||
|
(!parts.is_empty()).then(|| parts.join("; "))
|
||||||
|
}
|
||||||
|
|
||||||
|
#[derive(Debug, Clone, Copy, PartialEq)]
|
||||||
|
struct GrokRateLimitSnapshot {
|
||||||
|
remaining: f64,
|
||||||
|
total: f64,
|
||||||
|
window_seconds: u64,
|
||||||
|
wait_time_seconds: Option<u64>,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl GrokRateLimitSnapshot {
|
||||||
|
fn reset_after_seconds(self) -> u64 {
|
||||||
|
self.wait_time_seconds.unwrap_or(self.window_seconds)
|
||||||
|
}
|
||||||
|
|
||||||
|
fn reset_at_source(self) -> &'static str {
|
||||||
|
if self.wait_time_seconds.is_some() {
|
||||||
|
"grok_rate_limits_wait_time"
|
||||||
|
} else {
|
||||||
|
"grok_rate_limits_window"
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fn parse_grok_rate_limits(body: &serde_json::Value) -> Option<GrokRateLimitSnapshot> {
|
||||||
|
let remaining = body
|
||||||
|
.get("remainingQueries")
|
||||||
|
.and_then(serde_json::Value::as_f64)?;
|
||||||
|
let total = body
|
||||||
|
.get("totalQueries")
|
||||||
|
.and_then(serde_json::Value::as_f64)
|
||||||
|
.unwrap_or(remaining.max(0.0));
|
||||||
|
let window_seconds = body
|
||||||
|
.get("windowSizeSeconds")
|
||||||
|
.and_then(serde_json::Value::as_u64)
|
||||||
|
.unwrap_or(72_000);
|
||||||
|
let wait_time_seconds = body
|
||||||
|
.get("waitTimeSeconds")
|
||||||
|
.and_then(serde_json::Value::as_u64);
|
||||||
|
Some(GrokRateLimitSnapshot {
|
||||||
|
remaining,
|
||||||
|
total,
|
||||||
|
window_seconds,
|
||||||
|
wait_time_seconds,
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
fn grok_pool_tier_hint_for_refresh(
|
||||||
|
key: &StoredProviderCatalogKey,
|
||||||
|
auth_config: Option<&serde_json::Value>,
|
||||||
|
) -> Option<&'static str> {
|
||||||
|
key.status_snapshot
|
||||||
|
.as_ref()
|
||||||
|
.and_then(serde_json::Value::as_object)
|
||||||
|
.and_then(|snapshot| snapshot.get("quota"))
|
||||||
|
.and_then(serde_json::Value::as_object)
|
||||||
|
.and_then(grok_pool_tier_from_quota_bucket)
|
||||||
|
.or_else(|| {
|
||||||
|
key.upstream_metadata
|
||||||
|
.as_ref()
|
||||||
|
.and_then(serde_json::Value::as_object)
|
||||||
|
.and_then(|metadata| metadata.get("grok"))
|
||||||
|
.and_then(serde_json::Value::as_object)
|
||||||
|
.and_then(grok_pool_tier_from_quota_bucket)
|
||||||
|
})
|
||||||
|
.or_else(|| {
|
||||||
|
auth_config
|
||||||
|
.and_then(serde_json::Value::as_object)
|
||||||
|
.and_then(grok_pool_tier_from_quota_bucket)
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn execute_grok_quota_plan(
|
||||||
|
state: &AdminAppState<'_>,
|
||||||
|
transport: &AdminGatewayProviderTransportSnapshot,
|
||||||
|
endpoint: &StoredProviderCatalogEndpoint,
|
||||||
|
body: serde_json::Value,
|
||||||
|
proxy_override: Option<&ProxySnapshot>,
|
||||||
|
) -> Result<ProviderQuotaExecutionOutcome, GatewayError> {
|
||||||
|
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(
|
||||||
|
proxy.as_ref(),
|
||||||
|
)));
|
||||||
|
let transport_profile = state.resolve_transport_profile(transport);
|
||||||
|
let base_url = grok_base_url(endpoint);
|
||||||
|
let headers = build_grok_quota_headers(
|
||||||
|
grok_auth_config(transport).as_ref(),
|
||||||
|
transport_profile.as_ref(),
|
||||||
|
&base_url,
|
||||||
|
)
|
||||||
|
.ok_or_else(|| {
|
||||||
|
GatewayError::Internal("unsupported Grok browser transport profile".to_string())
|
||||||
|
})?;
|
||||||
|
let plan = ExecutionPlan {
|
||||||
|
request_id: format!("grok-quota:{}", transport.key.id),
|
||||||
|
candidate_id: None,
|
||||||
|
provider_name: Some("grok".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,
|
||||||
|
GROK_RATE_LIMITS_PATH.trim_start_matches('/')
|
||||||
|
),
|
||||||
|
headers,
|
||||||
|
content_type: Some("application/json".to_string()),
|
||||||
|
content_encoding: None,
|
||||||
|
body: RequestBody::from_json(body),
|
||||||
|
stream: false,
|
||||||
|
client_api_format: "openai:responses".to_string(),
|
||||||
|
provider_api_format: "grok:rate_limits".to_string(),
|
||||||
|
model_name: Some("grok-quota".to_string()),
|
||||||
|
proxy,
|
||||||
|
transport_profile,
|
||||||
|
timeouts,
|
||||||
|
};
|
||||||
|
|
||||||
|
execute_provider_quota_plan(state, transport, plan, "grok").await
|
||||||
|
}
|
||||||
|
|
||||||
|
fn grok_quota_error_detail(result: &ExecutionResult) -> Option<String> {
|
||||||
|
extract_execution_error_message(result).or_else(|| {
|
||||||
|
let body = result.body.as_ref()?.body_bytes_b64.as_deref()?;
|
||||||
|
let decoded = base64::engine::general_purpose::STANDARD
|
||||||
|
.decode(body)
|
||||||
|
.ok()?;
|
||||||
|
let text = String::from_utf8_lossy(&decoded).trim().to_string();
|
||||||
|
(!text.is_empty()).then_some(text)
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
fn grok_is_cloudflare_challenge(message: &str) -> bool {
|
||||||
|
let lowered = message.to_ascii_lowercase();
|
||||||
|
lowered.contains("cloudflare")
|
||||||
|
|| lowered.contains("just a moment")
|
||||||
|
|| lowered.contains("__cf_chl")
|
||||||
|
|| lowered.contains("cf-ray")
|
||||||
|
}
|
||||||
|
|
||||||
|
fn grok_quota_invalid_reason(status_code: u16, upstream_message: Option<&str>) -> String {
|
||||||
|
let message = upstream_message.unwrap_or_default().trim();
|
||||||
|
if status_code == 403 && grok_is_cloudflare_challenge(message) {
|
||||||
|
return format!(
|
||||||
|
"{OAUTH_REFRESH_FAILED_PREFIX}Grok Cloudflare 验证失败,请重新从同一浏览器复制最新 Cookie 和 User-Agent,或配置可通过 Cloudflare 的代理运行时"
|
||||||
|
);
|
||||||
|
}
|
||||||
|
let detail = if message.is_empty() {
|
||||||
|
match status_code {
|
||||||
|
401 => "Grok Token 无效或已过期",
|
||||||
|
403 => "Grok 账户访问受限",
|
||||||
|
_ => "Grok 请求失败",
|
||||||
|
}
|
||||||
|
} else {
|
||||||
|
message
|
||||||
|
};
|
||||||
|
match status_code {
|
||||||
|
401 => format!("{OAUTH_EXPIRED_PREFIX}{detail}"),
|
||||||
|
403 => format!("{OAUTH_ACCOUNT_BLOCK_PREFIX}{detail}"),
|
||||||
|
_ => detail.to_string(),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fn grok_quota_result_message(reason: &str) -> String {
|
||||||
|
for prefix in [
|
||||||
|
OAUTH_REFRESH_FAILED_PREFIX,
|
||||||
|
OAUTH_EXPIRED_PREFIX,
|
||||||
|
OAUTH_ACCOUNT_BLOCK_PREFIX,
|
||||||
|
] {
|
||||||
|
if let Some(message) = reason.strip_prefix(prefix) {
|
||||||
|
return message.trim().to_string();
|
||||||
|
}
|
||||||
|
}
|
||||||
|
reason.trim().to_string()
|
||||||
|
}
|
||||||
|
|
||||||
|
pub(crate) async fn refresh_grok_provider_quota_locally(
|
||||||
|
state: &AdminAppState<'_>,
|
||||||
|
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;
|
||||||
|
let mut failed_count = 0usize;
|
||||||
|
|
||||||
|
for key in keys {
|
||||||
|
let transport = match state
|
||||||
|
.read_provider_transport_snapshot(&provider.id, &endpoint.id, &key.id)
|
||||||
|
.await?
|
||||||
|
{
|
||||||
|
Some(transport) => transport,
|
||||||
|
None => {
|
||||||
|
failed_count += 1;
|
||||||
|
results.push(json!({
|
||||||
|
"key_id": key.id,
|
||||||
|
"key_name": key.name,
|
||||||
|
"status": "error",
|
||||||
|
"message": "Provider transport snapshot unavailable",
|
||||||
|
}));
|
||||||
|
continue;
|
||||||
|
}
|
||||||
|
};
|
||||||
|
|
||||||
|
if grok_auth_config(&transport).is_none() {
|
||||||
|
failed_count += 1;
|
||||||
|
results.push(json!({
|
||||||
|
"key_id": key.id,
|
||||||
|
"key_name": key.name,
|
||||||
|
"status": "error",
|
||||||
|
"message": "缺少 Grok 账号会话信息,请先导入 Token",
|
||||||
|
}));
|
||||||
|
continue;
|
||||||
|
}
|
||||||
|
|
||||||
|
let auth_config = grok_auth_config(&transport);
|
||||||
|
let mut quota_by_model = serde_json::Map::new();
|
||||||
|
let mut refreshed = false;
|
||||||
|
let mut invalid_reason = None::<String>;
|
||||||
|
let mut invalid_at = key.oauth_invalid_at_unix_secs;
|
||||||
|
let mut last_status_code = None::<u16>;
|
||||||
|
let mut last_error_message = None::<String>;
|
||||||
|
let mut metadata_update = serde_json::Map::new();
|
||||||
|
let base_url = grok_base_url(endpoint);
|
||||||
|
|
||||||
|
let supported_windows = grok_supported_quota_windows_for_tier(
|
||||||
|
grok_pool_tier_hint_for_refresh(&key, auth_config.as_ref()),
|
||||||
|
);
|
||||||
|
for (quota_key, mode_name) in supported_windows.iter().copied() {
|
||||||
|
let result = match execute_grok_quota_plan(
|
||||||
|
state,
|
||||||
|
&transport,
|
||||||
|
endpoint,
|
||||||
|
json!({ "modelName": mode_name }),
|
||||||
|
proxy_override.as_ref(),
|
||||||
|
)
|
||||||
|
.await?
|
||||||
|
{
|
||||||
|
ProviderQuotaExecutionOutcome::Response(result) => result,
|
||||||
|
ProviderQuotaExecutionOutcome::Failure(detail) => {
|
||||||
|
last_error_message = Some(format!("rate-limits 请求执行失败: {detail}"));
|
||||||
|
continue;
|
||||||
|
}
|
||||||
|
};
|
||||||
|
last_status_code = Some(result.status_code);
|
||||||
|
|
||||||
|
if result.status_code == 200 {
|
||||||
|
if let Some(body_json) = result
|
||||||
|
.body
|
||||||
|
.as_ref()
|
||||||
|
.and_then(|body| body.json_body.as_ref())
|
||||||
|
{
|
||||||
|
if let Some(rate_limit) = parse_grok_rate_limits(body_json) {
|
||||||
|
refreshed = true;
|
||||||
|
let now_unix_secs = SystemTime::now()
|
||||||
|
.duration_since(UNIX_EPOCH)
|
||||||
|
.ok()
|
||||||
|
.map(|duration| duration.as_secs())
|
||||||
|
.unwrap_or(0);
|
||||||
|
let reset_after_seconds = rate_limit.reset_after_seconds();
|
||||||
|
let reset_at = now_unix_secs.saturating_add(reset_after_seconds);
|
||||||
|
quota_by_model.insert(
|
||||||
|
(*quota_key).to_string(),
|
||||||
|
json!({
|
||||||
|
"display_name": *mode_name,
|
||||||
|
"remaining_fraction": if rate_limit.total > 0.0 { Some((rate_limit.remaining / rate_limit.total).clamp(0.0, 1.0)) } else { None::<f64> },
|
||||||
|
"used_percent": if rate_limit.total > 0.0 { Some(((rate_limit.total - rate_limit.remaining).max(0.0) / rate_limit.total * 100.0).clamp(0.0, 100.0)) } else { None::<f64> },
|
||||||
|
"remaining": rate_limit.remaining,
|
||||||
|
"total": rate_limit.total,
|
||||||
|
"window_seconds": rate_limit.window_seconds,
|
||||||
|
"wait_time_seconds": rate_limit.wait_time_seconds,
|
||||||
|
"reset_after_seconds": reset_after_seconds,
|
||||||
|
"reset_at": reset_at,
|
||||||
|
"next_reset_at": reset_at,
|
||||||
|
"reset_at_source": rate_limit.reset_at_source(),
|
||||||
|
"is_exhausted": rate_limit.remaining <= 0.0,
|
||||||
|
}),
|
||||||
|
);
|
||||||
|
} else {
|
||||||
|
last_error_message = Some(
|
||||||
|
"Grok rate-limits 未返回 remainingQueries/totalQueries".to_string(),
|
||||||
|
);
|
||||||
|
}
|
||||||
|
} else {
|
||||||
|
last_error_message = Some("Grok rate-limits 未返回 JSON 数据".to_string());
|
||||||
|
}
|
||||||
|
} else if matches!(result.status_code, 401 | 403) {
|
||||||
|
let now_unix_secs = SystemTime::now()
|
||||||
|
.duration_since(UNIX_EPOCH)
|
||||||
|
.ok()
|
||||||
|
.map(|duration| duration.as_secs())
|
||||||
|
.unwrap_or(0);
|
||||||
|
invalid_at = Some(now_unix_secs);
|
||||||
|
let error_detail = grok_quota_error_detail(&result);
|
||||||
|
invalid_reason = Some(grok_quota_invalid_reason(
|
||||||
|
result.status_code,
|
||||||
|
error_detail.as_deref(),
|
||||||
|
));
|
||||||
|
last_error_message = invalid_reason.as_deref().map(grok_quota_result_message);
|
||||||
|
} else {
|
||||||
|
let error_detail =
|
||||||
|
grok_quota_error_detail(&result).unwrap_or_else(|| "Grok 请求失败".to_string());
|
||||||
|
last_error_message = Some(format!(
|
||||||
|
"Grok rate-limits 请求失败({}): {error_detail}",
|
||||||
|
result.status_code
|
||||||
|
));
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
if refreshed {
|
||||||
|
if let Some(pool_tier) = grok_pool_tier_from_quota_bucket("a_by_model)
|
||||||
|
.or_else(|| grok_pool_tier_hint_for_refresh(&key, auth_config.as_ref()))
|
||||||
|
{
|
||||||
|
let pool_tier_value = json!(pool_tier);
|
||||||
|
metadata_update.insert("pool_tier".to_string(), pool_tier_value.clone());
|
||||||
|
metadata_update
|
||||||
|
.entry("plan_type".to_string())
|
||||||
|
.or_insert(pool_tier_value);
|
||||||
|
}
|
||||||
|
metadata_update.insert(
|
||||||
|
"updated_at".to_string(),
|
||||||
|
json!(SystemTime::now()
|
||||||
|
.duration_since(UNIX_EPOCH)
|
||||||
|
.ok()
|
||||||
|
.map(|duration| duration.as_secs())
|
||||||
|
.unwrap_or(0)),
|
||||||
|
);
|
||||||
|
metadata_update.insert("base_url".to_string(), json!(base_url));
|
||||||
|
metadata_update.insert("quota_by_model".to_string(), json!(quota_by_model));
|
||||||
|
}
|
||||||
|
|
||||||
|
let metadata_update_value = if metadata_update.is_empty() {
|
||||||
|
None
|
||||||
|
} else {
|
||||||
|
Some(serde_json::Value::Object({
|
||||||
|
let mut map = serde_json::Map::new();
|
||||||
|
map.insert(
|
||||||
|
"grok".to_string(),
|
||||||
|
serde_json::Value::Object(metadata_update.clone()),
|
||||||
|
);
|
||||||
|
map
|
||||||
|
}))
|
||||||
|
};
|
||||||
|
|
||||||
|
if !persist_provider_quota_refresh_state(
|
||||||
|
state,
|
||||||
|
&key.id,
|
||||||
|
metadata_update_value.as_ref(),
|
||||||
|
invalid_at,
|
||||||
|
invalid_reason,
|
||||||
|
None,
|
||||||
|
)
|
||||||
|
.await?
|
||||||
|
{
|
||||||
|
failed_count += 1;
|
||||||
|
results.push(json!({
|
||||||
|
"key_id": key.id,
|
||||||
|
"key_name": key.name,
|
||||||
|
"status": "error",
|
||||||
|
"message": "Key 状态写入失败",
|
||||||
|
}));
|
||||||
|
continue;
|
||||||
|
}
|
||||||
|
|
||||||
|
if refreshed {
|
||||||
|
success_count += 1;
|
||||||
|
} else {
|
||||||
|
failed_count += 1;
|
||||||
|
}
|
||||||
|
|
||||||
|
let mut payload = serde_json::Map::new();
|
||||||
|
payload.insert("key_id".to_string(), json!(key.id));
|
||||||
|
payload.insert("key_name".to_string(), json!(key.name));
|
||||||
|
payload.insert(
|
||||||
|
"status".to_string(),
|
||||||
|
json!(if refreshed { "success" } else { "error" }),
|
||||||
|
);
|
||||||
|
if let Some(metadata) = metadata_update.get("quota_by_model").cloned() {
|
||||||
|
payload.insert("metadata".to_string(), metadata);
|
||||||
|
}
|
||||||
|
if let Some(quota_snapshot) = build_quota_snapshot_payload(
|
||||||
|
"grok",
|
||||||
|
key.status_snapshot.as_ref(),
|
||||||
|
metadata_update_value.as_ref(),
|
||||||
|
) {
|
||||||
|
payload.insert("quota_snapshot".to_string(), quota_snapshot);
|
||||||
|
}
|
||||||
|
if !refreshed {
|
||||||
|
payload.insert(
|
||||||
|
"message".to_string(),
|
||||||
|
json!(last_error_message.unwrap_or_else(|| {
|
||||||
|
"Grok rate-limits 未返回可用配额数据".to_string()
|
||||||
|
})),
|
||||||
|
);
|
||||||
|
if let Some(status_code) = last_status_code {
|
||||||
|
payload.insert("status_code".to_string(), json!(status_code));
|
||||||
|
}
|
||||||
|
}
|
||||||
|
results.push(serde_json::Value::Object(payload));
|
||||||
|
}
|
||||||
|
|
||||||
|
Ok(Some(json!({
|
||||||
|
"success": success_count,
|
||||||
|
"failed": failed_count,
|
||||||
|
"total": success_count + failed_count,
|
||||||
|
"results": results,
|
||||||
|
"message": format!("已处理 {} 个 Key", success_count + failed_count),
|
||||||
|
"auto_removed": 0,
|
||||||
|
})))
|
||||||
|
}
|
||||||
|
|
||||||
|
#[cfg(test)]
|
||||||
|
mod tests {
|
||||||
|
use super::{
|
||||||
|
build_grok_quota_cookie, build_grok_quota_headers, grok_pool_tier_hint_for_refresh,
|
||||||
|
grok_quota_error_detail, grok_quota_invalid_reason, grok_quota_result_message,
|
||||||
|
parse_grok_rate_limits,
|
||||||
|
};
|
||||||
|
use crate::handlers::admin::provider::shared::payloads::OAUTH_REFRESH_FAILED_PREFIX;
|
||||||
|
use aether_contracts::{ExecutionResult, ResponseBody};
|
||||||
|
use aether_data_contracts::repository::provider_catalog::StoredProviderCatalogKey;
|
||||||
|
use base64::Engine as _;
|
||||||
|
use serde_json::json;
|
||||||
|
use std::collections::BTreeMap;
|
||||||
|
|
||||||
|
fn sample_key(
|
||||||
|
status_snapshot: Option<serde_json::Value>,
|
||||||
|
upstream_metadata: Option<serde_json::Value>,
|
||||||
|
) -> StoredProviderCatalogKey {
|
||||||
|
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.status_snapshot = status_snapshot;
|
||||||
|
key.upstream_metadata = upstream_metadata;
|
||||||
|
key
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn quota_cookie_preserves_grok_session_and_clearance() {
|
||||||
|
let auth_config = json!({
|
||||||
|
"sso_token": "sso=abc",
|
||||||
|
"sso_rw_token": "sso-rw=rw",
|
||||||
|
"cf_clearance": "cf"
|
||||||
|
});
|
||||||
|
|
||||||
|
let cookie = build_grok_quota_cookie(Some(&auth_config)).expect("cookie should build");
|
||||||
|
|
||||||
|
assert_eq!(cookie, "sso=abc; sso-rw=rw; cf_clearance=cf");
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn quota_cookie_removes_duplicate_session_cookies_from_cf_profile() {
|
||||||
|
let auth_config = json!({
|
||||||
|
"sso_token": "abc",
|
||||||
|
"sso_rw_token": "rw",
|
||||||
|
"cf_cookies": "i18nextLng=zh; sso=ignored; sso-rw=ignored-rw; cf_clearance=cf"
|
||||||
|
});
|
||||||
|
|
||||||
|
let cookie = build_grok_quota_cookie(Some(&auth_config)).expect("cookie should build");
|
||||||
|
|
||||||
|
assert_eq!(cookie, "sso=abc; sso-rw=rw; i18nextLng=zh; cf_clearance=cf");
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn quota_headers_use_resolved_transport_profile_user_agent() {
|
||||||
|
let auth_config = json!({
|
||||||
|
"sso_token": "abc",
|
||||||
|
"user_agent": "Mozilla/5.0 custom"
|
||||||
|
});
|
||||||
|
let transport_profile = aether_provider_transport::grok_browser_resolved_transport_profile(
|
||||||
|
Some("chrome137"),
|
||||||
|
"test",
|
||||||
|
)
|
||||||
|
.expect("profile should resolve");
|
||||||
|
|
||||||
|
let headers = build_grok_quota_headers(
|
||||||
|
Some(&auth_config),
|
||||||
|
Some(&transport_profile),
|
||||||
|
"https://grok.com",
|
||||||
|
)
|
||||||
|
.expect("headers should build");
|
||||||
|
|
||||||
|
assert!(headers
|
||||||
|
.get("user-agent")
|
||||||
|
.is_some_and(|value| value.contains("Chrome/137.0.0.0")));
|
||||||
|
assert_eq!(
|
||||||
|
headers.get("sec-ch-ua"),
|
||||||
|
Some(
|
||||||
|
&r#""Google Chrome";v="137", "Chromium";v="137", "Not(A:Brand";v="24""#.to_string()
|
||||||
|
)
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn quota_headers_default_to_chrome136_clearance_profile() {
|
||||||
|
let auth_config = json!({
|
||||||
|
"sso_token": "abc"
|
||||||
|
});
|
||||||
|
|
||||||
|
let transport_profile =
|
||||||
|
aether_provider_transport::grok_browser_resolved_transport_profile(None, "test")
|
||||||
|
.expect("profile should resolve");
|
||||||
|
let headers = build_grok_quota_headers(
|
||||||
|
Some(&auth_config),
|
||||||
|
Some(&transport_profile),
|
||||||
|
"https://grok.com",
|
||||||
|
)
|
||||||
|
.expect("headers should build");
|
||||||
|
|
||||||
|
assert!(headers
|
||||||
|
.get("user-agent")
|
||||||
|
.is_some_and(|value| value.contains("Chrome/136.0.0.0")));
|
||||||
|
assert_eq!(
|
||||||
|
headers.get("sec-ch-ua"),
|
||||||
|
Some(
|
||||||
|
&r#""Google Chrome";v="136", "Chromium";v="136", "Not(A:Brand";v="24""#.to_string()
|
||||||
|
)
|
||||||
|
);
|
||||||
|
assert_eq!(
|
||||||
|
headers.get("sec-ch-ua-platform"),
|
||||||
|
Some(&r#""macOS""#.to_string())
|
||||||
|
);
|
||||||
|
assert!(headers.contains_key("x-statsig-id"));
|
||||||
|
assert!(headers.contains_key("x-xai-request-id"));
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn quota_headers_do_not_mark_rate_limits_as_grok_app_chat_runtime() {
|
||||||
|
let auth_config = json!({
|
||||||
|
"sso_token": "abc"
|
||||||
|
});
|
||||||
|
|
||||||
|
let transport_profile =
|
||||||
|
aether_provider_transport::grok_browser_resolved_transport_profile(None, "test")
|
||||||
|
.expect("profile should resolve");
|
||||||
|
let headers = build_grok_quota_headers(
|
||||||
|
Some(&auth_config),
|
||||||
|
Some(&transport_profile),
|
||||||
|
"https://grok.com",
|
||||||
|
)
|
||||||
|
.expect("headers should build");
|
||||||
|
|
||||||
|
assert!(!headers.contains_key(aether_provider_transport::GROK_INTERNAL_HEADER));
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn parses_grok_wait_time_seconds_as_authoritative_reset_delay() {
|
||||||
|
let body = json!({
|
||||||
|
"windowSizeSeconds": 86_400,
|
||||||
|
"remainingQueries": 0,
|
||||||
|
"waitTimeSeconds": 12_648,
|
||||||
|
"totalQueries": 30,
|
||||||
|
"lowEffortRateLimits": null,
|
||||||
|
"highEffortRateLimits": null
|
||||||
|
});
|
||||||
|
|
||||||
|
let rate_limits = parse_grok_rate_limits(&body).expect("rate limits should parse");
|
||||||
|
|
||||||
|
assert_eq!(rate_limits.remaining, 0.0);
|
||||||
|
assert_eq!(rate_limits.total, 30.0);
|
||||||
|
assert_eq!(rate_limits.window_seconds, 86_400);
|
||||||
|
assert_eq!(rate_limits.wait_time_seconds, Some(12_648));
|
||||||
|
assert_eq!(rate_limits.reset_after_seconds(), 12_648);
|
||||||
|
assert_eq!(rate_limits.reset_at_source(), "grok_rate_limits_wait_time");
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn parses_grok_rate_limits_falls_back_to_window_when_wait_time_is_absent() {
|
||||||
|
let body = json!({
|
||||||
|
"windowSizeSeconds": 86_400,
|
||||||
|
"remainingQueries": 12,
|
||||||
|
"totalQueries": 30
|
||||||
|
});
|
||||||
|
|
||||||
|
let rate_limits = parse_grok_rate_limits(&body).expect("rate limits should parse");
|
||||||
|
|
||||||
|
assert_eq!(rate_limits.remaining, 12.0);
|
||||||
|
assert_eq!(rate_limits.total, 30.0);
|
||||||
|
assert_eq!(rate_limits.window_seconds, 86_400);
|
||||||
|
assert_eq!(rate_limits.wait_time_seconds, None);
|
||||||
|
assert_eq!(rate_limits.reset_after_seconds(), 86_400);
|
||||||
|
assert_eq!(rate_limits.reset_at_source(), "grok_rate_limits_window");
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn infers_grok_pool_tier_from_live_quota_totals() {
|
||||||
|
let key = sample_key(
|
||||||
|
Some(json!({
|
||||||
|
"quota": {
|
||||||
|
"pool_tier": "heavy"
|
||||||
|
}
|
||||||
|
})),
|
||||||
|
None,
|
||||||
|
);
|
||||||
|
|
||||||
|
assert_eq!(
|
||||||
|
grok_pool_tier_hint_for_refresh(&key, Some(&json!({}))),
|
||||||
|
Some("heavy")
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn infers_basic_grok_pool_tier_from_fast_quota_when_auto_is_absent() {
|
||||||
|
let key = sample_key(
|
||||||
|
None,
|
||||||
|
Some(json!({
|
||||||
|
"grok": {
|
||||||
|
"plan_type": "basic"
|
||||||
|
}
|
||||||
|
})),
|
||||||
|
);
|
||||||
|
|
||||||
|
assert_eq!(grok_pool_tier_hint_for_refresh(&key, None), Some("basic"));
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn cloudflare_challenge_403_is_not_account_block() {
|
||||||
|
let body = "<!DOCTYPE html><html><head><title>Just a moment...</title></head><body>Cloudflare</body></html>";
|
||||||
|
let result = ExecutionResult {
|
||||||
|
request_id: "grok-quota:test".to_string(),
|
||||||
|
candidate_id: None,
|
||||||
|
status_code: 403,
|
||||||
|
headers: BTreeMap::new(),
|
||||||
|
body: Some(ResponseBody {
|
||||||
|
json_body: None,
|
||||||
|
body_bytes_b64: Some(base64::engine::general_purpose::STANDARD.encode(body)),
|
||||||
|
}),
|
||||||
|
telemetry: None,
|
||||||
|
error: None,
|
||||||
|
};
|
||||||
|
|
||||||
|
let detail = grok_quota_error_detail(&result).expect("html body should be decoded");
|
||||||
|
let reason = grok_quota_invalid_reason(result.status_code, Some(&detail));
|
||||||
|
|
||||||
|
assert!(reason.starts_with("[REFRESH_FAILED] "));
|
||||||
|
assert!(!reason.starts_with("[ACCOUNT_BLOCK] "));
|
||||||
|
assert!(reason.contains("Cloudflare"));
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn quota_result_message_removes_status_prefix() {
|
||||||
|
let reason = format!("{OAUTH_REFRESH_FAILED_PREFIX}Grok Cloudflare 验证失败");
|
||||||
|
|
||||||
|
assert_eq!(
|
||||||
|
grok_quota_result_message(&reason),
|
||||||
|
"Grok Cloudflare 验证失败"
|
||||||
|
);
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -2,5 +2,6 @@ pub(crate) mod antigravity;
|
|||||||
pub(crate) mod chatgpt_web;
|
pub(crate) mod chatgpt_web;
|
||||||
pub(crate) mod codex;
|
pub(crate) mod codex;
|
||||||
pub(crate) mod dispatch;
|
pub(crate) mod dispatch;
|
||||||
|
pub(crate) mod grok;
|
||||||
pub(crate) mod kiro;
|
pub(crate) mod kiro;
|
||||||
pub(crate) mod shared;
|
pub(crate) mod shared;
|
||||||
|
|||||||
@@ -63,6 +63,13 @@ fn select_provider_oauth_runtime_endpoint(
|
|||||||
.trim()
|
.trim()
|
||||||
.eq_ignore_ascii_case("openai:image")
|
.eq_ignore_ascii_case("openai:image")
|
||||||
}),
|
}),
|
||||||
|
"grok" => matching_endpoint(endpoints, include_inactive, |endpoint| {
|
||||||
|
endpoint
|
||||||
|
.api_format
|
||||||
|
.trim()
|
||||||
|
.eq_ignore_ascii_case("openai:chat")
|
||||||
|
})
|
||||||
|
.or_else(|| matching_endpoint(endpoints, include_inactive, |_| true)),
|
||||||
"antigravity" => matching_endpoint(endpoints, include_inactive, |endpoint| {
|
"antigravity" => matching_endpoint(endpoints, include_inactive, |endpoint| {
|
||||||
endpoint
|
endpoint
|
||||||
.api_format
|
.api_format
|
||||||
|
|||||||
@@ -40,7 +40,7 @@ pub(super) async fn admin_provider_ops_sub2api_balance_payload(
|
|||||||
"query_balance",
|
"query_balance",
|
||||||
message,
|
message,
|
||||||
None,
|
None,
|
||||||
)
|
);
|
||||||
}
|
}
|
||||||
};
|
};
|
||||||
|
|
||||||
|
|||||||
@@ -57,7 +57,7 @@ pub(super) async fn handle_admin_provider_ops_action(
|
|||||||
Err(_) => {
|
Err(_) => {
|
||||||
return Ok(Some(bad_request_detail_response(
|
return Ok(Some(bad_request_detail_response(
|
||||||
"请求体必须是合法的 JSON 对象",
|
"请求体必须是合法的 JSON 对象",
|
||||||
)))
|
)));
|
||||||
}
|
}
|
||||||
};
|
};
|
||||||
let payload =
|
let payload =
|
||||||
@@ -67,7 +67,7 @@ pub(super) async fn handle_admin_provider_ops_action(
|
|||||||
Err(_) => {
|
Err(_) => {
|
||||||
return Ok(Some(bad_request_detail_response(
|
return Ok(Some(bad_request_detail_response(
|
||||||
"请求体必须是合法的 JSON 对象",
|
"请求体必须是合法的 JSON 对象",
|
||||||
)))
|
)));
|
||||||
}
|
}
|
||||||
};
|
};
|
||||||
payload.config
|
payload.config
|
||||||
|
|||||||
@@ -3,7 +3,9 @@ use crate::handlers::admin::provider::shared::support::{
|
|||||||
};
|
};
|
||||||
use crate::handlers::admin::request::AdminAppState;
|
use crate::handlers::admin::request::AdminAppState;
|
||||||
use crate::handlers::admin::shared::{provider_key_status_snapshot_payload, unix_secs_to_rfc3339};
|
use crate::handlers::admin::shared::{provider_key_status_snapshot_payload, unix_secs_to_rfc3339};
|
||||||
use crate::provider_key_auth::{provider_key_auth_semantics, provider_key_effective_api_formats};
|
use crate::provider_key_auth::{
|
||||||
|
provider_key_auth_semantics, provider_key_can_refresh_oauth, provider_key_effective_api_formats,
|
||||||
|
};
|
||||||
use aether_admin::provider::pool as admin_provider_pool_pure;
|
use aether_admin::provider::pool as admin_provider_pool_pure;
|
||||||
use aether_admin::provider::quota as admin_provider_quota_pure;
|
use aether_admin::provider::quota as admin_provider_quota_pure;
|
||||||
use aether_data_contracts::repository::pool_scores::StoredPoolMemberScore;
|
use aether_data_contracts::repository::pool_scores::StoredPoolMemberScore;
|
||||||
@@ -597,6 +599,185 @@ fn admin_pool_build_antigravity_account_quota_from_snapshot(
|
|||||||
))
|
))
|
||||||
}
|
}
|
||||||
|
|
||||||
|
fn admin_pool_grok_quota_window_label(
|
||||||
|
window: &serde_json::Map<String, serde_json::Value>,
|
||||||
|
) -> String {
|
||||||
|
let raw_code = window
|
||||||
|
.get("code")
|
||||||
|
.and_then(serde_json::Value::as_str)
|
||||||
|
.map(str::trim)
|
||||||
|
.unwrap_or_default()
|
||||||
|
.trim_start_matches("model:")
|
||||||
|
.to_ascii_lowercase();
|
||||||
|
let raw_label = window
|
||||||
|
.get("label")
|
||||||
|
.or_else(|| window.get("model"))
|
||||||
|
.and_then(serde_json::Value::as_str)
|
||||||
|
.map(str::trim)
|
||||||
|
.filter(|value| !value.is_empty())
|
||||||
|
.unwrap_or(raw_code.as_str())
|
||||||
|
.to_ascii_lowercase();
|
||||||
|
match raw_label.as_str() {
|
||||||
|
"quota_auto" | "auto" => "Auto".to_string(),
|
||||||
|
"quota_fast" | "fast" => "Fast".to_string(),
|
||||||
|
"quota_expert" | "expert" => "Expert".to_string(),
|
||||||
|
"quota_heavy" | "heavy" => "Heavy".to_string(),
|
||||||
|
"quota_grok_4_3" | "grok-420-computer-use-sa" => "Grok 4.3".to_string(),
|
||||||
|
_ => match raw_code.as_str() {
|
||||||
|
"quota_auto" | "auto" => "Auto".to_string(),
|
||||||
|
"quota_fast" | "fast" => "Fast".to_string(),
|
||||||
|
"quota_expert" | "expert" => "Expert".to_string(),
|
||||||
|
"quota_heavy" | "heavy" => "Heavy".to_string(),
|
||||||
|
"quota_grok_4_3" | "grok-420-computer-use-sa" => "Grok 4.3".to_string(),
|
||||||
|
_ => window
|
||||||
|
.get("label")
|
||||||
|
.or_else(|| window.get("model"))
|
||||||
|
.and_then(serde_json::Value::as_str)
|
||||||
|
.map(str::trim)
|
||||||
|
.filter(|value| !value.is_empty())
|
||||||
|
.unwrap_or("模式")
|
||||||
|
.to_string(),
|
||||||
|
},
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fn admin_pool_quota_window_remaining_percent(
|
||||||
|
window: &serde_json::Map<String, serde_json::Value>,
|
||||||
|
) -> Option<f64> {
|
||||||
|
admin_pool_json_to_f64(window.get("remaining_ratio"))
|
||||||
|
.map(|value| (value * 100.0).clamp(0.0, 100.0))
|
||||||
|
.or_else(|| {
|
||||||
|
admin_pool_json_to_f64(window.get("used_ratio"))
|
||||||
|
.map(|value| ((1.0 - value) * 100.0).clamp(0.0, 100.0))
|
||||||
|
})
|
||||||
|
.or_else(|| {
|
||||||
|
admin_pool_json_to_f64(window.get("remaining_value"))
|
||||||
|
.zip(admin_pool_json_to_f64(window.get("limit_value")))
|
||||||
|
.and_then(|(remaining, limit)| {
|
||||||
|
(limit > 0.0).then_some((remaining / limit * 100.0).clamp(0.0, 100.0))
|
||||||
|
})
|
||||||
|
})
|
||||||
|
.or_else(|| {
|
||||||
|
admin_pool_json_to_f64(window.get("used_value"))
|
||||||
|
.zip(admin_pool_json_to_f64(window.get("limit_value")))
|
||||||
|
.and_then(|(used, limit)| {
|
||||||
|
(limit > 0.0).then_some(((1.0 - used / limit) * 100.0).clamp(0.0, 100.0))
|
||||||
|
})
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
fn admin_pool_quota_window_value_text(
|
||||||
|
window: &serde_json::Map<String, serde_json::Value>,
|
||||||
|
) -> Option<String> {
|
||||||
|
let limit_value =
|
||||||
|
admin_pool_json_to_f64(window.get("limit_value")).filter(|value| *value > 0.0)?;
|
||||||
|
if let Some(remaining_value) = admin_pool_json_to_f64(window.get("remaining_value")) {
|
||||||
|
return Some(format!(
|
||||||
|
"{}/{}",
|
||||||
|
admin_pool_format_quota_value(remaining_value),
|
||||||
|
admin_pool_format_quota_value(limit_value),
|
||||||
|
));
|
||||||
|
}
|
||||||
|
admin_pool_json_to_f64(window.get("used_value")).map(|used_value| {
|
||||||
|
format!(
|
||||||
|
"{}/{}",
|
||||||
|
admin_pool_format_quota_value((limit_value - used_value).max(0.0)),
|
||||||
|
admin_pool_format_quota_value(limit_value),
|
||||||
|
)
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
fn admin_pool_build_grok_account_quota_from_snapshot(
|
||||||
|
quota_snapshot: &serde_json::Map<String, serde_json::Value>,
|
||||||
|
) -> Option<String> {
|
||||||
|
let code = quota_snapshot
|
||||||
|
.get("code")
|
||||||
|
.and_then(serde_json::Value::as_str)
|
||||||
|
.map(str::trim)
|
||||||
|
.unwrap_or_default();
|
||||||
|
if code.eq_ignore_ascii_case("banned") {
|
||||||
|
return quota_snapshot
|
||||||
|
.get("label")
|
||||||
|
.and_then(serde_json::Value::as_str)
|
||||||
|
.map(ToOwned::to_owned)
|
||||||
|
.or_else(|| Some("账号已封禁".to_string()));
|
||||||
|
}
|
||||||
|
if code.eq_ignore_ascii_case("forbidden") {
|
||||||
|
return quota_snapshot
|
||||||
|
.get("label")
|
||||||
|
.and_then(serde_json::Value::as_str)
|
||||||
|
.map(ToOwned::to_owned)
|
||||||
|
.or_else(|| Some("访问受限".to_string()));
|
||||||
|
}
|
||||||
|
|
||||||
|
let model_parts = admin_pool_quota_windows(quota_snapshot)
|
||||||
|
.into_iter()
|
||||||
|
.filter(|window| {
|
||||||
|
window
|
||||||
|
.get("scope")
|
||||||
|
.and_then(serde_json::Value::as_str)
|
||||||
|
.is_some_and(|scope| scope.eq_ignore_ascii_case("model"))
|
||||||
|
})
|
||||||
|
.filter_map(|window| {
|
||||||
|
let remaining_percent = admin_pool_quota_window_remaining_percent(window)?;
|
||||||
|
let mut part = format!(
|
||||||
|
"{}剩余 {}",
|
||||||
|
admin_pool_grok_quota_window_label(window),
|
||||||
|
admin_pool_format_percent(remaining_percent),
|
||||||
|
);
|
||||||
|
if let Some(value_text) = admin_pool_quota_window_value_text(window) {
|
||||||
|
part.push_str(&format!(" ({value_text})"));
|
||||||
|
}
|
||||||
|
Some(part)
|
||||||
|
})
|
||||||
|
.collect::<Vec<_>>();
|
||||||
|
|
||||||
|
if !model_parts.is_empty() {
|
||||||
|
return Some(model_parts.join(" | "));
|
||||||
|
}
|
||||||
|
|
||||||
|
let window = admin_pool_quota_window(quota_snapshot, "usage")
|
||||||
|
.or_else(|| admin_pool_quota_windows(quota_snapshot).into_iter().next())?;
|
||||||
|
let remaining_value = admin_pool_json_to_f64(window.get("remaining_value"));
|
||||||
|
let limit_value = admin_pool_json_to_f64(window.get("limit_value"));
|
||||||
|
if let (Some(remaining_value), Some(limit_value)) = (remaining_value, limit_value) {
|
||||||
|
if limit_value > 0.0 && remaining_value <= 0.0 {
|
||||||
|
return Some(format!(
|
||||||
|
"剩余 {}/{}",
|
||||||
|
admin_pool_format_quota_value(remaining_value),
|
||||||
|
admin_pool_format_quota_value(limit_value),
|
||||||
|
));
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
if let Some(remaining_percent) = admin_pool_quota_window_remaining_percent(window) {
|
||||||
|
if let Some(value_text) = admin_pool_quota_window_value_text(window) {
|
||||||
|
return Some(format!(
|
||||||
|
"剩余 {} ({value_text})",
|
||||||
|
admin_pool_format_percent(remaining_percent),
|
||||||
|
));
|
||||||
|
}
|
||||||
|
return Some(format!(
|
||||||
|
"剩余 {}",
|
||||||
|
admin_pool_format_percent(remaining_percent),
|
||||||
|
));
|
||||||
|
}
|
||||||
|
|
||||||
|
match (remaining_value, limit_value) {
|
||||||
|
(Some(remaining_value), Some(limit_value)) if limit_value > 0.0 => Some(format!(
|
||||||
|
"剩余 {}/{}",
|
||||||
|
admin_pool_format_quota_value(remaining_value),
|
||||||
|
admin_pool_format_quota_value(limit_value),
|
||||||
|
)),
|
||||||
|
_ => quota_snapshot
|
||||||
|
.get("label")
|
||||||
|
.and_then(serde_json::Value::as_str)
|
||||||
|
.map(str::trim)
|
||||||
|
.filter(|value| !value.is_empty())
|
||||||
|
.map(ToOwned::to_owned),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
fn admin_pool_build_gemini_cli_account_quota_from_snapshot(
|
fn admin_pool_build_gemini_cli_account_quota_from_snapshot(
|
||||||
quota_snapshot: &serde_json::Map<String, serde_json::Value>,
|
quota_snapshot: &serde_json::Map<String, serde_json::Value>,
|
||||||
) -> Option<String> {
|
) -> Option<String> {
|
||||||
@@ -699,6 +880,13 @@ fn admin_pool_build_account_quota(
|
|||||||
return Some(account_quota);
|
return Some(account_quota);
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
"grok" => {
|
||||||
|
if let Some(account_quota) =
|
||||||
|
admin_pool_build_grok_account_quota_from_snapshot(quota_snapshot)
|
||||||
|
{
|
||||||
|
return Some(account_quota);
|
||||||
|
}
|
||||||
|
}
|
||||||
"gemini_cli" => {
|
"gemini_cli" => {
|
||||||
if let Some(account_quota) =
|
if let Some(account_quota) =
|
||||||
admin_pool_build_gemini_cli_account_quota_from_snapshot(quota_snapshot)
|
admin_pool_build_gemini_cli_account_quota_from_snapshot(quota_snapshot)
|
||||||
@@ -962,7 +1150,10 @@ pub(super) fn build_admin_pool_key_payload(
|
|||||||
);
|
);
|
||||||
payload.insert(
|
payload.insert(
|
||||||
"can_refresh_oauth".to_string(),
|
"can_refresh_oauth".to_string(),
|
||||||
json!(auth_semantics.can_refresh_oauth()),
|
json!(provider_key_can_refresh_oauth(
|
||||||
|
auth_semantics,
|
||||||
|
auth_config.as_ref()
|
||||||
|
)),
|
||||||
);
|
);
|
||||||
payload.insert(
|
payload.insert(
|
||||||
"can_export_oauth".to_string(),
|
"can_export_oauth".to_string(),
|
||||||
@@ -1225,4 +1416,39 @@ mod tests {
|
|||||||
assert_eq!(usage["total_tokens"], json!(375));
|
assert_eq!(usage["total_tokens"], json!(375));
|
||||||
assert_eq!(usage["total_cost_usd"], json!("0.60000000"));
|
assert_eq!(usage["total_cost_usd"], json!("0.60000000"));
|
||||||
}
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn grok_model_quota_is_rendered_for_pool_rows() {
|
||||||
|
let quota_snapshot = json!({
|
||||||
|
"provider_type": "grok",
|
||||||
|
"code": "ok",
|
||||||
|
"exhausted": false,
|
||||||
|
"plan_type": "heavy",
|
||||||
|
"pool_tier": "heavy",
|
||||||
|
"windows": [
|
||||||
|
{
|
||||||
|
"code": "model:quota_auto",
|
||||||
|
"label": "auto",
|
||||||
|
"scope": "model",
|
||||||
|
"remaining_ratio": 0.4,
|
||||||
|
"used_value": 90,
|
||||||
|
"limit_value": 150
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"code": "model:quota_heavy",
|
||||||
|
"label": "heavy",
|
||||||
|
"scope": "model",
|
||||||
|
"remaining_ratio": 0.0,
|
||||||
|
"used_value": 20,
|
||||||
|
"limit_value": 20
|
||||||
|
}
|
||||||
|
]
|
||||||
|
});
|
||||||
|
let quota_snapshot = quota_snapshot.as_object().unwrap();
|
||||||
|
|
||||||
|
assert_eq!(
|
||||||
|
admin_pool_build_account_quota("grok", Some(quota_snapshot)),
|
||||||
|
Some("Auto剩余 40.0% (60/150) | Heavy剩余 0.0% (0/20)".to_string())
|
||||||
|
);
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -3,7 +3,7 @@ use super::{
|
|||||||
AdminPoolResolveSelectionRequest, ADMIN_POOL_PROVIDER_CATALOG_READER_UNAVAILABLE_DETAIL,
|
AdminPoolResolveSelectionRequest, ADMIN_POOL_PROVIDER_CATALOG_READER_UNAVAILABLE_DETAIL,
|
||||||
};
|
};
|
||||||
use crate::handlers::admin::request::{AdminAppState, AdminRequestContext};
|
use crate::handlers::admin::request::{AdminAppState, AdminRequestContext};
|
||||||
use crate::provider_key_auth::provider_key_auth_semantics;
|
use crate::provider_key_auth::{provider_key_auth_semantics, provider_key_can_refresh_oauth};
|
||||||
use crate::GatewayError;
|
use crate::GatewayError;
|
||||||
use aether_admin::provider::pool as admin_provider_pool_pure;
|
use aether_admin::provider::pool as admin_provider_pool_pure;
|
||||||
use axum::{
|
use axum::{
|
||||||
@@ -94,6 +94,7 @@ pub(super) async fn build_admin_pool_resolve_selection_response(
|
|||||||
.iter()
|
.iter()
|
||||||
.map(|key| {
|
.map(|key| {
|
||||||
let auth_semantics = provider_key_auth_semantics(key, &provider_type);
|
let auth_semantics = provider_key_auth_semantics(key, &provider_type);
|
||||||
|
let auth_config = state.parse_catalog_auth_config_json(key);
|
||||||
json!({
|
json!({
|
||||||
"key_id": key.id,
|
"key_id": key.id,
|
||||||
"key_name": key.name,
|
"key_name": key.name,
|
||||||
@@ -102,7 +103,7 @@ pub(super) async fn build_admin_pool_resolve_selection_response(
|
|||||||
"credential_kind": auth_semantics.credential_kind().as_str(),
|
"credential_kind": auth_semantics.credential_kind().as_str(),
|
||||||
"runtime_auth_kind": auth_semantics.runtime_auth_kind().as_str(),
|
"runtime_auth_kind": auth_semantics.runtime_auth_kind().as_str(),
|
||||||
"oauth_managed": auth_semantics.oauth_managed(),
|
"oauth_managed": auth_semantics.oauth_managed(),
|
||||||
"can_refresh_oauth": auth_semantics.can_refresh_oauth(),
|
"can_refresh_oauth": provider_key_can_refresh_oauth(auth_semantics, auth_config.as_ref()),
|
||||||
"can_export_oauth": auth_semantics.can_export_oauth(),
|
"can_export_oauth": auth_semantics.can_export_oauth(),
|
||||||
"can_edit_oauth": auth_semantics.can_edit_oauth(),
|
"can_edit_oauth": auth_semantics.can_edit_oauth(),
|
||||||
})
|
})
|
||||||
|
|||||||
@@ -17,6 +17,9 @@ use crate::ai_serving::{
|
|||||||
};
|
};
|
||||||
use crate::clock::current_unix_ms;
|
use crate::clock::current_unix_ms;
|
||||||
use crate::execution_runtime;
|
use crate::execution_runtime;
|
||||||
|
use crate::handlers::admin::provider::shared::model_test_capabilities::{
|
||||||
|
admin_provider_model_supports_image_generation, admin_provider_model_test_capabilities_payload,
|
||||||
|
};
|
||||||
use crate::handlers::admin::request::{AdminAppState, AdminGatewayProviderTransportSnapshot};
|
use crate::handlers::admin::request::{AdminAppState, AdminGatewayProviderTransportSnapshot};
|
||||||
use crate::handlers::shared::provider_pool::{
|
use crate::handlers::shared::provider_pool::{
|
||||||
admin_provider_pool_config_from_config_value, read_admin_provider_pool_runtime_state,
|
admin_provider_pool_config_from_config_value, read_admin_provider_pool_runtime_state,
|
||||||
@@ -100,6 +103,181 @@ struct ProviderQueryKeyFetchResult {
|
|||||||
has_success: bool,
|
has_success: bool,
|
||||||
}
|
}
|
||||||
|
|
||||||
|
fn provider_query_model_id(model: &Value) -> Option<&str> {
|
||||||
|
model
|
||||||
|
.get("id")
|
||||||
|
.and_then(Value::as_str)
|
||||||
|
.map(str::trim)
|
||||||
|
.filter(|value| !value.is_empty())
|
||||||
|
}
|
||||||
|
|
||||||
|
fn provider_query_grok_required_tier_rank(model_id: &str) -> Option<u8> {
|
||||||
|
match model_id.trim() {
|
||||||
|
"grok-4.20-0309-non-reasoning" | "grok-4.20-fast" | "grok-imagine-image-lite" => Some(0),
|
||||||
|
"grok-4.20-0309"
|
||||||
|
| "grok-4.20-0309-reasoning"
|
||||||
|
| "grok-4.20-0309-non-reasoning-super"
|
||||||
|
| "grok-4.20-0309-super"
|
||||||
|
| "grok-4.20-0309-reasoning-super"
|
||||||
|
| "grok-4.20-auto"
|
||||||
|
| "grok-4.20-expert"
|
||||||
|
| "grok-4.3-beta"
|
||||||
|
| "grok-imagine-image"
|
||||||
|
| "grok-imagine-image-pro"
|
||||||
|
| "grok-imagine-image-edit" => Some(1),
|
||||||
|
"grok-4.20-0309-non-reasoning-heavy"
|
||||||
|
| "grok-4.20-0309-heavy"
|
||||||
|
| "grok-4.20-0309-reasoning-heavy"
|
||||||
|
| "grok-4.20-multi-agent-0309"
|
||||||
|
| "grok-4.20-heavy" => Some(2),
|
||||||
|
_ => None,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fn provider_query_normalize_grok_pool_tier(value: Option<&str>) -> Option<&'static str> {
|
||||||
|
match value?.trim().to_ascii_lowercase().as_str() {
|
||||||
|
"basic" => Some("basic"),
|
||||||
|
"super" => Some("super"),
|
||||||
|
"heavy" => Some("heavy"),
|
||||||
|
_ => None,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fn provider_query_grok_pool_tier_rank(value: Option<&str>) -> u8 {
|
||||||
|
match provider_query_normalize_grok_pool_tier(value).unwrap_or("basic") {
|
||||||
|
"heavy" => 2,
|
||||||
|
"super" => 1,
|
||||||
|
_ => 0,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fn provider_query_grok_quota_string(quota: &Map<String, Value>, fields: &[&str]) -> Option<String> {
|
||||||
|
fields.iter().find_map(|field| {
|
||||||
|
quota
|
||||||
|
.get(*field)
|
||||||
|
.and_then(Value::as_str)
|
||||||
|
.map(str::trim)
|
||||||
|
.filter(|value| !value.is_empty())
|
||||||
|
.map(ToOwned::to_owned)
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
fn provider_query_grok_window_limit(quota: &Map<String, Value>, model_name: &str) -> Option<f64> {
|
||||||
|
quota
|
||||||
|
.get("windows")
|
||||||
|
.and_then(Value::as_array)?
|
||||||
|
.iter()
|
||||||
|
.filter_map(Value::as_object)
|
||||||
|
.find(|window| {
|
||||||
|
window
|
||||||
|
.get("model")
|
||||||
|
.and_then(Value::as_str)
|
||||||
|
.is_some_and(|value| value.trim() == model_name)
|
||||||
|
})
|
||||||
|
.and_then(|window| window.get("limit_value"))
|
||||||
|
.and_then(Value::as_f64)
|
||||||
|
.filter(|value| value.is_finite() && *value > 0.0)
|
||||||
|
}
|
||||||
|
|
||||||
|
fn provider_query_grok_pool_tier_from_quota(quota: &Map<String, Value>) -> Option<&'static str> {
|
||||||
|
if let Some(tier) =
|
||||||
|
provider_query_grok_quota_string(quota, &["pool_tier", "tier", "plan_type", "plan"])
|
||||||
|
.and_then(|value| provider_query_normalize_grok_pool_tier(Some(&value)))
|
||||||
|
{
|
||||||
|
return Some(tier);
|
||||||
|
}
|
||||||
|
|
||||||
|
if let Some(auto_total) = provider_query_grok_window_limit(quota, "quota_auto") {
|
||||||
|
if (auto_total - 150.0).abs() < f64::EPSILON {
|
||||||
|
return Some("heavy");
|
||||||
|
}
|
||||||
|
if (auto_total - 50.0).abs() < f64::EPSILON {
|
||||||
|
return Some("super");
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
if let Some(fast_total) = provider_query_grok_window_limit(quota, "quota_fast") {
|
||||||
|
if (fast_total - 400.0).abs() < f64::EPSILON {
|
||||||
|
return Some("heavy");
|
||||||
|
}
|
||||||
|
if (fast_total - 140.0).abs() < f64::EPSILON {
|
||||||
|
return Some("super");
|
||||||
|
}
|
||||||
|
if (fast_total - 30.0).abs() < f64::EPSILON {
|
||||||
|
return Some("basic");
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
None
|
||||||
|
}
|
||||||
|
|
||||||
|
fn provider_query_grok_key_pool_tier(key: &StoredProviderCatalogKey) -> Option<&'static str> {
|
||||||
|
key.status_snapshot
|
||||||
|
.as_ref()
|
||||||
|
.and_then(Value::as_object)
|
||||||
|
.and_then(|snapshot| snapshot.get("quota"))
|
||||||
|
.and_then(Value::as_object)
|
||||||
|
.and_then(provider_query_grok_pool_tier_from_quota)
|
||||||
|
}
|
||||||
|
|
||||||
|
fn provider_query_filter_models_for_key(
|
||||||
|
provider: &StoredProviderCatalogProvider,
|
||||||
|
key: &StoredProviderCatalogKey,
|
||||||
|
models: Vec<Value>,
|
||||||
|
) -> Vec<Value> {
|
||||||
|
if !provider.provider_type.trim().eq_ignore_ascii_case("grok") {
|
||||||
|
return models;
|
||||||
|
}
|
||||||
|
|
||||||
|
let allowed_rank = provider_query_grok_pool_tier_rank(provider_query_grok_key_pool_tier(key));
|
||||||
|
models
|
||||||
|
.into_iter()
|
||||||
|
.filter(|model| {
|
||||||
|
provider_query_model_id(model)
|
||||||
|
.and_then(provider_query_grok_required_tier_rank)
|
||||||
|
.is_some_and(|required_rank| required_rank <= allowed_rank)
|
||||||
|
})
|
||||||
|
.collect()
|
||||||
|
}
|
||||||
|
|
||||||
|
fn provider_query_attach_model_test_capabilities(
|
||||||
|
provider: &StoredProviderCatalogProvider,
|
||||||
|
models: Vec<Value>,
|
||||||
|
) -> Vec<Value> {
|
||||||
|
models
|
||||||
|
.into_iter()
|
||||||
|
.map(|mut model| {
|
||||||
|
let Some(object) = model.as_object_mut() else {
|
||||||
|
return model;
|
||||||
|
};
|
||||||
|
let model_id = object
|
||||||
|
.get("id")
|
||||||
|
.and_then(Value::as_str)
|
||||||
|
.map(str::trim)
|
||||||
|
.unwrap_or_default()
|
||||||
|
.to_string();
|
||||||
|
let supports_image_generation = admin_provider_model_supports_image_generation(
|
||||||
|
&provider.provider_type,
|
||||||
|
&model_id,
|
||||||
|
object
|
||||||
|
.get("supports_image_generation")
|
||||||
|
.or_else(|| object.get("effective_supports_image_generation"))
|
||||||
|
.and_then(Value::as_bool)
|
||||||
|
.unwrap_or(false),
|
||||||
|
);
|
||||||
|
object.insert(
|
||||||
|
"model_test_capabilities".to_string(),
|
||||||
|
admin_provider_model_test_capabilities_payload(
|
||||||
|
&provider.provider_type,
|
||||||
|
&model_id,
|
||||||
|
supports_image_generation,
|
||||||
|
),
|
||||||
|
);
|
||||||
|
model
|
||||||
|
})
|
||||||
|
.collect()
|
||||||
|
}
|
||||||
|
|
||||||
fn provider_query_codex_preset_fallback(
|
fn provider_query_codex_preset_fallback(
|
||||||
provider: &StoredProviderCatalogProvider,
|
provider: &StoredProviderCatalogProvider,
|
||||||
) -> Option<ProviderQueryKeyFetchResult> {
|
) -> Option<ProviderQueryKeyFetchResult> {
|
||||||
@@ -245,8 +423,9 @@ async fn provider_query_fetch_models_for_key(
|
|||||||
if let Some(cached_models) =
|
if let Some(cached_models) =
|
||||||
provider_query_read_cached_models(state, &provider.id, &key.id).await
|
provider_query_read_cached_models(state, &provider.id, &key.id).await
|
||||||
{
|
{
|
||||||
|
let models = provider_query_filter_models_for_key(provider, key, cached_models);
|
||||||
return Ok(ProviderQueryKeyFetchResult {
|
return Ok(ProviderQueryKeyFetchResult {
|
||||||
models: cached_models,
|
models,
|
||||||
error: None,
|
error: None,
|
||||||
from_cache: true,
|
from_cache: true,
|
||||||
has_success: true,
|
has_success: true,
|
||||||
@@ -257,8 +436,13 @@ async fn provider_query_fetch_models_for_key(
|
|||||||
let selected_endpoints = selected_models_fetch_endpoints(endpoints, key);
|
let selected_endpoints = selected_models_fetch_endpoints(endpoints, key);
|
||||||
if selected_endpoints.is_empty() {
|
if selected_endpoints.is_empty() {
|
||||||
if let Some(models) = preset_models_for_provider(&provider.provider_type) {
|
if let Some(models) = preset_models_for_provider(&provider.provider_type) {
|
||||||
|
let models = provider_query_filter_models_for_key(
|
||||||
|
provider,
|
||||||
|
key,
|
||||||
|
aggregate_models_for_cache(&models),
|
||||||
|
);
|
||||||
return Ok(ProviderQueryKeyFetchResult {
|
return Ok(ProviderQueryKeyFetchResult {
|
||||||
models: aggregate_models_for_cache(&models),
|
models,
|
||||||
error: None,
|
error: None,
|
||||||
from_cache: false,
|
from_cache: false,
|
||||||
has_success: true,
|
has_success: true,
|
||||||
@@ -342,7 +526,7 @@ async fn provider_query_fetch_models_for_key(
|
|||||||
}
|
}
|
||||||
|
|
||||||
Ok(ProviderQueryKeyFetchResult {
|
Ok(ProviderQueryKeyFetchResult {
|
||||||
models: unique_models,
|
models: provider_query_filter_models_for_key(provider, key, unique_models),
|
||||||
error,
|
error,
|
||||||
from_cache: false,
|
from_cache: false,
|
||||||
has_success: outcome.has_success,
|
has_success: outcome.has_success,
|
||||||
@@ -397,11 +581,12 @@ pub(crate) async fn build_admin_provider_query_models_response(
|
|||||||
force_refresh,
|
force_refresh,
|
||||||
)
|
)
|
||||||
.await?;
|
.await?;
|
||||||
let success = !result.models.is_empty();
|
let models = provider_query_attach_model_test_capabilities(&provider, result.models);
|
||||||
|
let success = !models.is_empty();
|
||||||
return Ok(Json(json!({
|
return Ok(Json(json!({
|
||||||
"success": success,
|
"success": success,
|
||||||
"data": {
|
"data": {
|
||||||
"models": result.models,
|
"models": models,
|
||||||
"error": result.error,
|
"error": result.error,
|
||||||
"from_cache": result.from_cache,
|
"from_cache": result.from_cache,
|
||||||
},
|
},
|
||||||
@@ -429,6 +614,7 @@ pub(crate) async fn build_admin_provider_query_models_response(
|
|||||||
{
|
{
|
||||||
if let Some(models) = provider_query_read_provider_cached_models(state, &provider.id).await
|
if let Some(models) = provider_query_read_provider_cached_models(state, &provider.id).await
|
||||||
{
|
{
|
||||||
|
let models = provider_query_attach_model_test_capabilities(&provider, models);
|
||||||
return Ok(Json(json!({
|
return Ok(Json(json!({
|
||||||
"success": !models.is_empty(),
|
"success": !models.is_empty(),
|
||||||
"data": {
|
"data": {
|
||||||
@@ -504,6 +690,7 @@ pub(crate) async fn build_admin_provider_query_models_response(
|
|||||||
if !success && error.is_none() {
|
if !success && error.is_none() {
|
||||||
error = Some(ADMIN_PROVIDER_QUERY_NO_MODELS_FROM_KEY_DETAIL.to_string());
|
error = Some(ADMIN_PROVIDER_QUERY_NO_MODELS_FROM_KEY_DETAIL.to_string());
|
||||||
}
|
}
|
||||||
|
let models = provider_query_attach_model_test_capabilities(&provider, models);
|
||||||
|
|
||||||
Ok(Json(json!({
|
Ok(Json(json!({
|
||||||
"success": success,
|
"success": success,
|
||||||
@@ -519,3 +706,152 @@ pub(crate) async fn build_admin_provider_query_models_response(
|
|||||||
}))
|
}))
|
||||||
.into_response())
|
.into_response())
|
||||||
}
|
}
|
||||||
|
|
||||||
|
#[cfg(test)]
|
||||||
|
mod tests {
|
||||||
|
use super::*;
|
||||||
|
use aether_data_contracts::repository::provider_catalog::{
|
||||||
|
StoredProviderCatalogKey, StoredProviderCatalogProvider,
|
||||||
|
};
|
||||||
|
|
||||||
|
fn grok_provider() -> StoredProviderCatalogProvider {
|
||||||
|
let mut provider = StoredProviderCatalogProvider::new(
|
||||||
|
"provider-1".to_string(),
|
||||||
|
"Grok".to_string(),
|
||||||
|
None,
|
||||||
|
"grok".to_string(),
|
||||||
|
)
|
||||||
|
.expect("provider should build");
|
||||||
|
provider.provider_type = "grok".to_string();
|
||||||
|
provider
|
||||||
|
}
|
||||||
|
|
||||||
|
fn grok_key_with_quota(quota: Value) -> StoredProviderCatalogKey {
|
||||||
|
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.status_snapshot = Some(json!({ "quota": quota }));
|
||||||
|
key
|
||||||
|
}
|
||||||
|
|
||||||
|
fn model(id: &str) -> Value {
|
||||||
|
json!({ "id": id })
|
||||||
|
}
|
||||||
|
|
||||||
|
fn filtered_ids(key: &StoredProviderCatalogKey) -> Vec<String> {
|
||||||
|
provider_query_filter_models_for_key(
|
||||||
|
&grok_provider(),
|
||||||
|
key,
|
||||||
|
vec![
|
||||||
|
model("grok-4.20-0309-non-reasoning"),
|
||||||
|
model("grok-4.20-auto"),
|
||||||
|
model("grok-4.20-heavy"),
|
||||||
|
model("grok-imagine-image-lite"),
|
||||||
|
model("grok-imagine-image"),
|
||||||
|
model("grok-imagine-image-edit"),
|
||||||
|
],
|
||||||
|
)
|
||||||
|
.into_iter()
|
||||||
|
.filter_map(|item| item.get("id").and_then(Value::as_str).map(str::to_string))
|
||||||
|
.collect()
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn provider_query_grok_basic_tier_hides_super_and_heavy_models() {
|
||||||
|
let key = grok_key_with_quota(json!({ "pool_tier": "basic" }));
|
||||||
|
|
||||||
|
assert_eq!(
|
||||||
|
filtered_ids(&key),
|
||||||
|
["grok-4.20-0309-non-reasoning", "grok-imagine-image-lite"]
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn provider_query_grok_super_tier_hides_heavy_models() {
|
||||||
|
let key = grok_key_with_quota(json!({ "plan_type": "super" }));
|
||||||
|
|
||||||
|
assert_eq!(
|
||||||
|
filtered_ids(&key),
|
||||||
|
[
|
||||||
|
"grok-4.20-0309-non-reasoning",
|
||||||
|
"grok-4.20-auto",
|
||||||
|
"grok-imagine-image-lite",
|
||||||
|
"grok-imagine-image",
|
||||||
|
"grok-imagine-image-edit"
|
||||||
|
]
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn provider_query_grok_heavy_tier_keeps_full_non_video_catalog() {
|
||||||
|
let key = grok_key_with_quota(json!({ "pool_tier": "heavy" }));
|
||||||
|
|
||||||
|
assert_eq!(
|
||||||
|
filtered_ids(&key),
|
||||||
|
[
|
||||||
|
"grok-4.20-0309-non-reasoning",
|
||||||
|
"grok-4.20-auto",
|
||||||
|
"grok-4.20-heavy",
|
||||||
|
"grok-imagine-image-lite",
|
||||||
|
"grok-imagine-image",
|
||||||
|
"grok-imagine-image-edit"
|
||||||
|
]
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn provider_query_grok_tier_falls_back_to_live_quota_windows() {
|
||||||
|
let key = grok_key_with_quota(json!({
|
||||||
|
"windows": [
|
||||||
|
{ "model": "quota_fast", "limit_value": 140.0 }
|
||||||
|
]
|
||||||
|
}));
|
||||||
|
|
||||||
|
assert_eq!(
|
||||||
|
filtered_ids(&key),
|
||||||
|
[
|
||||||
|
"grok-4.20-0309-non-reasoning",
|
||||||
|
"grok-4.20-auto",
|
||||||
|
"grok-imagine-image-lite",
|
||||||
|
"grok-imagine-image",
|
||||||
|
"grok-imagine-image-edit"
|
||||||
|
]
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn provider_query_attaches_model_test_capabilities_to_models() {
|
||||||
|
let models = provider_query_attach_model_test_capabilities(
|
||||||
|
&grok_provider(),
|
||||||
|
vec![
|
||||||
|
model("grok-4.20-fast"),
|
||||||
|
model("grok-imagine-image"),
|
||||||
|
model("grok-imagine-image-edit"),
|
||||||
|
],
|
||||||
|
);
|
||||||
|
|
||||||
|
assert!(models[0]["model_test_capabilities"]["openai:image"].is_null());
|
||||||
|
assert_eq!(
|
||||||
|
models[1]["model_test_capabilities"]["openai:image"]["max_generation_count"],
|
||||||
|
json!(4)
|
||||||
|
);
|
||||||
|
assert_eq!(
|
||||||
|
models[1]["model_test_capabilities"]["openai:image"]["supports_generation"],
|
||||||
|
json!(true)
|
||||||
|
);
|
||||||
|
assert_eq!(
|
||||||
|
models[2]["model_test_capabilities"]["openai:image"]["supports_generation"],
|
||||||
|
json!(false)
|
||||||
|
);
|
||||||
|
assert_eq!(
|
||||||
|
models[2]["model_test_capabilities"]["openai:image"]["supports_edit"],
|
||||||
|
json!(true)
|
||||||
|
);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|||||||
@@ -81,6 +81,7 @@ use tracing::{debug, warn};
|
|||||||
use uuid::Uuid;
|
use uuid::Uuid;
|
||||||
|
|
||||||
mod adapter;
|
mod adapter;
|
||||||
|
mod capabilities;
|
||||||
mod model_mapping;
|
mod model_mapping;
|
||||||
mod summary;
|
mod summary;
|
||||||
|
|
||||||
@@ -88,13 +89,17 @@ use self::adapter::{
|
|||||||
provider_query_antigravity_test_unsupported_reason,
|
provider_query_antigravity_test_unsupported_reason,
|
||||||
provider_query_antigravity_unsupported_reason,
|
provider_query_antigravity_unsupported_reason,
|
||||||
provider_query_default_antigravity_endpoint_test_body,
|
provider_query_default_antigravity_endpoint_test_body,
|
||||||
provider_query_model_test_endpoint_priority, provider_query_normalize_api_format_alias,
|
provider_query_grok_test_unsupported_reason, provider_query_model_test_endpoint_priority,
|
||||||
provider_query_standard_test_client_api_format,
|
provider_query_normalize_api_format_alias, provider_query_standard_test_client_api_format,
|
||||||
provider_query_standard_test_unsupported_reason,
|
provider_query_standard_test_unsupported_reason,
|
||||||
provider_query_test_adapter_for_provider_api_format,
|
provider_query_test_adapter_for_provider_api_format,
|
||||||
provider_query_transport_supports_model_test_execution,
|
provider_query_transport_supports_model_test_execution,
|
||||||
provider_query_unsupported_test_api_format_message, ProviderQueryTestAdapter,
|
provider_query_unsupported_test_api_format_message, ProviderQueryTestAdapter,
|
||||||
};
|
};
|
||||||
|
use self::capabilities::{
|
||||||
|
provider_query_openai_image_normalize_failure_message,
|
||||||
|
provider_query_openai_image_normalize_options,
|
||||||
|
};
|
||||||
use self::model_mapping::{
|
use self::model_mapping::{
|
||||||
provider_query_resolve_explicit_mapped_effective_model,
|
provider_query_resolve_explicit_mapped_effective_model,
|
||||||
provider_query_resolve_global_effective_model,
|
provider_query_resolve_global_effective_model,
|
||||||
@@ -530,12 +535,20 @@ fn provider_query_build_test_request_body_for_route(
|
|||||||
provider_query_build_test_request_body_with_model_policy(payload, model, override_custom_model)
|
provider_query_build_test_request_body_with_model_policy(payload, model, override_custom_model)
|
||||||
}
|
}
|
||||||
|
|
||||||
fn provider_query_build_test_request_body_with_model_policy(
|
fn provider_query_build_test_request_body_for_api_format(
|
||||||
payload: &Value,
|
payload: &Value,
|
||||||
model: &str,
|
model: &str,
|
||||||
override_custom_model: bool,
|
route_path: &str,
|
||||||
|
client_api_format: &str,
|
||||||
) -> Value {
|
) -> Value {
|
||||||
|
let client_api_format = provider_query_normalize_api_format_alias(client_api_format);
|
||||||
|
let override_custom_model = route_path.ends_with("/test-model-failover")
|
||||||
|
|| provider_query_extract_mapped_model_name(payload).is_some();
|
||||||
if let Some(mut body) = provider_query_extract_request_body(payload) {
|
if let Some(mut body) = provider_query_extract_request_body(payload) {
|
||||||
|
let has_conversation = provider_query_request_body_has_conversation_for_api_format(
|
||||||
|
&body,
|
||||||
|
client_api_format.as_str(),
|
||||||
|
);
|
||||||
if let Some(object) = body.as_object_mut() {
|
if let Some(object) = body.as_object_mut() {
|
||||||
if override_custom_model {
|
if override_custom_model {
|
||||||
object.insert("model".to_string(), Value::String(model.to_string()));
|
object.insert("model".to_string(), Value::String(model.to_string()));
|
||||||
@@ -544,6 +557,123 @@ fn provider_query_build_test_request_body_with_model_policy(
|
|||||||
.entry("model".to_string())
|
.entry("model".to_string())
|
||||||
.or_insert_with(|| Value::String(model.to_string()));
|
.or_insert_with(|| Value::String(model.to_string()));
|
||||||
}
|
}
|
||||||
|
if !has_conversation {
|
||||||
|
provider_query_insert_default_test_conversation(
|
||||||
|
object,
|
||||||
|
client_api_format.as_str(),
|
||||||
|
payload,
|
||||||
|
);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return body;
|
||||||
|
}
|
||||||
|
|
||||||
|
let message = provider_query_extract_message(payload)
|
||||||
|
.unwrap_or_else(|| DEFAULT_PROVIDER_QUERY_TEST_MESSAGE.to_string());
|
||||||
|
match client_api_format.as_str() {
|
||||||
|
"openai:responses" | "openai:responses:compact" => json!({
|
||||||
|
"model": model,
|
||||||
|
"input": message,
|
||||||
|
"max_output_tokens": 30,
|
||||||
|
"temperature": 0.7,
|
||||||
|
"stream": true,
|
||||||
|
}),
|
||||||
|
"claude:messages" => json!({
|
||||||
|
"model": model,
|
||||||
|
"messages": [{
|
||||||
|
"role": "user",
|
||||||
|
"content": message
|
||||||
|
}],
|
||||||
|
"max_tokens": 30,
|
||||||
|
"temperature": 0.7,
|
||||||
|
"stream": true,
|
||||||
|
}),
|
||||||
|
_ => json!({
|
||||||
|
"model": model,
|
||||||
|
"messages": [{
|
||||||
|
"role": "user",
|
||||||
|
"content": message
|
||||||
|
}],
|
||||||
|
"max_tokens": 30,
|
||||||
|
"temperature": 0.7,
|
||||||
|
"stream": true,
|
||||||
|
}),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fn provider_query_build_grok_test_request_body_for_api_format(
|
||||||
|
payload: &Value,
|
||||||
|
model: &str,
|
||||||
|
route_path: &str,
|
||||||
|
client_api_format: &str,
|
||||||
|
) -> Value {
|
||||||
|
provider_query_build_test_request_body_for_api_format(
|
||||||
|
payload,
|
||||||
|
model,
|
||||||
|
route_path,
|
||||||
|
client_api_format,
|
||||||
|
)
|
||||||
|
}
|
||||||
|
|
||||||
|
fn provider_query_insert_default_test_conversation(
|
||||||
|
object: &mut Map<String, Value>,
|
||||||
|
client_api_format: &str,
|
||||||
|
payload: &Value,
|
||||||
|
) {
|
||||||
|
let message = provider_query_extract_message(payload)
|
||||||
|
.unwrap_or_else(|| DEFAULT_PROVIDER_QUERY_TEST_MESSAGE.to_string());
|
||||||
|
match client_api_format {
|
||||||
|
"openai:responses" | "openai:responses:compact" => {
|
||||||
|
object.insert("input".to_string(), Value::String(message));
|
||||||
|
}
|
||||||
|
"claude:messages" => {
|
||||||
|
object.insert(
|
||||||
|
"messages".to_string(),
|
||||||
|
json!([{ "role": "user", "content": message }]),
|
||||||
|
);
|
||||||
|
}
|
||||||
|
_ => {
|
||||||
|
object.insert(
|
||||||
|
"messages".to_string(),
|
||||||
|
json!([{ "role": "user", "content": message }]),
|
||||||
|
);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fn provider_query_grok_test_client_api_format(provider_api_format: &str) -> &'static str {
|
||||||
|
match provider_query_normalize_api_format_alias(provider_api_format).as_str() {
|
||||||
|
"openai:responses" | "openai:responses:compact" => "openai:responses",
|
||||||
|
"claude:messages" => "claude:messages",
|
||||||
|
_ => "openai:chat",
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fn provider_query_build_test_request_body_with_model_policy(
|
||||||
|
payload: &Value,
|
||||||
|
model: &str,
|
||||||
|
override_custom_model: bool,
|
||||||
|
) -> Value {
|
||||||
|
if let Some(mut body) = provider_query_extract_request_body(payload) {
|
||||||
|
let has_conversation = provider_query_request_body_has_conversation(&body);
|
||||||
|
if let Some(object) = body.as_object_mut() {
|
||||||
|
if override_custom_model {
|
||||||
|
object.insert("model".to_string(), Value::String(model.to_string()));
|
||||||
|
} else {
|
||||||
|
object
|
||||||
|
.entry("model".to_string())
|
||||||
|
.or_insert_with(|| Value::String(model.to_string()));
|
||||||
|
}
|
||||||
|
if !has_conversation {
|
||||||
|
object.insert(
|
||||||
|
"messages".to_string(),
|
||||||
|
json!([{
|
||||||
|
"role": "user",
|
||||||
|
"content": provider_query_extract_message(payload)
|
||||||
|
.unwrap_or_else(|| DEFAULT_PROVIDER_QUERY_TEST_MESSAGE.to_string())
|
||||||
|
}]),
|
||||||
|
);
|
||||||
|
}
|
||||||
}
|
}
|
||||||
return body;
|
return body;
|
||||||
}
|
}
|
||||||
@@ -561,6 +691,73 @@ fn provider_query_build_test_request_body_with_model_policy(
|
|||||||
})
|
})
|
||||||
}
|
}
|
||||||
|
|
||||||
|
fn provider_query_request_body_has_conversation(body: &Value) -> bool {
|
||||||
|
body.get("messages")
|
||||||
|
.and_then(Value::as_array)
|
||||||
|
.map(|messages| {
|
||||||
|
messages
|
||||||
|
.iter()
|
||||||
|
.any(|message| value_has_non_empty_text(message.get("content")))
|
||||||
|
})
|
||||||
|
.unwrap_or(false)
|
||||||
|
|| value_has_non_empty_text(body.get("input"))
|
||||||
|
|| value_has_non_empty_text(body.get("prompt"))
|
||||||
|
|| value_has_non_empty_text(body.get("query"))
|
||||||
|
|| value_has_non_empty_text(body.get("system"))
|
||||||
|
}
|
||||||
|
|
||||||
|
fn provider_query_request_body_has_conversation_for_api_format(
|
||||||
|
body: &Value,
|
||||||
|
client_api_format: &str,
|
||||||
|
) -> bool {
|
||||||
|
match provider_query_normalize_api_format_alias(client_api_format).as_str() {
|
||||||
|
"openai:responses" | "openai:responses:compact" => {
|
||||||
|
value_has_non_empty_text(body.get("input"))
|
||||||
|
|| value_has_non_empty_text(body.get("prompt"))
|
||||||
|
}
|
||||||
|
"claude:messages" => {
|
||||||
|
body.get("messages")
|
||||||
|
.and_then(Value::as_array)
|
||||||
|
.map(|messages| {
|
||||||
|
messages
|
||||||
|
.iter()
|
||||||
|
.any(|message| value_has_non_empty_text(message.get("content")))
|
||||||
|
})
|
||||||
|
.unwrap_or(false)
|
||||||
|
|| value_has_non_empty_text(body.get("system"))
|
||||||
|
}
|
||||||
|
_ => provider_query_request_body_has_conversation(body),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fn provider_query_request_body_is_openai_responses_shape(body: &Value) -> bool {
|
||||||
|
let Some(object) = body.as_object() else {
|
||||||
|
return false;
|
||||||
|
};
|
||||||
|
[
|
||||||
|
"input",
|
||||||
|
"tools",
|
||||||
|
"tool_choice",
|
||||||
|
"instructions",
|
||||||
|
"previous_response_id",
|
||||||
|
]
|
||||||
|
.iter()
|
||||||
|
.any(|key| object.contains_key(*key))
|
||||||
|
}
|
||||||
|
|
||||||
|
fn value_has_non_empty_text(value: Option<&Value>) -> bool {
|
||||||
|
match value {
|
||||||
|
Some(Value::String(value)) => !value.trim().is_empty(),
|
||||||
|
Some(Value::Array(values)) => values
|
||||||
|
.iter()
|
||||||
|
.any(|value| value_has_non_empty_text(Some(value))),
|
||||||
|
Some(Value::Object(values)) => values
|
||||||
|
.values()
|
||||||
|
.any(|value| value_has_non_empty_text(Some(value))),
|
||||||
|
_ => false,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
fn provider_query_request_body_model<'a>(request_body: &'a Value, fallback: &'a str) -> &'a str {
|
fn provider_query_request_body_model<'a>(request_body: &'a Value, fallback: &'a str) -> &'a str {
|
||||||
request_body
|
request_body
|
||||||
.get("model")
|
.get("model")
|
||||||
@@ -1562,6 +1759,32 @@ fn provider_query_chatgpt_web_image_internal_url(base_url: &str) -> String {
|
|||||||
format!("{base_url}/__aether/chatgpt-web-image")
|
format!("{base_url}/__aether/chatgpt-web-image")
|
||||||
}
|
}
|
||||||
|
|
||||||
|
fn provider_query_openai_image_test_upstream_url(
|
||||||
|
transport: &AdminGatewayProviderTransportSnapshot,
|
||||||
|
request_query: Option<&str>,
|
||||||
|
) -> String {
|
||||||
|
if transport
|
||||||
|
.provider
|
||||||
|
.provider_type
|
||||||
|
.trim()
|
||||||
|
.eq_ignore_ascii_case("chatgpt_web")
|
||||||
|
{
|
||||||
|
provider_query_chatgpt_web_image_internal_url(&transport.endpoint.base_url)
|
||||||
|
} else if transport
|
||||||
|
.provider
|
||||||
|
.provider_type
|
||||||
|
.trim()
|
||||||
|
.eq_ignore_ascii_case("grok")
|
||||||
|
{
|
||||||
|
crate::provider_transport::build_grok_upstream_url(
|
||||||
|
transport,
|
||||||
|
crate::provider_transport::GROK_CHAT_PATH,
|
||||||
|
)
|
||||||
|
} else {
|
||||||
|
crate::provider_transport::build_openai_image_upstream_url(transport, request_query)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
async fn provider_query_finalize_openai_image_result(
|
async fn provider_query_finalize_openai_image_result(
|
||||||
route_path: &str,
|
route_path: &str,
|
||||||
trace_id: &str,
|
trace_id: &str,
|
||||||
@@ -1661,12 +1884,16 @@ async fn provider_query_execute_openai_image_test_candidate(
|
|||||||
*synthetic_request.headers_mut() = incoming_request_headers;
|
*synthetic_request.headers_mut() = incoming_request_headers;
|
||||||
let (parts, _) = synthetic_request.into_parts();
|
let (parts, _) = synthetic_request.into_parts();
|
||||||
|
|
||||||
let Some(normalized_request) =
|
let provider_type = transport.provider.provider_type.as_str();
|
||||||
crate::ai_serving::normalize_openai_image_request(&parts, &request_body, None)
|
let Some(normalized_request) = crate::ai_serving::normalize_openai_image_request_with_options(
|
||||||
else {
|
&parts,
|
||||||
|
&request_body,
|
||||||
|
None,
|
||||||
|
provider_query_openai_image_normalize_options(provider_type),
|
||||||
|
) else {
|
||||||
return Ok(provider_query_skipped_execution_outcome(
|
return Ok(provider_query_skipped_execution_outcome(
|
||||||
request_body.clone(),
|
request_body.clone(),
|
||||||
"Provider request body could not be normalized for openai:image",
|
provider_query_openai_image_normalize_failure_message(provider_type, &request_body),
|
||||||
));
|
));
|
||||||
};
|
};
|
||||||
|
|
||||||
@@ -1675,6 +1902,11 @@ async fn provider_query_execute_openai_image_test_candidate(
|
|||||||
.provider_type
|
.provider_type
|
||||||
.trim()
|
.trim()
|
||||||
.eq_ignore_ascii_case("chatgpt_web");
|
.eq_ignore_ascii_case("chatgpt_web");
|
||||||
|
let is_grok = transport
|
||||||
|
.provider
|
||||||
|
.provider_type
|
||||||
|
.trim()
|
||||||
|
.eq_ignore_ascii_case("grok");
|
||||||
let mut provider_request_body = if is_chatgpt_web {
|
let mut provider_request_body = if is_chatgpt_web {
|
||||||
match crate::ai_serving::build_chatgpt_web_image_request_body(&parts, &request_body, None) {
|
match crate::ai_serving::build_chatgpt_web_image_request_body(&parts, &request_body, None) {
|
||||||
Ok(body) => body,
|
Ok(body) => body,
|
||||||
@@ -1702,17 +1934,33 @@ async fn provider_query_execute_openai_image_test_candidate(
|
|||||||
"Provider auth is unavailable for openai:image",
|
"Provider auth is unavailable for openai:image",
|
||||||
));
|
));
|
||||||
};
|
};
|
||||||
|
let transport_profile = state.resolve_transport_profile(&transport);
|
||||||
|
|
||||||
let Some(mut request_headers) = crate::provider_transport::build_openai_image_headers(
|
let Some(mut request_headers) = (if is_grok {
|
||||||
crate::provider_transport::ProviderOpenAiImageHeadersInput {
|
crate::provider_transport::build_grok_browser_headers(
|
||||||
headers: &parts.headers,
|
crate::provider_transport::GrokHeaderInput {
|
||||||
auth_header: &auth_header,
|
transport: &transport,
|
||||||
auth_value: &auth_value,
|
transport_profile: transport_profile.as_ref(),
|
||||||
header_rules: transport.endpoint.header_rules.as_ref(),
|
request_headers: Some(&parts.headers),
|
||||||
provider_request_body: &provider_request_body,
|
content_type: "application/json",
|
||||||
original_request_body: &request_body,
|
accept: "*/*",
|
||||||
},
|
header_rules: transport.endpoint.header_rules.as_ref(),
|
||||||
) else {
|
provider_request_body: &provider_request_body,
|
||||||
|
original_request_body: &request_body,
|
||||||
|
},
|
||||||
|
)
|
||||||
|
} else {
|
||||||
|
crate::provider_transport::build_openai_image_headers(
|
||||||
|
crate::provider_transport::ProviderOpenAiImageHeadersInput {
|
||||||
|
headers: &parts.headers,
|
||||||
|
auth_header: &auth_header,
|
||||||
|
auth_value: &auth_value,
|
||||||
|
header_rules: transport.endpoint.header_rules.as_ref(),
|
||||||
|
provider_request_body: &provider_request_body,
|
||||||
|
original_request_body: &request_body,
|
||||||
|
},
|
||||||
|
)
|
||||||
|
}) else {
|
||||||
return Ok(ProviderQueryExecutionOutcome {
|
return Ok(ProviderQueryExecutionOutcome {
|
||||||
status: "failed",
|
status: "failed",
|
||||||
skip_reason: None,
|
skip_reason: None,
|
||||||
@@ -1728,6 +1976,7 @@ async fn provider_query_execute_openai_image_test_candidate(
|
|||||||
};
|
};
|
||||||
if is_chatgpt_web {
|
if is_chatgpt_web {
|
||||||
request_headers.insert("x-aether-chatgpt-web-image".to_string(), "1".to_string());
|
request_headers.insert("x-aether-chatgpt-web-image".to_string(), "1".to_string());
|
||||||
|
} else if is_grok {
|
||||||
} else {
|
} else {
|
||||||
crate::ai_serving::apply_codex_openai_responses_special_headers(
|
crate::ai_serving::apply_codex_openai_responses_special_headers(
|
||||||
&mut request_headers,
|
&mut request_headers,
|
||||||
@@ -1761,16 +2010,12 @@ async fn provider_query_execute_openai_image_test_candidate(
|
|||||||
.filter(|value| !value.is_empty())
|
.filter(|value| !value.is_empty())
|
||||||
.unwrap_or(request_model.as_str())
|
.unwrap_or(request_model.as_str())
|
||||||
.to_string();
|
.to_string();
|
||||||
let image_request = if is_chatgpt_web {
|
let image_request = if is_chatgpt_web || is_grok {
|
||||||
provider_request_body.clone()
|
provider_request_body.clone()
|
||||||
} else {
|
} else {
|
||||||
normalized_request.summary_json.clone()
|
normalized_request.summary_json.clone()
|
||||||
};
|
};
|
||||||
let request_url = if is_chatgpt_web {
|
let request_url = provider_query_openai_image_test_upstream_url(&transport, parts.uri.query());
|
||||||
provider_query_chatgpt_web_image_internal_url(&transport.endpoint.base_url)
|
|
||||||
} else {
|
|
||||||
crate::provider_transport::build_openai_image_upstream_url(&transport, parts.uri.query())
|
|
||||||
};
|
|
||||||
let upstream_is_stream = provider_request_body
|
let upstream_is_stream = provider_request_body
|
||||||
.get("stream")
|
.get("stream")
|
||||||
.and_then(Value::as_bool)
|
.and_then(Value::as_bool)
|
||||||
@@ -1796,11 +2041,27 @@ async fn provider_query_execute_openai_image_test_candidate(
|
|||||||
proxy: state
|
proxy: state
|
||||||
.resolve_transport_proxy_snapshot_with_tunnel_affinity(&transport)
|
.resolve_transport_proxy_snapshot_with_tunnel_affinity(&transport)
|
||||||
.await,
|
.await,
|
||||||
transport_profile: state.resolve_transport_profile(&transport),
|
transport_profile: transport_profile.clone(),
|
||||||
timeouts: state.resolve_transport_execution_timeouts(&transport),
|
timeouts: state.resolve_transport_execution_timeouts(&transport),
|
||||||
};
|
};
|
||||||
|
|
||||||
let result = if is_chatgpt_web {
|
let result = if is_grok {
|
||||||
|
let report_context = json!({
|
||||||
|
"client_api_format": "openai:image",
|
||||||
|
"provider_api_format": "openai:image",
|
||||||
|
"provider_type": "grok",
|
||||||
|
"model": request_model,
|
||||||
|
"mapped_model": mapped_model,
|
||||||
|
"image_request": image_request.clone(),
|
||||||
|
});
|
||||||
|
state
|
||||||
|
.execute_execution_runtime_sync_plan_with_report_context(
|
||||||
|
Some(trace_id),
|
||||||
|
&plan,
|
||||||
|
Some(&report_context),
|
||||||
|
)
|
||||||
|
.await?
|
||||||
|
} else if is_chatgpt_web {
|
||||||
let report_context = json!({
|
let report_context = json!({
|
||||||
"client_api_format": "openai:image",
|
"client_api_format": "openai:image",
|
||||||
"provider_api_format": "openai:image",
|
"provider_api_format": "openai:image",
|
||||||
@@ -2115,6 +2376,177 @@ async fn provider_query_execute_antigravity_test_candidate(
|
|||||||
})
|
})
|
||||||
}
|
}
|
||||||
|
|
||||||
|
async fn provider_query_execute_grok_test_candidate(
|
||||||
|
state: &AdminAppState<'_>,
|
||||||
|
provider: &StoredProviderCatalogProvider,
|
||||||
|
candidate: &ProviderQueryTestCandidate,
|
||||||
|
payload: &Value,
|
||||||
|
route_path: &str,
|
||||||
|
trace_id: &str,
|
||||||
|
) -> Result<ProviderQueryExecutionOutcome, GatewayError> {
|
||||||
|
let Some(transport) = state
|
||||||
|
.read_provider_transport_snapshot(&provider.id, &candidate.endpoint.id, &candidate.key.id)
|
||||||
|
.await?
|
||||||
|
else {
|
||||||
|
return Ok(provider_query_skipped_execution_outcome(
|
||||||
|
Value::Null,
|
||||||
|
"Provider transport snapshot is unavailable",
|
||||||
|
));
|
||||||
|
};
|
||||||
|
|
||||||
|
let provider_api_format =
|
||||||
|
provider_query_normalize_api_format_alias(&candidate.endpoint.api_format);
|
||||||
|
let client_api_format = provider_query_grok_test_client_api_format(&provider_api_format);
|
||||||
|
let request_body = provider_query_build_grok_test_request_body_for_api_format(
|
||||||
|
payload,
|
||||||
|
&candidate.effective_model,
|
||||||
|
route_path,
|
||||||
|
client_api_format,
|
||||||
|
);
|
||||||
|
if let Some(reason) =
|
||||||
|
provider_query_grok_test_unsupported_reason(&transport, &provider_api_format)
|
||||||
|
{
|
||||||
|
return Ok(provider_query_skipped_execution_outcome(
|
||||||
|
request_body,
|
||||||
|
format!(
|
||||||
|
"{} ({reason})",
|
||||||
|
provider_query_unsupported_test_api_format_message(&candidate.endpoint.api_format)
|
||||||
|
),
|
||||||
|
));
|
||||||
|
}
|
||||||
|
|
||||||
|
let incoming_request_headers = provider_query_extract_request_headers(payload);
|
||||||
|
let mut synthetic_request = http::Request::builder()
|
||||||
|
.uri(route_path)
|
||||||
|
.body(())
|
||||||
|
.map_err(|err| GatewayError::Internal(err.to_string()))?;
|
||||||
|
*synthetic_request.headers_mut() = incoming_request_headers;
|
||||||
|
let (parts, _) = synthetic_request.into_parts();
|
||||||
|
|
||||||
|
let request_model =
|
||||||
|
provider_query_request_body_model(&request_body, &candidate.effective_model);
|
||||||
|
let request_url = crate::provider_transport::build_grok_upstream_url(
|
||||||
|
&transport,
|
||||||
|
crate::provider_transport::GROK_CHAT_PATH,
|
||||||
|
);
|
||||||
|
let provider_request_body = crate::provider_transport::build_grok_app_chat_body(
|
||||||
|
client_api_format,
|
||||||
|
Some(request_model),
|
||||||
|
&request_body,
|
||||||
|
);
|
||||||
|
let report_context = json!({
|
||||||
|
"provider_type": provider.provider_type,
|
||||||
|
"provider_api_format": provider_api_format,
|
||||||
|
"client_api_format": client_api_format,
|
||||||
|
"model": request_model,
|
||||||
|
"mapped_model": candidate.effective_model,
|
||||||
|
"request_path": route_path,
|
||||||
|
"request_body": request_body,
|
||||||
|
});
|
||||||
|
let transport_profile = state.resolve_transport_profile(&transport);
|
||||||
|
let Some(request_headers) = crate::provider_transport::build_grok_browser_headers(
|
||||||
|
crate::provider_transport::GrokHeaderInput {
|
||||||
|
transport: &transport,
|
||||||
|
transport_profile: transport_profile.as_ref(),
|
||||||
|
request_headers: Some(&parts.headers),
|
||||||
|
content_type: "application/json",
|
||||||
|
accept: "text/event-stream",
|
||||||
|
header_rules: transport.endpoint.header_rules.as_ref(),
|
||||||
|
provider_request_body: &provider_request_body,
|
||||||
|
original_request_body: &request_body,
|
||||||
|
},
|
||||||
|
) else {
|
||||||
|
return Ok(ProviderQueryExecutionOutcome {
|
||||||
|
status: "failed",
|
||||||
|
skip_reason: None,
|
||||||
|
error_message: Some("provider request headers build failed".to_string()),
|
||||||
|
status_code: None,
|
||||||
|
latency_ms: None,
|
||||||
|
request_url,
|
||||||
|
request_headers: BTreeMap::new(),
|
||||||
|
request_body: provider_request_body,
|
||||||
|
response_headers: BTreeMap::new(),
|
||||||
|
response_body: None,
|
||||||
|
});
|
||||||
|
};
|
||||||
|
|
||||||
|
let plan = ExecutionPlan {
|
||||||
|
request_id: trace_id.to_string(),
|
||||||
|
candidate_id: Some(format!("provider-query-{}", candidate.key.id)),
|
||||||
|
provider_name: Some(provider.name.clone()),
|
||||||
|
provider_id: provider.id.clone(),
|
||||||
|
endpoint_id: candidate.endpoint.id.clone(),
|
||||||
|
key_id: candidate.key.id.clone(),
|
||||||
|
method: "POST".to_string(),
|
||||||
|
url: request_url.clone(),
|
||||||
|
headers: request_headers.clone(),
|
||||||
|
content_type: Some("application/json".to_string()),
|
||||||
|
content_encoding: None,
|
||||||
|
body: RequestBody::from_json(request_body.clone()),
|
||||||
|
stream: true,
|
||||||
|
client_api_format: client_api_format.to_string(),
|
||||||
|
provider_api_format: provider_api_format.clone(),
|
||||||
|
model_name: Some(request_model.to_string()),
|
||||||
|
proxy: state
|
||||||
|
.resolve_transport_proxy_snapshot_with_tunnel_affinity(&transport)
|
||||||
|
.await,
|
||||||
|
transport_profile,
|
||||||
|
timeouts: state.resolve_transport_execution_timeouts(&transport),
|
||||||
|
};
|
||||||
|
|
||||||
|
let result = match state
|
||||||
|
.execute_execution_runtime_sync_plan_with_report_context(
|
||||||
|
Some(trace_id),
|
||||||
|
&plan,
|
||||||
|
Some(&report_context),
|
||||||
|
)
|
||||||
|
.await
|
||||||
|
{
|
||||||
|
Ok(result) => result,
|
||||||
|
Err(err) => {
|
||||||
|
return Ok(ProviderQueryExecutionOutcome {
|
||||||
|
status: "failed",
|
||||||
|
skip_reason: None,
|
||||||
|
error_message: Some(format!("model test execution failed: {err:?}")),
|
||||||
|
status_code: None,
|
||||||
|
latency_ms: None,
|
||||||
|
request_url,
|
||||||
|
request_headers,
|
||||||
|
request_body: provider_request_body,
|
||||||
|
response_headers: BTreeMap::new(),
|
||||||
|
response_body: None,
|
||||||
|
});
|
||||||
|
}
|
||||||
|
};
|
||||||
|
let response_body = result.body.as_ref().and_then(|body| body.json_body.clone());
|
||||||
|
let did_fail = result.status_code >= 400 || response_body.is_none();
|
||||||
|
let error_message = if did_fail {
|
||||||
|
provider_query_extract_error_message(&result).or_else(|| {
|
||||||
|
response_body.is_none().then(|| {
|
||||||
|
format!(
|
||||||
|
"Provider returned HTTP {} without a model-test response body",
|
||||||
|
result.status_code
|
||||||
|
)
|
||||||
|
})
|
||||||
|
})
|
||||||
|
} else {
|
||||||
|
None
|
||||||
|
};
|
||||||
|
|
||||||
|
Ok(ProviderQueryExecutionOutcome {
|
||||||
|
status: if did_fail { "failed" } else { "success" },
|
||||||
|
skip_reason: None,
|
||||||
|
error_message,
|
||||||
|
status_code: Some(result.status_code),
|
||||||
|
latency_ms: result.telemetry.as_ref().and_then(|value| value.elapsed_ms),
|
||||||
|
request_url,
|
||||||
|
request_headers,
|
||||||
|
request_body: provider_request_body,
|
||||||
|
response_headers: result.headers,
|
||||||
|
response_body,
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
async fn provider_query_execute_standard_test_candidate(
|
async fn provider_query_execute_standard_test_candidate(
|
||||||
state: &AdminAppState<'_>,
|
state: &AdminAppState<'_>,
|
||||||
provider: &StoredProviderCatalogProvider,
|
provider: &StoredProviderCatalogProvider,
|
||||||
@@ -2132,22 +2564,25 @@ async fn provider_query_execute_standard_test_candidate(
|
|||||||
"Provider transport snapshot is unavailable",
|
"Provider transport snapshot is unavailable",
|
||||||
));
|
));
|
||||||
};
|
};
|
||||||
let original_request_body = provider_query_build_test_request_body_for_route(
|
let provider_api_format = candidate.endpoint.api_format.as_str();
|
||||||
|
let normalized_provider_api_format =
|
||||||
|
crate::ai_serving::normalize_api_format_alias(provider_api_format);
|
||||||
|
let client_api_format =
|
||||||
|
provider_query_standard_test_client_api_format(normalized_provider_api_format.as_str());
|
||||||
|
let original_request_body = provider_query_build_test_request_body_for_api_format(
|
||||||
payload,
|
payload,
|
||||||
&candidate.effective_model,
|
&candidate.effective_model,
|
||||||
route_path,
|
route_path,
|
||||||
|
client_api_format,
|
||||||
);
|
);
|
||||||
if !provider_query_transport_supports_model_test_execution(
|
if !provider_query_transport_supports_model_test_execution(
|
||||||
state,
|
state,
|
||||||
&transport,
|
&transport,
|
||||||
candidate.endpoint.api_format.as_str(),
|
provider_api_format,
|
||||||
) {
|
) {
|
||||||
return Ok(provider_query_skipped_execution_outcome(
|
return Ok(provider_query_skipped_execution_outcome(
|
||||||
original_request_body,
|
original_request_body,
|
||||||
provider_query_standard_test_unsupported_reason(
|
provider_query_standard_test_unsupported_reason(&transport, provider_api_format),
|
||||||
&transport,
|
|
||||||
candidate.endpoint.api_format.as_str(),
|
|
||||||
),
|
|
||||||
));
|
));
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -2159,11 +2594,6 @@ async fn provider_query_execute_standard_test_candidate(
|
|||||||
let request_model =
|
let request_model =
|
||||||
provider_query_request_body_model(&request_body, &candidate.effective_model);
|
provider_query_request_body_model(&request_body, &candidate.effective_model);
|
||||||
|
|
||||||
let provider_api_format = candidate.endpoint.api_format.as_str();
|
|
||||||
let normalized_provider_api_format =
|
|
||||||
crate::ai_serving::normalize_api_format_alias(provider_api_format);
|
|
||||||
let client_api_format =
|
|
||||||
provider_query_standard_test_client_api_format(normalized_provider_api_format.as_str());
|
|
||||||
let upstream_is_stream = provider_query_resolve_standard_test_upstream_is_stream(
|
let upstream_is_stream = provider_query_resolve_standard_test_upstream_is_stream(
|
||||||
transport.endpoint.config.as_ref(),
|
transport.endpoint.config.as_ref(),
|
||||||
transport.provider.provider_type.as_str(),
|
transport.provider.provider_type.as_str(),
|
||||||
@@ -2229,12 +2659,20 @@ async fn provider_query_execute_standard_test_candidate(
|
|||||||
}
|
}
|
||||||
"openai:responses" | "openai:responses:compact" => {
|
"openai:responses" | "openai:responses:compact" => {
|
||||||
let Some(mut provider_request_body) =
|
let Some(mut provider_request_body) =
|
||||||
crate::ai_serving::build_cross_format_openai_chat_request_body(
|
(if provider_query_request_body_is_openai_responses_shape(&request_body) {
|
||||||
&request_body,
|
crate::ai_serving::build_local_openai_responses_request_body(
|
||||||
request_model,
|
&request_body,
|
||||||
normalized_provider_api_format.as_str(),
|
request_model,
|
||||||
upstream_is_stream,
|
upstream_is_stream,
|
||||||
)
|
)
|
||||||
|
} else {
|
||||||
|
crate::ai_serving::build_cross_format_openai_chat_request_body(
|
||||||
|
&request_body,
|
||||||
|
request_model,
|
||||||
|
normalized_provider_api_format.as_str(),
|
||||||
|
upstream_is_stream,
|
||||||
|
)
|
||||||
|
})
|
||||||
else {
|
else {
|
||||||
return Ok(provider_query_skipped_execution_outcome(
|
return Ok(provider_query_skipped_execution_outcome(
|
||||||
request_body.clone(),
|
request_body.clone(),
|
||||||
@@ -2659,6 +3097,12 @@ async fn build_admin_provider_query_kiro_failover_response(
|
|||||||
)
|
)
|
||||||
.await
|
.await
|
||||||
}
|
}
|
||||||
|
Some(ProviderQueryTestAdapter::Grok) => {
|
||||||
|
provider_query_execute_grok_test_candidate(
|
||||||
|
state, &provider, candidate, payload, route_path, &trace_id,
|
||||||
|
)
|
||||||
|
.await
|
||||||
|
}
|
||||||
Some(ProviderQueryTestAdapter::Standard) => {
|
Some(ProviderQueryTestAdapter::Standard) => {
|
||||||
provider_query_execute_standard_test_candidate(
|
provider_query_execute_standard_test_candidate(
|
||||||
state, &provider, candidate, payload, route_path, &trace_id,
|
state, &provider, candidate, payload, route_path, &trace_id,
|
||||||
@@ -2718,7 +3162,10 @@ async fn build_admin_provider_query_kiro_failover_response(
|
|||||||
));
|
));
|
||||||
if is_success {
|
if is_success {
|
||||||
success_body = response_body;
|
success_body = response_body;
|
||||||
success_stream = matches!(adapter, Some(ProviderQueryTestAdapter::Kiro));
|
success_stream = matches!(
|
||||||
|
adapter,
|
||||||
|
Some(ProviderQueryTestAdapter::Kiro | ProviderQueryTestAdapter::Grok)
|
||||||
|
);
|
||||||
winning_candidate_index = Some(candidate_index);
|
winning_candidate_index = Some(candidate_index);
|
||||||
break;
|
break;
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -10,6 +10,7 @@ use serde_json::{json, Value};
|
|||||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||||
pub(super) enum ProviderQueryTestAdapter {
|
pub(super) enum ProviderQueryTestAdapter {
|
||||||
Standard,
|
Standard,
|
||||||
|
Grok,
|
||||||
Kiro,
|
Kiro,
|
||||||
OpenAiImage,
|
OpenAiImage,
|
||||||
Antigravity,
|
Antigravity,
|
||||||
@@ -134,6 +135,65 @@ pub(super) fn provider_query_antigravity_test_unsupported_reason(
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
pub(super) fn provider_query_grok_test_unsupported_reason(
|
||||||
|
transport: &AdminGatewayProviderTransportSnapshot,
|
||||||
|
api_format: &str,
|
||||||
|
) -> Option<&'static str> {
|
||||||
|
if !transport.provider.is_active {
|
||||||
|
return Some("provider_inactive");
|
||||||
|
}
|
||||||
|
if !transport.endpoint.is_active {
|
||||||
|
return Some("endpoint_inactive");
|
||||||
|
}
|
||||||
|
if !transport.key.is_active {
|
||||||
|
return Some("key_inactive");
|
||||||
|
}
|
||||||
|
if !transport
|
||||||
|
.provider
|
||||||
|
.provider_type
|
||||||
|
.trim()
|
||||||
|
.eq_ignore_ascii_case("grok")
|
||||||
|
{
|
||||||
|
return Some("transport_provider_type_unsupported");
|
||||||
|
}
|
||||||
|
let normalized_api_format = provider_query_normalize_api_format_alias(api_format);
|
||||||
|
if !matches!(
|
||||||
|
normalized_api_format.as_str(),
|
||||||
|
"openai:chat" | "openai:responses" | "openai:responses:compact" | "claude:messages"
|
||||||
|
) {
|
||||||
|
return Some("transport_api_format_mismatch");
|
||||||
|
}
|
||||||
|
if provider_query_normalize_api_format_alias(&transport.endpoint.api_format)
|
||||||
|
!= normalized_api_format
|
||||||
|
{
|
||||||
|
return Some("transport_api_format_mismatch");
|
||||||
|
}
|
||||||
|
if crate::provider_transport::resolve_grok_session_auth(transport).is_none() {
|
||||||
|
return Some("transport_oauth_resolution_unsupported");
|
||||||
|
}
|
||||||
|
if !crate::provider_transport::header_rules_are_locally_supported(
|
||||||
|
transport.endpoint.header_rules.as_ref(),
|
||||||
|
) {
|
||||||
|
return Some("transport_header_rules_unsupported");
|
||||||
|
}
|
||||||
|
if !crate::provider_transport::body_rules_are_locally_supported(
|
||||||
|
transport.endpoint.body_rules.as_ref(),
|
||||||
|
) {
|
||||||
|
return Some("transport_body_rules_unsupported");
|
||||||
|
}
|
||||||
|
if !crate::provider_transport::transport_proxy_is_locally_supported(transport) {
|
||||||
|
return Some("transport_proxy_unsupported");
|
||||||
|
}
|
||||||
|
if crate::provider_transport::transport_profile_is_configured(transport)
|
||||||
|
&& crate::provider_transport::resolve_transport_profile(transport).is_none()
|
||||||
|
{
|
||||||
|
return Some("transport_profile_unsupported");
|
||||||
|
}
|
||||||
|
|
||||||
|
None
|
||||||
|
}
|
||||||
|
|
||||||
pub(super) fn provider_query_normalize_api_format_alias(value: &str) -> String {
|
pub(super) fn provider_query_normalize_api_format_alias(value: &str) -> String {
|
||||||
crate::ai_serving::normalize_api_format_alias(value)
|
crate::ai_serving::normalize_api_format_alias(value)
|
||||||
}
|
}
|
||||||
@@ -147,6 +207,15 @@ pub(super) fn provider_query_test_adapter_for_provider_api_format(
|
|||||||
}
|
}
|
||||||
|
|
||||||
let normalized_api_format = provider_query_normalize_api_format_alias(api_format);
|
let normalized_api_format = provider_query_normalize_api_format_alias(api_format);
|
||||||
|
if provider_type.trim().eq_ignore_ascii_case("grok") {
|
||||||
|
return match normalized_api_format.as_str() {
|
||||||
|
"openai:chat" | "openai:responses" | "openai:responses:compact" | "claude:messages" => {
|
||||||
|
Some(ProviderQueryTestAdapter::Grok)
|
||||||
|
}
|
||||||
|
"openai:image" => Some(ProviderQueryTestAdapter::OpenAiImage),
|
||||||
|
_ => None,
|
||||||
|
};
|
||||||
|
}
|
||||||
if normalized_api_format == "openai:image" {
|
if normalized_api_format == "openai:image" {
|
||||||
return Some(ProviderQueryTestAdapter::OpenAiImage);
|
return Some(ProviderQueryTestAdapter::OpenAiImage);
|
||||||
}
|
}
|
||||||
@@ -182,6 +251,16 @@ pub(super) fn provider_query_model_test_endpoint_priority(
|
|||||||
let normalized_api_format = provider_query_normalize_api_format_alias(api_format);
|
let normalized_api_format = provider_query_normalize_api_format_alias(api_format);
|
||||||
match provider_query_test_adapter_for_provider_api_format(provider_type, api_format)? {
|
match provider_query_test_adapter_for_provider_api_format(provider_type, api_format)? {
|
||||||
ProviderQueryTestAdapter::Kiro => Some(0),
|
ProviderQueryTestAdapter::Kiro => Some(0),
|
||||||
|
ProviderQueryTestAdapter::Grok => {
|
||||||
|
if matches!(
|
||||||
|
normalized_api_format.as_str(),
|
||||||
|
"openai:chat" | "openai:responses" | "openai:responses:compact" | "claude:messages"
|
||||||
|
) {
|
||||||
|
Some(0)
|
||||||
|
} else {
|
||||||
|
Some(2)
|
||||||
|
}
|
||||||
|
}
|
||||||
ProviderQueryTestAdapter::Antigravity => Some(1),
|
ProviderQueryTestAdapter::Antigravity => Some(1),
|
||||||
ProviderQueryTestAdapter::OpenAiImage => Some(2),
|
ProviderQueryTestAdapter::OpenAiImage => Some(2),
|
||||||
ProviderQueryTestAdapter::Standard => {
|
ProviderQueryTestAdapter::Standard => {
|
||||||
@@ -234,6 +313,9 @@ pub(super) fn provider_query_transport_supports_model_test_execution(
|
|||||||
)
|
)
|
||||||
.is_none()
|
.is_none()
|
||||||
}
|
}
|
||||||
|
Some(ProviderQueryTestAdapter::Grok) => {
|
||||||
|
provider_query_grok_test_unsupported_reason(transport, api_format).is_none()
|
||||||
|
}
|
||||||
Some(ProviderQueryTestAdapter::Standard) => match crate::ai_serving::normalize_api_format_alias(api_format).as_str() {
|
Some(ProviderQueryTestAdapter::Standard) => match crate::ai_serving::normalize_api_format_alias(api_format).as_str() {
|
||||||
"openai:chat" => {
|
"openai:chat" => {
|
||||||
crate::provider_transport::policy::supports_local_openai_chat_transport(transport)
|
crate::provider_transport::policy::supports_local_openai_chat_transport(transport)
|
||||||
|
|||||||
@@ -0,0 +1,50 @@
|
|||||||
|
use crate::handlers::admin::provider::shared::model_test_capabilities::{
|
||||||
|
admin_provider_openai_image_normalize_options, admin_provider_openai_image_test_capability,
|
||||||
|
AdminProviderOpenAiImageTestCapability,
|
||||||
|
};
|
||||||
|
use serde_json::Value;
|
||||||
|
|
||||||
|
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||||
|
pub(super) struct ProviderQueryOpenAiImageTestCapability(AdminProviderOpenAiImageTestCapability);
|
||||||
|
|
||||||
|
pub(super) fn provider_query_openai_image_test_capability(
|
||||||
|
provider_type: &str,
|
||||||
|
) -> ProviderQueryOpenAiImageTestCapability {
|
||||||
|
ProviderQueryOpenAiImageTestCapability(admin_provider_openai_image_test_capability(
|
||||||
|
provider_type,
|
||||||
|
))
|
||||||
|
}
|
||||||
|
|
||||||
|
pub(super) fn provider_query_openai_image_normalize_options(
|
||||||
|
provider_type: &str,
|
||||||
|
) -> crate::ai_serving::OpenAiImageNormalizeOptions {
|
||||||
|
admin_provider_openai_image_normalize_options(provider_type)
|
||||||
|
}
|
||||||
|
|
||||||
|
pub(super) fn provider_query_openai_image_requested_count(request_body: &Value) -> Option<u64> {
|
||||||
|
request_body.get("n").and_then(|value| {
|
||||||
|
value.as_u64().or_else(|| {
|
||||||
|
value
|
||||||
|
.as_str()
|
||||||
|
.map(str::trim)
|
||||||
|
.filter(|value| !value.is_empty())
|
||||||
|
.and_then(|value| value.parse::<u64>().ok())
|
||||||
|
})
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
pub(super) fn provider_query_openai_image_normalize_failure_message(
|
||||||
|
provider_type: &str,
|
||||||
|
request_body: &Value,
|
||||||
|
) -> String {
|
||||||
|
let capability = provider_query_openai_image_test_capability(provider_type);
|
||||||
|
if provider_query_openai_image_requested_count(request_body)
|
||||||
|
.is_some_and(|value| !capability.0.supports_generation_count(value))
|
||||||
|
{
|
||||||
|
return format!(
|
||||||
|
"Provider request body could not be normalized for openai:image: selected provider supports n=1..{} for generation",
|
||||||
|
capability.0.max_generation_count
|
||||||
|
);
|
||||||
|
}
|
||||||
|
"Provider request body could not be normalized for openai:image".to_string()
|
||||||
|
}
|
||||||
@@ -1,3 +1,5 @@
|
|||||||
|
use std::collections::BTreeMap;
|
||||||
|
|
||||||
use super::super::provider_query_key_display_name;
|
use super::super::provider_query_key_display_name;
|
||||||
use super::{ProviderQueryExecutionOutcome, ProviderQueryTestCandidate};
|
use super::{ProviderQueryExecutionOutcome, ProviderQueryTestCandidate};
|
||||||
use serde_json::{json, Value};
|
use serde_json::{json, Value};
|
||||||
@@ -22,13 +24,43 @@ pub(super) fn provider_query_test_attempt_payload(
|
|||||||
"status_code": execution.status_code,
|
"status_code": execution.status_code,
|
||||||
"latency_ms": execution.latency_ms,
|
"latency_ms": execution.latency_ms,
|
||||||
"request_url": execution.request_url,
|
"request_url": execution.request_url,
|
||||||
"request_headers": execution.request_headers,
|
"request_headers": provider_query_redact_diagnostic_headers(&execution.request_headers),
|
||||||
"request_body": execution.request_body,
|
"request_body": execution.request_body,
|
||||||
"response_headers": execution.response_headers,
|
"response_headers": provider_query_redact_diagnostic_headers(&execution.response_headers),
|
||||||
"response_body": execution.response_body,
|
"response_body": execution.response_body,
|
||||||
})
|
})
|
||||||
}
|
}
|
||||||
|
|
||||||
|
fn provider_query_redact_diagnostic_headers(
|
||||||
|
headers: &BTreeMap<String, String>,
|
||||||
|
) -> BTreeMap<String, String> {
|
||||||
|
headers
|
||||||
|
.iter()
|
||||||
|
.map(|(name, value)| {
|
||||||
|
if provider_query_header_is_sensitive(name) {
|
||||||
|
(name.clone(), "<redacted>".to_string())
|
||||||
|
} else {
|
||||||
|
(name.clone(), value.clone())
|
||||||
|
}
|
||||||
|
})
|
||||||
|
.collect()
|
||||||
|
}
|
||||||
|
|
||||||
|
fn provider_query_header_is_sensitive(name: &str) -> bool {
|
||||||
|
matches!(
|
||||||
|
name.trim().to_ascii_lowercase().as_str(),
|
||||||
|
"authorization"
|
||||||
|
| "proxy-authorization"
|
||||||
|
| "cookie"
|
||||||
|
| "set-cookie"
|
||||||
|
| "x-api-key"
|
||||||
|
| "api-key"
|
||||||
|
| "x-goog-api-key"
|
||||||
|
| "anthropic-api-key"
|
||||||
|
| "openai-api-key"
|
||||||
|
)
|
||||||
|
}
|
||||||
|
|
||||||
pub(super) fn provider_query_candidate_summary_payload(
|
pub(super) fn provider_query_candidate_summary_payload(
|
||||||
total_candidates: usize,
|
total_candidates: usize,
|
||||||
total_attempts: usize,
|
total_attempts: usize,
|
||||||
@@ -133,3 +165,37 @@ pub(super) fn provider_query_candidate_summary_payload(
|
|||||||
.unwrap_or(Value::Null),
|
.unwrap_or(Value::Null),
|
||||||
})
|
})
|
||||||
}
|
}
|
||||||
|
|
||||||
|
#[cfg(test)]
|
||||||
|
mod tests {
|
||||||
|
use super::*;
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn provider_query_diagnostic_headers_redact_credentials() {
|
||||||
|
let headers = BTreeMap::from([
|
||||||
|
("cookie".to_string(), "sso=secret".to_string()),
|
||||||
|
("authorization".to_string(), "Bearer secret".to_string()),
|
||||||
|
("x-goog-api-key".to_string(), "secret".to_string()),
|
||||||
|
("content-type".to_string(), "application/json".to_string()),
|
||||||
|
]);
|
||||||
|
|
||||||
|
let redacted = provider_query_redact_diagnostic_headers(&headers);
|
||||||
|
|
||||||
|
assert_eq!(
|
||||||
|
redacted.get("cookie").map(String::as_str),
|
||||||
|
Some("<redacted>")
|
||||||
|
);
|
||||||
|
assert_eq!(
|
||||||
|
redacted.get("authorization").map(String::as_str),
|
||||||
|
Some("<redacted>")
|
||||||
|
);
|
||||||
|
assert_eq!(
|
||||||
|
redacted.get("x-goog-api-key").map(String::as_str),
|
||||||
|
Some("<redacted>")
|
||||||
|
);
|
||||||
|
assert_eq!(
|
||||||
|
redacted.get("content-type").map(String::as_str),
|
||||||
|
Some("application/json")
|
||||||
|
);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|||||||
@@ -1,6 +1,68 @@
|
|||||||
use super::*;
|
use super::*;
|
||||||
|
use crate::handlers::admin::request::AdminGatewayProviderTransportSnapshot;
|
||||||
use serde_json::json;
|
use serde_json::json;
|
||||||
|
|
||||||
|
fn sample_openai_image_transport(provider_type: &str) -> AdminGatewayProviderTransportSnapshot {
|
||||||
|
AdminGatewayProviderTransportSnapshot {
|
||||||
|
provider: crate::provider_transport::snapshot::GatewayProviderTransportProvider {
|
||||||
|
id: "provider-1".to_string(),
|
||||||
|
name: "Provider".to_string(),
|
||||||
|
provider_type: provider_type.to_string(),
|
||||||
|
website: None,
|
||||||
|
is_active: true,
|
||||||
|
keep_priority_on_conversion: false,
|
||||||
|
enable_format_conversion: false,
|
||||||
|
concurrent_limit: None,
|
||||||
|
max_retries: None,
|
||||||
|
proxy: None,
|
||||||
|
request_timeout_secs: None,
|
||||||
|
stream_first_byte_timeout_secs: None,
|
||||||
|
config: None,
|
||||||
|
},
|
||||||
|
endpoint: crate::provider_transport::snapshot::GatewayProviderTransportEndpoint {
|
||||||
|
id: "endpoint-1".to_string(),
|
||||||
|
provider_id: "provider-1".to_string(),
|
||||||
|
api_format: "openai:image".to_string(),
|
||||||
|
api_family: None,
|
||||||
|
endpoint_kind: None,
|
||||||
|
is_active: true,
|
||||||
|
base_url: "https://grok.com/".to_string(),
|
||||||
|
header_rules: None,
|
||||||
|
body_rules: None,
|
||||||
|
max_retries: None,
|
||||||
|
custom_path: None,
|
||||||
|
config: None,
|
||||||
|
format_acceptance_config: None,
|
||||||
|
proxy: None,
|
||||||
|
},
|
||||||
|
key: crate::provider_transport::snapshot::GatewayProviderTransportKey {
|
||||||
|
id: "key-1".to_string(),
|
||||||
|
provider_id: "provider-1".to_string(),
|
||||||
|
name: "key".to_string(),
|
||||||
|
auth_type: "oauth".to_string(),
|
||||||
|
is_active: true,
|
||||||
|
api_formats: None,
|
||||||
|
auth_type_by_format: None,
|
||||||
|
allow_auth_channel_mismatch_formats: None,
|
||||||
|
allowed_models: None,
|
||||||
|
capabilities: None,
|
||||||
|
rate_multipliers: None,
|
||||||
|
global_priority_by_format: None,
|
||||||
|
expires_at_unix_secs: None,
|
||||||
|
proxy: None,
|
||||||
|
fingerprint: None,
|
||||||
|
decrypted_api_key: String::new(),
|
||||||
|
decrypted_auth_config: Some(
|
||||||
|
json!({
|
||||||
|
"sso_token": "abc",
|
||||||
|
"sso_rw_token": "rw"
|
||||||
|
})
|
||||||
|
.to_string(),
|
||||||
|
),
|
||||||
|
},
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
#[test]
|
#[test]
|
||||||
fn provider_query_test_request_body_preserves_custom_model() {
|
fn provider_query_test_request_body_preserves_custom_model() {
|
||||||
let payload = json!({
|
let payload = json!({
|
||||||
@@ -28,6 +90,41 @@ fn provider_query_test_request_body_defaults_missing_model() {
|
|||||||
assert_eq!(body["model"], json!("fallback-model"));
|
assert_eq!(body["model"], json!("fallback-model"));
|
||||||
}
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn provider_query_test_request_body_fills_empty_conversation() {
|
||||||
|
let payload = json!({
|
||||||
|
"request_body": {
|
||||||
|
"model": "custom-upstream-model",
|
||||||
|
"messages": []
|
||||||
|
}
|
||||||
|
});
|
||||||
|
|
||||||
|
let body = provider_query_build_test_request_body(&payload, "fallback-model");
|
||||||
|
|
||||||
|
assert_eq!(body["model"], json!("custom-upstream-model"));
|
||||||
|
assert_eq!(
|
||||||
|
body["messages"],
|
||||||
|
json!([{ "role": "user", "content": DEFAULT_PROVIDER_QUERY_TEST_MESSAGE }])
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn provider_query_test_request_body_keeps_non_empty_conversation() {
|
||||||
|
let payload = json!({
|
||||||
|
"request_body": {
|
||||||
|
"model": "custom-upstream-model",
|
||||||
|
"messages": [{ "role": "user", "content": "custom prompt" }]
|
||||||
|
}
|
||||||
|
});
|
||||||
|
|
||||||
|
let body = provider_query_build_test_request_body(&payload, "fallback-model");
|
||||||
|
|
||||||
|
assert_eq!(
|
||||||
|
body["messages"],
|
||||||
|
json!([{ "role": "user", "content": "custom prompt" }])
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
#[test]
|
#[test]
|
||||||
fn provider_query_failover_request_body_overrides_custom_model() {
|
fn provider_query_failover_request_body_overrides_custom_model() {
|
||||||
let payload = json!({
|
let payload = json!({
|
||||||
@@ -168,6 +265,54 @@ fn provider_query_standard_test_aggregates_responses_stream_body() {
|
|||||||
assert_eq!(body["output"][0]["content"][0]["text"], json!("Hello"));
|
assert_eq!(body["output"][0]["content"][0]["text"], json!("Hello"));
|
||||||
}
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn provider_query_standard_test_aggregates_responses_image_generation_call() {
|
||||||
|
let stream_body = concat!(
|
||||||
|
"event: response.created\n",
|
||||||
|
"data: {\"type\":\"response.created\",\"response\":{\"id\":\"resp_img_123\",\"object\":\"response\",\"model\":\"gpt-5.4-mini\",\"status\":\"in_progress\",\"output\":[]}}\n\n",
|
||||||
|
"event: response.output_item.done\n",
|
||||||
|
"data: {\"type\":\"response.output_item.done\",\"output_index\":0,\"item\":{\"id\":\"ig_123\",\"type\":\"image_generation_call\",\"status\":\"completed\",\"output_format\":\"png\",\"result\":\"aGVsbG8=\"}}\n\n",
|
||||||
|
"event: response.completed\n",
|
||||||
|
"data: {\"type\":\"response.completed\",\"response\":{\"id\":\"resp_img_123\",\"object\":\"response\",\"model\":\"gpt-5.4-mini\",\"status\":\"completed\",\"output\":[]}}\n\n",
|
||||||
|
);
|
||||||
|
let result = aether_contracts::ExecutionResult {
|
||||||
|
request_id: "provider-test".to_string(),
|
||||||
|
candidate_id: Some("candidate-0".to_string()),
|
||||||
|
status_code: 200,
|
||||||
|
headers: BTreeMap::new(),
|
||||||
|
body: Some(aether_contracts::ResponseBody {
|
||||||
|
json_body: None,
|
||||||
|
body_bytes_b64: Some(
|
||||||
|
base64::engine::general_purpose::STANDARD.encode(stream_body.as_bytes()),
|
||||||
|
),
|
||||||
|
}),
|
||||||
|
telemetry: None,
|
||||||
|
error: None,
|
||||||
|
};
|
||||||
|
|
||||||
|
let body = provider_query_standard_execution_response_body("openai:responses", &result)
|
||||||
|
.expect("responses image stream body should aggregate");
|
||||||
|
|
||||||
|
assert_eq!(body["output"][0]["type"], json!("image_generation_call"));
|
||||||
|
assert_eq!(body["output"][0]["result"], json!("aGVsbG8="));
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn provider_query_responses_test_request_body_defaults_to_responses_input() {
|
||||||
|
let payload = json!({"message": "hello from responses"});
|
||||||
|
|
||||||
|
let body = provider_query_build_test_request_body_for_api_format(
|
||||||
|
&payload,
|
||||||
|
"gpt-5.4-mini",
|
||||||
|
"/api/admin/provider-query/test-model",
|
||||||
|
"openai:responses",
|
||||||
|
);
|
||||||
|
|
||||||
|
assert_eq!(body["model"], json!("gpt-5.4-mini"));
|
||||||
|
assert_eq!(body["input"], json!("hello from responses"));
|
||||||
|
assert!(body.get("messages").is_none());
|
||||||
|
}
|
||||||
|
|
||||||
#[test]
|
#[test]
|
||||||
fn provider_query_test_adapter_routes_fixed_provider_endpoint_types() {
|
fn provider_query_test_adapter_routes_fixed_provider_endpoint_types() {
|
||||||
assert_eq!(
|
assert_eq!(
|
||||||
@@ -204,6 +349,22 @@ fn provider_query_test_adapter_routes_fixed_provider_endpoint_types() {
|
|||||||
),
|
),
|
||||||
Some(ProviderQueryTestAdapter::Antigravity)
|
Some(ProviderQueryTestAdapter::Antigravity)
|
||||||
);
|
);
|
||||||
|
assert_eq!(
|
||||||
|
provider_query_test_adapter_for_provider_api_format("grok", "openai:chat"),
|
||||||
|
Some(ProviderQueryTestAdapter::Grok)
|
||||||
|
);
|
||||||
|
assert_eq!(
|
||||||
|
provider_query_test_adapter_for_provider_api_format("grok", "openai:responses"),
|
||||||
|
Some(ProviderQueryTestAdapter::Grok)
|
||||||
|
);
|
||||||
|
assert_eq!(
|
||||||
|
provider_query_test_adapter_for_provider_api_format("grok", "claude:messages"),
|
||||||
|
Some(ProviderQueryTestAdapter::Grok)
|
||||||
|
);
|
||||||
|
assert_eq!(
|
||||||
|
provider_query_test_adapter_for_provider_api_format("grok", "openai:image"),
|
||||||
|
Some(ProviderQueryTestAdapter::OpenAiImage)
|
||||||
|
);
|
||||||
assert_eq!(
|
assert_eq!(
|
||||||
provider_query_test_adapter_for_provider_api_format("custom", "openai:embedding"),
|
provider_query_test_adapter_for_provider_api_format("custom", "openai:embedding"),
|
||||||
Some(ProviderQueryTestAdapter::Standard)
|
Some(ProviderQueryTestAdapter::Standard)
|
||||||
@@ -244,12 +405,154 @@ fn provider_query_endpoint_priority_prefers_text_before_cli_and_image() {
|
|||||||
provider_query_model_test_endpoint_priority("chatgpt_web", "openai:image"),
|
provider_query_model_test_endpoint_priority("chatgpt_web", "openai:image"),
|
||||||
Some(2)
|
Some(2)
|
||||||
);
|
);
|
||||||
|
assert_eq!(
|
||||||
|
provider_query_model_test_endpoint_priority("grok", "openai:chat"),
|
||||||
|
Some(0)
|
||||||
|
);
|
||||||
|
assert_eq!(
|
||||||
|
provider_query_model_test_endpoint_priority("grok", "openai:responses"),
|
||||||
|
Some(0)
|
||||||
|
);
|
||||||
assert_eq!(
|
assert_eq!(
|
||||||
provider_query_model_test_endpoint_priority("antigravity", "gemini:generate_content"),
|
provider_query_model_test_endpoint_priority("antigravity", "gemini:generate_content"),
|
||||||
Some(1)
|
Some(1)
|
||||||
);
|
);
|
||||||
}
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn provider_query_grok_model_test_body_maps_non_reasoning_model_to_fast_mode() {
|
||||||
|
let payload = json!({
|
||||||
|
"request_body": {
|
||||||
|
"model": "grok-4.20-0309-non-reasoning",
|
||||||
|
"messages": [
|
||||||
|
{"role": "system", "content": "be concise"},
|
||||||
|
{"role": "user", "content": "hello"}
|
||||||
|
]
|
||||||
|
}
|
||||||
|
});
|
||||||
|
let request_body = provider_query_build_test_request_body_for_route(
|
||||||
|
&payload,
|
||||||
|
"grok-4.20-0309-non-reasoning",
|
||||||
|
"/api/admin/provider-query/test-model",
|
||||||
|
);
|
||||||
|
|
||||||
|
let upstream_body = crate::provider_transport::build_grok_app_chat_body(
|
||||||
|
"openai:chat",
|
||||||
|
Some(provider_query_request_body_model(
|
||||||
|
&request_body,
|
||||||
|
"grok-4.20-0309-non-reasoning",
|
||||||
|
)),
|
||||||
|
&request_body,
|
||||||
|
);
|
||||||
|
|
||||||
|
assert_eq!(upstream_body["modeId"], json!("fast"));
|
||||||
|
assert_eq!(
|
||||||
|
upstream_body["message"],
|
||||||
|
json!("[system]: be concise\n\n[user]: hello")
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn provider_query_grok_model_test_uses_responses_client_body_for_responses_endpoint() {
|
||||||
|
let payload = json!({
|
||||||
|
"request_body": {
|
||||||
|
"model": "grok-4.20-0309-non-reasoning",
|
||||||
|
"input": "hello from responses body"
|
||||||
|
}
|
||||||
|
});
|
||||||
|
let request_body = provider_query_build_grok_test_request_body_for_api_format(
|
||||||
|
&payload,
|
||||||
|
"grok-4.20-0309-non-reasoning",
|
||||||
|
"/api/admin/provider-query/test-model",
|
||||||
|
"openai:responses",
|
||||||
|
);
|
||||||
|
|
||||||
|
let upstream_body = crate::provider_transport::build_grok_app_chat_body(
|
||||||
|
provider_query_grok_test_client_api_format("openai:responses"),
|
||||||
|
Some(provider_query_request_body_model(
|
||||||
|
&request_body,
|
||||||
|
"grok-4.20-0309-non-reasoning",
|
||||||
|
)),
|
||||||
|
&request_body,
|
||||||
|
);
|
||||||
|
|
||||||
|
assert_eq!(upstream_body["modeId"], json!("fast"));
|
||||||
|
assert_eq!(upstream_body["message"], json!("hello from responses body"));
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn provider_query_grok_model_test_uses_responses_input_when_existing_body_has_messages() {
|
||||||
|
let payload = json!({
|
||||||
|
"request_body": {
|
||||||
|
"model": "grok-4.20-0309-non-reasoning",
|
||||||
|
"messages": [{
|
||||||
|
"role": "user",
|
||||||
|
"content": "hello from stale chat body"
|
||||||
|
}]
|
||||||
|
}
|
||||||
|
});
|
||||||
|
let request_body = provider_query_build_grok_test_request_body_for_api_format(
|
||||||
|
&payload,
|
||||||
|
"grok-4.20-0309-non-reasoning",
|
||||||
|
"/api/admin/provider-query/test-model",
|
||||||
|
"openai:responses",
|
||||||
|
);
|
||||||
|
|
||||||
|
assert_eq!(request_body["model"], json!("grok-4.20-0309-non-reasoning"));
|
||||||
|
assert_eq!(
|
||||||
|
request_body["input"],
|
||||||
|
json!("Hello! This is a test message.")
|
||||||
|
);
|
||||||
|
assert!(request_body.get("messages").is_some());
|
||||||
|
|
||||||
|
let upstream_body = crate::provider_transport::build_grok_app_chat_body(
|
||||||
|
provider_query_grok_test_client_api_format("openai:responses"),
|
||||||
|
Some(provider_query_request_body_model(
|
||||||
|
&request_body,
|
||||||
|
"grok-4.20-0309-non-reasoning",
|
||||||
|
)),
|
||||||
|
&request_body,
|
||||||
|
);
|
||||||
|
|
||||||
|
assert_eq!(upstream_body["modeId"], json!("fast"));
|
||||||
|
assert_eq!(
|
||||||
|
upstream_body["message"],
|
||||||
|
json!("Hello! This is a test message.")
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn provider_query_grok_model_test_defaults_claude_messages_body_for_claude_endpoint() {
|
||||||
|
let payload = json!({});
|
||||||
|
let request_body = provider_query_build_grok_test_request_body_for_api_format(
|
||||||
|
&payload,
|
||||||
|
"grok-4.20-0309-non-reasoning",
|
||||||
|
"/api/admin/provider-query/test-model",
|
||||||
|
"claude:messages",
|
||||||
|
);
|
||||||
|
|
||||||
|
assert_eq!(request_body["model"], json!("grok-4.20-0309-non-reasoning"));
|
||||||
|
assert_eq!(
|
||||||
|
request_body["messages"],
|
||||||
|
json!([{ "role": "user", "content": DEFAULT_PROVIDER_QUERY_TEST_MESSAGE }])
|
||||||
|
);
|
||||||
|
|
||||||
|
let upstream_body = crate::provider_transport::build_grok_app_chat_body(
|
||||||
|
provider_query_grok_test_client_api_format("claude:messages"),
|
||||||
|
Some(provider_query_request_body_model(
|
||||||
|
&request_body,
|
||||||
|
"grok-4.20-0309-non-reasoning",
|
||||||
|
)),
|
||||||
|
&request_body,
|
||||||
|
);
|
||||||
|
|
||||||
|
assert_eq!(upstream_body["modeId"], json!("fast"));
|
||||||
|
assert_eq!(
|
||||||
|
upstream_body["message"],
|
||||||
|
json!("[user]: Hello! This is a test message.")
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
#[test]
|
#[test]
|
||||||
fn provider_query_candidate_summary_marks_unused_after_first_success() {
|
fn provider_query_candidate_summary_marks_unused_after_first_success() {
|
||||||
let attempts = vec![json!({
|
let attempts = vec![json!({
|
||||||
@@ -360,3 +663,76 @@ fn provider_query_failover_image_test_request_body_overrides_model() {
|
|||||||
|
|
||||||
assert_eq!(body["model"], json!("new-image-model"));
|
assert_eq!(body["model"], json!("new-image-model"));
|
||||||
}
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn provider_query_grok_image_test_allows_multi_generation_count() {
|
||||||
|
let request = http::Request::builder()
|
||||||
|
.uri("/v1/images/generations")
|
||||||
|
.body(())
|
||||||
|
.expect("request should build");
|
||||||
|
let (parts, _) = request.into_parts();
|
||||||
|
let body = json!({
|
||||||
|
"model": "grok-imagine-image",
|
||||||
|
"prompt": "draw",
|
||||||
|
"n": 2
|
||||||
|
});
|
||||||
|
|
||||||
|
let normalized = crate::ai_serving::normalize_openai_image_request_with_options(
|
||||||
|
&parts,
|
||||||
|
&body,
|
||||||
|
None,
|
||||||
|
provider_query_openai_image_normalize_options("grok"),
|
||||||
|
)
|
||||||
|
.expect("grok image model tests should allow multi-image generation");
|
||||||
|
let provider_body = crate::ai_serving::build_openai_image_provider_request_body(&normalized);
|
||||||
|
|
||||||
|
assert_eq!(provider_body["n"], json!(2));
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn provider_query_grok_image_test_uses_grok_app_chat_upstream_url() {
|
||||||
|
let transport = sample_openai_image_transport("grok");
|
||||||
|
|
||||||
|
assert_eq!(
|
||||||
|
provider_query_openai_image_test_upstream_url(&transport, Some("trace=1")),
|
||||||
|
"https://grok.com/rest/app-chat/conversations/new"
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn provider_query_chatgpt_web_image_test_uses_internal_upstream_url() {
|
||||||
|
let transport = sample_openai_image_transport("chatgpt_web");
|
||||||
|
|
||||||
|
assert_eq!(
|
||||||
|
provider_query_openai_image_test_upstream_url(&transport, Some("trace=1")),
|
||||||
|
"https://grok.com/__aether/chatgpt-web-image"
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn provider_query_non_grok_image_test_keeps_single_generation_boundary() {
|
||||||
|
let request = http::Request::builder()
|
||||||
|
.uri("/v1/images/generations")
|
||||||
|
.body(())
|
||||||
|
.expect("request should build");
|
||||||
|
let (parts, _) = request.into_parts();
|
||||||
|
let body = json!({
|
||||||
|
"model": "gpt-image-2",
|
||||||
|
"prompt": "draw",
|
||||||
|
"n": 2
|
||||||
|
});
|
||||||
|
|
||||||
|
assert!(
|
||||||
|
crate::ai_serving::normalize_openai_image_request_with_options(
|
||||||
|
&parts,
|
||||||
|
&body,
|
||||||
|
None,
|
||||||
|
provider_query_openai_image_normalize_options("chatgpt_web"),
|
||||||
|
)
|
||||||
|
.is_none()
|
||||||
|
);
|
||||||
|
assert_eq!(
|
||||||
|
provider_query_openai_image_normalize_failure_message("chatgpt_web", &body),
|
||||||
|
"Provider request body could not be normalized for openai:image: selected provider supports n=1..1 for generation"
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|||||||
@@ -1,3 +1,4 @@
|
|||||||
|
pub(crate) mod model_test_capabilities;
|
||||||
pub(crate) mod paths;
|
pub(crate) mod paths;
|
||||||
pub(crate) mod payloads;
|
pub(crate) mod payloads;
|
||||||
pub(crate) mod support;
|
pub(crate) mod support;
|
||||||
|
|||||||
@@ -0,0 +1,121 @@
|
|||||||
|
use crate::image_capabilities::{
|
||||||
|
openai_image_normalize_options_for_provider, openai_image_provider_max_generation_count,
|
||||||
|
};
|
||||||
|
use serde_json::{json, Value};
|
||||||
|
|
||||||
|
const GROK_IMAGE_MODEL_IDS: &[&str] = &[
|
||||||
|
"grok-imagine-image-lite",
|
||||||
|
"grok-imagine-image",
|
||||||
|
"grok-imagine-image-pro",
|
||||||
|
"grok-imagine-image-edit",
|
||||||
|
];
|
||||||
|
const GROK_IMAGE_EDIT_MODEL_ID: &str = "grok-imagine-image-edit";
|
||||||
|
|
||||||
|
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||||
|
pub(crate) struct AdminProviderOpenAiImageTestCapability {
|
||||||
|
pub(crate) max_generation_count: u64,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl AdminProviderOpenAiImageTestCapability {
|
||||||
|
pub(crate) fn supports_generation_count(self, count: u64) -> bool {
|
||||||
|
count >= 1 && count <= self.max_generation_count
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
pub(crate) fn admin_provider_openai_image_test_capability(
|
||||||
|
provider_type: &str,
|
||||||
|
) -> AdminProviderOpenAiImageTestCapability {
|
||||||
|
AdminProviderOpenAiImageTestCapability {
|
||||||
|
max_generation_count: openai_image_provider_max_generation_count(provider_type),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
pub(crate) fn admin_provider_openai_image_normalize_options(
|
||||||
|
provider_type: &str,
|
||||||
|
) -> crate::ai_serving::OpenAiImageNormalizeOptions {
|
||||||
|
openai_image_normalize_options_for_provider(provider_type)
|
||||||
|
}
|
||||||
|
|
||||||
|
pub(crate) fn admin_provider_model_test_capabilities_payload(
|
||||||
|
provider_type: &str,
|
||||||
|
model_id: &str,
|
||||||
|
supports_image_generation: bool,
|
||||||
|
) -> Value {
|
||||||
|
let provider_type = provider_type.trim();
|
||||||
|
let model_id = model_id.trim();
|
||||||
|
let is_grok_image_edit =
|
||||||
|
provider_type.eq_ignore_ascii_case("grok") && model_id == GROK_IMAGE_EDIT_MODEL_ID;
|
||||||
|
let openai_image = if supports_image_generation {
|
||||||
|
Some(json!({
|
||||||
|
"max_generation_count": admin_provider_openai_image_test_capability(provider_type).max_generation_count,
|
||||||
|
"supports_generation": !is_grok_image_edit,
|
||||||
|
"supports_edit": is_grok_image_edit,
|
||||||
|
}))
|
||||||
|
} else {
|
||||||
|
None
|
||||||
|
};
|
||||||
|
|
||||||
|
json!({
|
||||||
|
"openai:image": openai_image,
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
pub(crate) fn admin_provider_model_supports_image_generation(
|
||||||
|
provider_type: &str,
|
||||||
|
model_id: &str,
|
||||||
|
fallback_supports_image_generation: bool,
|
||||||
|
) -> bool {
|
||||||
|
if provider_type.trim().eq_ignore_ascii_case("grok") {
|
||||||
|
let model_id = model_id.trim();
|
||||||
|
return GROK_IMAGE_MODEL_IDS
|
||||||
|
.iter()
|
||||||
|
.any(|candidate| model_id.eq_ignore_ascii_case(candidate));
|
||||||
|
}
|
||||||
|
fallback_supports_image_generation
|
||||||
|
}
|
||||||
|
|
||||||
|
#[cfg(test)]
|
||||||
|
mod tests {
|
||||||
|
use super::*;
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn grok_image_generation_models_expose_multi_image_capability() {
|
||||||
|
let payload =
|
||||||
|
admin_provider_model_test_capabilities_payload("grok", "grok-imagine-image", true);
|
||||||
|
|
||||||
|
assert_eq!(payload["openai:image"]["max_generation_count"], 4);
|
||||||
|
assert_eq!(payload["openai:image"]["supports_generation"], true);
|
||||||
|
assert_eq!(payload["openai:image"]["supports_edit"], false);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn grok_image_edit_model_is_edit_only_for_generation_tests() {
|
||||||
|
let payload =
|
||||||
|
admin_provider_model_test_capabilities_payload("grok", "grok-imagine-image-edit", true);
|
||||||
|
|
||||||
|
assert_eq!(payload["openai:image"]["max_generation_count"], 4);
|
||||||
|
assert_eq!(payload["openai:image"]["supports_generation"], false);
|
||||||
|
assert_eq!(payload["openai:image"]["supports_edit"], true);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn non_image_models_report_null_image_test_capability() {
|
||||||
|
let payload = admin_provider_model_test_capabilities_payload("openai", "gpt-5.5", false);
|
||||||
|
|
||||||
|
assert!(payload["openai:image"].is_null());
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn grok_image_support_uses_catalog_model_ids_not_global_fallback() {
|
||||||
|
assert!(admin_provider_model_supports_image_generation(
|
||||||
|
"grok",
|
||||||
|
"grok-imagine-image-pro",
|
||||||
|
false,
|
||||||
|
));
|
||||||
|
assert!(!admin_provider_model_supports_image_generation(
|
||||||
|
"grok",
|
||||||
|
"grok-4.20-fast",
|
||||||
|
true,
|
||||||
|
));
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -4,9 +4,9 @@ pub(crate) fn normalize_provider_type_input(value: &str) -> Result<String, Strin
|
|||||||
let normalized = value.trim().to_ascii_lowercase();
|
let normalized = value.trim().to_ascii_lowercase();
|
||||||
match normalized.as_str() {
|
match normalized.as_str() {
|
||||||
"custom" | "claude_code" | "kiro" | "codex" | "chatgpt_web" | "gemini_cli"
|
"custom" | "claude_code" | "kiro" | "codex" | "chatgpt_web" | "gemini_cli"
|
||||||
| "antigravity" | "vertex_ai" => Ok(normalized),
|
| "antigravity" | "vertex_ai" | "grok" => Ok(normalized),
|
||||||
_ => Err(
|
_ => Err(
|
||||||
"provider_type 仅支持 custom / claude_code / kiro / codex / chatgpt_web / gemini_cli / antigravity / vertex_ai"
|
"provider_type 仅支持 custom / claude_code / kiro / codex / chatgpt_web / gemini_cli / antigravity / vertex_ai / grok"
|
||||||
.to_string(),
|
.to_string(),
|
||||||
),
|
),
|
||||||
}
|
}
|
||||||
@@ -260,6 +260,14 @@ mod tests {
|
|||||||
);
|
);
|
||||||
}
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn normalize_provider_type_supports_grok() {
|
||||||
|
assert_eq!(
|
||||||
|
normalize_provider_type_input(" Grok ").expect("type should normalize"),
|
||||||
|
"grok"
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
#[test]
|
#[test]
|
||||||
fn normalize_api_format_list_dedupes_canonical_formats() {
|
fn normalize_api_format_list_dedupes_canonical_formats() {
|
||||||
assert_eq!(
|
assert_eq!(
|
||||||
|
|||||||
@@ -205,7 +205,7 @@ pub(crate) async fn build_admin_update_provider_record(
|
|||||||
updated.stream_first_byte_timeout_secs = match payload.stream_first_byte_timeout {
|
updated.stream_first_byte_timeout_secs = match payload.stream_first_byte_timeout {
|
||||||
Some(value) if (1.0..=300.0).contains(&value) => Some(value),
|
Some(value) if (1.0..=300.0).contains(&value) => Some(value),
|
||||||
Some(_) => {
|
Some(_) => {
|
||||||
return Err("stream_first_byte_timeout 必须是 1 到 300 之间的数字".to_string())
|
return Err("stream_first_byte_timeout 必须是 1 到 300 之间的数字".to_string());
|
||||||
}
|
}
|
||||||
None => None,
|
None => None,
|
||||||
};
|
};
|
||||||
|
|||||||
@@ -250,4 +250,19 @@ impl<'a> AdminAppState<'a> {
|
|||||||
crate::execution_runtime::execute_execution_runtime_sync_plan(self.app, trace_id, plan)
|
crate::execution_runtime::execute_execution_runtime_sync_plan(self.app, trace_id, plan)
|
||||||
.await
|
.await
|
||||||
}
|
}
|
||||||
|
|
||||||
|
pub(crate) async fn execute_execution_runtime_sync_plan_with_report_context(
|
||||||
|
&self,
|
||||||
|
trace_id: Option<&str>,
|
||||||
|
plan: &aether_contracts::ExecutionPlan,
|
||||||
|
report_context: Option<&serde_json::Value>,
|
||||||
|
) -> Result<aether_contracts::ExecutionResult, GatewayError> {
|
||||||
|
crate::execution_runtime::execute_execution_runtime_sync_plan_with_report_context(
|
||||||
|
self.app,
|
||||||
|
trace_id,
|
||||||
|
plan,
|
||||||
|
report_context,
|
||||||
|
)
|
||||||
|
.await
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -18,8 +18,7 @@ use std::collections::BTreeMap;
|
|||||||
use std::io::Read;
|
use std::io::Read;
|
||||||
use url::Url;
|
use url::Url;
|
||||||
|
|
||||||
const KIRO_IDC_AMZ_USER_AGENT: &str =
|
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";
|
||||||
"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";
|
|
||||||
const ADMIN_PROVIDER_OAUTH_TIMEOUT_MS: u64 = 30_000;
|
const ADMIN_PROVIDER_OAUTH_TIMEOUT_MS: u64 = 30_000;
|
||||||
const ADMIN_PROVIDER_OAUTH_PROXY_TIMEOUT_MS: u64 = 60_000;
|
const ADMIN_PROVIDER_OAUTH_PROXY_TIMEOUT_MS: u64 = 60_000;
|
||||||
|
|
||||||
|
|||||||
@@ -84,7 +84,7 @@ impl<'a> AdminAppState<'a> {
|
|||||||
Json(json!({ "detail": "请求数据验证失败" })),
|
Json(json!({ "detail": "请求数据验证失败" })),
|
||||||
)
|
)
|
||||||
.into_response(),
|
.into_response(),
|
||||||
))
|
));
|
||||||
}
|
}
|
||||||
};
|
};
|
||||||
|
|
||||||
|
|||||||
@@ -1197,7 +1197,7 @@ impl<'a> AdminAppState<'a> {
|
|||||||
Err(_) => {
|
Err(_) => {
|
||||||
return Ok(Err(invalid_request(format!(
|
return Ok(Err(invalid_request(format!(
|
||||||
"Provider '{provider_name}' 配置格式无效"
|
"Provider '{provider_name}' 配置格式无效"
|
||||||
))))
|
))));
|
||||||
}
|
}
|
||||||
};
|
};
|
||||||
let mut updated = invalid!(
|
let mut updated = invalid!(
|
||||||
@@ -1230,7 +1230,7 @@ impl<'a> AdminAppState<'a> {
|
|||||||
Err(_) => {
|
Err(_) => {
|
||||||
return Ok(Err(invalid_request(format!(
|
return Ok(Err(invalid_request(format!(
|
||||||
"Provider '{provider_name}' 配置格式无效"
|
"Provider '{provider_name}' 配置格式无效"
|
||||||
))))
|
))));
|
||||||
}
|
}
|
||||||
};
|
};
|
||||||
let (mut record, shift_existing_priorities_from) =
|
let (mut record, shift_existing_priorities_from) =
|
||||||
@@ -1301,7 +1301,7 @@ impl<'a> AdminAppState<'a> {
|
|||||||
Err(_) => {
|
Err(_) => {
|
||||||
return Ok(Err(invalid_request(
|
return Ok(Err(invalid_request(
|
||||||
"Provider Endpoint 配置格式无效",
|
"Provider Endpoint 配置格式无效",
|
||||||
)))
|
)));
|
||||||
}
|
}
|
||||||
};
|
};
|
||||||
let (fields, payload) = patch.into_parts();
|
let (fields, payload) = patch.into_parts();
|
||||||
@@ -1504,7 +1504,7 @@ impl<'a> AdminAppState<'a> {
|
|||||||
) {
|
) {
|
||||||
Ok(patch) => patch,
|
Ok(patch) => patch,
|
||||||
Err(_) => {
|
Err(_) => {
|
||||||
return Ok(Err(invalid_request("Provider Key 配置格式无效")))
|
return Ok(Err(invalid_request("Provider Key 配置格式无效")));
|
||||||
}
|
}
|
||||||
};
|
};
|
||||||
let mut updated = invalid!(
|
let mut updated = invalid!(
|
||||||
@@ -1985,7 +1985,7 @@ impl<'a> AdminAppState<'a> {
|
|||||||
Err(_) => {
|
Err(_) => {
|
||||||
return Ok(Err(invalid_request(
|
return Ok(Err(invalid_request(
|
||||||
"merge_mode 仅支持 skip / overwrite / error",
|
"merge_mode 仅支持 skip / overwrite / error",
|
||||||
)))
|
)));
|
||||||
}
|
}
|
||||||
};
|
};
|
||||||
let empty = Vec::new();
|
let empty = Vec::new();
|
||||||
|
|||||||
@@ -1,7 +1,7 @@
|
|||||||
use crate::handlers::shared::{json_string_list, unix_secs_to_rfc3339};
|
use crate::handlers::shared::{json_string_list, unix_secs_to_rfc3339};
|
||||||
use crate::provider_key_auth::{
|
use crate::provider_key_auth::{
|
||||||
provider_key_auth_semantics, provider_key_configured_api_formats,
|
provider_key_auth_semantics, provider_key_can_refresh_oauth,
|
||||||
provider_key_inherits_provider_api_formats,
|
provider_key_configured_api_formats, provider_key_inherits_provider_api_formats,
|
||||||
};
|
};
|
||||||
use crate::AppState;
|
use crate::AppState;
|
||||||
use aether_admin::provider::quota as admin_provider_quota_pure;
|
use aether_admin::provider::quota as admin_provider_quota_pure;
|
||||||
@@ -10,6 +10,9 @@ use aether_admin::provider::status as admin_provider_status_pure;
|
|||||||
use aether_crypto::DEVELOPMENT_ENCRYPTION_KEY;
|
use aether_crypto::DEVELOPMENT_ENCRYPTION_KEY;
|
||||||
use aether_crypto::{decrypt_python_fernet_ciphertext, encrypt_python_fernet_plaintext};
|
use aether_crypto::{decrypt_python_fernet_ciphertext, encrypt_python_fernet_plaintext};
|
||||||
use aether_data_contracts::repository::provider_catalog::StoredProviderCatalogKey;
|
use aether_data_contracts::repository::provider_catalog::StoredProviderCatalogKey;
|
||||||
|
use aether_provider_pool::{
|
||||||
|
grok_pool_tier_from_quota_bucket, grok_supported_quota_windows_for_tier,
|
||||||
|
};
|
||||||
use serde_json::{json, Map, Value};
|
use serde_json::{json, Map, Value};
|
||||||
use std::borrow::Cow;
|
use std::borrow::Cow;
|
||||||
use std::time::{SystemTime, UNIX_EPOCH};
|
use std::time::{SystemTime, UNIX_EPOCH};
|
||||||
@@ -461,6 +464,24 @@ fn model_quota_window_snapshot(
|
|||||||
item: &Map<String, Value>,
|
item: &Map<String, Value>,
|
||||||
observed_at_unix_secs: Option<u64>,
|
observed_at_unix_secs: Option<u64>,
|
||||||
) -> Option<Value> {
|
) -> Option<Value> {
|
||||||
|
let remaining_value = item
|
||||||
|
.get("remaining")
|
||||||
|
.or_else(|| item.get("remaining_value"))
|
||||||
|
.and_then(admin_provider_quota_pure::coerce_json_f64);
|
||||||
|
let limit_value = item
|
||||||
|
.get("total")
|
||||||
|
.or_else(|| item.get("limit_value"))
|
||||||
|
.and_then(admin_provider_quota_pure::coerce_json_f64)
|
||||||
|
.filter(|value| *value > 0.0);
|
||||||
|
let used_value = item
|
||||||
|
.get("used")
|
||||||
|
.or_else(|| item.get("used_value"))
|
||||||
|
.and_then(admin_provider_quota_pure::coerce_json_f64)
|
||||||
|
.or_else(|| {
|
||||||
|
remaining_value
|
||||||
|
.zip(limit_value)
|
||||||
|
.map(|(remaining, limit)| (limit - remaining).max(0.0))
|
||||||
|
});
|
||||||
let used_ratio = item
|
let used_ratio = item
|
||||||
.get("used_percent")
|
.get("used_percent")
|
||||||
.and_then(admin_provider_quota_pure::coerce_json_f64)
|
.and_then(admin_provider_quota_pure::coerce_json_f64)
|
||||||
@@ -489,6 +510,8 @@ fn model_quota_window_snapshot(
|
|||||||
&& reset_at.is_none()
|
&& reset_at.is_none()
|
||||||
&& reset_seconds.is_none()
|
&& reset_seconds.is_none()
|
||||||
&& is_exhausted.is_none()
|
&& is_exhausted.is_none()
|
||||||
|
&& remaining_value.is_none()
|
||||||
|
&& limit_value.is_none()
|
||||||
{
|
{
|
||||||
return None;
|
return None;
|
||||||
}
|
}
|
||||||
@@ -507,12 +530,29 @@ fn model_quota_window_snapshot(
|
|||||||
window.insert("model".to_string(), json!(model_name));
|
window.insert("model".to_string(), json!(model_name));
|
||||||
window.insert("used_ratio".to_string(), json!(used_ratio));
|
window.insert("used_ratio".to_string(), json!(used_ratio));
|
||||||
window.insert("remaining_ratio".to_string(), json!(remaining_ratio));
|
window.insert("remaining_ratio".to_string(), json!(remaining_ratio));
|
||||||
|
window.insert("used_value".to_string(), json!(used_value));
|
||||||
|
window.insert("remaining_value".to_string(), json!(remaining_value));
|
||||||
|
window.insert("limit_value".to_string(), json!(limit_value));
|
||||||
window.insert("reset_at".to_string(), json!(reset_at));
|
window.insert("reset_at".to_string(), json!(reset_at));
|
||||||
window.insert("reset_seconds".to_string(), json!(reset_seconds));
|
window.insert("reset_seconds".to_string(), json!(reset_seconds));
|
||||||
window.insert("is_exhausted".to_string(), json!(is_exhausted));
|
window.insert("is_exhausted".to_string(), json!(is_exhausted));
|
||||||
Some(Value::Object(window))
|
Some(Value::Object(window))
|
||||||
}
|
}
|
||||||
|
|
||||||
|
fn provider_quota_metadata_string(
|
||||||
|
metadata: &Map<String, Value>,
|
||||||
|
fields: &[&str],
|
||||||
|
) -> Option<String> {
|
||||||
|
fields.iter().find_map(|field| {
|
||||||
|
metadata
|
||||||
|
.get(*field)
|
||||||
|
.and_then(Value::as_str)
|
||||||
|
.map(str::trim)
|
||||||
|
.filter(|value| !value.is_empty())
|
||||||
|
.map(ToOwned::to_owned)
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
fn quota_windows_usage_ratio(windows: &[Value]) -> Option<f64> {
|
fn quota_windows_usage_ratio(windows: &[Value]) -> Option<f64> {
|
||||||
windows
|
windows
|
||||||
.iter()
|
.iter()
|
||||||
@@ -1126,6 +1166,72 @@ fn build_antigravity_quota_status_snapshot(
|
|||||||
}))
|
}))
|
||||||
}
|
}
|
||||||
|
|
||||||
|
fn build_grok_quota_status_snapshot(
|
||||||
|
upstream_metadata: Option<&Value>,
|
||||||
|
source: &str,
|
||||||
|
) -> Option<Value> {
|
||||||
|
let metadata = provider_quota_metadata_bucket(upstream_metadata, "grok")?;
|
||||||
|
let observed_at_unix_secs = provider_quota_timestamp_unix_secs(metadata.get("updated_at"));
|
||||||
|
let inferred_pool_tier = grok_pool_tier_from_quota_bucket(metadata);
|
||||||
|
let pool_tier = provider_quota_metadata_string(metadata, &["pool_tier", "tier"])
|
||||||
|
.or_else(|| inferred_pool_tier.map(ToOwned::to_owned));
|
||||||
|
let plan_type = provider_quota_metadata_string(metadata, &["plan_type", "plan"])
|
||||||
|
.or_else(|| pool_tier.clone());
|
||||||
|
let supported_windows = grok_supported_quota_windows_for_tier(pool_tier.as_deref());
|
||||||
|
let windows = provider_quota_model_bucket(metadata)
|
||||||
|
.map(|models| {
|
||||||
|
models
|
||||||
|
.iter()
|
||||||
|
.filter_map(|(model_name, item)| {
|
||||||
|
if !supported_windows
|
||||||
|
.iter()
|
||||||
|
.any(|(quota_key, _)| *quota_key == model_name.as_str())
|
||||||
|
{
|
||||||
|
return None;
|
||||||
|
}
|
||||||
|
model_quota_window_snapshot(
|
||||||
|
model_name,
|
||||||
|
item.as_object()?,
|
||||||
|
observed_at_unix_secs,
|
||||||
|
)
|
||||||
|
})
|
||||||
|
.collect::<Vec<_>>()
|
||||||
|
})
|
||||||
|
.unwrap_or_default();
|
||||||
|
|
||||||
|
if windows.is_empty() && observed_at_unix_secs.is_none() {
|
||||||
|
return None;
|
||||||
|
}
|
||||||
|
|
||||||
|
let usage_ratio = quota_windows_usage_ratio(&windows);
|
||||||
|
let reset_seconds = quota_windows_min_reset_seconds(&windows);
|
||||||
|
let reset_at = quota_windows_min_reset_at(&windows);
|
||||||
|
let exhausted = quota_windows_all_exhausted(&windows);
|
||||||
|
|
||||||
|
Some(json!({
|
||||||
|
"version": 2,
|
||||||
|
"provider_type": "grok",
|
||||||
|
"code": if exhausted { "exhausted" } else { "ok" },
|
||||||
|
"label": if exhausted { Some("额度耗尽") } else { None::<&str> },
|
||||||
|
"reason": if exhausted {
|
||||||
|
Some("所有 Grok 模式额度已耗尽")
|
||||||
|
} else {
|
||||||
|
None::<&str>
|
||||||
|
},
|
||||||
|
"freshness": "fresh",
|
||||||
|
"source": source,
|
||||||
|
"observed_at": observed_at_unix_secs,
|
||||||
|
"exhausted": exhausted,
|
||||||
|
"usage_ratio": usage_ratio,
|
||||||
|
"updated_at": observed_at_unix_secs,
|
||||||
|
"reset_at": reset_at,
|
||||||
|
"reset_seconds": reset_seconds,
|
||||||
|
"plan_type": plan_type,
|
||||||
|
"pool_tier": pool_tier,
|
||||||
|
"windows": windows,
|
||||||
|
}))
|
||||||
|
}
|
||||||
|
|
||||||
fn build_gemini_cli_quota_status_snapshot(
|
fn build_gemini_cli_quota_status_snapshot(
|
||||||
upstream_metadata: Option<&Value>,
|
upstream_metadata: Option<&Value>,
|
||||||
source: &str,
|
source: &str,
|
||||||
@@ -1228,6 +1334,7 @@ pub(crate) fn sync_provider_key_quota_status_snapshot(
|
|||||||
"kiro" => build_kiro_quota_status_snapshot(upstream_metadata, source),
|
"kiro" => build_kiro_quota_status_snapshot(upstream_metadata, source),
|
||||||
"chatgpt_web" => build_chatgpt_web_quota_status_snapshot(upstream_metadata, source),
|
"chatgpt_web" => build_chatgpt_web_quota_status_snapshot(upstream_metadata, source),
|
||||||
"antigravity" => build_antigravity_quota_status_snapshot(upstream_metadata, source),
|
"antigravity" => build_antigravity_quota_status_snapshot(upstream_metadata, source),
|
||||||
|
"grok" => build_grok_quota_status_snapshot(upstream_metadata, source),
|
||||||
"gemini_cli" => build_gemini_cli_quota_status_snapshot(upstream_metadata, source),
|
"gemini_cli" => build_gemini_cli_quota_status_snapshot(upstream_metadata, source),
|
||||||
_ => None,
|
_ => None,
|
||||||
}?;
|
}?;
|
||||||
@@ -1590,7 +1697,10 @@ pub(crate) fn build_admin_provider_key_response(
|
|||||||
);
|
);
|
||||||
payload.insert(
|
payload.insert(
|
||||||
"can_refresh_oauth".to_string(),
|
"can_refresh_oauth".to_string(),
|
||||||
json!(auth_semantics.can_refresh_oauth()),
|
json!(provider_key_can_refresh_oauth(
|
||||||
|
auth_semantics,
|
||||||
|
auth_config.as_ref()
|
||||||
|
)),
|
||||||
);
|
);
|
||||||
payload.insert(
|
payload.insert(
|
||||||
"can_export_oauth".to_string(),
|
"can_export_oauth".to_string(),
|
||||||
@@ -2099,6 +2209,69 @@ mod tests {
|
|||||||
assert_eq!(window.get("remaining_ratio"), Some(&json!(0.96)));
|
assert_eq!(window.get("remaining_ratio"), Some(&json!(0.96)));
|
||||||
}
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn provider_key_status_snapshot_payload_backfills_grok_model_quota() {
|
||||||
|
let mut key = sample_catalog_key();
|
||||||
|
key.upstream_metadata = Some(json!({
|
||||||
|
"grok": {
|
||||||
|
"updated_at": 1_778_067_246u64,
|
||||||
|
"pool_tier": "heavy",
|
||||||
|
"plan_type": "heavy",
|
||||||
|
"quota_by_model": {
|
||||||
|
"quota_auto": {
|
||||||
|
"display_name": "auto",
|
||||||
|
"remaining_fraction": 0.4,
|
||||||
|
"used_percent": 60.0,
|
||||||
|
"remaining": 60.0,
|
||||||
|
"total": 150.0,
|
||||||
|
"reset_at": 1_778_157_172u64,
|
||||||
|
"is_exhausted": false
|
||||||
|
},
|
||||||
|
"quota_heavy": {
|
||||||
|
"display_name": "heavy",
|
||||||
|
"remaining_fraction": 0.0,
|
||||||
|
"used_percent": 100.0,
|
||||||
|
"reset_at": 1_778_157_172u64,
|
||||||
|
"is_exhausted": true
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}));
|
||||||
|
|
||||||
|
let payload = provider_key_status_snapshot_payload(&key, "grok");
|
||||||
|
let quota = payload
|
||||||
|
.get("quota")
|
||||||
|
.and_then(Value::as_object)
|
||||||
|
.expect("quota snapshot should be object");
|
||||||
|
let windows = quota
|
||||||
|
.get("windows")
|
||||||
|
.and_then(Value::as_array)
|
||||||
|
.expect("grok quota windows should exist");
|
||||||
|
|
||||||
|
assert_eq!(quota.get("provider_type"), Some(&json!("grok")));
|
||||||
|
assert_eq!(quota.get("code"), Some(&json!("ok")));
|
||||||
|
assert_eq!(quota.get("plan_type"), Some(&json!("heavy")));
|
||||||
|
assert_eq!(quota.get("pool_tier"), Some(&json!("heavy")));
|
||||||
|
assert_eq!(quota.get("exhausted"), Some(&json!(false)));
|
||||||
|
assert_eq!(quota.get("usage_ratio"), Some(&json!(1.0)));
|
||||||
|
assert_eq!(quota.get("reset_at"), Some(&json!(1_778_157_172u64)));
|
||||||
|
assert_eq!(windows.len(), 2);
|
||||||
|
assert!(windows.iter().any(|window| {
|
||||||
|
window
|
||||||
|
.get("code")
|
||||||
|
.and_then(Value::as_str)
|
||||||
|
.is_some_and(|code| code == "model:quota_auto")
|
||||||
|
}));
|
||||||
|
let auto = windows
|
||||||
|
.iter()
|
||||||
|
.filter_map(Value::as_object)
|
||||||
|
.find(|window| window.get("code") == Some(&json!("model:quota_auto")))
|
||||||
|
.expect("auto quota window should exist");
|
||||||
|
assert_eq!(auto.get("remaining_value"), Some(&json!(60.0)));
|
||||||
|
assert_eq!(auto.get("limit_value"), Some(&json!(150.0)));
|
||||||
|
assert_eq!(auto.get("used_value"), Some(&json!(90.0)));
|
||||||
|
}
|
||||||
|
|
||||||
#[test]
|
#[test]
|
||||||
fn provider_key_status_snapshot_payload_preserves_existing_materialized_quota_snapshot() {
|
fn provider_key_status_snapshot_payload_preserves_existing_materialized_quota_snapshot() {
|
||||||
let mut key = sample_catalog_key();
|
let mut key = sample_catalog_key();
|
||||||
|
|||||||
@@ -4,6 +4,7 @@ use aether_data_contracts::repository::provider_catalog::{
|
|||||||
use aether_provider_transport::provider_types::{
|
use aether_provider_transport::provider_types::{
|
||||||
fixed_provider_key_inherits_api_formats, provider_type_is_fixed,
|
fixed_provider_key_inherits_api_formats, provider_type_is_fixed,
|
||||||
};
|
};
|
||||||
|
use serde_json::{Map, Value};
|
||||||
use std::collections::BTreeSet;
|
use std::collections::BTreeSet;
|
||||||
|
|
||||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||||
@@ -81,6 +82,18 @@ impl ProviderKeyAuthSemantics {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
pub(crate) fn provider_key_can_refresh_oauth(
|
||||||
|
auth_semantics: ProviderKeyAuthSemantics,
|
||||||
|
auth_config: Option<&Map<String, Value>>,
|
||||||
|
) -> bool {
|
||||||
|
auth_semantics.can_refresh_oauth()
|
||||||
|
&& auth_config
|
||||||
|
.and_then(|config| config.get("refresh_token"))
|
||||||
|
.and_then(Value::as_str)
|
||||||
|
.map(str::trim)
|
||||||
|
.is_some_and(|value| !value.is_empty())
|
||||||
|
}
|
||||||
|
|
||||||
fn normalized_auth_type(key: &StoredProviderCatalogKey) -> String {
|
fn normalized_auth_type(key: &StoredProviderCatalogKey) -> String {
|
||||||
key.auth_type.trim().to_ascii_lowercase()
|
key.auth_type.trim().to_ascii_lowercase()
|
||||||
}
|
}
|
||||||
@@ -106,6 +119,10 @@ fn provider_uses_bearer_oauth_runtime(provider_type: &str) -> bool {
|
|||||||
)
|
)
|
||||||
}
|
}
|
||||||
|
|
||||||
|
fn provider_uses_grok_session_runtime(provider_type: &str) -> bool {
|
||||||
|
provider_type.trim().eq_ignore_ascii_case("grok")
|
||||||
|
}
|
||||||
|
|
||||||
fn provider_key_is_legacy_kiro_oauth_session(
|
fn provider_key_is_legacy_kiro_oauth_session(
|
||||||
key: &StoredProviderCatalogKey,
|
key: &StoredProviderCatalogKey,
|
||||||
provider_type: &str,
|
provider_type: &str,
|
||||||
@@ -122,7 +139,8 @@ pub(crate) fn provider_key_auth_semantics(
|
|||||||
) -> ProviderKeyAuthSemantics {
|
) -> ProviderKeyAuthSemantics {
|
||||||
let auth_type = normalized_auth_type(key);
|
let auth_type = normalized_auth_type(key);
|
||||||
let oauth_managed = auth_type == "oauth"
|
let oauth_managed = auth_type == "oauth"
|
||||||
|| provider_key_is_legacy_kiro_oauth_session(key, provider_type, &auth_type);
|
|| provider_key_is_legacy_kiro_oauth_session(key, provider_type, &auth_type)
|
||||||
|
|| (provider_uses_grok_session_runtime(provider_type) && key_has_auth_config(key));
|
||||||
let credential_kind = if oauth_managed {
|
let credential_kind = if oauth_managed {
|
||||||
ProviderKeyCredentialKind::OAuthSession
|
ProviderKeyCredentialKind::OAuthSession
|
||||||
} else if matches!(auth_type.as_str(), "service_account" | "vertex_ai") {
|
} else if matches!(auth_type.as_str(), "service_account" | "vertex_ai") {
|
||||||
@@ -135,6 +153,8 @@ pub(crate) fn provider_key_auth_semantics(
|
|||||||
ProviderKeyCredentialKind::OAuthSession => {
|
ProviderKeyCredentialKind::OAuthSession => {
|
||||||
if provider_uses_bearer_oauth_runtime(provider_type) {
|
if provider_uses_bearer_oauth_runtime(provider_type) {
|
||||||
ProviderKeyRuntimeAuthKind::Bearer
|
ProviderKeyRuntimeAuthKind::Bearer
|
||||||
|
} else if provider_uses_grok_session_runtime(provider_type) {
|
||||||
|
ProviderKeyRuntimeAuthKind::Unknown
|
||||||
} else {
|
} else {
|
||||||
ProviderKeyRuntimeAuthKind::Unknown
|
ProviderKeyRuntimeAuthKind::Unknown
|
||||||
}
|
}
|
||||||
@@ -226,7 +246,7 @@ pub(crate) fn provider_key_effective_api_formats(
|
|||||||
#[cfg(test)]
|
#[cfg(test)]
|
||||||
mod tests {
|
mod tests {
|
||||||
use super::{
|
use super::{
|
||||||
provider_active_api_formats, provider_key_auth_semantics,
|
provider_active_api_formats, provider_key_auth_semantics, provider_key_can_refresh_oauth,
|
||||||
provider_key_configured_api_formats, provider_key_effective_api_formats,
|
provider_key_configured_api_formats, provider_key_effective_api_formats,
|
||||||
provider_key_inherits_provider_api_formats, ProviderKeyCredentialKind,
|
provider_key_inherits_provider_api_formats, ProviderKeyCredentialKind,
|
||||||
ProviderKeyRuntimeAuthKind,
|
ProviderKeyRuntimeAuthKind,
|
||||||
@@ -293,6 +313,46 @@ mod tests {
|
|||||||
);
|
);
|
||||||
}
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn recognizes_grok_oauth_session_as_managed_without_bearer_runtime() {
|
||||||
|
let mut key = sample_key("oauth");
|
||||||
|
key.encrypted_auth_config = Some(r#"{"sso_token":"abc"}"#.to_string());
|
||||||
|
|
||||||
|
let semantics = provider_key_auth_semantics(&key, "grok");
|
||||||
|
|
||||||
|
assert!(semantics.oauth_managed());
|
||||||
|
assert_eq!(
|
||||||
|
semantics.credential_kind(),
|
||||||
|
ProviderKeyCredentialKind::OAuthSession
|
||||||
|
);
|
||||||
|
assert_eq!(
|
||||||
|
semantics.runtime_auth_kind(),
|
||||||
|
ProviderKeyRuntimeAuthKind::Unknown
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn refresh_capability_requires_stored_refresh_token() {
|
||||||
|
let semantics = provider_key_auth_semantics(&sample_key("oauth"), "codex");
|
||||||
|
|
||||||
|
assert!(!provider_key_can_refresh_oauth(
|
||||||
|
semantics,
|
||||||
|
json!({
|
||||||
|
"access_token": "access-token",
|
||||||
|
"access_token_import_temporary": true
|
||||||
|
})
|
||||||
|
.as_object()
|
||||||
|
));
|
||||||
|
assert!(!provider_key_can_refresh_oauth(
|
||||||
|
semantics,
|
||||||
|
json!({ "refresh_token": " " }).as_object()
|
||||||
|
));
|
||||||
|
assert!(provider_key_can_refresh_oauth(
|
||||||
|
semantics,
|
||||||
|
json!({ "refresh_token": "refresh-token" }).as_object()
|
||||||
|
));
|
||||||
|
}
|
||||||
|
|
||||||
#[test]
|
#[test]
|
||||||
fn recognizes_legacy_kiro_bearer_key_with_auth_config_as_oauth_managed() {
|
fn recognizes_legacy_kiro_bearer_key_with_auth_config_as_oauth_managed() {
|
||||||
let mut key = sample_key("bearer");
|
let mut key = sample_key("bearer");
|
||||||
|
|||||||
@@ -94,6 +94,7 @@ async fn gateway_handles_admin_provider_models_locally_with_trusted_admin_princi
|
|||||||
assert_eq!(items[0]["effective_input_price"], 3.0);
|
assert_eq!(items[0]["effective_input_price"], 3.0);
|
||||||
assert_eq!(items[0]["effective_output_price"], 15.0);
|
assert_eq!(items[0]["effective_output_price"], 15.0);
|
||||||
assert_eq!(items[0]["effective_supports_streaming"], true);
|
assert_eq!(items[0]["effective_supports_streaming"], true);
|
||||||
|
assert!(items[0]["model_test_capabilities"]["openai:image"].is_null());
|
||||||
assert_eq!(items[0]["created_at"], "2024-03-21T05:46:40Z");
|
assert_eq!(items[0]["created_at"], "2024-03-21T05:46:40Z");
|
||||||
assert_eq!(items[0]["updated_at"], "2024-03-21T05:48:20Z");
|
assert_eq!(items[0]["updated_at"], "2024-03-21T05:48:20Z");
|
||||||
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
|
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
|
||||||
|
|||||||
@@ -1967,6 +1967,144 @@ async fn gateway_streams_codex_openai_responses_upstream_for_admin_pool_model_te
|
|||||||
execution_runtime_handle.abort();
|
execution_runtime_handle.abort();
|
||||||
}
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn gateway_routes_grok_responses_admin_pool_model_test_through_grok_runtime() {
|
||||||
|
let execution_runtime = Router::new().route(
|
||||||
|
"/v1/execute/sync",
|
||||||
|
any(move |Json(plan): Json<ExecutionPlan>| async move {
|
||||||
|
assert_eq!(plan.provider_id, "provider-grok");
|
||||||
|
assert_eq!(plan.endpoint_id, "endpoint-grok-responses");
|
||||||
|
assert_eq!(plan.key_id, "key-grok-oauth");
|
||||||
|
assert_eq!(plan.client_api_format, "openai:responses");
|
||||||
|
assert_eq!(plan.provider_api_format, "openai:responses");
|
||||||
|
assert_eq!(plan.url, "https://grok.com/rest/app-chat/conversations/new");
|
||||||
|
assert_eq!(plan.model_name.as_deref(), Some("grok-4.20-fast"));
|
||||||
|
assert!(plan.stream, "Grok model test should request a stream");
|
||||||
|
assert_eq!(
|
||||||
|
plan.headers
|
||||||
|
.get(aether_provider_transport::GROK_INTERNAL_HEADER)
|
||||||
|
.map(String::as_str),
|
||||||
|
Some("1")
|
||||||
|
);
|
||||||
|
assert_eq!(
|
||||||
|
plan.headers.get("cookie").map(String::as_str),
|
||||||
|
Some("sso=grok-sso; sso-rw=grok-rw")
|
||||||
|
);
|
||||||
|
let body = plan.body.json_body.as_ref().expect("json body");
|
||||||
|
assert_eq!(body["model"], json!("grok-4.20-fast"));
|
||||||
|
assert_eq!(body["input"], json!("Hello! This is a test message."));
|
||||||
|
assert_eq!(
|
||||||
|
body["messages"][0]["content"],
|
||||||
|
json!("stale chat-shaped frontend body")
|
||||||
|
);
|
||||||
|
Json(json!({
|
||||||
|
"request_id": plan.request_id,
|
||||||
|
"candidate_id": plan.candidate_id,
|
||||||
|
"status_code": 200,
|
||||||
|
"headers": {
|
||||||
|
"content-type": "application/json"
|
||||||
|
},
|
||||||
|
"body": {
|
||||||
|
"json_body": {
|
||||||
|
"id": "resp-grok-model-test",
|
||||||
|
"model": "grok-4.20-fast",
|
||||||
|
"output_text": "ok"
|
||||||
|
}
|
||||||
|
},
|
||||||
|
"telemetry": {
|
||||||
|
"elapsed_ms": 18
|
||||||
|
}
|
||||||
|
}))
|
||||||
|
}),
|
||||||
|
);
|
||||||
|
|
||||||
|
let (execution_runtime_url, execution_runtime_handle) = start_server(execution_runtime).await;
|
||||||
|
let mut provider = sample_provider("provider-grok", "Grok", 10);
|
||||||
|
provider.provider_type = "grok".to_string();
|
||||||
|
provider.config = Some(json!({"pool_advanced": {}}));
|
||||||
|
let mut key = sample_key(
|
||||||
|
"key-grok-oauth",
|
||||||
|
"provider-grok",
|
||||||
|
"openai:responses",
|
||||||
|
"__placeholder__",
|
||||||
|
);
|
||||||
|
key.auth_type = "oauth".to_string();
|
||||||
|
key.encrypted_auth_config = Some(
|
||||||
|
aether_crypto::encrypt_python_fernet_plaintext(
|
||||||
|
DEVELOPMENT_ENCRYPTION_KEY,
|
||||||
|
r#"{
|
||||||
|
"provider_type":"grok",
|
||||||
|
"sso_token":"grok-sso",
|
||||||
|
"sso_rw_token":"grok-rw"
|
||||||
|
}"#,
|
||||||
|
)
|
||||||
|
.expect("auth config should encrypt"),
|
||||||
|
);
|
||||||
|
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
|
||||||
|
vec![provider],
|
||||||
|
vec![sample_endpoint(
|
||||||
|
"endpoint-grok-responses",
|
||||||
|
"provider-grok",
|
||||||
|
"openai:responses",
|
||||||
|
"https://grok.com",
|
||||||
|
)],
|
||||||
|
vec![key],
|
||||||
|
));
|
||||||
|
|
||||||
|
let gateway = build_router_with_state(
|
||||||
|
build_state_with_execution_runtime_override(execution_runtime_url)
|
||||||
|
.with_data_state_for_tests(GatewayDataState::with_provider_transport_reader_for_tests(
|
||||||
|
provider_catalog_repository,
|
||||||
|
DEVELOPMENT_ENCRYPTION_KEY.to_string(),
|
||||||
|
)),
|
||||||
|
);
|
||||||
|
let (gateway_url, gateway_handle) = start_server(gateway).await;
|
||||||
|
|
||||||
|
let response = reqwest::Client::new()
|
||||||
|
.post(format!(
|
||||||
|
"{gateway_url}/api/admin/provider-query/test-model-failover"
|
||||||
|
))
|
||||||
|
.header(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")
|
||||||
|
.json(&json!({
|
||||||
|
"provider_id": "provider-grok",
|
||||||
|
"mode": "pool",
|
||||||
|
"model": "grok-4.20-fast",
|
||||||
|
"failover_models": ["grok-4.20-fast"],
|
||||||
|
"api_format": "openai:responses",
|
||||||
|
"endpoint_id": "endpoint-grok-responses",
|
||||||
|
"request_body": {
|
||||||
|
"model": "grok-4.20-fast",
|
||||||
|
"messages": [{
|
||||||
|
"role": "user",
|
||||||
|
"content": "stale chat-shaped frontend body"
|
||||||
|
}],
|
||||||
|
"stream": true
|
||||||
|
}
|
||||||
|
}))
|
||||||
|
.send()
|
||||||
|
.await
|
||||||
|
.expect("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["success"], json!(true));
|
||||||
|
assert_eq!(payload["attempts"][0]["status"], json!("success"));
|
||||||
|
assert_eq!(
|
||||||
|
payload["attempts"][0]["request_body"]["message"],
|
||||||
|
json!("Hello! This is a test message.")
|
||||||
|
);
|
||||||
|
assert_eq!(
|
||||||
|
payload["attempts"][0]["request_headers"][aether_provider_transport::GROK_INTERNAL_HEADER],
|
||||||
|
json!("1")
|
||||||
|
);
|
||||||
|
|
||||||
|
gateway_handle.abort();
|
||||||
|
execution_runtime_handle.abort();
|
||||||
|
}
|
||||||
|
|
||||||
#[tokio::test]
|
#[tokio::test]
|
||||||
async fn gateway_uses_pool_scheduler_order_for_admin_pool_model_test() {
|
async fn gateway_uses_pool_scheduler_order_for_admin_pool_model_test() {
|
||||||
let execution_runtime = Router::new().route(
|
let execution_runtime = Router::new().route(
|
||||||
|
|||||||
@@ -1,4 +1,5 @@
|
|||||||
import apiClient from './client'
|
import apiClient from './client'
|
||||||
|
import type { ModelTestCapabilities } from './endpoints/types'
|
||||||
import axios from 'axios'
|
import axios from 'axios'
|
||||||
import { cachedRequest, buildCacheKey } from '@/utils/cache'
|
import { cachedRequest, buildCacheKey } from '@/utils/cache'
|
||||||
import type { BillingSummary } from './auth'
|
import type { BillingSummary } from './auth'
|
||||||
@@ -420,6 +421,8 @@ export interface ProviderModelsQueryResponse {
|
|||||||
owned_by?: string
|
owned_by?: string
|
||||||
display_name?: string
|
display_name?: string
|
||||||
api_format?: string
|
api_format?: string
|
||||||
|
api_formats?: string[]
|
||||||
|
model_test_capabilities?: ModelTestCapabilities | null
|
||||||
}>
|
}>
|
||||||
error?: string
|
error?: string
|
||||||
from_cache?: boolean
|
from_cache?: boolean
|
||||||
|
|||||||
@@ -112,6 +112,7 @@ export interface PoolPresetMeta {
|
|||||||
export interface PoolKeyDetail {
|
export interface PoolKeyDetail {
|
||||||
key_id: string
|
key_id: string
|
||||||
key_name: string
|
key_name: string
|
||||||
|
provider_type?: string | null
|
||||||
is_active: boolean
|
is_active: boolean
|
||||||
auth_type: string
|
auth_type: string
|
||||||
auth_type_by_format?: Record<string, 'api_key' | 'bearer'> | null
|
auth_type_by_format?: Record<string, 'api_key' | 'bearer'> | null
|
||||||
|
|||||||
@@ -111,6 +111,12 @@ export async function importProviderRefreshToken(
|
|||||||
account_id?: string
|
account_id?: string
|
||||||
account_user_id?: string
|
account_user_id?: string
|
||||||
plan_type?: string
|
plan_type?: string
|
||||||
|
pool_tier?: string
|
||||||
|
sso_rw_token?: string
|
||||||
|
cf_cookies?: string
|
||||||
|
cf_clearance?: string
|
||||||
|
user_agent?: string
|
||||||
|
browser_profile?: string
|
||||||
user_id?: string
|
user_id?: string
|
||||||
account_name?: string
|
account_name?: string
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -91,7 +91,7 @@ export async function updateProvider(
|
|||||||
providerId: string,
|
providerId: string,
|
||||||
data: Partial<{
|
data: Partial<{
|
||||||
name: string
|
name: string
|
||||||
provider_type: 'custom' | 'vertex_ai' | 'claude_code' | 'codex' | 'chatgpt_web' | 'gemini_cli' | 'antigravity' | 'kiro'
|
provider_type: 'custom' | 'vertex_ai' | 'claude_code' | 'codex' | 'chatgpt_web' | 'gemini_cli' | 'antigravity' | 'kiro' | 'grok'
|
||||||
description: string | null
|
description: string | null
|
||||||
website: string
|
website: string
|
||||||
provider_priority: number
|
provider_priority: number
|
||||||
@@ -126,7 +126,7 @@ export async function updateProvider(
|
|||||||
export async function createProvider(
|
export async function createProvider(
|
||||||
data: {
|
data: {
|
||||||
name: string
|
name: string
|
||||||
provider_type?: 'custom' | 'vertex_ai' | 'claude_code' | 'codex' | 'chatgpt_web' | 'gemini_cli' | 'antigravity' | 'kiro'
|
provider_type?: 'custom' | 'vertex_ai' | 'claude_code' | 'codex' | 'chatgpt_web' | 'gemini_cli' | 'antigravity' | 'kiro' | 'grok'
|
||||||
description?: string
|
description?: string
|
||||||
website?: string
|
website?: string
|
||||||
billing_type?: 'monthly_quota' | 'pay_as_you_go' | 'free_tier'
|
billing_type?: 'monthly_quota' | 'pay_as_you_go' | 'free_tier'
|
||||||
|
|||||||
@@ -389,11 +389,25 @@ export interface ChatGPTWebUpstreamMetadata {
|
|||||||
user_id?: string | null
|
user_id?: string | null
|
||||||
}
|
}
|
||||||
|
|
||||||
|
export interface GrokUpstreamMetadata {
|
||||||
|
updated_at?: number // Unix 时间戳(秒)
|
||||||
|
plan_type?: string | null
|
||||||
|
pool_tier?: string | null
|
||||||
|
is_banned?: boolean | null
|
||||||
|
ban_reason?: string | null
|
||||||
|
last_rate_limit_probe_at?: number | null
|
||||||
|
clearance_state?: string | null
|
||||||
|
email?: string | null
|
||||||
|
account_id?: string | null
|
||||||
|
account_user_id?: string | null
|
||||||
|
}
|
||||||
|
|
||||||
export interface UpstreamMetadata {
|
export interface UpstreamMetadata {
|
||||||
codex?: CodexUpstreamMetadata
|
codex?: CodexUpstreamMetadata
|
||||||
antigravity?: AntigravityUpstreamMetadata
|
antigravity?: AntigravityUpstreamMetadata
|
||||||
kiro?: KiroUpstreamMetadata
|
kiro?: KiroUpstreamMetadata
|
||||||
chatgpt_web?: ChatGPTWebUpstreamMetadata
|
chatgpt_web?: ChatGPTWebUpstreamMetadata
|
||||||
|
grok?: GrokUpstreamMetadata
|
||||||
}
|
}
|
||||||
|
|
||||||
// 按格式的健康度数据
|
// 按格式的健康度数据
|
||||||
@@ -512,7 +526,7 @@ export interface PublicEndpointStatusMonitorResponse {
|
|||||||
formats: PublicEndpointStatusMonitor[]
|
formats: PublicEndpointStatusMonitor[]
|
||||||
}
|
}
|
||||||
|
|
||||||
export type ProviderType = 'custom' | 'claude_code' | 'codex' | 'chatgpt_web' | 'gemini_cli' | 'antigravity' | 'kiro' | 'vertex_ai'
|
export type ProviderType = 'custom' | 'claude_code' | 'codex' | 'chatgpt_web' | 'gemini_cli' | 'antigravity' | 'kiro' | 'grok' | 'vertex_ai'
|
||||||
|
|
||||||
export interface ClaudeCodeAdvancedConfig {
|
export interface ClaudeCodeAdvancedConfig {
|
||||||
// 会话数量控制:null/undefined 表示不限制
|
// 会话数量控制:null/undefined 表示不限制
|
||||||
|
|||||||
@@ -64,6 +64,7 @@ export interface QuotaStatusSnapshot {
|
|||||||
reset_at?: number | null
|
reset_at?: number | null
|
||||||
reset_seconds?: number | null
|
reset_seconds?: number | null
|
||||||
plan_type?: string | null
|
plan_type?: string | null
|
||||||
|
pool_tier?: string | null
|
||||||
credits?: QuotaCreditsSnapshot | null
|
credits?: QuotaCreditsSnapshot | null
|
||||||
windows?: QuotaWindowSnapshot[] | null
|
windows?: QuotaWindowSnapshot[] | null
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -51,7 +51,7 @@
|
|||||||
@update:model-value="(v: string) => { selectedProxyNodeId = v; proxyPopoverOpen = false }"
|
@update:model-value="(v: string) => { selectedProxyNodeId = v; proxyPopoverOpen = false }"
|
||||||
/>
|
/>
|
||||||
<p class="text-[10px] text-muted-foreground">
|
<p class="text-[10px] text-muted-foreground">
|
||||||
{{ selectedProxyNodeId ? '授权、刷新、额度查询均走此代理' : '未设置,依次回退到提供商代理 → 系统代理' }}
|
{{ selectedProxyNodeId ? `${providerCredentialActionLabel}、刷新、额度查询均走此代理` : '未设置,依次回退到提供商代理 → 系统代理' }}
|
||||||
</p>
|
</p>
|
||||||
</div>
|
</div>
|
||||||
</PopoverContent>
|
</PopoverContent>
|
||||||
@@ -60,7 +60,10 @@
|
|||||||
|
|
||||||
<div class="space-y-4">
|
<div class="space-y-4">
|
||||||
<!-- Tab 切换 -->
|
<!-- Tab 切换 -->
|
||||||
<div class="flex rounded-lg border border-border p-0.5 bg-muted/30">
|
<div
|
||||||
|
v-if="showAuthorizationMode"
|
||||||
|
class="flex rounded-lg border border-border p-0.5 bg-muted/30"
|
||||||
|
>
|
||||||
<button
|
<button
|
||||||
class="flex-1 px-3 py-1.5 text-xs font-medium rounded-md transition-all"
|
class="flex-1 px-3 py-1.5 text-xs font-medium rounded-md transition-all"
|
||||||
:class="[
|
:class="[
|
||||||
@@ -79,7 +82,7 @@
|
|||||||
: 'text-muted-foreground hover:text-foreground'"
|
: 'text-muted-foreground hover:text-foreground'"
|
||||||
@click="switchMode('import')"
|
@click="switchMode('import')"
|
||||||
>
|
>
|
||||||
导入授权
|
{{ importModeLabel }}
|
||||||
</button>
|
</button>
|
||||||
</div>
|
</div>
|
||||||
|
|
||||||
@@ -458,11 +461,12 @@
|
|||||||
v-model="importText"
|
v-model="importText"
|
||||||
:disabled="importing"
|
:disabled="importing"
|
||||||
:reset-key="importInputResetKey"
|
:reset-key="importInputResetKey"
|
||||||
drop-title="拖入授权文件或点击选择"
|
:drop-title="importDropTitle"
|
||||||
drop-hint="支持 .json / .txt,可多选"
|
:drop-hint="importDropHint"
|
||||||
manual-placeholder="粘贴 Refresh Token / Access Token 或 JSON 内容"
|
:manual-placeholder="importManualPlaceholder"
|
||||||
paste-toggle-text="或手动粘贴 Token"
|
:manual-description="importManualDescription"
|
||||||
file-toggle-text="或选择 JSON 文件导入"
|
:paste-toggle-text="importPasteToggleText"
|
||||||
|
:file-toggle-text="importFileToggleText"
|
||||||
textarea-class="min-h-[200px] text-xs font-mono break-all !rounded-xl"
|
textarea-class="min-h-[200px] text-xs font-mono break-all !rounded-xl"
|
||||||
@error="handleImportInputError"
|
@error="handleImportInputError"
|
||||||
/>
|
/>
|
||||||
@@ -523,7 +527,7 @@
|
|||||||
取消
|
取消
|
||||||
</Button>
|
</Button>
|
||||||
<Button
|
<Button
|
||||||
v-if="mode === 'oauth' && !isKiroProvider"
|
v-if="mode === 'oauth' && showAuthorizationMode && !isKiroProvider"
|
||||||
:disabled="!canCompleteOAuth"
|
:disabled="!canCompleteOAuth"
|
||||||
@click="handleCompleteOAuth"
|
@click="handleCompleteOAuth"
|
||||||
>
|
>
|
||||||
@@ -541,7 +545,7 @@
|
|||||||
:disabled="!canImport"
|
:disabled="!canImport"
|
||||||
@click="handleImport"
|
@click="handleImport"
|
||||||
>
|
>
|
||||||
{{ importing ? (importTask ? `导入中 ${importTask.progress_percent}%` : '导入中...') : '导入' }}
|
{{ importing ? (importTask ? `导入中 ${importTask.progress_percent}%` : '导入中...') : importButtonLabel }}
|
||||||
</Button>
|
</Button>
|
||||||
</template>
|
</template>
|
||||||
</Dialog>
|
</Dialog>
|
||||||
@@ -644,7 +648,7 @@ function getSelectedNodeLabel(): string {
|
|||||||
|
|
||||||
// 模式
|
// 模式
|
||||||
type DialogMode = 'oauth' | 'import'
|
type DialogMode = 'oauth' | 'import'
|
||||||
const mode = ref<DialogMode>('oauth')
|
const mode = ref<DialogMode>((props.providerType || '').toLowerCase() === 'grok' ? 'import' : 'oauth')
|
||||||
|
|
||||||
// OAuth 状态
|
// OAuth 状态
|
||||||
interface OAuthState {
|
interface OAuthState {
|
||||||
@@ -736,6 +740,9 @@ const importPolling = ref(false)
|
|||||||
const isOpen = computed(() => props.open)
|
const isOpen = computed(() => props.open)
|
||||||
|
|
||||||
const isKiroProvider = computed(() => (props.providerType || '').toLowerCase() === 'kiro')
|
const isKiroProvider = computed(() => (props.providerType || '').toLowerCase() === 'kiro')
|
||||||
|
const isGrokProvider = computed(() => (props.providerType || '').toLowerCase() === 'grok')
|
||||||
|
const showAuthorizationMode = computed(() => !isGrokProvider.value)
|
||||||
|
const defaultMode = computed<DialogMode>(() => (isGrokProvider.value ? 'import' : 'oauth'))
|
||||||
|
|
||||||
const isSocialDeviceAuth = computed(() =>
|
const isSocialDeviceAuth = computed(() =>
|
||||||
device.value.auth_type === 'google' || device.value.auth_type === 'github'
|
device.value.auth_type === 'google' || device.value.auth_type === 'github'
|
||||||
@@ -782,6 +789,32 @@ const canImport = computed(() => {
|
|||||||
return importText.value.trim().length > 0 && !importing.value
|
return importText.value.trim().length > 0 && !importing.value
|
||||||
})
|
})
|
||||||
|
|
||||||
|
const importModeLabel = computed(() => (isGrokProvider.value ? '导入账号' : '导入授权'))
|
||||||
|
const importButtonLabel = computed(() => (isGrokProvider.value ? '导入账号' : '导入'))
|
||||||
|
const importDropTitle = computed(() => (
|
||||||
|
isGrokProvider.value ? '拖入 Grok 账号文件或点击选择' : '拖入授权文件或点击选择'
|
||||||
|
))
|
||||||
|
const importDropHint = computed(() => (
|
||||||
|
isGrokProvider.value ? '支持 .json / .txt,可多选、批量导入' : '支持 .json / .txt,可多选'
|
||||||
|
))
|
||||||
|
const importManualPlaceholder = computed(() => (
|
||||||
|
isGrokProvider.value
|
||||||
|
? '粘贴 Grok sso/session token,支持每行一个;或粘贴包含 token、sso_token、access_token、plan_type、pool_tier 的 JSON'
|
||||||
|
: '粘贴 Refresh Token / Access Token 或 JSON 内容'
|
||||||
|
))
|
||||||
|
const importManualDescription = computed(() => (
|
||||||
|
isGrokProvider.value
|
||||||
|
? 'plan_type / pool_tier 会作为账号套餐与能力特征保存,不是路由池选择。'
|
||||||
|
: ''
|
||||||
|
))
|
||||||
|
const importPasteToggleText = computed(() => (
|
||||||
|
isGrokProvider.value ? '或手动粘贴 Grok Token' : '或手动粘贴 Token'
|
||||||
|
))
|
||||||
|
const importFileToggleText = computed(() => (
|
||||||
|
isGrokProvider.value ? '或选择 Grok Token 文件导入' : '或选择 JSON 文件导入'
|
||||||
|
))
|
||||||
|
const providerCredentialActionLabel = computed(() => (isGrokProvider.value ? '导入' : '授权'))
|
||||||
|
|
||||||
function stopImportPolling() {
|
function stopImportPolling() {
|
||||||
if (importPollTimer) {
|
if (importPollTimer) {
|
||||||
clearTimeout(importPollTimer)
|
clearTimeout(importPollTimer)
|
||||||
@@ -923,7 +956,7 @@ function resetDeviceRuntimeState() {
|
|||||||
device.value.error = ''
|
device.value.error = ''
|
||||||
}
|
}
|
||||||
|
|
||||||
function isKiroDeviceAuthOptionDisabled(authType: DeviceAuthType): boolean {
|
function isKiroDeviceAuthOptionDisabled(_authType: DeviceAuthType): boolean {
|
||||||
if (device.value.starting) {
|
if (device.value.starting) {
|
||||||
return !isSocialDeviceAuth.value
|
return !isSocialDeviceAuth.value
|
||||||
}
|
}
|
||||||
@@ -976,11 +1009,12 @@ function resetForm() {
|
|||||||
importInputResetKey.value += 1
|
importInputResetKey.value += 1
|
||||||
proxyPopoverOpen.value = false
|
proxyPopoverOpen.value = false
|
||||||
selectedProxyNodeId.value = ''
|
selectedProxyNodeId.value = ''
|
||||||
mode.value = 'oauth'
|
mode.value = defaultMode.value
|
||||||
}
|
}
|
||||||
|
|
||||||
function switchMode(newMode: DialogMode) {
|
function switchMode(newMode: DialogMode) {
|
||||||
if (mode.value === newMode) return
|
if (mode.value === newMode) return
|
||||||
|
if (newMode === 'oauth' && !showAuthorizationMode.value) return
|
||||||
|
|
||||||
mode.value = newMode
|
mode.value = newMode
|
||||||
if (newMode === 'oauth') {
|
if (newMode === 'oauth') {
|
||||||
@@ -1011,6 +1045,7 @@ function openAuthorizationUrl() {
|
|||||||
|
|
||||||
async function initOAuth() {
|
async function initOAuth() {
|
||||||
if (!props.providerId) return
|
if (!props.providerId) return
|
||||||
|
if (!showAuthorizationMode.value) return
|
||||||
if (isKiroProvider.value) return
|
if (isKiroProvider.value) return
|
||||||
if (oauth.value.starting) return
|
if (oauth.value.starting) return
|
||||||
|
|
||||||
@@ -1095,6 +1130,12 @@ function parseImportText(text: string): {
|
|||||||
account_id?: string
|
account_id?: string
|
||||||
account_user_id?: string
|
account_user_id?: string
|
||||||
plan_type?: string
|
plan_type?: string
|
||||||
|
pool_tier?: string
|
||||||
|
sso_rw_token?: string
|
||||||
|
cf_cookies?: string
|
||||||
|
cf_clearance?: string
|
||||||
|
user_agent?: string
|
||||||
|
browser_profile?: string
|
||||||
user_id?: string
|
user_id?: string
|
||||||
account_name?: string
|
account_name?: string
|
||||||
} | null {
|
} | null {
|
||||||
@@ -1106,30 +1147,50 @@ function parseImportText(text: string): {
|
|||||||
return { refresh_token: trimmed }
|
return { refresh_token: trimmed }
|
||||||
}
|
}
|
||||||
|
|
||||||
|
if (isGrokProvider.value) {
|
||||||
|
const cookieImport = parseGrokCookieImport(trimmed)
|
||||||
|
if (cookieImport) {
|
||||||
|
return cookieImport
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
try {
|
try {
|
||||||
const parsed: unknown = JSON.parse(trimmed)
|
const parsed: unknown = JSON.parse(trimmed)
|
||||||
if (typeof parsed === 'object' && parsed !== null) {
|
if (typeof parsed === 'object' && parsed !== null) {
|
||||||
const obj = parsed as Record<string, unknown>
|
const obj = parsed as Record<string, unknown>
|
||||||
|
const grokCookieImport = isGrokProvider.value
|
||||||
|
? parseGrokCookieImport(normalizeStringField(obj.cookie) ?? normalizeStringField(obj.cookieHeader) ?? '')
|
||||||
|
: null
|
||||||
const refreshToken = obj.refresh_token
|
const refreshToken = obj.refresh_token
|
||||||
const refreshTokenCamel = obj.refreshToken
|
const refreshTokenCamel = obj.refreshToken
|
||||||
const accessToken = obj.access_token
|
const accessToken = obj.access_token
|
||||||
const accessTokenCamel = obj.accessToken
|
const accessTokenCamel = obj.accessToken
|
||||||
|
const grokSsoToken = isGrokProvider.value
|
||||||
|
? normalizeStringField(obj.sso_token) ?? normalizeStringField(obj.ssoToken) ?? normalizeStringField(obj.token) ?? grokCookieImport?.access_token
|
||||||
|
: undefined
|
||||||
const normalizedRefreshToken = typeof refreshToken === 'string' && refreshToken.trim()
|
const normalizedRefreshToken = typeof refreshToken === 'string' && refreshToken.trim()
|
||||||
? refreshToken.trim()
|
? refreshToken.trim()
|
||||||
: (typeof refreshTokenCamel === 'string' && refreshTokenCamel.trim() ? refreshTokenCamel.trim() : undefined)
|
: (typeof refreshTokenCamel === 'string' && refreshTokenCamel.trim() ? refreshTokenCamel.trim() : undefined)
|
||||||
const normalizedAccessToken = typeof accessToken === 'string' && accessToken.trim()
|
const normalizedAccessToken = typeof accessToken === 'string' && accessToken.trim()
|
||||||
? accessToken.trim()
|
? accessToken.trim()
|
||||||
: (typeof accessTokenCamel === 'string' && accessTokenCamel.trim() ? accessTokenCamel.trim() : undefined)
|
: (typeof accessTokenCamel === 'string' && accessTokenCamel.trim() ? accessTokenCamel.trim() : undefined)
|
||||||
if (normalizedRefreshToken || normalizedAccessToken) {
|
const importedAccessToken = normalizedAccessToken ?? grokSsoToken
|
||||||
|
if (normalizedRefreshToken || importedAccessToken) {
|
||||||
return {
|
return {
|
||||||
refresh_token: normalizedRefreshToken,
|
refresh_token: normalizedRefreshToken,
|
||||||
access_token: normalizedAccessToken,
|
access_token: importedAccessToken,
|
||||||
expires_at: normalizeNumberField(obj.expires_at) ?? normalizeNumberField(obj.expiresAt),
|
expires_at: normalizeNumberField(obj.expires_at) ?? normalizeNumberField(obj.expiresAt),
|
||||||
name: (typeof obj.name === 'string' ? obj.name : undefined) || (typeof obj.oauth_email === 'string' ? obj.oauth_email : undefined),
|
name: (typeof obj.name === 'string' ? obj.name : undefined) || (typeof obj.oauth_email === 'string' ? obj.oauth_email : undefined),
|
||||||
email: normalizeStringField(obj.email) ?? normalizeStringField(obj.oauth_email),
|
email: normalizeStringField(obj.email) ?? normalizeStringField(obj.oauth_email),
|
||||||
account_id: normalizeStringField(obj.account_id) ?? normalizeStringField(obj.accountId) ?? normalizeStringField(obj.chatgpt_account_id) ?? normalizeStringField(obj.chatgptAccountId),
|
account_id: normalizeStringField(obj.account_id) ?? normalizeStringField(obj.accountId) ?? normalizeStringField(obj.chatgpt_account_id) ?? normalizeStringField(obj.chatgptAccountId),
|
||||||
account_user_id: normalizeStringField(obj.account_user_id) ?? normalizeStringField(obj.accountUserId) ?? normalizeStringField(obj.chatgpt_account_user_id) ?? normalizeStringField(obj.chatgptAccountUserId),
|
account_user_id: normalizeStringField(obj.account_user_id) ?? normalizeStringField(obj.accountUserId) ?? normalizeStringField(obj.chatgpt_account_user_id) ?? normalizeStringField(obj.chatgptAccountUserId),
|
||||||
plan_type: normalizeStringField(obj.plan_type) ?? normalizeStringField(obj.planType) ?? normalizeStringField(obj.chatgpt_plan_type) ?? normalizeStringField(obj.chatgptPlanType),
|
plan_type: normalizeStringField(obj.plan_type) ?? normalizeStringField(obj.planType) ?? normalizeStringField(obj.chatgpt_plan_type) ?? normalizeStringField(obj.chatgptPlanType),
|
||||||
|
pool_tier: isGrokProvider.value ? normalizeStringField(obj.pool_tier) ?? normalizeStringField(obj.poolTier) ?? normalizeStringField(obj.tier) : undefined,
|
||||||
|
sso_rw_token: isGrokProvider.value ? normalizeStringField(obj.sso_rw_token) ?? normalizeStringField(obj.ssoRwToken) ?? grokCookieImport?.sso_rw_token : undefined,
|
||||||
|
cf_cookies: isGrokProvider.value ? normalizeStringField(obj.cf_cookies) ?? normalizeStringField(obj.cfCookies) ?? grokCookieImport?.cf_cookies : undefined,
|
||||||
|
cf_clearance: isGrokProvider.value ? normalizeStringField(obj.cf_clearance) ?? normalizeStringField(obj.cfClearance) ?? grokCookieImport?.cf_clearance : undefined,
|
||||||
|
user_agent: isGrokProvider.value ? normalizeStringField(obj.user_agent) ?? normalizeStringField(obj.userAgent) ?? grokCookieImport?.user_agent : undefined,
|
||||||
|
browser_profile: isGrokProvider.value ? normalizeStringField(obj.browser_profile) ?? normalizeStringField(obj.browserProfile) ?? normalizeStringField(obj.browser) ?? normalizeStringField(obj.impersonate) ?? grokCookieImport?.browser_profile : undefined,
|
||||||
user_id: normalizeStringField(obj.user_id) ?? normalizeStringField(obj.userId) ?? normalizeStringField(obj.chatgpt_user_id) ?? normalizeStringField(obj.chatgptUserId),
|
user_id: normalizeStringField(obj.user_id) ?? normalizeStringField(obj.userId) ?? normalizeStringField(obj.chatgpt_user_id) ?? normalizeStringField(obj.chatgptUserId),
|
||||||
account_name: normalizeStringField(obj.account_name) ?? normalizeStringField(obj.accountName),
|
account_name: normalizeStringField(obj.account_name) ?? normalizeStringField(obj.accountName),
|
||||||
}
|
}
|
||||||
@@ -1147,6 +1208,72 @@ function parseImportText(text: string): {
|
|||||||
return { refresh_token: trimmed }
|
return { refresh_token: trimmed }
|
||||||
}
|
}
|
||||||
|
|
||||||
|
function parseGrokCookieImport(text: string): {
|
||||||
|
access_token: string
|
||||||
|
sso_rw_token?: string
|
||||||
|
cf_cookies?: string
|
||||||
|
cf_clearance?: string
|
||||||
|
user_agent?: string
|
||||||
|
browser_profile?: string
|
||||||
|
user_id?: string
|
||||||
|
} | null {
|
||||||
|
const cookies = parseCookieHeader(text)
|
||||||
|
const sso = cookies.get('sso')
|
||||||
|
if (!sso) return null
|
||||||
|
const userAgent = currentBrowserUserAgent()
|
||||||
|
|
||||||
|
return {
|
||||||
|
access_token: sso,
|
||||||
|
sso_rw_token: cookies.get('sso-rw'),
|
||||||
|
cf_cookies: buildGrokCookieProfile(cookies),
|
||||||
|
cf_clearance: cookies.get('cf_clearance'),
|
||||||
|
user_agent: userAgent,
|
||||||
|
browser_profile: inferGrokBrowserProfile(userAgent),
|
||||||
|
user_id: cookies.get('x-userid'),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
function currentBrowserUserAgent(): string | undefined {
|
||||||
|
const value = typeof navigator !== 'undefined' ? navigator.userAgent?.trim() : ''
|
||||||
|
return value || undefined
|
||||||
|
}
|
||||||
|
|
||||||
|
function inferGrokBrowserProfile(userAgent: string | undefined): string | undefined {
|
||||||
|
const value = (userAgent || '').toLowerCase()
|
||||||
|
if (!value) return 'chrome136'
|
||||||
|
if (value.includes('firefox/')) return 'firefox'
|
||||||
|
if (value.includes('safari/') && !value.includes('chrome/') && !value.includes('chromium/')) {
|
||||||
|
return value.includes('iphone') || value.includes('ipad') ? 'safari_ios' : 'safari'
|
||||||
|
}
|
||||||
|
return 'chrome136'
|
||||||
|
}
|
||||||
|
|
||||||
|
function buildGrokCookieProfile(cookies: Map<string, string>): string | undefined {
|
||||||
|
const parts: string[] = []
|
||||||
|
for (const [name, value] of cookies) {
|
||||||
|
if (name === 'sso' || name === 'sso-rw') continue
|
||||||
|
parts.push(`${name}=${value}`)
|
||||||
|
}
|
||||||
|
return parts.length > 0 ? parts.join('; ') : undefined
|
||||||
|
}
|
||||||
|
|
||||||
|
function parseCookieHeader(text: string): Map<string, string> {
|
||||||
|
const normalized = text.trim().replace(/^cookie:\s*/i, '')
|
||||||
|
const cookies = new Map<string, string>()
|
||||||
|
for (const segment of normalized.split(';')) {
|
||||||
|
const part = segment.trim()
|
||||||
|
if (!part) continue
|
||||||
|
const separator = part.indexOf('=')
|
||||||
|
if (separator <= 0) continue
|
||||||
|
const name = part.slice(0, separator).trim().toLowerCase()
|
||||||
|
const value = part.slice(separator + 1).trim()
|
||||||
|
if (name && value) {
|
||||||
|
cookies.set(name, value)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return cookies
|
||||||
|
}
|
||||||
|
|
||||||
function normalizeStringField(value: unknown): string | undefined {
|
function normalizeStringField(value: unknown): string | undefined {
|
||||||
return typeof value === 'string' && value.trim() ? value.trim() : undefined
|
return typeof value === 'string' && value.trim() ? value.trim() : undefined
|
||||||
}
|
}
|
||||||
@@ -1404,6 +1531,10 @@ onBeforeUnmount(() => {
|
|||||||
watch(() => props.open, (newOpen) => {
|
watch(() => props.open, (newOpen) => {
|
||||||
if (newOpen) {
|
if (newOpen) {
|
||||||
proxyNodesStore.ensureLoaded()
|
proxyNodesStore.ensureLoaded()
|
||||||
|
mode.value = defaultMode.value
|
||||||
|
if (!showAuthorizationMode.value) {
|
||||||
|
return
|
||||||
|
}
|
||||||
if (isKiroProvider.value) {
|
if (isKiroProvider.value) {
|
||||||
void ensureKiroSocialDeviceAuth()
|
void ensureKiroSocialDeviceAuth()
|
||||||
} else {
|
} else {
|
||||||
@@ -1417,6 +1548,10 @@ watch(() => props.open, (newOpen) => {
|
|||||||
watch(
|
watch(
|
||||||
() => [props.open, props.providerId, props.providerType] as const,
|
() => [props.open, props.providerId, props.providerType] as const,
|
||||||
() => {
|
() => {
|
||||||
|
if (props.open && !showAuthorizationMode.value) {
|
||||||
|
mode.value = 'import'
|
||||||
|
return
|
||||||
|
}
|
||||||
if (props.open && isKiroProvider.value && mode.value === 'oauth') {
|
if (props.open && isKiroProvider.value && mode.value === 'oauth') {
|
||||||
void ensureKiroSocialDeviceAuth()
|
void ensureKiroSocialDeviceAuth()
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -858,6 +858,7 @@ const PROVIDER_TYPE_LABELS: Record<string, string> = {
|
|||||||
gemini_cli: 'Gemini CLI',
|
gemini_cli: 'Gemini CLI',
|
||||||
antigravity: 'Antigravity',
|
antigravity: 'Antigravity',
|
||||||
kiro: 'Kiro',
|
kiro: 'Kiro',
|
||||||
|
grok: 'Grok',
|
||||||
}
|
}
|
||||||
|
|
||||||
function formatProviderType(type?: string): string {
|
function formatProviderType(type?: string): string {
|
||||||
|
|||||||
@@ -346,7 +346,7 @@
|
|||||||
<Copy class="w-2.5 h-2.5" />
|
<Copy class="w-2.5 h-2.5" />
|
||||||
</Button>
|
</Button>
|
||||||
<!-- OAuth 状态(失效/过期/倒计时)和刷新按钮 -->
|
<!-- OAuth 状态(失效/过期/倒计时)和刷新按钮 -->
|
||||||
<template v-if="shouldShowOAuthRefreshControl(key)">
|
<template v-if="shouldShowOAuthRefreshControl(key, provider.provider_type)">
|
||||||
<!-- 账号级别异常:醒目提示 + 清除按钮 -->
|
<!-- 账号级别异常:醒目提示 + 清除按钮 -->
|
||||||
<template v-if="isAccountLevelBlock(key)">
|
<template v-if="isAccountLevelBlock(key)">
|
||||||
<Badge
|
<Badge
|
||||||
@@ -1293,6 +1293,7 @@ import type {
|
|||||||
AntigravityModelQuota,
|
AntigravityModelQuota,
|
||||||
CodexUpstreamMetadata,
|
CodexUpstreamMetadata,
|
||||||
ChatGPTWebUpstreamMetadata,
|
ChatGPTWebUpstreamMetadata,
|
||||||
|
GrokUpstreamMetadata,
|
||||||
KiroUpstreamMetadata,
|
KiroUpstreamMetadata,
|
||||||
QuotaStatusSnapshot,
|
QuotaStatusSnapshot,
|
||||||
QuotaWindowSnapshot,
|
QuotaWindowSnapshot,
|
||||||
@@ -1964,7 +1965,7 @@ function quotaSnapshotHasDisplayData(quota: QuotaStatusSnapshot | null | undefin
|
|||||||
|
|
||||||
function getQuotaSnapshotForProvider(
|
function getQuotaSnapshotForProvider(
|
||||||
key: EndpointAPIKey,
|
key: EndpointAPIKey,
|
||||||
providerType: 'codex' | 'kiro' | 'antigravity' | 'chatgpt_web' | 'gemini_cli',
|
providerType: 'codex' | 'kiro' | 'antigravity' | 'chatgpt_web' | 'gemini_cli' | 'grok',
|
||||||
): QuotaStatusSnapshot | null {
|
): QuotaStatusSnapshot | null {
|
||||||
const quota = key.status_snapshot?.quota
|
const quota = key.status_snapshot?.quota
|
||||||
if (!quota) return null
|
if (!quota) return null
|
||||||
@@ -2168,6 +2169,66 @@ function hasKiroQuotaDisplayData(key: EndpointAPIKey): boolean {
|
|||||||
return !!kiro && (kiro.usage_percentage !== undefined || kiro.usage_limit !== undefined)
|
return !!kiro && (kiro.usage_percentage !== undefined || kiro.usage_limit !== undefined)
|
||||||
}
|
}
|
||||||
|
|
||||||
|
type GrokQuotaDisplay = GrokUpstreamMetadata & {
|
||||||
|
usage_percentage?: number
|
||||||
|
usage_limit?: number
|
||||||
|
current_usage?: number
|
||||||
|
remaining?: number
|
||||||
|
next_reset_at?: number
|
||||||
|
}
|
||||||
|
|
||||||
|
function getGrokQuotaDisplay(key: EndpointAPIKey): GrokQuotaDisplay | null {
|
||||||
|
const quota = getQuotaSnapshotForProvider(key, 'grok')
|
||||||
|
if (!quota) return null
|
||||||
|
|
||||||
|
const display: GrokQuotaDisplay = {}
|
||||||
|
const updatedAt = getQuotaSnapshotUpdatedAt(quota)
|
||||||
|
if (updatedAt !== undefined) display.updated_at = updatedAt
|
||||||
|
if (quota.plan_type) display.plan_type = quota.plan_type
|
||||||
|
if (quota.pool_tier) display.pool_tier = quota.pool_tier
|
||||||
|
|
||||||
|
const code = String(quota.code || '').trim().toLowerCase()
|
||||||
|
if (code === 'banned' || code === 'forbidden') {
|
||||||
|
display.is_banned = true
|
||||||
|
if (quota.reason) display.ban_reason = quota.reason
|
||||||
|
}
|
||||||
|
|
||||||
|
const usageWindow =
|
||||||
|
getQuotaWindow(quota, 'usage')
|
||||||
|
?? getQuotaWindowByScope(quota, 'account')[0]
|
||||||
|
?? getQuotaWindowByScope(quota, 'model')
|
||||||
|
.map(window => ({
|
||||||
|
window,
|
||||||
|
remainingPercent: getQuotaWindowRemainingPercent(window),
|
||||||
|
}))
|
||||||
|
.filter((item): item is { window: QuotaWindowSnapshot, remainingPercent: number } => item.remainingPercent !== undefined)
|
||||||
|
.sort((a, b) => a.remainingPercent - b.remainingPercent)[0]?.window
|
||||||
|
?? null
|
||||||
|
if (usageWindow) {
|
||||||
|
const usedPercent = getQuotaWindowUsedPercent(usageWindow)
|
||||||
|
if (usedPercent !== undefined) display.usage_percentage = usedPercent
|
||||||
|
if (typeof usageWindow.used_value === 'number') display.current_usage = usageWindow.used_value
|
||||||
|
if (typeof usageWindow.limit_value === 'number') display.usage_limit = usageWindow.limit_value
|
||||||
|
if (typeof usageWindow.remaining_value === 'number') display.remaining = usageWindow.remaining_value
|
||||||
|
|
||||||
|
const nextResetAt =
|
||||||
|
getQuotaWindowResetAt(usageWindow)
|
||||||
|
?? (() => {
|
||||||
|
const resetSeconds = getQuotaWindowResetSeconds(usageWindow)
|
||||||
|
if (updatedAt === undefined || resetSeconds === undefined) return undefined
|
||||||
|
return updatedAt + resetSeconds
|
||||||
|
})()
|
||||||
|
if (nextResetAt !== undefined) display.next_reset_at = nextResetAt
|
||||||
|
}
|
||||||
|
|
||||||
|
return Object.keys(display).length > 0 ? display : null
|
||||||
|
}
|
||||||
|
|
||||||
|
function hasGrokQuotaDisplayData(key: EndpointAPIKey): boolean {
|
||||||
|
const grok = getGrokQuotaDisplay(key)
|
||||||
|
return !!grok && (grok.usage_percentage !== undefined || grok.usage_limit !== undefined)
|
||||||
|
}
|
||||||
|
|
||||||
type ChatGPTWebQuotaDisplay = ChatGPTWebUpstreamMetadata & {
|
type ChatGPTWebQuotaDisplay = ChatGPTWebUpstreamMetadata & {
|
||||||
image_quota_remaining_percent?: number
|
image_quota_remaining_percent?: number
|
||||||
image_quota_used_percent?: number
|
image_quota_used_percent?: number
|
||||||
@@ -2435,6 +2496,28 @@ function shouldAutoRefreshKiroQuota(): boolean {
|
|||||||
return false
|
return false
|
||||||
}
|
}
|
||||||
|
|
||||||
|
function shouldAutoRefreshGrokQuota(): boolean {
|
||||||
|
if (provider.value?.provider_type !== 'grok') return false
|
||||||
|
const now = Math.floor(Date.now() / 1000)
|
||||||
|
|
||||||
|
for (const { key } of allKeys.value) {
|
||||||
|
if (!key.is_active) continue
|
||||||
|
|
||||||
|
if (isTokenExpiringSoon(key, now)) return true
|
||||||
|
|
||||||
|
if (!hasGrokQuotaDisplayData(key)) {
|
||||||
|
return true
|
||||||
|
}
|
||||||
|
|
||||||
|
const updatedAt = getGrokQuotaDisplay(key)?.updated_at
|
||||||
|
if (typeof updatedAt !== 'number' || (now - updatedAt) > AUTO_QUOTA_REFRESH_STALE_SECONDS) {
|
||||||
|
return true
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
return false
|
||||||
|
}
|
||||||
|
|
||||||
function shouldAutoRefreshChatGPTWebQuota(): boolean {
|
function shouldAutoRefreshChatGPTWebQuota(): boolean {
|
||||||
if (provider.value?.provider_type !== 'chatgpt_web') return false
|
if (provider.value?.provider_type !== 'chatgpt_web') return false
|
||||||
const now = Math.floor(Date.now() / 1000)
|
const now = Math.floor(Date.now() / 1000)
|
||||||
@@ -2541,7 +2624,7 @@ async function autoRefreshQuotaInBackground(options: { ignoreCooldown?: boolean
|
|||||||
if (refreshingQuota.value) return
|
if (refreshingQuota.value) return
|
||||||
|
|
||||||
const providerType = provider.value?.provider_type
|
const providerType = provider.value?.provider_type
|
||||||
if (providerType !== 'codex' && providerType !== 'antigravity' && providerType !== 'kiro' && providerType !== 'chatgpt_web') return
|
if (providerType !== 'codex' && providerType !== 'antigravity' && providerType !== 'kiro' && providerType !== 'chatgpt_web' && providerType !== 'grok') return
|
||||||
|
|
||||||
// 检查是否需要刷新
|
// 检查是否需要刷新
|
||||||
let shouldRefresh = false
|
let shouldRefresh = false
|
||||||
@@ -2551,6 +2634,8 @@ async function autoRefreshQuotaInBackground(options: { ignoreCooldown?: boolean
|
|||||||
shouldRefresh = shouldAutoRefreshAntigravityQuota()
|
shouldRefresh = shouldAutoRefreshAntigravityQuota()
|
||||||
} else if (providerType === 'kiro') {
|
} else if (providerType === 'kiro') {
|
||||||
shouldRefresh = shouldAutoRefreshKiroQuota()
|
shouldRefresh = shouldAutoRefreshKiroQuota()
|
||||||
|
} else if (providerType === 'grok') {
|
||||||
|
shouldRefresh = shouldAutoRefreshGrokQuota()
|
||||||
} else if (providerType === 'chatgpt_web') {
|
} else if (providerType === 'chatgpt_web') {
|
||||||
shouldRefresh = shouldAutoRefreshChatGPTWebQuota()
|
shouldRefresh = shouldAutoRefreshChatGPTWebQuota()
|
||||||
}
|
}
|
||||||
@@ -2564,6 +2649,8 @@ async function autoRefreshQuotaInBackground(options: { ignoreCooldown?: boolean
|
|||||||
hadCachedQuota = allKeys.value.some(({ key }) => key.is_active && hasAntigravityQuotaDisplayData(key))
|
hadCachedQuota = allKeys.value.some(({ key }) => key.is_active && hasAntigravityQuotaDisplayData(key))
|
||||||
} else if (providerType === 'kiro') {
|
} else if (providerType === 'kiro') {
|
||||||
hadCachedQuota = allKeys.value.some(({ key }) => key.is_active && hasKiroQuotaDisplayData(key))
|
hadCachedQuota = allKeys.value.some(({ key }) => key.is_active && hasKiroQuotaDisplayData(key))
|
||||||
|
} else if (providerType === 'grok') {
|
||||||
|
hadCachedQuota = allKeys.value.some(({ key }) => key.is_active && hasGrokQuotaDisplayData(key))
|
||||||
} else if (providerType === 'chatgpt_web') {
|
} else if (providerType === 'chatgpt_web') {
|
||||||
hadCachedQuota = allKeys.value.some(({ key }) => key.is_active && hasChatGPTWebQuotaDisplayData(key))
|
hadCachedQuota = allKeys.value.some(({ key }) => key.is_active && hasChatGPTWebQuotaDisplayData(key))
|
||||||
}
|
}
|
||||||
@@ -3030,6 +3117,9 @@ function formatOAuthPlanType(planType: string): string {
|
|||||||
team: 'Team',
|
team: 'Team',
|
||||||
enterprise: 'Enterprise',
|
enterprise: 'Enterprise',
|
||||||
ultra: 'Ultra',
|
ultra: 'Ultra',
|
||||||
|
basic: 'Basic',
|
||||||
|
super: 'Super',
|
||||||
|
heavy: 'Heavy',
|
||||||
}
|
}
|
||||||
return labels[planType.toLowerCase()] || planType
|
return labels[planType.toLowerCase()] || planType
|
||||||
}
|
}
|
||||||
@@ -3377,6 +3467,9 @@ function getOAuthPlanTypeClass(planType: string): string {
|
|||||||
ultra: 'border-amber-500/50 text-amber-600 dark:text-amber-400',
|
ultra: 'border-amber-500/50 text-amber-600 dark:text-amber-400',
|
||||||
'pro+': 'border-purple-500/50 text-purple-600 dark:text-purple-400',
|
'pro+': 'border-purple-500/50 text-purple-600 dark:text-purple-400',
|
||||||
power: 'border-amber-500/50 text-amber-600 dark:text-amber-400',
|
power: 'border-amber-500/50 text-amber-600 dark:text-amber-400',
|
||||||
|
basic: 'border-primary/50 text-primary',
|
||||||
|
super: 'border-green-500/50 text-green-600 dark:text-green-400',
|
||||||
|
heavy: 'border-amber-500/50 text-amber-600 dark:text-amber-400',
|
||||||
}
|
}
|
||||||
return classes[planType.toLowerCase()] || ''
|
return classes[planType.toLowerCase()] || ''
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -60,6 +60,9 @@
|
|||||||
<SelectItem value="gemini_cli">
|
<SelectItem value="gemini_cli">
|
||||||
Gemini CLI
|
Gemini CLI
|
||||||
</SelectItem>
|
</SelectItem>
|
||||||
|
<SelectItem value="grok">
|
||||||
|
Grok
|
||||||
|
</SelectItem>
|
||||||
<SelectItem value="kiro">
|
<SelectItem value="kiro">
|
||||||
Kiro
|
Kiro
|
||||||
</SelectItem>
|
</SelectItem>
|
||||||
@@ -87,6 +90,9 @@
|
|||||||
<SelectItem value="gemini_cli">
|
<SelectItem value="gemini_cli">
|
||||||
Gemini CLI
|
Gemini CLI
|
||||||
</SelectItem>
|
</SelectItem>
|
||||||
|
<SelectItem value="grok">
|
||||||
|
Grok
|
||||||
|
</SelectItem>
|
||||||
<SelectItem value="kiro">
|
<SelectItem value="kiro">
|
||||||
Kiro
|
Kiro
|
||||||
</SelectItem>
|
</SelectItem>
|
||||||
@@ -355,7 +361,7 @@ const defaultPriority = computed(() => {
|
|||||||
// 表单数据
|
// 表单数据
|
||||||
const form = ref({
|
const form = ref({
|
||||||
name: '',
|
name: '',
|
||||||
provider_type: 'custom' as 'custom' | 'vertex_ai' | 'claude_code' | 'codex' | 'chatgpt_web' | 'gemini_cli' | 'antigravity' | 'kiro',
|
provider_type: 'custom' as 'custom' | 'vertex_ai' | 'claude_code' | 'codex' | 'chatgpt_web' | 'gemini_cli' | 'antigravity' | 'kiro' | 'grok',
|
||||||
description: '',
|
description: '',
|
||||||
website: '',
|
website: '',
|
||||||
// 计费配置
|
// 计费配置
|
||||||
|
|||||||
@@ -0,0 +1,389 @@
|
|||||||
|
/* eslint-disable vue/one-component-per-file, vue/require-default-prop */
|
||||||
|
import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest'
|
||||||
|
import { createApp, nextTick, type App } from 'vue'
|
||||||
|
import OAuthAccountDialog from '@/features/providers/components/OAuthAccountDialog.vue'
|
||||||
|
|
||||||
|
const endpointMocks = vi.hoisted(() => ({
|
||||||
|
startProviderLevelOAuth: vi.fn(),
|
||||||
|
completeProviderLevelOAuth: vi.fn(),
|
||||||
|
importProviderRefreshToken: vi.fn(),
|
||||||
|
startBatchImportOAuthTask: vi.fn(),
|
||||||
|
getBatchImportOAuthTaskStatus: vi.fn(),
|
||||||
|
startDeviceAuthorize: vi.fn(),
|
||||||
|
pollDeviceAuthorize: vi.fn(),
|
||||||
|
getAwsRegions: vi.fn(),
|
||||||
|
}))
|
||||||
|
|
||||||
|
vi.mock('@/api/endpoints', () => endpointMocks)
|
||||||
|
|
||||||
|
vi.mock('@/components/ui', async () => {
|
||||||
|
const { defineComponent, h } = await import('vue')
|
||||||
|
|
||||||
|
const passthrough = (name: string, tag = 'div') => defineComponent({
|
||||||
|
name,
|
||||||
|
setup(_, { slots }) {
|
||||||
|
return () => h(tag, slots.default?.())
|
||||||
|
},
|
||||||
|
})
|
||||||
|
|
||||||
|
const Dialog = defineComponent({
|
||||||
|
name: 'DialogStub',
|
||||||
|
props: {
|
||||||
|
modelValue: Boolean,
|
||||||
|
},
|
||||||
|
setup(props, { slots }) {
|
||||||
|
return () => props.modelValue
|
||||||
|
? h('section', [slots.headerActions?.(), slots.default?.(), slots.footer?.()])
|
||||||
|
: null
|
||||||
|
},
|
||||||
|
})
|
||||||
|
|
||||||
|
const Button = defineComponent({
|
||||||
|
name: 'ButtonStub',
|
||||||
|
inheritAttrs: false,
|
||||||
|
props: {
|
||||||
|
disabled: Boolean,
|
||||||
|
variant: String,
|
||||||
|
size: String,
|
||||||
|
},
|
||||||
|
setup(props, { attrs, slots }) {
|
||||||
|
return () => h('button', {
|
||||||
|
...attrs,
|
||||||
|
disabled: props.disabled,
|
||||||
|
type: attrs.type ?? 'button',
|
||||||
|
}, slots.default?.())
|
||||||
|
},
|
||||||
|
})
|
||||||
|
|
||||||
|
const Textarea = defineComponent({
|
||||||
|
name: 'TextareaStub',
|
||||||
|
inheritAttrs: false,
|
||||||
|
props: {
|
||||||
|
modelValue: {
|
||||||
|
type: String,
|
||||||
|
default: '',
|
||||||
|
},
|
||||||
|
},
|
||||||
|
emits: ['update:modelValue'],
|
||||||
|
setup(props, { attrs, emit }) {
|
||||||
|
return () => h('textarea', {
|
||||||
|
...attrs,
|
||||||
|
value: props.modelValue,
|
||||||
|
onInput: (event: Event) => emit('update:modelValue', (event.target as HTMLTextAreaElement).value),
|
||||||
|
})
|
||||||
|
},
|
||||||
|
})
|
||||||
|
|
||||||
|
return {
|
||||||
|
Dialog,
|
||||||
|
Button,
|
||||||
|
Textarea,
|
||||||
|
Popover: passthrough('PopoverStub'),
|
||||||
|
PopoverTrigger: passthrough('PopoverTriggerStub'),
|
||||||
|
PopoverContent: passthrough('PopoverContentStub'),
|
||||||
|
}
|
||||||
|
})
|
||||||
|
|
||||||
|
vi.mock('radix-vue', async () => {
|
||||||
|
const { defineComponent, h } = await import('vue')
|
||||||
|
const passthrough = (name: string) => defineComponent({
|
||||||
|
name,
|
||||||
|
setup(_, { slots }) {
|
||||||
|
return () => h('div', slots.default?.())
|
||||||
|
},
|
||||||
|
})
|
||||||
|
|
||||||
|
return {
|
||||||
|
ComboboxAnchor: passthrough('ComboboxAnchorStub'),
|
||||||
|
ComboboxContent: passthrough('ComboboxContentStub'),
|
||||||
|
ComboboxEmpty: passthrough('ComboboxEmptyStub'),
|
||||||
|
ComboboxInput: passthrough('ComboboxInputStub'),
|
||||||
|
ComboboxItem: passthrough('ComboboxItemStub'),
|
||||||
|
ComboboxRoot: passthrough('ComboboxRootStub'),
|
||||||
|
ComboboxTrigger: passthrough('ComboboxTriggerStub'),
|
||||||
|
ComboboxViewport: passthrough('ComboboxViewportStub'),
|
||||||
|
}
|
||||||
|
})
|
||||||
|
|
||||||
|
vi.mock('@/components/common/JsonImportInput.vue', async () => {
|
||||||
|
const { defineComponent, h } = await import('vue')
|
||||||
|
|
||||||
|
return {
|
||||||
|
default: defineComponent({
|
||||||
|
name: 'JsonImportInputStub',
|
||||||
|
props: {
|
||||||
|
modelValue: {
|
||||||
|
type: String,
|
||||||
|
default: '',
|
||||||
|
},
|
||||||
|
dropTitle: {
|
||||||
|
type: String,
|
||||||
|
default: '',
|
||||||
|
},
|
||||||
|
dropHint: {
|
||||||
|
type: String,
|
||||||
|
default: '',
|
||||||
|
},
|
||||||
|
manualPlaceholder: {
|
||||||
|
type: String,
|
||||||
|
default: '',
|
||||||
|
},
|
||||||
|
manualDescription: {
|
||||||
|
type: String,
|
||||||
|
default: '',
|
||||||
|
},
|
||||||
|
pasteToggleText: {
|
||||||
|
type: String,
|
||||||
|
default: '',
|
||||||
|
},
|
||||||
|
fileToggleText: {
|
||||||
|
type: String,
|
||||||
|
default: '',
|
||||||
|
},
|
||||||
|
},
|
||||||
|
emits: ['update:modelValue'],
|
||||||
|
setup(props, { emit }) {
|
||||||
|
return () => h('div', [
|
||||||
|
h('p', { 'data-testid': 'drop-title' }, props.dropTitle),
|
||||||
|
h('p', { 'data-testid': 'drop-hint' }, props.dropHint),
|
||||||
|
h('p', { 'data-testid': 'manual-description' }, props.manualDescription),
|
||||||
|
h('p', props.pasteToggleText),
|
||||||
|
h('p', props.fileToggleText),
|
||||||
|
h('textarea', {
|
||||||
|
placeholder: props.manualPlaceholder,
|
||||||
|
value: props.modelValue,
|
||||||
|
onInput: (event: Event) => emit('update:modelValue', (event.target as HTMLTextAreaElement).value),
|
||||||
|
}),
|
||||||
|
])
|
||||||
|
},
|
||||||
|
}),
|
||||||
|
}
|
||||||
|
})
|
||||||
|
|
||||||
|
vi.mock('@/components/ui/Label.vue', () => ({}))
|
||||||
|
vi.mock('./ProxyNodeSelect.vue', () => ({}))
|
||||||
|
vi.mock('@/features/providers/components/ProxyNodeSelect.vue', async () => {
|
||||||
|
const { defineComponent, h } = await import('vue')
|
||||||
|
return {
|
||||||
|
default: defineComponent({
|
||||||
|
name: 'ProxyNodeSelectStub',
|
||||||
|
setup() {
|
||||||
|
return () => h('div')
|
||||||
|
},
|
||||||
|
}),
|
||||||
|
}
|
||||||
|
})
|
||||||
|
|
||||||
|
vi.mock('@/stores/proxy-nodes', () => ({
|
||||||
|
useProxyNodesStore: () => ({
|
||||||
|
nodes: [],
|
||||||
|
onlineNodes: [],
|
||||||
|
loading: false,
|
||||||
|
ensureLoaded: vi.fn(),
|
||||||
|
}),
|
||||||
|
}))
|
||||||
|
|
||||||
|
vi.mock('@/composables/useToast', () => ({
|
||||||
|
useToast: () => ({
|
||||||
|
success: vi.fn(),
|
||||||
|
error: vi.fn(),
|
||||||
|
}),
|
||||||
|
}))
|
||||||
|
|
||||||
|
vi.mock('@/composables/useClipboard', () => ({
|
||||||
|
useClipboard: () => ({
|
||||||
|
copyToClipboard: vi.fn(),
|
||||||
|
}),
|
||||||
|
}))
|
||||||
|
|
||||||
|
vi.mock('@/composables/useTotp', () => ({
|
||||||
|
useTotp: () => ({
|
||||||
|
code: { value: '' },
|
||||||
|
remaining: { value: 0 },
|
||||||
|
start: vi.fn(),
|
||||||
|
stop: vi.fn(),
|
||||||
|
}),
|
||||||
|
}))
|
||||||
|
|
||||||
|
vi.mock('lucide-vue-next', async () => {
|
||||||
|
const { defineComponent, h } = await import('vue')
|
||||||
|
const Icon = defineComponent({
|
||||||
|
name: 'IconStub',
|
||||||
|
setup() {
|
||||||
|
return () => h('span')
|
||||||
|
},
|
||||||
|
})
|
||||||
|
|
||||||
|
return {
|
||||||
|
UserPlus: Icon,
|
||||||
|
Copy: Icon,
|
||||||
|
ExternalLink: Icon,
|
||||||
|
Globe: Icon,
|
||||||
|
AlertCircle: Icon,
|
||||||
|
ShieldCheck: Icon,
|
||||||
|
ChevronsUpDown: Icon,
|
||||||
|
Check: Icon,
|
||||||
|
}
|
||||||
|
})
|
||||||
|
|
||||||
|
const mountedApps: Array<{ app: App, root: HTMLElement }> = []
|
||||||
|
|
||||||
|
function mountDialog(providerType = 'grok') {
|
||||||
|
const root = document.createElement('div')
|
||||||
|
document.body.appendChild(root)
|
||||||
|
const app = createApp(OAuthAccountDialog, {
|
||||||
|
open: true,
|
||||||
|
providerId: 'provider-1',
|
||||||
|
providerType,
|
||||||
|
})
|
||||||
|
app.mount(root)
|
||||||
|
mountedApps.push({ app, root })
|
||||||
|
return root
|
||||||
|
}
|
||||||
|
|
||||||
|
async function settle() {
|
||||||
|
await nextTick()
|
||||||
|
await Promise.resolve()
|
||||||
|
}
|
||||||
|
|
||||||
|
function getButton(root: HTMLElement, text: string) {
|
||||||
|
return Array.from(root.querySelectorAll('button'))
|
||||||
|
.find(button => button.textContent?.includes(text))
|
||||||
|
}
|
||||||
|
|
||||||
|
function getImportTextarea(root: HTMLElement) {
|
||||||
|
const textarea = root.querySelector('textarea')
|
||||||
|
if (!(textarea instanceof HTMLTextAreaElement)) {
|
||||||
|
throw new Error('Expected import textarea to exist')
|
||||||
|
}
|
||||||
|
return textarea
|
||||||
|
}
|
||||||
|
|
||||||
|
describe('OAuthAccountDialog Grok import', () => {
|
||||||
|
beforeEach(() => {
|
||||||
|
endpointMocks.startProviderLevelOAuth.mockReset()
|
||||||
|
endpointMocks.completeProviderLevelOAuth.mockReset()
|
||||||
|
endpointMocks.importProviderRefreshToken.mockReset()
|
||||||
|
endpointMocks.startBatchImportOAuthTask.mockReset()
|
||||||
|
endpointMocks.getBatchImportOAuthTaskStatus.mockReset()
|
||||||
|
endpointMocks.startDeviceAuthorize.mockReset()
|
||||||
|
endpointMocks.pollDeviceAuthorize.mockReset()
|
||||||
|
endpointMocks.getAwsRegions.mockReset()
|
||||||
|
|
||||||
|
endpointMocks.importProviderRefreshToken.mockResolvedValue({
|
||||||
|
provider_type: 'grok',
|
||||||
|
has_refresh_token: false,
|
||||||
|
email: 'grok@example.com',
|
||||||
|
replaced: false,
|
||||||
|
})
|
||||||
|
endpointMocks.startBatchImportOAuthTask.mockResolvedValue({
|
||||||
|
task_id: 'task-1',
|
||||||
|
status: 'submitted',
|
||||||
|
total: 2,
|
||||||
|
processed: 0,
|
||||||
|
success: 0,
|
||||||
|
failed: 0,
|
||||||
|
progress_percent: 0,
|
||||||
|
})
|
||||||
|
})
|
||||||
|
|
||||||
|
afterEach(() => {
|
||||||
|
for (const { app, root } of mountedApps.splice(0)) {
|
||||||
|
app.unmount()
|
||||||
|
root.remove()
|
||||||
|
}
|
||||||
|
})
|
||||||
|
|
||||||
|
it('opens Grok in import mode without starting unsupported OAuth', async () => {
|
||||||
|
const root = mountDialog('grok')
|
||||||
|
await settle()
|
||||||
|
|
||||||
|
expect(endpointMocks.startProviderLevelOAuth).not.toHaveBeenCalled()
|
||||||
|
expect(root.textContent).not.toContain('获取授权')
|
||||||
|
expect(root.querySelector('textarea')?.getAttribute('placeholder')).toContain('Grok sso/session token')
|
||||||
|
expect(root.textContent).toContain('plan_type / pool_tier')
|
||||||
|
expect(getButton(root, '导入账号')).toBeTruthy()
|
||||||
|
})
|
||||||
|
|
||||||
|
it('maps a single Grok JSON token into account metadata import payload', async () => {
|
||||||
|
const root = mountDialog('grok')
|
||||||
|
await settle()
|
||||||
|
|
||||||
|
const textarea = getImportTextarea(root)
|
||||||
|
textarea.value = JSON.stringify({
|
||||||
|
token: 'sso-1',
|
||||||
|
planType: 'super',
|
||||||
|
tier: 'heavy',
|
||||||
|
email: 'grok@example.com',
|
||||||
|
accountName: 'Grok Heavy',
|
||||||
|
})
|
||||||
|
textarea.dispatchEvent(new Event('input'))
|
||||||
|
await settle()
|
||||||
|
|
||||||
|
getButton(root, '导入账号')?.click()
|
||||||
|
await settle()
|
||||||
|
|
||||||
|
expect(endpointMocks.importProviderRefreshToken).toHaveBeenCalledWith('provider-1', {
|
||||||
|
access_token: 'sso-1',
|
||||||
|
account_name: 'Grok Heavy',
|
||||||
|
email: 'grok@example.com',
|
||||||
|
plan_type: 'super',
|
||||||
|
pool_tier: 'heavy',
|
||||||
|
sso_rw_token: undefined,
|
||||||
|
cf_cookies: undefined,
|
||||||
|
cf_clearance: undefined,
|
||||||
|
user_agent: undefined,
|
||||||
|
browser_profile: undefined,
|
||||||
|
proxy_node_id: undefined,
|
||||||
|
refresh_token: undefined,
|
||||||
|
expires_at: undefined,
|
||||||
|
name: undefined,
|
||||||
|
account_id: undefined,
|
||||||
|
account_user_id: undefined,
|
||||||
|
user_id: undefined,
|
||||||
|
})
|
||||||
|
})
|
||||||
|
|
||||||
|
it('keeps Grok multiline token import on the batch task path', async () => {
|
||||||
|
const root = mountDialog('grok')
|
||||||
|
await settle()
|
||||||
|
|
||||||
|
const textarea = getImportTextarea(root)
|
||||||
|
textarea.value = 'sso-1\nsso-2'
|
||||||
|
textarea.dispatchEvent(new Event('input'))
|
||||||
|
await settle()
|
||||||
|
|
||||||
|
getButton(root, '导入账号')?.click()
|
||||||
|
await settle()
|
||||||
|
|
||||||
|
expect(endpointMocks.startBatchImportOAuthTask).toHaveBeenCalledWith(
|
||||||
|
'provider-1',
|
||||||
|
'sso-1\nsso-2',
|
||||||
|
undefined,
|
||||||
|
)
|
||||||
|
expect(endpointMocks.importProviderRefreshToken).not.toHaveBeenCalled()
|
||||||
|
})
|
||||||
|
|
||||||
|
it('extracts Grok account fields from a pasted browser cookie header', async () => {
|
||||||
|
const root = mountDialog('grok')
|
||||||
|
await settle()
|
||||||
|
|
||||||
|
const textarea = getImportTextarea(root)
|
||||||
|
textarea.value = 'i18nextLng=zh; cf_clearance=cf-1; sso-rw=rw-1; sso=sso-1; x-userid=user-1'
|
||||||
|
textarea.dispatchEvent(new Event('input'))
|
||||||
|
await settle()
|
||||||
|
|
||||||
|
getButton(root, '导入账号')?.click()
|
||||||
|
await settle()
|
||||||
|
|
||||||
|
expect(endpointMocks.importProviderRefreshToken).toHaveBeenCalledWith('provider-1', expect.objectContaining({
|
||||||
|
access_token: 'sso-1',
|
||||||
|
sso_rw_token: 'rw-1',
|
||||||
|
cf_cookies: 'i18nextlng=zh; cf_clearance=cf-1; x-userid=user-1',
|
||||||
|
cf_clearance: 'cf-1',
|
||||||
|
user_agent: expect.any(String),
|
||||||
|
browser_profile: 'chrome136',
|
||||||
|
user_id: 'user-1',
|
||||||
|
}))
|
||||||
|
})
|
||||||
|
})
|
||||||
@@ -8,4 +8,10 @@ describe('providerTypeUtils', () => {
|
|||||||
expect(isOAuthAccountProviderType('ChatGPT_Web')).toBe(true)
|
expect(isOAuthAccountProviderType('ChatGPT_Web')).toBe(true)
|
||||||
expect(isKeyManagedProviderType('chatgpt_web')).toBe(false)
|
expect(isKeyManagedProviderType('chatgpt_web')).toBe(false)
|
||||||
})
|
})
|
||||||
|
|
||||||
|
it('treats Grok as an OAuth account provider', () => {
|
||||||
|
expect(isOAuthAccountProviderType('grok')).toBe(true)
|
||||||
|
expect(isOAuthAccountProviderType('GROK')).toBe(true)
|
||||||
|
expect(isKeyManagedProviderType('grok')).toBe(false)
|
||||||
|
})
|
||||||
})
|
})
|
||||||
|
|||||||
@@ -11,6 +11,7 @@ const oauthAccountProviderTypes = new Set([
|
|||||||
'gemini_cli',
|
'gemini_cli',
|
||||||
'antigravity',
|
'antigravity',
|
||||||
'kiro',
|
'kiro',
|
||||||
|
'grok',
|
||||||
])
|
])
|
||||||
|
|
||||||
export const isOAuthAccountProviderType = (providerType?: string | null): boolean =>
|
export const isOAuthAccountProviderType = (providerType?: string | null): boolean =>
|
||||||
|
|||||||
30
frontend/src/utils/__tests__/providerKeyAuth.spec.ts
Normal file
30
frontend/src/utils/__tests__/providerKeyAuth.spec.ts
Normal file
@@ -0,0 +1,30 @@
|
|||||||
|
import { describe, expect, it } from 'vitest'
|
||||||
|
|
||||||
|
import {
|
||||||
|
getProviderMaskedSecretLabel,
|
||||||
|
shouldShowOAuthRefreshControl,
|
||||||
|
} from '@/utils/providerKeyAuth'
|
||||||
|
|
||||||
|
describe('providerKeyAuth', () => {
|
||||||
|
it('renders Grok OAuth-managed cookies as sessions without OAuth refresh controls', () => {
|
||||||
|
const key = {
|
||||||
|
auth_type: 'oauth',
|
||||||
|
oauth_managed: true,
|
||||||
|
can_refresh_oauth: false,
|
||||||
|
}
|
||||||
|
|
||||||
|
expect(getProviderMaskedSecretLabel(key, 'grok')).toBe('[Session Cookie]')
|
||||||
|
expect(shouldShowOAuthRefreshControl(key, 'grok')).toBe(false)
|
||||||
|
})
|
||||||
|
|
||||||
|
it('keeps standard OAuth providers on OAuth token semantics', () => {
|
||||||
|
const key = {
|
||||||
|
auth_type: 'oauth',
|
||||||
|
oauth_managed: true,
|
||||||
|
can_refresh_oauth: false,
|
||||||
|
}
|
||||||
|
|
||||||
|
expect(getProviderMaskedSecretLabel(key, 'codex')).toBe('[OAuth Token]')
|
||||||
|
expect(shouldShowOAuthRefreshControl(key, 'codex')).toBe(true)
|
||||||
|
})
|
||||||
|
})
|
||||||
@@ -39,4 +39,69 @@ describe('providerKeyQuota', () => {
|
|||||||
},
|
},
|
||||||
}, 'codex')).toBe('周剩余 90.0% | 5H剩余 80.0% | Spark5H剩余 60.0% | Spark周剩余 95.0%')
|
}, 'codex')).toBe('周剩余 90.0% | 5H剩余 80.0% | Spark5H剩余 60.0% | Spark周剩余 95.0%')
|
||||||
})
|
})
|
||||||
|
|
||||||
|
it('formats Grok account quota from structured quota windows', () => {
|
||||||
|
expect(getQuotaDisplayText({
|
||||||
|
status_snapshot: {
|
||||||
|
oauth: {
|
||||||
|
code: 'valid',
|
||||||
|
},
|
||||||
|
account: {
|
||||||
|
code: 'ok',
|
||||||
|
blocked: false,
|
||||||
|
},
|
||||||
|
quota: {
|
||||||
|
provider_type: 'grok',
|
||||||
|
code: 'ok',
|
||||||
|
exhausted: false,
|
||||||
|
windows: [
|
||||||
|
{
|
||||||
|
scope: 'account',
|
||||||
|
used_value: 2,
|
||||||
|
limit_value: 10,
|
||||||
|
remaining_ratio: 0.8,
|
||||||
|
},
|
||||||
|
],
|
||||||
|
},
|
||||||
|
},
|
||||||
|
}, 'grok')).toBe('剩余 80.0% (8/10)')
|
||||||
|
})
|
||||||
|
|
||||||
|
it('formats Grok mode quota from model-scoped windows', () => {
|
||||||
|
expect(getQuotaDisplayText({
|
||||||
|
status_snapshot: {
|
||||||
|
oauth: {
|
||||||
|
code: 'valid',
|
||||||
|
},
|
||||||
|
account: {
|
||||||
|
code: 'ok',
|
||||||
|
blocked: false,
|
||||||
|
},
|
||||||
|
quota: {
|
||||||
|
provider_type: 'grok',
|
||||||
|
code: 'ok',
|
||||||
|
exhausted: false,
|
||||||
|
plan_type: 'heavy',
|
||||||
|
windows: [
|
||||||
|
{
|
||||||
|
code: 'model:quota_auto',
|
||||||
|
label: 'auto',
|
||||||
|
scope: 'model',
|
||||||
|
remaining_ratio: 0.4,
|
||||||
|
used_value: 90,
|
||||||
|
limit_value: 150,
|
||||||
|
},
|
||||||
|
{
|
||||||
|
code: 'model:quota_heavy',
|
||||||
|
label: 'heavy',
|
||||||
|
scope: 'model',
|
||||||
|
remaining_ratio: 0,
|
||||||
|
used_value: 20,
|
||||||
|
limit_value: 20,
|
||||||
|
},
|
||||||
|
],
|
||||||
|
},
|
||||||
|
},
|
||||||
|
}, 'grok')).toBe('Auto剩余 40.0% (60/150) | Heavy剩余 0.0% (0/20)')
|
||||||
|
})
|
||||||
})
|
})
|
||||||
|
|||||||
@@ -4,6 +4,7 @@ export const OAUTH_ICONS: Record<string, string> = {
|
|||||||
github: `<svg viewBox="0 0 24 24" fill="currentColor"><path d="M12 0c-6.626 0-12 5.373-12 12 0 5.302 3.438 9.8 8.207 11.387.599.111.793-.261.793-.577v-2.234c-3.338.726-4.033-1.416-4.033-1.416-.546-1.387-1.333-1.756-1.333-1.756-1.089-.745.083-.729.083-.729 1.205.084 1.839 1.237 1.839 1.237 1.07 1.834 2.807 1.304 3.492.997.107-.775.418-1.305.762-1.604-2.665-.305-5.467-1.334-5.467-5.931 0-1.311.469-2.381 1.236-3.221-.124-.303-.535-1.524.117-3.176 0 0 1.008-.322 3.301 1.23.957-.266 1.983-.399 3.003-.404 1.02.005 2.047.138 3.006.404 2.291-1.552 3.297-1.23 3.297-1.23.653 1.653.242 2.874.118 3.176.77.84 1.235 1.911 1.235 3.221 0 4.609-2.807 5.624-5.479 5.921.43.372.823 1.102.823 2.222v3.293c0 .319.192.694.801.576 4.765-1.589 8.199-6.086 8.199-11.386 0-6.627-5.373-12-12-12z"/></svg>`,
|
github: `<svg viewBox="0 0 24 24" fill="currentColor"><path d="M12 0c-6.626 0-12 5.373-12 12 0 5.302 3.438 9.8 8.207 11.387.599.111.793-.261.793-.577v-2.234c-3.338.726-4.033-1.416-4.033-1.416-.546-1.387-1.333-1.756-1.333-1.756-1.089-.745.083-.729.083-.729 1.205.084 1.839 1.237 1.839 1.237 1.07 1.834 2.807 1.304 3.492.997.107-.775.418-1.305.762-1.604-2.665-.305-5.467-1.334-5.467-5.931 0-1.311.469-2.381 1.236-3.221-.124-.303-.535-1.524.117-3.176 0 0 1.008-.322 3.301 1.23.957-.266 1.983-.399 3.003-.404 1.02.005 2.047.138 3.006.404 2.291-1.552 3.297-1.23 3.297-1.23.653 1.653.242 2.874.118 3.176.77.84 1.235 1.911 1.235 3.221 0 4.609-2.807 5.624-5.479 5.921.43.372.823 1.102.823 2.222v3.293c0 .319.192.694.801.576 4.765-1.589 8.199-6.086 8.199-11.386 0-6.627-5.373-12-12-12z"/></svg>`,
|
||||||
google: `<svg viewBox="0 0 24 24"><path fill="#4285F4" d="M22.56 12.25c0-.78-.07-1.53-.2-2.25H12v4.26h5.92c-.26 1.37-1.04 2.53-2.21 3.31v2.77h3.57c2.08-1.92 3.28-4.74 3.28-8.09z"/><path fill="#34A853" d="M12 23c2.97 0 5.46-.98 7.28-2.66l-3.57-2.77c-.98.66-2.23 1.06-3.71 1.06-2.86 0-5.29-1.93-6.16-4.53H2.18v2.84C3.99 20.53 7.7 23 12 23z"/><path fill="#FBBC05" d="M5.84 14.09c-.22-.66-.35-1.36-.35-2.09s.13-1.43.35-2.09V7.07H2.18C1.43 8.55 1 10.22 1 12s.43 3.45 1.18 4.93l2.85-2.22.81-.62z"/><path fill="#EA4335" d="M12 5.38c1.62 0 3.06.56 4.21 1.64l3.15-3.15C17.45 2.09 14.97 1 12 1 7.7 1 3.99 3.47 2.18 7.07l3.66 2.84c.87-2.6 3.3-4.53 6.16-4.53z"/></svg>`,
|
google: `<svg viewBox="0 0 24 24"><path fill="#4285F4" d="M22.56 12.25c0-.78-.07-1.53-.2-2.25H12v4.26h5.92c-.26 1.37-1.04 2.53-2.21 3.31v2.77h3.57c2.08-1.92 3.28-4.74 3.28-8.09z"/><path fill="#34A853" d="M12 23c2.97 0 5.46-.98 7.28-2.66l-3.57-2.77c-.98.66-2.23 1.06-3.71 1.06-2.86 0-5.29-1.93-6.16-4.53H2.18v2.84C3.99 20.53 7.7 23 12 23z"/><path fill="#FBBC05" d="M5.84 14.09c-.22-.66-.35-1.36-.35-2.09s.13-1.43.35-2.09V7.07H2.18C1.43 8.55 1 10.22 1 12s.43 3.45 1.18 4.93l2.85-2.22.81-.62z"/><path fill="#EA4335" d="M12 5.38c1.62 0 3.06.56 4.21 1.64l3.15-3.15C17.45 2.09 14.97 1 12 1 7.7 1 3.99 3.47 2.18 7.07l3.66 2.84c.87-2.6 3.3-4.53 6.16-4.53z"/></svg>`,
|
||||||
gemini_cli: `<svg viewBox="0 0 24 24"><path fill="#4285F4" d="M22.56 12.25c0-.78-.07-1.53-.2-2.25H12v4.26h5.92c-.26 1.37-1.04 2.53-2.21 3.31v2.77h3.57c2.08-1.92 3.28-4.74 3.28-8.09z"/><path fill="#34A853" d="M12 23c2.97 0 5.46-.98 7.28-2.66l-3.57-2.77c-.98.66-2.23 1.06-3.71 1.06-2.86 0-5.29-1.93-6.16-4.53H2.18v2.84C3.99 20.53 7.7 23 12 23z"/><path fill="#FBBC05" d="M5.84 14.09c-.22-.66-.35-1.36-.35-2.09s.13-1.43.35-2.09V7.07H2.18C1.43 8.55 1 10.22 1 12s.43 3.45 1.18 4.93l2.85-2.22.81-.62z"/><path fill="#EA4335" d="M12 5.38c1.62 0 3.06.56 4.21 1.64l3.15-3.15C17.45 2.09 14.97 1 12 1 7.7 1 3.99 3.47 2.18 7.07l3.66 2.84c.87-2.6 3.3-4.53 6.16-4.53z"/></svg>`,
|
gemini_cli: `<svg viewBox="0 0 24 24"><path fill="#4285F4" d="M22.56 12.25c0-.78-.07-1.53-.2-2.25H12v4.26h5.92c-.26 1.37-1.04 2.53-2.21 3.31v2.77h3.57c2.08-1.92 3.28-4.74 3.28-8.09z"/><path fill="#34A853" d="M12 23c2.97 0 5.46-.98 7.28-2.66l-3.57-2.77c-.98.66-2.23 1.06-3.71 1.06-2.86 0-5.29-1.93-6.16-4.53H2.18v2.84C3.99 20.53 7.7 23 12 23z"/><path fill="#FBBC05" d="M5.84 14.09c-.22-.66-.35-1.36-.35-2.09s.13-1.43.35-2.09V7.07H2.18C1.43 8.55 1 10.22 1 12s.43 3.45 1.18 4.93l2.85-2.22.81-.62z"/><path fill="#EA4335" d="M12 5.38c1.62 0 3.06.56 4.21 1.64l3.15-3.15C17.45 2.09 14.97 1 12 1 7.7 1 3.99 3.47 2.18 7.07l3.66 2.84c.87-2.6 3.3-4.53 6.16-4.53z"/></svg>`,
|
||||||
|
grok: `<svg viewBox="0 0 24 24" fill="none" xmlns="http://www.w3.org/2000/svg"><circle cx="12" cy="12" r="10" fill="#111827"/><path d="M7 7L17 17M17 7L7 17" stroke="#FFFFFF" stroke-width="2" stroke-linecap="round"/></svg>`,
|
||||||
}
|
}
|
||||||
|
|
||||||
// Default icon when provider type is not found
|
// Default icon when provider type is not found
|
||||||
|
|||||||
@@ -8,6 +8,9 @@ const PLAN_TYPE_LABELS: Record<string, string> = {
|
|||||||
'pro+': 'Pro+',
|
'pro+': 'Pro+',
|
||||||
power: 'Power',
|
power: 'Power',
|
||||||
ultra: 'Ultra',
|
ultra: 'Ultra',
|
||||||
|
basic: 'Basic',
|
||||||
|
super: 'Super',
|
||||||
|
heavy: 'Heavy',
|
||||||
}
|
}
|
||||||
|
|
||||||
const PLAN_TYPE_CLASS_NAMES: Record<string, string> = {
|
const PLAN_TYPE_CLASS_NAMES: Record<string, string> = {
|
||||||
@@ -20,6 +23,9 @@ const PLAN_TYPE_CLASS_NAMES: Record<string, string> = {
|
|||||||
ultra: 'border-amber-500/50 text-amber-600 dark:text-amber-400',
|
ultra: 'border-amber-500/50 text-amber-600 dark:text-amber-400',
|
||||||
'pro+': 'border-purple-500/50 text-purple-600 dark:text-purple-400',
|
'pro+': 'border-purple-500/50 text-purple-600 dark:text-purple-400',
|
||||||
power: 'border-amber-500/50 text-amber-600 dark:text-amber-400',
|
power: 'border-amber-500/50 text-amber-600 dark:text-amber-400',
|
||||||
|
basic: 'border-primary/50 text-primary',
|
||||||
|
super: 'border-green-500/50 text-green-600 dark:text-green-400',
|
||||||
|
heavy: 'border-amber-500/50 text-amber-600 dark:text-amber-400',
|
||||||
}
|
}
|
||||||
|
|
||||||
export function normalizeOAuthPlanType(planType?: string | null): string | null {
|
export function normalizeOAuthPlanType(planType?: string | null): string | null {
|
||||||
|
|||||||
@@ -1,4 +1,5 @@
|
|||||||
export interface ProviderKeyAuthCarrier {
|
export interface ProviderKeyAuthCarrier {
|
||||||
|
provider_type?: string | null
|
||||||
auth_type?: string | null
|
auth_type?: string | null
|
||||||
credential_kind?: string | null
|
credential_kind?: string | null
|
||||||
runtime_auth_kind?: string | null
|
runtime_auth_kind?: string | null
|
||||||
@@ -15,6 +16,14 @@ function normalizeText(value: unknown): string | null {
|
|||||||
return text || null
|
return text || null
|
||||||
}
|
}
|
||||||
|
|
||||||
|
function resolveProviderType(input: ProviderKeyAuthCarrier, providerType?: string | null): string | null {
|
||||||
|
return normalizeText(providerType) ?? normalizeText(input.provider_type)
|
||||||
|
}
|
||||||
|
|
||||||
|
function isGrokSessionCredential(input: ProviderKeyAuthCarrier, providerType?: string | null): boolean {
|
||||||
|
return resolveProviderType(input, providerType) === 'grok' && isOAuthManagedCredential(input)
|
||||||
|
}
|
||||||
|
|
||||||
export function getProviderCredentialKind(
|
export function getProviderCredentialKind(
|
||||||
input: ProviderKeyAuthCarrier,
|
input: ProviderKeyAuthCarrier,
|
||||||
): 'raw_secret' | 'oauth_session' | 'service_account' {
|
): 'raw_secret' | 'oauth_session' | 'service_account' {
|
||||||
@@ -78,7 +87,11 @@ export function canRefreshOAuthCredential(input: ProviderKeyAuthCarrier): boolea
|
|||||||
return isOAuthManagedCredential(input)
|
return isOAuthManagedCredential(input)
|
||||||
}
|
}
|
||||||
|
|
||||||
export function shouldShowOAuthRefreshControl(input: ProviderKeyAuthCarrier): boolean {
|
export function shouldShowOAuthRefreshControl(
|
||||||
|
input: ProviderKeyAuthCarrier,
|
||||||
|
providerType?: string | null,
|
||||||
|
): boolean {
|
||||||
|
if (isGrokSessionCredential(input, providerType)) return false
|
||||||
return isOAuthManagedCredential(input)
|
return isOAuthManagedCredential(input)
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -103,7 +116,11 @@ export function getProviderAuthLabel(input: ProviderKeyAuthCarrier): string {
|
|||||||
return getProviderRuntimeAuthKind(input) === 'bearer' ? 'Bearer' : 'API Key'
|
return getProviderRuntimeAuthKind(input) === 'bearer' ? 'Bearer' : 'API Key'
|
||||||
}
|
}
|
||||||
|
|
||||||
export function getProviderMaskedSecretLabel(input: ProviderKeyAuthCarrier): string {
|
export function getProviderMaskedSecretLabel(
|
||||||
|
input: ProviderKeyAuthCarrier,
|
||||||
|
providerType?: string | null,
|
||||||
|
): string {
|
||||||
|
if (isGrokSessionCredential(input, providerType)) return '[Session Cookie]'
|
||||||
if (isOAuthManagedCredential(input)) return '[OAuth Token]'
|
if (isOAuthManagedCredential(input)) return '[OAuth Token]'
|
||||||
if (isServiceAccountCredential(input)) return '[Service Account]'
|
if (isServiceAccountCredential(input)) return '[Service Account]'
|
||||||
if (getProviderRuntimeAuthKind(input) === 'mixed') return '[Key]'
|
if (getProviderRuntimeAuthKind(input) === 'mixed') return '[Key]'
|
||||||
|
|||||||
@@ -94,6 +94,37 @@ function formatQuotaValue(value: number | null | undefined): string {
|
|||||||
return normalized.toFixed(1)
|
return normalized.toFixed(1)
|
||||||
}
|
}
|
||||||
|
|
||||||
|
function getQuotaWindowValueText(window: QuotaWindowSnapshot | null | undefined): string | null {
|
||||||
|
if (!window || typeof window.limit_value !== 'number' || window.limit_value <= 0) return null
|
||||||
|
if (typeof window.remaining_value === 'number') {
|
||||||
|
return `${formatQuotaValue(window.remaining_value)}/${formatQuotaValue(window.limit_value)}`
|
||||||
|
}
|
||||||
|
if (typeof window.used_value === 'number') {
|
||||||
|
return `${formatQuotaValue(Math.max(window.limit_value - window.used_value, 0))}/${formatQuotaValue(window.limit_value)}`
|
||||||
|
}
|
||||||
|
return null
|
||||||
|
}
|
||||||
|
|
||||||
|
const GROK_QUOTA_MODE_LABELS: Record<string, string> = {
|
||||||
|
quota_auto: 'Auto',
|
||||||
|
auto: 'Auto',
|
||||||
|
quota_fast: 'Fast',
|
||||||
|
fast: 'Fast',
|
||||||
|
quota_expert: 'Expert',
|
||||||
|
expert: 'Expert',
|
||||||
|
quota_heavy: 'Heavy',
|
||||||
|
heavy: 'Heavy',
|
||||||
|
quota_grok_4_3: 'Grok 4.3',
|
||||||
|
'grok-420-computer-use-sa': 'Grok 4.3',
|
||||||
|
}
|
||||||
|
|
||||||
|
function getGrokQuotaWindowLabel(window: QuotaWindowSnapshot): string {
|
||||||
|
const rawCode = normalizeText(window.code)?.replace(/^model:/i, '') || ''
|
||||||
|
const rawLabel = normalizeText(window.label) || normalizeText(window.model) || rawCode
|
||||||
|
const normalized = (rawLabel || rawCode).trim().toLowerCase()
|
||||||
|
return GROK_QUOTA_MODE_LABELS[normalized] || GROK_QUOTA_MODE_LABELS[rawCode.toLowerCase()] || rawLabel || rawCode || '模式'
|
||||||
|
}
|
||||||
|
|
||||||
function getCodexQuotaText(quota: QuotaStatusSnapshot): string | null {
|
function getCodexQuotaText(quota: QuotaStatusSnapshot): string | null {
|
||||||
const parts: string[] = []
|
const parts: string[] = []
|
||||||
for (const [label, code] of [
|
for (const [label, code] of [
|
||||||
@@ -142,6 +173,47 @@ function getKiroQuotaText(quota: QuotaStatusSnapshot): string | null {
|
|||||||
return normalizeText(quota.label)
|
return normalizeText(quota.label)
|
||||||
}
|
}
|
||||||
|
|
||||||
|
function getGrokQuotaText(quota: QuotaStatusSnapshot): string | null {
|
||||||
|
const code = normalizeText(quota.code)?.toLowerCase()
|
||||||
|
if (code === 'banned') {
|
||||||
|
return normalizeText(quota.label) || '账号已封禁'
|
||||||
|
}
|
||||||
|
if (code === 'forbidden') {
|
||||||
|
return normalizeText(quota.label) || '访问受限'
|
||||||
|
}
|
||||||
|
|
||||||
|
const modelWindows = getQuotaWindowsByScope(quota, 'model')
|
||||||
|
const modelParts = modelWindows
|
||||||
|
.map((window) => {
|
||||||
|
const remainingPercent = getQuotaWindowRemainingPercent(window)
|
||||||
|
if (remainingPercent == null) return null
|
||||||
|
const valueText = getQuotaWindowValueText(window)
|
||||||
|
return `${getGrokQuotaWindowLabel(window)}剩余 ${formatPercent(remainingPercent)}${valueText ? ` (${valueText})` : ''}`
|
||||||
|
})
|
||||||
|
.filter((value): value is string => value != null)
|
||||||
|
|
||||||
|
if (modelParts.length > 0) return modelParts.join(' | ')
|
||||||
|
|
||||||
|
const window = getQuotaWindow(quota, 'usage') ?? getQuotaWindowsByScope(quota, 'account')[0] ?? null
|
||||||
|
const remainingPercent = getQuotaWindowRemainingPercent(window)
|
||||||
|
if (typeof window?.remaining_value === 'number' && typeof window.limit_value === 'number' && window.limit_value > 0 && window.remaining_value <= 0) {
|
||||||
|
return `剩余 ${formatQuotaValue(window.remaining_value)}/${formatQuotaValue(window.limit_value)}`
|
||||||
|
}
|
||||||
|
if (remainingPercent != null) {
|
||||||
|
const valueText = getQuotaWindowValueText(window)
|
||||||
|
if (valueText) {
|
||||||
|
return `剩余 ${formatPercent(remainingPercent)} (${valueText})`
|
||||||
|
}
|
||||||
|
return `剩余 ${formatPercent(remainingPercent)}`
|
||||||
|
}
|
||||||
|
|
||||||
|
if (typeof window?.remaining_value === 'number' && typeof window.limit_value === 'number' && window.limit_value > 0) {
|
||||||
|
return `剩余 ${formatQuotaValue(window.remaining_value)}/${formatQuotaValue(window.limit_value)}`
|
||||||
|
}
|
||||||
|
|
||||||
|
return normalizeText(quota.label)
|
||||||
|
}
|
||||||
|
|
||||||
function getAntigravityQuotaText(quota: QuotaStatusSnapshot): string | null {
|
function getAntigravityQuotaText(quota: QuotaStatusSnapshot): string | null {
|
||||||
const code = normalizeText(quota.code)?.toLowerCase()
|
const code = normalizeText(quota.code)?.toLowerCase()
|
||||||
if (code === 'forbidden') {
|
if (code === 'forbidden') {
|
||||||
@@ -238,6 +310,8 @@ export function getQuotaSnapshotFallbackText(
|
|||||||
return getCodexQuotaText(quota)
|
return getCodexQuotaText(quota)
|
||||||
case 'kiro':
|
case 'kiro':
|
||||||
return getKiroQuotaText(quota)
|
return getKiroQuotaText(quota)
|
||||||
|
case 'grok':
|
||||||
|
return getGrokQuotaText(quota)
|
||||||
case 'antigravity':
|
case 'antigravity':
|
||||||
return getAntigravityQuotaText(quota)
|
return getAntigravityQuotaText(quota)
|
||||||
case 'gemini_cli':
|
case 'gemini_cli':
|
||||||
|
|||||||
@@ -573,7 +573,7 @@
|
|||||||
<Copy class="w-2.5 h-2.5" />
|
<Copy class="w-2.5 h-2.5" />
|
||||||
</Button>
|
</Button>
|
||||||
<span class="font-mono">
|
<span class="font-mono">
|
||||||
{{ getProviderMaskedSecretLabel(key) }}
|
{{ getProviderMaskedSecretLabel(key, selectedProviderType) }}
|
||||||
</span>
|
</span>
|
||||||
<template v-if="keyUiStateMap[key.key_id]?.showOAuthRefreshControl">
|
<template v-if="keyUiStateMap[key.key_id]?.showOAuthRefreshControl">
|
||||||
<Button
|
<Button
|
||||||
@@ -603,7 +603,7 @@
|
|||||||
</span>
|
</span>
|
||||||
</template>
|
</template>
|
||||||
<Badge
|
<Badge
|
||||||
v-if="key.oauth_plan_type"
|
v-if="keyUiStateMap[key.key_id]?.planLabel"
|
||||||
variant="outline"
|
variant="outline"
|
||||||
class="text-[9px] px-1 py-0 h-4 shrink-0"
|
class="text-[9px] px-1 py-0 h-4 shrink-0"
|
||||||
:class="keyUiStateMap[key.key_id]?.planClass || ''"
|
:class="keyUiStateMap[key.key_id]?.planClass || ''"
|
||||||
@@ -637,10 +637,11 @@
|
|||||||
<div class="flex items-center justify-between text-[10px] leading-none">
|
<div class="flex items-center justify-between text-[10px] leading-none">
|
||||||
<span class="text-muted-foreground font-medium shrink-0">{{ getQuotaProgressLabel(item.label) }}</span>
|
<span class="text-muted-foreground font-medium shrink-0">{{ getQuotaProgressLabel(item.label) }}</span>
|
||||||
<span
|
<span
|
||||||
v-if="getQuotaProgressDisplayText(item)"
|
v-if="getQuotaProgressResetDisplayText(item)"
|
||||||
|
data-testid="pool-quota-reset-text"
|
||||||
class="text-muted-foreground/80 tabular-nums truncate"
|
class="text-muted-foreground/80 tabular-nums truncate"
|
||||||
:title="item.detail"
|
:title="getQuotaProgressResetDisplayText(item)"
|
||||||
>{{ getQuotaProgressDisplayText(item) }}</span>
|
>{{ getQuotaProgressResetDisplayText(item) }}</span>
|
||||||
</div>
|
</div>
|
||||||
<div class="flex items-center gap-1.5">
|
<div class="flex items-center gap-1.5">
|
||||||
<div class="relative flex-1 h-1.5 rounded-full bg-border overflow-hidden">
|
<div class="relative flex-1 h-1.5 rounded-full bg-border overflow-hidden">
|
||||||
@@ -651,9 +652,10 @@
|
|||||||
/>
|
/>
|
||||||
</div>
|
</div>
|
||||||
<span
|
<span
|
||||||
|
data-testid="pool-quota-meter-text"
|
||||||
class="shrink-0 text-[10px] font-medium tabular-nums leading-none"
|
class="shrink-0 text-[10px] font-medium tabular-nums leading-none"
|
||||||
:class="getQuotaRemainingClassByRemaining(item.remainingPercent)"
|
:class="getQuotaRemainingClassByRemaining(item.remainingPercent)"
|
||||||
>{{ item.remainingPercent.toFixed(1) }}%</span>
|
>{{ getQuotaProgressMeterDisplayText(item) }}</span>
|
||||||
</div>
|
</div>
|
||||||
</div>
|
</div>
|
||||||
</div>
|
</div>
|
||||||
@@ -1153,11 +1155,12 @@
|
|||||||
>
|
>
|
||||||
<div class="flex items-center justify-between text-[10px] leading-none">
|
<div class="flex items-center justify-between text-[10px] leading-none">
|
||||||
<span class="text-muted-foreground font-medium shrink-0">{{ getQuotaProgressLabel(item.label) }}</span>
|
<span class="text-muted-foreground font-medium shrink-0">{{ getQuotaProgressLabel(item.label) }}</span>
|
||||||
<span
|
<span
|
||||||
v-if="getQuotaProgressDisplayText(item)"
|
v-if="getQuotaProgressResetDisplayText(item)"
|
||||||
class="text-muted-foreground/80 tabular-nums truncate"
|
data-testid="pool-quota-reset-text"
|
||||||
:title="item.detail"
|
class="text-muted-foreground/80 tabular-nums truncate"
|
||||||
>{{ getQuotaProgressDisplayText(item) }}</span>
|
:title="getQuotaProgressResetDisplayText(item)"
|
||||||
|
>{{ getQuotaProgressResetDisplayText(item) }}</span>
|
||||||
</div>
|
</div>
|
||||||
<div class="flex items-center gap-1.5">
|
<div class="flex items-center gap-1.5">
|
||||||
<div class="relative flex-1 h-1.5 rounded-full bg-border overflow-hidden">
|
<div class="relative flex-1 h-1.5 rounded-full bg-border overflow-hidden">
|
||||||
@@ -1168,9 +1171,10 @@
|
|||||||
/>
|
/>
|
||||||
</div>
|
</div>
|
||||||
<span
|
<span
|
||||||
|
data-testid="pool-quota-meter-text"
|
||||||
class="shrink-0 text-[10px] font-medium tabular-nums leading-none"
|
class="shrink-0 text-[10px] font-medium tabular-nums leading-none"
|
||||||
:class="getQuotaRemainingClassByRemaining(item.remainingPercent)"
|
:class="getQuotaRemainingClassByRemaining(item.remainingPercent)"
|
||||||
>{{ item.remainingPercent.toFixed(1) }}%</span>
|
>{{ getQuotaProgressMeterDisplayText(item) }}</span>
|
||||||
</div>
|
</div>
|
||||||
</div>
|
</div>
|
||||||
</div>
|
</div>
|
||||||
@@ -2073,6 +2077,7 @@ const showAccountQuotaColumn = computed(() => {
|
|||||||
|| selectedProviderType.value === 'gemini_cli'
|
|| selectedProviderType.value === 'gemini_cli'
|
||||||
|| selectedProviderType.value === 'kiro'
|
|| selectedProviderType.value === 'kiro'
|
||||||
|| selectedProviderType.value === 'antigravity'
|
|| selectedProviderType.value === 'antigravity'
|
||||||
|
|| selectedProviderType.value === 'grok'
|
||||||
|| selectedProviderType.value === 'chatgpt_web'
|
|| selectedProviderType.value === 'chatgpt_web'
|
||||||
})
|
})
|
||||||
|
|
||||||
@@ -2377,8 +2382,9 @@ const keyUiStateMap = computed<Record<string, PoolKeyUiState>>(() => {
|
|||||||
const visibleOAuthState = getVisibleOAuthState(key)
|
const visibleOAuthState = getVisibleOAuthState(key)
|
||||||
const oauthOrgBadge = getOAuthOrgBadge(key)
|
const oauthOrgBadge = getOAuthOrgBadge(key)
|
||||||
const quotaFallbackText = getQuotaFallbackText(key)
|
const quotaFallbackText = getQuotaFallbackText(key)
|
||||||
|
const planType = resolvePoolKeyPlanType(key)
|
||||||
const canRefreshToken = canRefreshOAuthCredential(key)
|
const canRefreshToken = canRefreshOAuthCredential(key)
|
||||||
const showOAuthRefreshControl = shouldShowOAuthRefreshControl(key)
|
const showOAuthRefreshControl = shouldShowOAuthRefreshControl(key, selectedProviderType.value)
|
||||||
|
|
||||||
map[key.key_id] = {
|
map[key.key_id] = {
|
||||||
rowClass: getRowClass(key),
|
rowClass: getRowClass(key),
|
||||||
@@ -2391,8 +2397,8 @@ const keyUiStateMap = computed<Record<string, PoolKeyUiState>>(() => {
|
|||||||
oauthRefreshButtonTitle: showOAuthRefreshControl ? getOAuthRefreshButtonTitle(key) : '',
|
oauthRefreshButtonTitle: showOAuthRefreshControl ? getOAuthRefreshButtonTitle(key) : '',
|
||||||
showOAuthRefreshControl,
|
showOAuthRefreshControl,
|
||||||
canRefreshToken,
|
canRefreshToken,
|
||||||
planLabel: key.oauth_plan_type ? formatOAuthPlanType(key.oauth_plan_type) : '',
|
planLabel: planType ? formatOAuthPlanType(planType) : '',
|
||||||
planClass: key.oauth_plan_type ? getOAuthPlanTypeClass(key.oauth_plan_type) : '',
|
planClass: planType ? getOAuthPlanTypeClass(planType) : '',
|
||||||
quotaFallbackText,
|
quotaFallbackText,
|
||||||
quotaTextClass: quotaFallbackText ? getQuotaTextClass(quotaFallbackText) : '',
|
quotaTextClass: quotaFallbackText ? getQuotaTextClass(quotaFallbackText) : '',
|
||||||
importedAtRelative: formatPoolKeyImportedAt(key),
|
importedAtRelative: formatPoolKeyImportedAt(key),
|
||||||
@@ -2470,6 +2476,7 @@ const quotaRefreshSupported = computed(() => {
|
|||||||
return selectedProviderType.value === 'codex'
|
return selectedProviderType.value === 'codex'
|
||||||
|| selectedProviderType.value === 'kiro'
|
|| selectedProviderType.value === 'kiro'
|
||||||
|| selectedProviderType.value === 'antigravity'
|
|| selectedProviderType.value === 'antigravity'
|
||||||
|
|| selectedProviderType.value === 'grok'
|
||||||
|| selectedProviderType.value === 'chatgpt_web'
|
|| selectedProviderType.value === 'chatgpt_web'
|
||||||
})
|
})
|
||||||
|
|
||||||
@@ -2589,7 +2596,12 @@ async function refreshCurrentPageQuotaInBackground(
|
|||||||
|
|
||||||
if (!options.silent) {
|
if (!options.silent) {
|
||||||
const skippedText = skippedCount > 0 ? `,冷却跳过 ${skippedCount}` : ''
|
const skippedText = skippedCount > 0 ? `,冷却跳过 ${skippedCount}` : ''
|
||||||
success(`当前页额度刷新完成:成功 ${successCount},失败 ${failedCount}${skippedText}`)
|
const firstFailureMessage = result.results.find(item => item.status !== 'success')?.message?.trim()
|
||||||
|
if (successCount === 0 && failedCount > 0 && firstFailureMessage) {
|
||||||
|
showError(`当前页额度刷新失败:${firstFailureMessage}${skippedText}`)
|
||||||
|
} else {
|
||||||
|
success(`当前页额度刷新完成:成功 ${successCount},失败 ${failedCount}${skippedText}`)
|
||||||
|
}
|
||||||
}
|
}
|
||||||
return true
|
return true
|
||||||
} catch (err) {
|
} catch (err) {
|
||||||
@@ -2714,7 +2726,7 @@ function toEndpointApiKey(key: PoolKeyDetail): EndpointAPIKey {
|
|||||||
id: key.key_id,
|
id: key.key_id,
|
||||||
provider_id: selectedProviderId.value || '',
|
provider_id: selectedProviderId.value || '',
|
||||||
api_formats: key.api_formats || [],
|
api_formats: key.api_formats || [],
|
||||||
api_key_masked: getProviderMaskedSecretLabel(key),
|
api_key_masked: getProviderMaskedSecretLabel(key, selectedProviderType.value),
|
||||||
auth_type: normalizeAuthTypeForEdit(key),
|
auth_type: normalizeAuthTypeForEdit(key),
|
||||||
auth_type_by_format: key.auth_type_by_format ?? null,
|
auth_type_by_format: key.auth_type_by_format ?? null,
|
||||||
credential_kind: key.credential_kind ?? null,
|
credential_kind: key.credential_kind ?? null,
|
||||||
@@ -3524,6 +3536,7 @@ function getMobileTagItems(key: PoolKeyDetail): PoolMobileTagItem[] {
|
|||||||
const accountAlert = getAccountAlertLabel(key)
|
const accountAlert = getAccountAlertLabel(key)
|
||||||
const oauthState = getVisibleOAuthState(key)
|
const oauthState = getVisibleOAuthState(key)
|
||||||
const orgBadge = getOAuthOrgBadge(key)
|
const orgBadge = getOAuthOrgBadge(key)
|
||||||
|
const planType = resolvePoolKeyPlanType(key)
|
||||||
|
|
||||||
return buildPoolMobileTagItems({
|
return buildPoolMobileTagItems({
|
||||||
accountStatusLabel: compactPoolStatusLabel(accountAlert),
|
accountStatusLabel: compactPoolStatusLabel(accountAlert),
|
||||||
@@ -3532,7 +3545,7 @@ function getMobileTagItems(key: PoolKeyDetail): PoolMobileTagItem[] {
|
|||||||
oauthStatusTone: getMobileOAuthTone(key),
|
oauthStatusTone: getMobileOAuthTone(key),
|
||||||
priorityLabel: `P${key.internal_priority ?? 50}`,
|
priorityLabel: `P${key.internal_priority ?? 50}`,
|
||||||
authLabel: getAuthTypeChipLabel(key),
|
authLabel: getAuthTypeChipLabel(key),
|
||||||
planLabel: key.oauth_plan_type ? formatOAuthPlanType(key.oauth_plan_type) : null,
|
planLabel: planType ? formatOAuthPlanType(planType) : null,
|
||||||
orgLabel: orgBadge?.label ?? null,
|
orgLabel: orgBadge?.label ?? null,
|
||||||
proxyLabel: key.proxy?.node_id ? '独立代理' : null,
|
proxyLabel: key.proxy?.node_id ? '独立代理' : null,
|
||||||
})
|
})
|
||||||
@@ -3565,6 +3578,9 @@ function formatOAuthPlanType(planType: string): string {
|
|||||||
ultra: 'Ultra',
|
ultra: 'Ultra',
|
||||||
'pro+': 'Pro+',
|
'pro+': 'Pro+',
|
||||||
power: 'Power',
|
power: 'Power',
|
||||||
|
basic: 'Basic',
|
||||||
|
super: 'Super',
|
||||||
|
heavy: 'Heavy',
|
||||||
}
|
}
|
||||||
return labelMap[planType.toLowerCase()] || planType
|
return labelMap[planType.toLowerCase()] || planType
|
||||||
}
|
}
|
||||||
@@ -3580,6 +3596,9 @@ function getOAuthPlanTypeClass(planType: string): string {
|
|||||||
ultra: 'border-amber-500/50 text-amber-600 dark:text-amber-400',
|
ultra: 'border-amber-500/50 text-amber-600 dark:text-amber-400',
|
||||||
'pro+': 'border-purple-500/50 text-purple-600 dark:text-purple-400',
|
'pro+': 'border-purple-500/50 text-purple-600 dark:text-purple-400',
|
||||||
power: 'border-amber-500/50 text-amber-600 dark:text-amber-400',
|
power: 'border-amber-500/50 text-amber-600 dark:text-amber-400',
|
||||||
|
basic: 'border-primary/50 text-primary',
|
||||||
|
super: 'border-green-500/50 text-green-600 dark:text-green-400',
|
||||||
|
heavy: 'border-amber-500/50 text-amber-600 dark:text-amber-400',
|
||||||
}
|
}
|
||||||
return classes[planType.toLowerCase()] || ''
|
return classes[planType.toLowerCase()] || ''
|
||||||
}
|
}
|
||||||
@@ -3674,7 +3693,7 @@ function getQuotaProgressLabel(label: string): string {
|
|||||||
}
|
}
|
||||||
|
|
||||||
function getQuotaProgressCountdown(item: QuotaProgressItem) {
|
function getQuotaProgressCountdown(item: QuotaProgressItem) {
|
||||||
if (!['5H', '周', 'Spark5H', 'Spark周'].includes(item.label)) return null
|
if (!['5H', '周', 'Spark5H', 'Spark周', 'Auto', 'Fast', 'Expert', 'Heavy', 'Grok 4.3'].includes(item.label)) return null
|
||||||
if (item.resetAtSeconds == null && item.resetSeconds == null) return null
|
if (item.resetAtSeconds == null && item.resetSeconds == null) return null
|
||||||
return getCodexResetCountdown(
|
return getCodexResetCountdown(
|
||||||
item.resetAtSeconds,
|
item.resetAtSeconds,
|
||||||
@@ -3704,11 +3723,16 @@ function shouldHideQuotaProgressDetailText(text: string | null | undefined): boo
|
|||||||
return (text ?? '').trim().includes('已重置')
|
return (text ?? '').trim().includes('已重置')
|
||||||
}
|
}
|
||||||
|
|
||||||
function getQuotaProgressDisplayText(item: QuotaProgressItem): string {
|
function getQuotaProgressResetDisplayText(item: QuotaProgressItem): string {
|
||||||
const countdownText = getQuotaProgressCountdownText(item)
|
const countdownText = getQuotaProgressCountdownText(item)
|
||||||
if (countdownText) return formatCompactQuotaCountdownText(countdownText)
|
if (countdownText) return formatCompactQuotaCountdownText(countdownText)
|
||||||
|
return ''
|
||||||
|
}
|
||||||
|
|
||||||
|
function getQuotaProgressMeterDisplayText(item: QuotaProgressItem): string {
|
||||||
const detail = item.detail?.trim() || ''
|
const detail = item.detail?.trim() || ''
|
||||||
return shouldHideQuotaProgressDetailText(detail) ? '' : detail
|
if (!shouldHideQuotaProgressDetailText(detail) && detail) return detail
|
||||||
|
return `${item.remainingPercent.toFixed(1)}%`
|
||||||
}
|
}
|
||||||
|
|
||||||
function getQuotaFallbackText(key: PoolKeyDetail): string | null {
|
function getQuotaFallbackText(key: PoolKeyDetail): string | null {
|
||||||
@@ -3718,6 +3742,11 @@ function getQuotaFallbackText(key: PoolKeyDetail): string | null {
|
|||||||
|
|
||||||
|
|
||||||
function getQuotaLabelOrder(label: string): number {
|
function getQuotaLabelOrder(label: string): number {
|
||||||
|
if (label === 'Auto') return 0
|
||||||
|
if (label === 'Fast') return 1
|
||||||
|
if (label === 'Expert') return 2
|
||||||
|
if (label === 'Heavy') return 3
|
||||||
|
if (label === 'Grok 4.3') return 4
|
||||||
if (label === '5H') return 0
|
if (label === '5H') return 0
|
||||||
if (label === '周') return 1
|
if (label === '周') return 1
|
||||||
if (label === 'Spark5H') return 2
|
if (label === 'Spark5H') return 2
|
||||||
@@ -3770,6 +3799,14 @@ function getQuotaSnapshotUpdatedAtSeconds(quota: QuotaStatusSnapshot | null | un
|
|||||||
return normalizeUnixSeconds(quota?.updated_at ?? quota?.observed_at ?? null)
|
return normalizeUnixSeconds(quota?.updated_at ?? quota?.observed_at ?? null)
|
||||||
}
|
}
|
||||||
|
|
||||||
|
function getQuotaSnapshotResetAtSeconds(quota: QuotaStatusSnapshot | null | undefined): number | null {
|
||||||
|
return normalizeUnixSeconds(quota?.reset_at ?? null)
|
||||||
|
}
|
||||||
|
|
||||||
|
function getQuotaSnapshotResetSeconds(quota: QuotaStatusSnapshot | null | undefined): number | null {
|
||||||
|
return normalizeRemainingSeconds(quota?.reset_seconds ?? null)
|
||||||
|
}
|
||||||
|
|
||||||
function getQuotaSnapshotWindow(
|
function getQuotaSnapshotWindow(
|
||||||
quota: QuotaStatusSnapshot | null | undefined,
|
quota: QuotaStatusSnapshot | null | undefined,
|
||||||
code: string,
|
code: string,
|
||||||
@@ -3830,6 +3867,47 @@ function formatQuotaValue(value: number | null | undefined): string {
|
|||||||
return normalized.toFixed(1)
|
return normalized.toFixed(1)
|
||||||
}
|
}
|
||||||
|
|
||||||
|
function getQuotaWindowValueText(window: QuotaWindowSnapshot | null | undefined): string | undefined {
|
||||||
|
if (!window || typeof window.limit_value !== 'number' || window.limit_value <= 0) return undefined
|
||||||
|
if (typeof window.remaining_value === 'number') {
|
||||||
|
return `${formatQuotaValue(window.remaining_value)}/${formatQuotaValue(window.limit_value)}`
|
||||||
|
}
|
||||||
|
if (typeof window.used_value === 'number') {
|
||||||
|
return `${formatQuotaValue(Math.max(window.limit_value - window.used_value, 0))}/${formatQuotaValue(window.limit_value)}`
|
||||||
|
}
|
||||||
|
return undefined
|
||||||
|
}
|
||||||
|
|
||||||
|
function resolvePoolKeyPlanType(key: PoolKeyDetail): string | null {
|
||||||
|
const direct = key.oauth_plan_type?.trim()
|
||||||
|
if (direct) return direct
|
||||||
|
const quota = getQuotaSnapshot(key)
|
||||||
|
const quotaPlan = quota?.plan_type?.trim()
|
||||||
|
if (quotaPlan) return quotaPlan
|
||||||
|
const quotaPoolTier = quota?.pool_tier?.trim()
|
||||||
|
return quotaPoolTier || null
|
||||||
|
}
|
||||||
|
|
||||||
|
const GROK_QUOTA_MODE_LABELS: Record<string, string> = {
|
||||||
|
quota_auto: 'Auto',
|
||||||
|
auto: 'Auto',
|
||||||
|
quota_fast: 'Fast',
|
||||||
|
fast: 'Fast',
|
||||||
|
quota_expert: 'Expert',
|
||||||
|
expert: 'Expert',
|
||||||
|
quota_heavy: 'Heavy',
|
||||||
|
heavy: 'Heavy',
|
||||||
|
quota_grok_4_3: 'Grok 4.3',
|
||||||
|
'grok-420-computer-use-sa': 'Grok 4.3',
|
||||||
|
}
|
||||||
|
|
||||||
|
function getGrokQuotaWindowLabel(window: QuotaWindowSnapshot): string {
|
||||||
|
const code = String(window.code || '').trim().replace(/^model:/i, '')
|
||||||
|
const label = String(window.label || window.model || code).trim()
|
||||||
|
const normalized = (label || code).toLowerCase()
|
||||||
|
return GROK_QUOTA_MODE_LABELS[normalized] || GROK_QUOTA_MODE_LABELS[code.toLowerCase()] || label || code || '模式'
|
||||||
|
}
|
||||||
|
|
||||||
function buildQuotaProgressItemsFromSnapshot(key: PoolKeyDetail): QuotaProgressItem[] {
|
function buildQuotaProgressItemsFromSnapshot(key: PoolKeyDetail): QuotaProgressItem[] {
|
||||||
const quota = getQuotaSnapshot(key)
|
const quota = getQuotaSnapshot(key)
|
||||||
if (!quota) return []
|
if (!quota) return []
|
||||||
@@ -3838,6 +3916,8 @@ function buildQuotaProgressItemsFromSnapshot(key: PoolKeyDetail): QuotaProgressI
|
|||||||
|
|
||||||
if (providerType === 'codex') {
|
if (providerType === 'codex') {
|
||||||
const items: QuotaProgressItem[] = []
|
const items: QuotaProgressItem[] = []
|
||||||
|
const quotaResetAtSeconds = getQuotaSnapshotResetAtSeconds(quota)
|
||||||
|
const quotaResetSeconds = getQuotaSnapshotResetSeconds(quota)
|
||||||
for (const [label, code] of [
|
for (const [label, code] of [
|
||||||
['5H', '5h'],
|
['5H', '5h'],
|
||||||
['周', 'weekly'],
|
['周', 'weekly'],
|
||||||
@@ -3850,8 +3930,8 @@ function buildQuotaProgressItemsFromSnapshot(key: PoolKeyDetail): QuotaProgressI
|
|||||||
items.push({
|
items.push({
|
||||||
label,
|
label,
|
||||||
remainingPercent,
|
remainingPercent,
|
||||||
resetAtSeconds: normalizeUnixSeconds(window?.reset_at ?? null),
|
resetAtSeconds: normalizeUnixSeconds(window?.reset_at ?? quotaResetAtSeconds ?? null),
|
||||||
resetSeconds: normalizeRemainingSeconds(window?.reset_seconds ?? null),
|
resetSeconds: normalizeRemainingSeconds(window?.reset_seconds ?? quotaResetSeconds ?? null),
|
||||||
updatedAtSeconds: getQuotaSnapshotUpdatedAtSeconds(quota),
|
updatedAtSeconds: getQuotaSnapshotUpdatedAtSeconds(quota),
|
||||||
})
|
})
|
||||||
}
|
}
|
||||||
@@ -3859,6 +3939,8 @@ function buildQuotaProgressItemsFromSnapshot(key: PoolKeyDetail): QuotaProgressI
|
|||||||
}
|
}
|
||||||
|
|
||||||
if (providerType === 'kiro') {
|
if (providerType === 'kiro') {
|
||||||
|
const quotaResetAtSeconds = getQuotaSnapshotResetAtSeconds(quota)
|
||||||
|
const quotaResetSeconds = getQuotaSnapshotResetSeconds(quota)
|
||||||
const window = getQuotaSnapshotWindow(quota, 'usage')
|
const window = getQuotaSnapshotWindow(quota, 'usage')
|
||||||
?? getQuotaSnapshotWindowsByScope(quota, 'account')[0]
|
?? getQuotaSnapshotWindowsByScope(quota, 'account')[0]
|
||||||
?? null
|
?? null
|
||||||
@@ -3873,8 +3955,45 @@ function buildQuotaProgressItemsFromSnapshot(key: PoolKeyDetail): QuotaProgressI
|
|||||||
label: '剩余',
|
label: '剩余',
|
||||||
remainingPercent,
|
remainingPercent,
|
||||||
detail,
|
detail,
|
||||||
resetAtSeconds: normalizeUnixSeconds(window?.reset_at ?? null),
|
resetAtSeconds: normalizeUnixSeconds(window?.reset_at ?? quotaResetAtSeconds ?? null),
|
||||||
resetSeconds: normalizeRemainingSeconds(window?.reset_seconds ?? null),
|
resetSeconds: normalizeRemainingSeconds(window?.reset_seconds ?? quotaResetSeconds ?? null),
|
||||||
|
updatedAtSeconds: getQuotaSnapshotUpdatedAtSeconds(quota),
|
||||||
|
}]
|
||||||
|
}
|
||||||
|
|
||||||
|
if (providerType === 'grok') {
|
||||||
|
const quotaResetAtSeconds = getQuotaSnapshotResetAtSeconds(quota)
|
||||||
|
const quotaResetSeconds = getQuotaSnapshotResetSeconds(quota)
|
||||||
|
const modelWindows = getQuotaSnapshotWindowsByScope(quota, 'model')
|
||||||
|
if (modelWindows.length > 0) {
|
||||||
|
return modelWindows
|
||||||
|
.map((window): QuotaProgressItem | null => {
|
||||||
|
const remainingPercent = getQuotaWindowRemainingPercent(window)
|
||||||
|
if (remainingPercent == null) return null
|
||||||
|
return {
|
||||||
|
label: getGrokQuotaWindowLabel(window),
|
||||||
|
remainingPercent,
|
||||||
|
detail: getQuotaWindowValueText(window),
|
||||||
|
resetAtSeconds: normalizeUnixSeconds(window?.reset_at ?? quotaResetAtSeconds ?? null),
|
||||||
|
resetSeconds: normalizeRemainingSeconds(window?.reset_seconds ?? quotaResetSeconds ?? null),
|
||||||
|
updatedAtSeconds: getQuotaSnapshotUpdatedAtSeconds(quota),
|
||||||
|
}
|
||||||
|
})
|
||||||
|
.filter((item): item is QuotaProgressItem => item != null)
|
||||||
|
}
|
||||||
|
|
||||||
|
const window = getQuotaSnapshotWindow(quota, 'usage')
|
||||||
|
?? getQuotaSnapshotWindowsByScope(quota, 'account')[0]
|
||||||
|
?? null
|
||||||
|
const remainingPercent = getQuotaWindowRemainingPercent(window)
|
||||||
|
if (remainingPercent == null) return []
|
||||||
|
|
||||||
|
return [{
|
||||||
|
label: '剩余',
|
||||||
|
remainingPercent,
|
||||||
|
detail: getQuotaWindowValueText(window),
|
||||||
|
resetAtSeconds: normalizeUnixSeconds(window?.reset_at ?? quotaResetAtSeconds ?? null),
|
||||||
|
resetSeconds: normalizeRemainingSeconds(window?.reset_seconds ?? quotaResetSeconds ?? null),
|
||||||
updatedAtSeconds: getQuotaSnapshotUpdatedAtSeconds(quota),
|
updatedAtSeconds: getQuotaSnapshotUpdatedAtSeconds(quota),
|
||||||
}]
|
}]
|
||||||
}
|
}
|
||||||
@@ -4075,29 +4194,6 @@ function getQuotaTextClass(quotaText: string): string {
|
|||||||
return 'text-[11px] text-foreground/90 leading-4'
|
return 'text-[11px] text-foreground/90 leading-4'
|
||||||
}
|
}
|
||||||
|
|
||||||
function formatStatInteger(value: number | null | undefined): string {
|
|
||||||
const n = Number(value ?? 0)
|
|
||||||
if (!Number.isFinite(n) || n <= 0) return '0'
|
|
||||||
return Math.round(n).toLocaleString('en-US')
|
|
||||||
}
|
|
||||||
|
|
||||||
function formatTokenCount(value: number | null | undefined): string {
|
|
||||||
const n = Number(value ?? 0)
|
|
||||||
if (!Number.isFinite(n) || n <= 0) return '0'
|
|
||||||
if (n >= 1_000_000) return `${(n / 1_000_000).toFixed(1)}M`
|
|
||||||
if (n >= 1_000) return `${(n / 1_000).toFixed(1)}K`
|
|
||||||
return String(Math.round(n))
|
|
||||||
}
|
|
||||||
|
|
||||||
function formatStatUsd(value: number | string | null | undefined): string {
|
|
||||||
const n = Number(value ?? 0)
|
|
||||||
if (!Number.isFinite(n) || n <= 0) return '$0.00'
|
|
||||||
if (n < 0.01) return `$${n.toFixed(4)}`
|
|
||||||
if (n < 1) return `$${n.toFixed(3)}`
|
|
||||||
if (n < 1000) return `$${n.toFixed(2)}`
|
|
||||||
return `$${n.toLocaleString('en-US', { minimumFractionDigits: 2, maximumFractionDigits: 2 })}`
|
|
||||||
}
|
|
||||||
|
|
||||||
function formatPoolScore(value: number | null | undefined): string {
|
function formatPoolScore(value: number | null | undefined): string {
|
||||||
const n = Number(value)
|
const n = Number(value)
|
||||||
if (!Number.isFinite(n)) return '-'
|
if (!Number.isFinite(n)) return '-'
|
||||||
|
|||||||
Reference in New Issue
Block a user