Merge origin/main into fix/gemini-cli-v1internal

This commit is contained in:
Mas0nShi
2026-05-28 11:58:00 +08:00
410 changed files with 38026 additions and 6621 deletions
@@ -1,6 +1,9 @@
use crate::handlers::admin::request::AdminAppState;
use crate::provider_key_auth::provider_key_effective_api_formats;
use aether_scheduler_core::count_recent_rpm_requests_for_provider_key_since;
use aether_scheduler_core::{
count_recent_rpm_requests_for_provider_key_since,
provider_key_circuit_payload_is_active_open_at,
};
use serde_json::json;
use std::time::{SystemTime, UNIX_EPOCH};
@@ -18,6 +21,10 @@ pub(crate) async fn build_admin_key_health_payload(
.await
.ok()
.and_then(|mut keys| keys.drain(..).next())?;
let now_unix_secs = SystemTime::now()
.duration_since(UNIX_EPOCH)
.map(|duration| duration.as_secs())
.unwrap_or_default();
let provider = state
.read_provider_catalog_providers_by_ids(std::slice::from_ref(&key.provider_id))
.await
@@ -81,10 +88,10 @@ pub(crate) async fn build_admin_key_health_payload(
.and_then(|value| value.get("last_failure_at"))
.cloned()
.unwrap_or(serde_json::Value::Null);
payload["circuit_breaker_open"] = json!(circuit_data
.and_then(|value| value.get("open"))
.and_then(serde_json::Value::as_bool)
.unwrap_or(false));
payload["circuit_breaker_open"] =
json!(circuit_data.is_some_and(
|value| provider_key_circuit_payload_is_active_open_at(value, now_unix_secs)
));
payload["circuit_breaker_open_at"] = circuit_data
.and_then(|value| value.get("open_at"))
.cloned()
@@ -107,15 +114,16 @@ pub(crate) async fn build_admin_key_health_payload(
.unwrap_or(0));
} else {
let mut formats_payload = serde_json::Map::new();
let mut any_circuit_open = false;
for format_name in
provider_key_effective_api_formats(&key, &provider.provider_type, &endpoints)
{
let health_data = health_by_format.and_then(|formats| formats.get(&format_name));
let circuit_data = circuit_by_format.and_then(|formats| formats.get(&format_name));
let is_open = circuit_data
.and_then(|value| value.get("open"))
.and_then(serde_json::Value::as_bool)
.unwrap_or(false);
let active_open = circuit_data.is_some_and(|value| {
provider_key_circuit_payload_is_active_open_at(value, now_unix_secs)
});
any_circuit_open |= active_open;
formats_payload.insert(
format_name.clone(),
json!({
@@ -134,8 +142,8 @@ pub(crate) async fn build_admin_key_health_payload(
.cloned()
.unwrap_or(serde_json::Value::Null),
"circuit_breaker": {
"state": if is_open { "open" } else { "closed" },
"open": is_open,
"state": if active_open { "open" } else { "closed" },
"open": active_open,
"open_at": circuit_data
.and_then(|value| value.get("open_at"))
.cloned()
@@ -167,13 +175,6 @@ pub(crate) async fn build_admin_key_health_payload(
.filter_map(serde_json::Value::as_f64)
.reduce(f64::min)
.unwrap_or(1.0);
let any_circuit_open = formats_payload.values().any(|value| {
value
.get("circuit_breaker")
.and_then(|circuit| circuit.get("open"))
.and_then(serde_json::Value::as_bool)
.unwrap_or(false)
});
payload["key_health_score"] = json!(key_health_score);
payload["any_circuit_open"] = json!(any_circuit_open);
@@ -4,7 +4,9 @@ use crate::handlers::public::{api_format_display_name, build_public_health_timel
use crate::handlers::shared::unix_ms_to_rfc3339;
use crate::provider_key_auth::provider_key_effective_api_formats;
use aether_data_contracts::repository::candidates::PublicHealthTimelineBucket;
use aether_scheduler_core::{is_provider_key_circuit_open, provider_key_health_score};
use aether_scheduler_core::{
any_provider_key_circuit_open_at, is_provider_key_circuit_open_at, provider_key_health_score,
};
use serde_json::json;
use std::collections::{BTreeMap, BTreeSet};
use std::time::{SystemTime, UNIX_EPOCH};
@@ -99,7 +101,9 @@ pub(crate) async fn build_admin_endpoint_health_status_payload(
.entry(api_format.clone())
.or_default()
.insert(key.id.clone());
if key.is_active && !is_provider_key_circuit_open(&key, &api_format) {
if key.is_active
&& !is_provider_key_circuit_open_at(&key, &api_format, now_unix_secs)
{
let key_health_score =
provider_key_health_score(&key, &api_format).unwrap_or(1.0);
active_keys_by_format
@@ -229,6 +233,10 @@ pub(crate) async fn build_admin_health_summary_payload(
return None;
}
let now_unix_secs = SystemTime::now()
.duration_since(UNIX_EPOCH)
.map(|duration| duration.as_secs())
.unwrap_or_default();
let providers = state
.list_provider_catalog_providers(false)
.await
@@ -286,20 +294,7 @@ pub(crate) async fn build_admin_health_summary_payload(
.count();
let circuit_open_keys = keys
.iter()
.filter(|key| {
key.circuit_breaker_by_format
.as_ref()
.and_then(serde_json::Value::as_object)
.map(|formats| {
formats.values().any(|circuit| {
circuit
.get("open")
.and_then(serde_json::Value::as_bool)
.unwrap_or(false)
})
})
.unwrap_or(false)
})
.filter(|key| any_provider_key_circuit_open_at(key, now_unix_secs))
.count();
Some(json!({
@@ -8,10 +8,12 @@ use aether_data_contracts::repository::provider_catalog::{
StoredProviderCatalogEndpoint, StoredProviderCatalogKey,
};
use aether_scheduler_core::{
is_provider_key_circuit_open, matches_model_mapping, provider_key_health_score,
is_provider_key_circuit_open_at, matches_model_mapping,
provider_key_circuit_payload_is_active_open_at, provider_key_health_score,
};
use serde_json::json;
use std::collections::BTreeMap;
use std::time::{SystemTime, UNIX_EPOCH};
use uuid::Uuid;
pub(crate) async fn build_admin_global_model_routing_payload(
@@ -86,6 +88,10 @@ pub(crate) async fn build_admin_global_model_routing_payload(
.flatten()
.and_then(|value| value.as_bool())
.unwrap_or(false);
let now_unix_secs = SystemTime::now()
.duration_since(UNIX_EPOCH)
.map(|duration| duration.as_secs())
.unwrap_or_default();
let global_model_mappings = global_model
.config
@@ -163,9 +169,11 @@ pub(crate) async fn build_admin_global_model_routing_payload(
entries
.iter()
.filter_map(|(api_format, value)| {
value.get("open")
.and_then(serde_json::Value::as_bool)
.filter(|is_open| *is_open)
provider_key_circuit_payload_is_active_open_at(
value,
now_unix_secs,
)
.then_some(())
.map(|_| api_format.clone())
})
.collect::<Vec<_>>()
@@ -188,7 +196,7 @@ pub(crate) async fn build_admin_global_model_routing_payload(
"effective_rpm": effective_rpm,
"allowed_models": allowed_models,
"health_score": provider_key_health_score(key, &endpoint.api_format),
"circuit_breaker_open": is_provider_key_circuit_open(key, &endpoint.api_format),
"circuit_breaker_open": is_provider_key_circuit_open_at(key, &endpoint.api_format, now_unix_secs),
"circuit_breaker_formats": circuit_breaker_formats,
"next_probe_at": next_probe_at,
});
@@ -1,6 +1,6 @@
use super::super::usage_helpers::admin_monitoring_usage_is_error;
use crate::handlers::admin::request::AdminAppState;
use crate::handlers::admin::shared::{provider_key_health_summary, unix_secs_to_rfc3339};
use crate::handlers::admin::shared::{provider_key_health_summary_at, unix_secs_to_rfc3339};
use crate::GatewayError;
use aether_data_contracts::repository::{
provider_catalog::StoredProviderCatalogKey, usage::UsageMonitoringErrorListQuery,
@@ -99,7 +99,7 @@ pub(super) async fn build_admin_monitoring_resilience_snapshot(
last_failure_at,
circuit_breaker_open,
circuit_by_format,
) = provider_key_health_summary(key);
) = provider_key_health_summary_at(key, now.timestamp().max(0) as u64);
if health_score < 0.8 {
degraded_keys += 1;
}
@@ -110,11 +110,12 @@ pub(super) async fn build_admin_monitoring_resilience_snapshot(
let open_formats = circuit_by_format
.iter()
.filter_map(|(api_format, value)| {
value
.get("open")
.and_then(serde_json::Value::as_bool)
.filter(|open| *open)
.map(|_| api_format.clone())
aether_scheduler_core::provider_key_circuit_payload_is_active_open_at(
value,
now.timestamp().max(0) as u64,
)
.then_some(())
.map(|_| api_format.clone())
})
.collect::<Vec<_>>();
@@ -1237,7 +1237,8 @@ async fn admin_monitoring_circuit_history_returns_local_payload() {
"openai:chat": {
"open": true,
"open_at": "2026-03-30T12:00:00+00:00",
"next_probe_at": "2026-03-30T12:05:00+00:00",
"next_probe_at": "2099-03-30T12:05:00+00:00",
"recovery_seconds": 300,
"reason": "错误率过高"
}
})),
@@ -81,6 +81,133 @@ async fn admin_monitoring_trace_request_returns_local_payload() {
assert_eq!(payload["candidates"][0]["status_code"], json!(502));
}
#[tokio::test]
async fn admin_monitoring_trace_request_resolves_usage_id_to_header_trace_id() {
let request_candidates = Arc::new(InMemoryRequestCandidateRepository::seed(vec![
sample_candidate(
"cand-used",
"trace-1",
0,
RequestCandidateStatus::Success,
Some(101),
Some(33),
Some(200),
),
]));
let provider_catalog = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![sample_provider()],
vec![sample_endpoint()],
vec![sample_key()],
));
let mut usage = sample_usage(
"usage-request-1",
"provider-1",
"OpenAI",
40,
0.02,
"completed",
Some(200),
100,
);
usage.id = "usage-row-1".to_string();
usage.candidate_id = Some("cand-used".to_string());
usage.request_headers = Some(json!({
"x-trace-id": "trace-1"
}));
let usage_repository = Arc::new(InMemoryUsageReadRepository::seed(vec![usage]));
let data_state =
crate::data::GatewayDataState::with_request_candidate_and_usage_repository_for_tests(
request_candidates,
usage_repository,
)
.with_provider_catalog_reader(provider_catalog);
let state = AppState::new()
.expect("state should build")
.with_data_state_for_tests(data_state);
let context = request_context(
http::Method::GET,
"/api/admin/monitoring/trace/usage-row-1?attempted_only=true",
);
let response = local_monitoring_response(&state, &context)
.await
.expect("handler should not error")
.expect("route should be handled locally");
assert_eq!(response.status(), http::StatusCode::OK);
let body = to_bytes(response.into_body(), usize::MAX)
.await
.expect("body should read");
let payload: serde_json::Value = serde_json::from_slice(&body).expect("json body should parse");
assert_eq!(payload["request_id"], json!("trace-1"));
assert_eq!(payload["candidates"][0]["id"], json!("cand-used"));
assert_eq!(
payload["candidates"][0]["extra_data"]["first_byte_time_ms"],
json!(30)
);
}
#[tokio::test]
async fn admin_monitoring_trace_request_resolves_usage_request_id_to_metadata_trace_id() {
let request_candidates = Arc::new(InMemoryRequestCandidateRepository::seed(vec![
sample_candidate(
"cand-used",
"trace-2",
0,
RequestCandidateStatus::Success,
Some(101),
Some(33),
Some(200),
),
]));
let provider_catalog = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![sample_provider()],
vec![sample_endpoint()],
vec![sample_key()],
));
let mut usage = sample_usage(
"usage-request-2",
"provider-1",
"OpenAI",
40,
0.02,
"completed",
Some(200),
100,
);
usage.candidate_id = Some("cand-used".to_string());
usage.request_metadata = Some(json!({
"trace_id": "trace-2"
}));
let usage_repository = Arc::new(InMemoryUsageReadRepository::seed(vec![usage]));
let data_state =
crate::data::GatewayDataState::with_request_candidate_and_usage_repository_for_tests(
request_candidates,
usage_repository,
)
.with_provider_catalog_reader(provider_catalog);
let state = AppState::new()
.expect("state should build")
.with_data_state_for_tests(data_state);
let context = request_context(
http::Method::GET,
"/api/admin/monitoring/trace/usage-request-2",
);
let response = local_monitoring_response(&state, &context)
.await
.expect("handler should not error")
.expect("route should be handled locally");
assert_eq!(response.status(), http::StatusCode::OK);
let body = to_bytes(response.into_body(), usize::MAX)
.await
.expect("body should read");
let payload: serde_json::Value = serde_json::from_slice(&body).expect("json body should parse");
assert_eq!(payload["request_id"], json!("trace-2"));
assert_eq!(payload["candidates"][0]["id"], json!("cand-used"));
}
#[tokio::test]
async fn admin_monitoring_trace_request_returns_oauth_account_label_from_auth_config() {
let request_candidates = Arc::new(InMemoryRequestCandidateRepository::seed(vec![
@@ -12,6 +12,7 @@ use aether_admin::observability::monitoring::{
use aether_data_contracts::repository::{
candidates::{DecisionTrace, RequestCandidateStatus},
provider_catalog::StoredProviderCatalogKey,
usage::StoredRequestUsageAudit,
};
use axum::{
body::Body,
@@ -21,12 +22,16 @@ use serde_json::{Map, Value};
use std::collections::BTreeMap;
use tracing::debug;
struct ResolvedAdminMonitoringTrace {
trace: DecisionTrace,
usage: Option<StoredRequestUsageAudit>,
}
pub(super) async fn build_admin_monitoring_trace_request_response(
state: &AdminAppState<'_>,
request_context: &AdminRequestContext<'_>,
) -> Result<Response<Body>, GatewayError> {
let admin_state = state;
let state = state.as_ref();
let Some(request_id) =
admin_monitoring_trace_request_id_from_path(&request_context.request_path)
else {
@@ -39,11 +44,8 @@ pub(super) async fn build_admin_monitoring_trace_request_response(
Err(detail) => return Ok(admin_monitoring_bad_request_response(detail)),
};
let Some(trace) = state
.data
.read_decision_trace(&request_id, attempted_only)
.await
.map_err(|err| GatewayError::Internal(err.to_string()))?
let Some(resolved) =
resolve_admin_monitoring_trace(admin_state, &request_id, attempted_only).await?
else {
debug!(
event_name = "admin_monitoring_request_trace_not_found",
@@ -58,22 +60,113 @@ pub(super) async fn build_admin_monitoring_trace_request_response(
attempted_only,
));
};
let usage = state
.data
.read_request_usage_audit(&request_id)
.await
.map_err(|err| GatewayError::Internal(err.to_string()))?;
let key_accounts = build_admin_monitoring_key_account_display_map(admin_state, &trace).await?;
let key_accounts =
build_admin_monitoring_key_account_display_map(admin_state, &resolved.trace).await?;
Ok(
build_admin_monitoring_trace_request_payload_response_with_key_accounts(
&trace,
usage.as_ref(),
&resolved.trace,
resolved.usage.as_ref(),
&key_accounts,
),
)
}
async fn resolve_admin_monitoring_trace(
state: &AdminAppState<'_>,
request_id: &str,
attempted_only: bool,
) -> Result<Option<ResolvedAdminMonitoringTrace>, GatewayError> {
let app = state.as_ref();
if let Some(trace) = app
.data
.read_decision_trace(request_id, attempted_only)
.await
.map_err(|err| GatewayError::Internal(err.to_string()))?
{
let usage = app
.data
.read_request_usage_audit(request_id)
.await
.map_err(|err| GatewayError::Internal(err.to_string()))?;
return Ok(Some(ResolvedAdminMonitoringTrace { trace, usage }));
}
let mut usage_candidates = Vec::new();
if let Some(usage) = app
.data
.read_request_usage_audit(request_id)
.await
.map_err(|err| GatewayError::Internal(err.to_string()))?
{
usage_candidates.push(usage);
}
if let Some(usage) = state.find_request_usage_by_id(request_id).await? {
if !usage_candidates.iter().any(|item| item.id == usage.id) {
usage_candidates.push(usage);
}
}
for usage in usage_candidates {
for trace_request_id in admin_monitoring_usage_trace_request_ids(&usage) {
if trace_request_id == request_id {
continue;
}
if let Some(trace) = app
.data
.read_decision_trace(&trace_request_id, attempted_only)
.await
.map_err(|err| GatewayError::Internal(err.to_string()))?
{
return Ok(Some(ResolvedAdminMonitoringTrace {
trace,
usage: Some(usage),
}));
}
}
}
Ok(None)
}
fn admin_monitoring_usage_trace_request_ids(usage: &StoredRequestUsageAudit) -> Vec<String> {
let mut ids = Vec::new();
push_non_empty_unique(&mut ids, usage.request_id.as_str());
if let Some(trace_id) = usage.trace_id() {
push_non_empty_unique(&mut ids, trace_id);
}
if let Some(trace_id) = usage_trace_id_from_headers(usage.request_headers.as_ref()) {
push_non_empty_unique(&mut ids, trace_id.as_str());
}
if let Some(trace_id) = usage_trace_id_from_headers(usage.provider_request_headers.as_ref()) {
push_non_empty_unique(&mut ids, trace_id.as_str());
}
ids
}
fn usage_trace_id_from_headers(headers: Option<&Value>) -> Option<String> {
let object = headers?.as_object()?;
object.iter().find_map(|(key, value)| {
key.eq_ignore_ascii_case(crate::constants::TRACE_ID_HEADER)
.then(|| {
value
.as_str()
.map(str::trim)
.filter(|value| !value.is_empty())
})
.flatten()
.map(ToOwned::to_owned)
})
}
fn push_non_empty_unique(values: &mut Vec<String>, value: &str) {
let value = value.trim();
if value.is_empty() || values.iter().any(|existing| existing == value) {
return;
}
values.push(value.to_string());
}
async fn build_admin_monitoring_key_account_display_map(
state: &AdminAppState<'_>,
trace: &DecisionTrace,
@@ -91,6 +91,26 @@ fn coerce_admin_provider_oauth_import_project_id(
}
}
fn json_import_expiry_value(value: Option<&serde_json::Value>) -> Option<u64> {
let value = value?;
json_u64_value(Some(value)).or_else(|| {
value
.as_str()
.map(str::trim)
.filter(|value| !value.is_empty())
.and_then(|value| chrono::DateTime::parse_from_rfc3339(value).ok())
.and_then(|value| u64::try_from(value.timestamp()).ok())
})
}
fn json_import_expiry_from_keys(
object: &serde_json::Map<String, serde_json::Value>,
keys: &[&str],
) -> Option<u64> {
keys.iter()
.find_map(|key| json_import_expiry_value(object.get(*key)))
}
fn grok_cookie_value(raw: &str, name: &str) -> Option<String> {
raw.trim()
.strip_prefix("Cookie:")
@@ -204,6 +224,8 @@ fn extract_admin_provider_oauth_batch_import_entry(
object
.get("sso_token")
.or_else(|| object.get("ssoToken"))
.or_else(|| object.get("session_token"))
.or_else(|| object.get("sessionToken"))
.or(grok_token_alias),
)
.or_else(|| {
@@ -254,7 +276,7 @@ fn extract_admin_provider_oauth_batch_import_entry(
refresh_token
};
let expires_at =
json_u64_value(object.get("expires_at").or_else(|| object.get("expiresAt")));
json_import_expiry_from_keys(object, &["expires_at", "expiresAt", "expired"]);
let account_id = coerce_admin_provider_oauth_import_str(
object
.get("account_id")
@@ -660,6 +682,21 @@ mod tests {
assert_eq!(entries[0].email.as_deref(), Some("[email protected]"));
}
#[test]
fn parses_common_chatgpt_web_json_aliases() {
let entries = parse_admin_provider_oauth_batch_import_entries(
"chatgpt_web",
r#"[{"session_token":"session-1","expired":"2030-01-01T00:00:00Z","chatgpt_account_id":"acc-1","chatgpt_plan_type":"plus"}]"#,
);
assert_eq!(entries.len(), 1);
assert_eq!(entries[0].refresh_token, None);
assert_eq!(entries[0].access_token.as_deref(), Some("session-1"));
assert_eq!(entries[0].expires_at, Some(1_893_456_000));
assert_eq!(entries[0].account_id.as_deref(), Some("acc-1"));
assert_eq!(entries[0].plan_type.as_deref(), Some("plus"));
}
#[test]
fn parses_plain_jwt_line_as_access_token() {
let token = unsigned_jwt(json!({
@@ -1,7 +1,9 @@
use super::super::super::errors::build_internal_control_error_response;
use super::super::super::provisioning::provider_oauth_token_payload_expires_at_unix_secs;
use super::super::super::quota::codex::refresh_codex_provider_quota_locally;
use super::super::super::runtime::resolve_provider_oauth_runtime_endpoints;
use super::super::super::runtime::{
resolve_provider_oauth_runtime_endpoints,
spawn_provider_oauth_account_state_refresh_after_update,
};
use super::super::super::state::{
admin_provider_oauth_template, enrich_admin_provider_oauth_auth_config,
is_fixed_provider_type_for_provider_oauth, json_non_empty_string,
@@ -219,50 +221,20 @@ pub(super) async fn handle_admin_provider_oauth_complete_key(
));
}
let mut account_state_recheck_attempted = false;
let mut account_state_recheck_error = None::<String>;
if provider_type == "codex" {
if let Some(endpoint) = runtime_endpoint {
let refreshed_key = state
.read_provider_catalog_keys_by_ids(std::slice::from_ref(&key_id))
.await?
.into_iter()
.next()
.unwrap_or_else(|| key.clone());
if let Some(result) = refresh_codex_provider_quota_locally(
state,
&provider,
&endpoint,
vec![refreshed_key],
request_proxy.clone(),
)
.await?
{
account_state_recheck_attempted = true;
let success = result
.get("success")
.and_then(serde_json::Value::as_u64)
.unwrap_or(0);
if success == 0 {
account_state_recheck_error = result
.get("results")
.and_then(serde_json::Value::as_array)
.and_then(|results| results.first())
.and_then(|value| value.get("message"))
.and_then(serde_json::Value::as_str)
.map(ToOwned::to_owned);
}
}
}
}
spawn_provider_oauth_account_state_refresh_after_update(
state.cloned_app(),
provider.clone(),
key_id.clone(),
request_proxy.clone(),
);
Ok(Json(json!({
"provider_type": provider_type,
"expires_at": expires_at,
"has_refresh_token": refresh_token.is_some(),
"email": auth_config.get("email").cloned().unwrap_or(serde_json::Value::Null),
"account_state_recheck_attempted": account_state_recheck_attempted,
"account_state_recheck_error": account_state_recheck_error,
"account_state_recheck_attempted": false,
"account_state_recheck_error": serde_json::Value::Null,
}))
.into_response())
}
@@ -828,13 +828,7 @@ async fn handle_admin_provider_oauth_windsurf_browser_device_poll(
};
callback_token.to_string()
} else {
let token = token.unwrap_or_default();
if windsurf_raw_api_key(token).is_none() {
return Ok(windsurf_browser_poll_error_response(
"浏览器授权请提交包含 state 的回调 URL;纯 token 请使用导入授权",
));
}
token.to_string()
token.unwrap_or_default().to_string()
};
let mut raw_credentials = serde_json::Map::new();
@@ -106,12 +106,21 @@ fn import_payload_string_any(
.map(ToOwned::to_owned)
}
fn import_payload_u64(
fn import_payload_u64_any(
payload: &serde_json::Map<String, serde_json::Value>,
snake_case: &str,
camel_case: &str,
keys: &[&str],
) -> Option<u64> {
json_u64_value(payload.get(snake_case).or_else(|| payload.get(camel_case)))
keys.iter().find_map(|key| {
let value = payload.get(*key)?;
json_u64_value(Some(value)).or_else(|| {
value
.as_str()
.map(str::trim)
.filter(|value| !value.is_empty())
.and_then(|value| chrono::DateTime::parse_from_rfc3339(value).ok())
.and_then(|value| u64::try_from(value.timestamp()).ok())
})
})
}
fn apply_single_import_hints(
@@ -409,9 +418,17 @@ pub(super) async fn handle_admin_provider_oauth_import_refresh_token(
let refresh_token_input = import_payload_string(&raw_payload, "refresh_token", "refreshToken");
let access_token_input = import_payload_string_any(
&raw_payload,
&["access_token", "accessToken", "sso_token", "ssoToken"],
&[
"access_token",
"accessToken",
"sso_token",
"ssoToken",
"session_token",
"sessionToken",
],
);
let imported_expires_at = import_payload_u64(&raw_payload, "expires_at", "expiresAt");
let imported_expires_at =
import_payload_u64_any(&raw_payload, &["expires_at", "expiresAt", "expired"]);
let name = raw_payload
.get("name")
.and_then(serde_json::Value::as_str)
@@ -622,8 +639,41 @@ pub(super) async fn handle_admin_provider_oauth_import_refresh_token(
#[cfg(test)]
mod tests {
use super::sanitize_windsurf_import_error;
use super::{
import_payload_string_any, import_payload_u64_any, sanitize_windsurf_import_error,
};
use aether_oauth::core::OAuthError;
use serde_json::json;
#[test]
fn single_import_accepts_session_token_alias() {
let payload = json!({
"session_token": "session-1",
})
.as_object()
.cloned()
.expect("payload should be an object");
assert_eq!(
import_payload_string_any(&payload, &["access_token", "session_token"]).as_deref(),
Some("session-1")
);
}
#[test]
fn single_import_accepts_iso_expired_alias() {
let payload = json!({
"expired": "2030-01-01T00:00:00Z",
})
.as_object()
.cloned()
.expect("payload should be an object");
assert_eq!(
import_payload_u64_any(&payload, &["expires_at", "expiresAt", "expired"]),
Some(1_893_456_000)
);
}
#[test]
fn windsurf_import_error_redacts_http_body() {
@@ -81,48 +81,42 @@ pub(super) async fn execute_admin_provider_oauth_refresh(
.await?;
if provider_auto_remove_banned_keys(provider.config.as_ref()) {
let now_unix_secs = helpers::unix_now_secs();
let latest_key = state
.read_provider_catalog_keys_by_ids(std::slice::from_ref(&key_id))
.await?
.into_iter()
.next();
if latest_key.as_ref().is_some_and(|latest_key| {
should_auto_remove_oauth_invalid_key(
latest_key,
None,
false,
now_unix_secs,
)
}) {
state
.clear_admin_provider_pool_cooldown(&provider.id, &key_id)
.await;
state
.reset_admin_provider_pool_cost(&provider.id, &key_id)
.await;
if state.delete_provider_catalog_key(&key_id).await? {
let deleted_key_ids = [key_id.clone()];
state
.cleanup_deleted_provider_catalog_refs(
&provider.id,
&[],
&deleted_key_ids,
let auto_removed = state
.cleanup_provider_catalog_key_if_current(
&provider,
&key_id,
|latest_key| {
should_auto_remove_oauth_invalid_key(
latest_key,
Some(&failure_reason),
false,
now_unix_secs,
)
.await?;
tracing::info!(
trace_id = %trace_id,
key_id = %key_id,
provider_id = %provider.id,
provider_type = %provider_type,
event_name = "auto_removed_oauth_refresh_failed",
"gateway manual provider oauth refresh auto-removed unusable key"
);
return Ok(RefreshDispatch::Respond(
response::oauth_refresh_auto_removed_response(&error_reason),
));
}
},
)
.await?;
if auto_removed {
tracing::info!(
trace_id = %trace_id,
key_id = %key_id,
provider_id = %provider.id,
provider_type = %provider_type,
event_name = "auto_removed_oauth_refresh_failed",
"gateway manual provider oauth refresh auto-removed unusable key"
);
return Ok(RefreshDispatch::Respond(
response::oauth_refresh_auto_removed_response(&error_reason),
));
}
}
tracing::info!(
trace_id = %trace_id,
key_id = %key_id,
provider_id = %provider.id,
provider_type = %provider_type,
event_name = "refresh_failed_retained",
"gateway manual provider oauth refresh failure retained key"
);
}
}
return Ok(RefreshDispatch::Respond(
@@ -171,9 +165,25 @@ pub(super) async fn execute_admin_provider_oauth_refresh(
};
if !helpers::key_is_account_blocked(&key, OAUTH_ACCOUNT_BLOCK_PREFIX) {
let _ = state
let previous_oauth_refresh_issue =
key.oauth_invalid_reason.as_deref().is_some_and(|reason| {
reason.lines().map(str::trim).any(|line| {
line.starts_with("[OAUTH_EXPIRED]") || line.starts_with("[REFRESH_FAILED]")
})
});
let cleared = state
.clear_provider_catalog_key_oauth_invalid_marker(&key_id)
.await?;
if cleared && previous_oauth_refresh_issue {
tracing::info!(
trace_id = %trace_id,
key_id = %key_id,
provider_id = %provider.id,
provider_type = %provider_type,
event_name = "refresh_fixed",
"gateway manual provider oauth refresh cleared oauth invalid marker"
);
}
}
let refreshed_key = state
@@ -201,10 +201,10 @@ pub(crate) async fn update_existing_provider_oauth_catalog_key(
.map(|duration| duration.as_secs())
.unwrap_or(0);
let mut updated = existing_key.clone();
updated.is_active = true;
updated.encrypted_api_key = Some(encrypted_api_key);
updated.encrypted_auth_config = Some(encrypted_auth_config);
updated.api_formats = provider_oauth_catalog_key_api_formats(provider_type, api_formats);
updated.is_active = true;
updated.expires_at_unix_secs = expires_at_unix_secs;
updated.oauth_invalid_at_unix_secs = None;
updated.oauth_invalid_reason = None;
@@ -1,9 +1,9 @@
use super::shared::{
build_provider_quota_execution_plan, build_quota_snapshot_payload, coerce_json_f64,
coerce_json_string, default_provider_quota_execution_timeouts, execute_provider_quota_plan,
extract_execution_error_message, oauth_refresh_auto_removed_result,
persist_provider_quota_refresh_state, quota_key_auto_removed,
quota_refresh_success_invalid_state, ProviderQuotaExecutionOutcome,
coerce_json_string, execute_provider_quota_plan, extract_execution_error_message,
oauth_refresh_auto_removed_result, persist_provider_quota_refresh_state,
quota_key_auto_removed, quota_refresh_success_invalid_state,
resolve_provider_quota_execution_timeouts, ProviderQuotaExecutionOutcome,
};
use crate::handlers::admin::request::{AdminAppState, AdminGatewayProviderTransportSnapshot};
use crate::GatewayError;
@@ -33,11 +33,10 @@ async fn execute_antigravity_quota_plan(
.await
}
};
let timeouts = state
.resolve_transport_execution_timeouts(transport)
.or(Some(default_provider_quota_execution_timeouts(
proxy.as_ref(),
)));
let timeouts = Some(resolve_provider_quota_execution_timeouts(
state.resolve_transport_execution_timeouts(transport),
proxy.as_ref(),
));
let spec = build_antigravity_pool_quota_request(
&transport.key.id,
&transport.endpoint.base_url,
@@ -1,16 +1,19 @@
use super::shared::{
build_quota_snapshot_payload, default_provider_quota_execution_timeouts,
execute_provider_quota_plan, extract_execution_error_message,
build_quota_snapshot_payload, execute_provider_quota_plan, extract_execution_error_message,
oauth_refresh_auto_removed_result, persist_provider_quota_refresh_state,
quota_key_auto_removed, quota_refresh_success_invalid_state, ProviderQuotaExecutionOutcome,
quota_key_auto_removed, quota_refresh_success_invalid_state,
resolve_provider_quota_execution_timeouts, ProviderQuotaExecutionOutcome,
};
use crate::handlers::admin::provider::shared::payloads::{
OAUTH_ACCOUNT_BLOCK_PREFIX, OAUTH_EXPIRED_PREFIX,
OAUTH_ACCOUNT_BLOCK_PREFIX, OAUTH_EXPIRED_PREFIX, OAUTH_REFRESH_FAILED_PREFIX,
};
use crate::handlers::admin::request::{AdminAppState, AdminGatewayProviderTransportSnapshot};
use crate::GatewayError;
use aether_admin::provider::quota::parse_chatgpt_web_conversation_init_response;
use aether_contracts::ProxySnapshot;
use aether_contracts::{
ExecutionResult, ProxySnapshot, ResolvedTransportProfile, TRANSPORT_BACKEND_BROWSER_WREQ,
TRANSPORT_HTTP_MODE_AUTO, TRANSPORT_POOL_SCOPE_KEY,
};
use aether_data_contracts::repository::provider_catalog::{
StoredProviderCatalogEndpoint, StoredProviderCatalogKey, StoredProviderCatalogProvider,
};
@@ -18,10 +21,12 @@ use aether_provider_pool::{
build_chatgpt_web_pool_quota_request, enrich_chatgpt_web_quota_metadata,
normalize_chatgpt_web_image_quota_limit,
};
use base64::Engine as _;
use serde_json::json;
use std::time::{SystemTime, UNIX_EPOCH};
const PLACEHOLDER_API_KEY: &str = "__placeholder__";
const CHATGPT_WEB_BROWSER_PROFILE: &str = "chrome143";
fn chatgpt_web_auth_config(
transport: &AdminGatewayProviderTransportSnapshot,
@@ -67,26 +72,105 @@ async fn execute_chatgpt_web_quota_plan(
.await
}
};
let timeouts = state
.resolve_transport_execution_timeouts(transport)
.or(Some(default_provider_quota_execution_timeouts(
proxy.as_ref(),
)));
let timeouts = Some(resolve_provider_quota_execution_timeouts(
state.resolve_transport_execution_timeouts(transport),
proxy.as_ref(),
));
let spec =
build_chatgpt_web_pool_quota_request(&transport.key.id, &endpoint.base_url, authorization);
let resolved_transport_profile = state.resolve_transport_profile(transport);
let plan = super::shared::build_provider_quota_execution_plan(
transport,
spec,
proxy,
state.resolve_transport_profile(transport),
chatgpt_web_quota_transport_profile(resolved_transport_profile.as_ref()),
timeouts,
);
execute_provider_quota_plan(state, transport, plan, "chatgpt_web").await
}
fn chatgpt_web_quota_transport_profile(
transport_profile: Option<&ResolvedTransportProfile>,
) -> Option<ResolvedTransportProfile> {
match transport_profile {
Some(profile)
if profile
.backend
.trim()
.eq_ignore_ascii_case(TRANSPORT_BACKEND_BROWSER_WREQ) =>
{
Some(profile.clone())
}
_ => Some(default_chatgpt_web_quota_transport_profile()),
}
}
fn default_chatgpt_web_quota_transport_profile() -> ResolvedTransportProfile {
ResolvedTransportProfile {
profile_id: CHATGPT_WEB_BROWSER_PROFILE.to_string(),
backend: TRANSPORT_BACKEND_BROWSER_WREQ.to_string(),
http_mode: TRANSPORT_HTTP_MODE_AUTO.to_string(),
pool_scope: TRANSPORT_POOL_SCOPE_KEY.to_string(),
header_fingerprint: None,
extra: Some(json!({
"browser_profile": CHATGPT_WEB_BROWSER_PROFILE,
"source": "chatgpt_web_quota_default",
})),
}
}
fn chatgpt_web_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 chatgpt_web_is_structured_account_block(message: &str) -> bool {
let lowered = message.to_ascii_lowercase();
[
"account has been disabled",
"account disabled",
"account has been deactivated",
"account_deactivated",
"account deactivated",
"organization has been disabled",
"organization_disabled",
"deactivated_workspace",
"account suspended",
"account banned",
"account_block",
"account blocked",
"访问被禁止",
"账户访问被禁止",
"账户已封禁",
"封禁",
"封号",
"被封",
]
.iter()
.any(|keyword| lowered.contains(keyword))
}
fn chatgpt_web_quota_403_refresh_failed_reason(message: Option<&str>) -> String {
let detail = message
.map(str::trim)
.filter(|value| !value.is_empty())
.filter(|value| !value.contains('<'))
.unwrap_or("ChatGPT Web 访问验证失败,请检查浏览器指纹、Cloudflare 验证或代理/地区限制");
format!("{OAUTH_REFRESH_FAILED_PREFIX}{detail}")
}
fn chatgpt_web_quota_invalid_reason(status_code: u16, upstream_message: Option<&str>) -> String {
let message = upstream_message.unwrap_or_default().trim();
if status_code == 403 && !chatgpt_web_is_structured_account_block(message) {
return chatgpt_web_quota_403_refresh_failed_reason(upstream_message);
}
let detail = if message.is_empty() {
match status_code {
401 => "ChatGPT Web Token 无效或已过期",
@@ -103,6 +187,19 @@ fn chatgpt_web_quota_invalid_reason(status_code: u16, upstream_message: Option<&
}
}
fn chatgpt_web_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_chatgpt_web_provider_quota_locally(
state: &AdminAppState<'_>,
provider: &StoredProviderCatalogProvider,
@@ -216,8 +313,20 @@ pub(crate) async fn refresh_chatgpt_web_provider_quota_locally(
message = Some("响应中未包含 ChatGPT Web 生图限额信息".to_string());
}
} else {
let err_msg = extract_execution_error_message(&result);
message = Some(match err_msg.as_deref() {
let err_msg = chatgpt_web_quota_error_detail(&result);
let invalid_reason = if matches!(result.status_code, 401 | 403) {
Some(chatgpt_web_quota_invalid_reason(
result.status_code,
err_msg.as_deref(),
))
} else {
None
};
let display_detail = invalid_reason
.as_deref()
.map(chatgpt_web_quota_result_message)
.or_else(|| err_msg.clone());
message = Some(match display_detail.as_deref() {
Some(detail) if !detail.is_empty() => {
format!(
"conversation/init 返回状态码 {}: {}",
@@ -229,12 +338,14 @@ pub(crate) async fn refresh_chatgpt_web_provider_quota_locally(
if matches!(result.status_code, 401 | 403) {
oauth_invalid_at_unix_secs = Some(now_unix_secs);
oauth_invalid_reason = Some(chatgpt_web_quota_invalid_reason(
result.status_code,
err_msg.as_deref(),
));
oauth_invalid_reason = invalid_reason;
status = if result.status_code == 401 {
"auth_invalid".to_string()
} else if oauth_invalid_reason
.as_deref()
.is_some_and(|reason| reason.starts_with(OAUTH_REFRESH_FAILED_PREFIX))
{
"refresh_failed".to_string()
} else {
"forbidden".to_string()
};
@@ -303,3 +414,81 @@ pub(crate) async fn refresh_chatgpt_web_provider_quota_locally(
"auto_removed": auto_removed_count,
})))
}
#[cfg(test)]
mod tests {
use super::*;
use aether_contracts::{ResponseBody, TRANSPORT_BACKEND_REQWEST_RUSTLS};
use base64::Engine as _;
use std::collections::BTreeMap;
#[test]
fn quota_refresh_defaults_to_browser_wreq_transport() {
let profile = chatgpt_web_quota_transport_profile(None).expect("transport profile");
assert_eq!(profile.backend, TRANSPORT_BACKEND_BROWSER_WREQ);
assert_eq!(profile.profile_id, CHATGPT_WEB_BROWSER_PROFILE);
assert_eq!(profile.http_mode, TRANSPORT_HTTP_MODE_AUTO);
assert_eq!(profile.pool_scope, TRANSPORT_POOL_SCOPE_KEY);
assert_eq!(
profile
.extra
.as_ref()
.and_then(|value| value.get("browser_profile"))
.and_then(serde_json::Value::as_str),
Some(CHATGPT_WEB_BROWSER_PROFILE)
);
}
#[test]
fn quota_refresh_overrides_non_browser_transport() {
let reqwest_profile = ResolvedTransportProfile {
profile_id: "chrome_136".to_string(),
backend: TRANSPORT_BACKEND_REQWEST_RUSTLS.to_string(),
http_mode: TRANSPORT_HTTP_MODE_AUTO.to_string(),
pool_scope: TRANSPORT_POOL_SCOPE_KEY.to_string(),
header_fingerprint: None,
extra: None,
};
let profile =
chatgpt_web_quota_transport_profile(Some(&reqwest_profile)).expect("transport profile");
assert_eq!(profile.backend, TRANSPORT_BACKEND_BROWSER_WREQ);
assert_eq!(profile.profile_id, CHATGPT_WEB_BROWSER_PROFILE);
}
#[test]
fn browser_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: "chatgpt-web-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 = chatgpt_web_quota_error_detail(&result).expect("html body should decode");
let reason = chatgpt_web_quota_invalid_reason(result.status_code, Some(&detail));
assert!(reason.starts_with(OAUTH_REFRESH_FAILED_PREFIX));
assert!(!reason.starts_with(OAUTH_ACCOUNT_BLOCK_PREFIX));
assert_eq!(
chatgpt_web_quota_result_message(&reason),
"ChatGPT Web 访问验证失败,请检查浏览器指纹、Cloudflare 验证或代理/地区限制"
);
}
#[test]
fn explicit_account_block_403_remains_account_block() {
let reason = chatgpt_web_quota_invalid_reason(403, Some("account has been deactivated"));
assert!(reason.starts_with(OAUTH_ACCOUNT_BLOCK_PREFIX));
}
}
@@ -44,6 +44,15 @@ fn merge_codex_quota_metadata(
serde_json::Value::Object(merged)
}
fn codex_oauth_refresh_issue_reason(reason: Option<&str>) -> bool {
reason.is_some_and(|reason| {
reason
.lines()
.map(str::trim)
.any(|line| line.starts_with("[OAUTH_EXPIRED]") || line.starts_with("[REFRESH_FAILED]"))
})
}
pub(crate) async fn refresh_codex_provider_quota_locally(
state: &AdminAppState<'_>,
provider: &StoredProviderCatalogProvider,
@@ -56,8 +65,13 @@ pub(crate) async fn refresh_codex_provider_quota_locally(
let mut success_count = 0usize;
let mut failed_count = 0usize;
let mut auto_removed_count = 0usize;
let mut refresh_fixed_count = 0usize;
let mut refresh_failed_retained_count = 0usize;
let mut auto_removed_hard_banned_count = 0usize;
for key in keys {
let had_oauth_refresh_issue =
codex_oauth_refresh_issue_reason(key.oauth_invalid_reason.as_deref());
let transport = match state
.read_provider_transport_snapshot(&provider.id, &endpoint.id, &key.id)
.await?
@@ -276,13 +290,9 @@ pub(crate) async fn refresh_codex_provider_quota_locally(
}
}
let auto_removed = auto_remove_abnormal_keys
let auto_remove_candidate = auto_remove_abnormal_keys
&& should_auto_remove_structured_reason(oauth_invalid_reason.as_deref());
if auto_removed {
if state.delete_provider_catalog_key(&key.id).await? {
auto_removed_count += 1;
}
} else if !persist_provider_quota_refresh_state(
let persisted = persist_provider_quota_refresh_state(
state,
&key.id,
metadata_update.as_ref(),
@@ -290,8 +300,8 @@ pub(crate) async fn refresh_codex_provider_quota_locally(
oauth_invalid_reason.clone(),
None,
)
.await?
{
.await?;
if !persisted {
failed_count += 1;
results.push(json!({
"key_id": key.id,
@@ -301,6 +311,29 @@ pub(crate) async fn refresh_codex_provider_quota_locally(
}));
continue;
}
let auto_removed = if auto_remove_candidate {
state
.cleanup_provider_catalog_key_if_current(provider, &key.id, |latest_key| {
should_auto_remove_structured_reason(latest_key.oauth_invalid_reason.as_deref())
})
.await?
} else {
false
};
if auto_removed {
auto_removed_count += 1;
auto_removed_hard_banned_count += 1;
}
let refresh_fixed =
status == "success" && had_oauth_refresh_issue && oauth_invalid_reason.is_none();
if refresh_fixed {
refresh_fixed_count += 1;
}
let refresh_failed_retained =
status != "success" && oauth_invalid_reason.is_some() && !auto_removed;
if refresh_failed_retained {
refresh_failed_retained_count += 1;
}
if status == "success" {
success_count += 1;
@@ -336,6 +369,13 @@ pub(crate) async fn refresh_codex_provider_quota_locally(
}
if auto_removed {
payload.insert("auto_removed".to_string(), json!(true));
payload.insert("auto_removed_hard_banned".to_string(), json!(true));
}
if refresh_fixed {
payload.insert("refresh_fixed".to_string(), json!(true));
}
if refresh_failed_retained {
payload.insert("refresh_failed_retained".to_string(), json!(true));
}
results.push(serde_json::Value::Object(payload));
}
@@ -346,5 +386,8 @@ pub(crate) async fn refresh_codex_provider_quota_locally(
"total": results.len(),
"results": results,
"auto_removed": auto_removed_count,
"refresh_fixed": refresh_fixed_count,
"refresh_failed_retained": refresh_failed_retained_count,
"auto_removed_hard_banned": auto_removed_hard_banned_count,
})))
}
@@ -1,6 +1,6 @@
use super::super::shared::{
build_provider_quota_execution_plan, default_provider_quota_execution_timeouts,
execute_provider_quota_plan, ProviderQuotaExecutionOutcome,
build_provider_quota_execution_plan, execute_provider_quota_plan,
resolve_provider_quota_execution_timeouts, ProviderQuotaExecutionOutcome,
};
use crate::handlers::admin::request::{AdminAppState, AdminGatewayProviderTransportSnapshot};
use crate::GatewayError;
@@ -38,11 +38,10 @@ pub(super) async fn execute_codex_quota_plan(
.await
}
};
let timeouts = state
.resolve_transport_execution_timeouts(transport)
.or(Some(default_provider_quota_execution_timeouts(
proxy.as_ref(),
)));
let timeouts = Some(resolve_provider_quota_execution_timeouts(
state.resolve_transport_execution_timeouts(transport),
proxy.as_ref(),
));
let plan = build_provider_quota_execution_plan(
transport,
spec,
@@ -1,8 +1,7 @@
use super::shared::{
build_quota_snapshot_payload, default_provider_quota_execution_timeouts,
execute_provider_quota_plan, extract_execution_error_message,
build_quota_snapshot_payload, execute_provider_quota_plan, extract_execution_error_message,
persist_provider_quota_refresh_state, quota_refresh_success_invalid_state,
ProviderQuotaExecutionOutcome,
resolve_provider_quota_execution_timeouts, ProviderQuotaExecutionOutcome,
};
use crate::handlers::admin::provider::shared::payloads::{
OAUTH_ACCOUNT_BLOCK_PREFIX, OAUTH_EXPIRED_PREFIX, OAUTH_REFRESH_FAILED_PREFIX,
@@ -245,11 +244,10 @@ async fn execute_grok_quota_plan(
.await
}
};
let timeouts = state
.resolve_transport_execution_timeouts(transport)
.or(Some(default_provider_quota_execution_timeouts(
proxy.as_ref(),
)));
let timeouts = Some(resolve_provider_quota_execution_timeouts(
state.resolve_transport_execution_timeouts(transport),
proxy.as_ref(),
));
let transport_profile = state.resolve_transport_profile(transport);
let base_url = grok_base_url(endpoint);
let headers = build_grok_quota_headers(
@@ -1,6 +1,6 @@
use super::super::shared::{
build_provider_quota_execution_plan, default_provider_quota_execution_timeouts,
execute_provider_quota_plan, ProviderQuotaExecutionOutcome,
build_provider_quota_execution_plan, execute_provider_quota_plan,
resolve_provider_quota_execution_timeouts, ProviderQuotaExecutionOutcome,
};
use crate::handlers::admin::request::{
AdminAppState, AdminGatewayProviderTransportSnapshot, AdminKiroRequestAuth,
@@ -23,11 +23,10 @@ pub(super) async fn execute_kiro_quota_plan(
.await
}
};
let timeouts = state
.resolve_transport_execution_timeouts(transport)
.or(Some(default_provider_quota_execution_timeouts(
proxy.as_ref(),
)));
let timeouts = Some(resolve_provider_quota_execution_timeouts(
state.resolve_transport_execution_timeouts(transport),
proxy.as_ref(),
));
let spec = build_kiro_pool_quota_request(
&transport.key.id,
&KiroPoolQuotaAuthInput {
@@ -46,6 +46,23 @@ pub(super) fn default_provider_quota_execution_timeouts(
}
}
pub(super) fn resolve_provider_quota_execution_timeouts(
configured: Option<ExecutionTimeouts>,
proxy: Option<&ProxySnapshot>,
) -> ExecutionTimeouts {
let defaults = default_provider_quota_execution_timeouts(proxy);
let Some(mut timeouts) = configured else {
return defaults;
};
timeouts.connect_ms = timeouts.connect_ms.or(defaults.connect_ms);
timeouts.read_ms = timeouts.read_ms.or(defaults.read_ms);
timeouts.write_ms = timeouts.write_ms.or(defaults.write_ms);
timeouts.pool_ms = timeouts.pool_ms.or(defaults.pool_ms);
timeouts.total_ms = timeouts.total_ms.or(defaults.total_ms);
timeouts.first_byte_ms = timeouts.first_byte_ms.or(defaults.first_byte_ms);
timeouts
}
pub(crate) fn provider_auto_remove_banned_keys(config: Option<&serde_json::Value>) -> bool {
admin_provider_quota_pure::provider_auto_remove_banned_keys(config)
}
@@ -317,12 +334,7 @@ pub(super) async fn execute_provider_quota_plan(
match state.execute_execution_runtime_sync_plan(None, &plan).await {
Ok(result) => Ok(ProviderQuotaExecutionOutcome::Response(result)),
Err(err) => {
let error = match err {
GatewayError::UpstreamUnavailable { message, .. }
| GatewayError::ControlUnavailable { message, .. }
| GatewayError::Client { message, .. }
| GatewayError::Internal(message) => message,
};
let error = err.into_message();
let proxy_node_id = plan
.proxy
.as_ref()
@@ -1,8 +1,8 @@
use super::shared::{
build_provider_quota_execution_plan, build_quota_snapshot_payload,
default_provider_quota_execution_timeouts, execute_provider_quota_plan,
build_provider_quota_execution_plan, build_quota_snapshot_payload, execute_provider_quota_plan,
extract_execution_error_message, persist_provider_quota_refresh_state,
quota_refresh_success_invalid_state, ProviderQuotaExecutionOutcome,
quota_refresh_success_invalid_state, resolve_provider_quota_execution_timeouts,
ProviderQuotaExecutionOutcome,
};
use crate::handlers::admin::request::{AdminAppState, AdminGatewayProviderTransportSnapshot};
use crate::GatewayError;
@@ -33,11 +33,10 @@ async fn execute_windsurf_probe_plan(
.await
}
};
let timeouts = state
.resolve_transport_execution_timeouts(transport)
.or(Some(default_provider_quota_execution_timeouts(
proxy.as_ref(),
)));
let timeouts = Some(resolve_provider_quota_execution_timeouts(
state.resolve_transport_execution_timeouts(transport),
proxy.as_ref(),
));
let plan = build_provider_quota_execution_plan(
transport,
spec,
@@ -301,12 +301,7 @@ fn admin_provider_ops_decode_response_bytes(
}
fn admin_provider_ops_gateway_error_message(error: GatewayError) -> String {
match error {
GatewayError::UpstreamUnavailable { message, .. }
| GatewayError::ControlUnavailable { message, .. }
| GatewayError::Client { message, .. }
| GatewayError::Internal(message) => message,
}
error.into_message()
}
pub(super) fn admin_provider_ops_verify_execution_error_message(error: &str) -> String {
@@ -16,7 +16,12 @@ use aether_runtime_state::{DataLayerError, RuntimeState};
use futures_util::future::join_all;
use std::collections::{BTreeMap, BTreeSet};
use std::time::{SystemTime, UNIX_EPOCH};
use tracing::warn;
use tracing::{info, warn};
const DEFAULT_POOL_RUNTIME_WINDOW_METRIC_KEY_LIMIT: usize = 512;
const MAX_POOL_RUNTIME_WINDOW_METRIC_KEY_LIMIT: usize = 10_000;
const POOL_RUNTIME_WINDOW_METRIC_KEY_LIMIT_ENV: &str =
"AETHER_GATEWAY_ADMIN_POOL_RUNTIME_WINDOW_METRIC_KEY_LIMIT";
fn current_unix_secs() -> u64 {
SystemTime::now()
@@ -29,6 +34,20 @@ fn should_load_active_probe_members(pool_config: &AdminProviderPoolConfig) -> bo
pool_config.probing_enabled
}
fn pool_runtime_window_metric_key_limit() -> usize {
std::env::var(POOL_RUNTIME_WINDOW_METRIC_KEY_LIMIT_ENV)
.ok()
.and_then(|value| value.trim().parse::<usize>().ok())
.filter(|value| *value > 0)
.unwrap_or(DEFAULT_POOL_RUNTIME_WINDOW_METRIC_KEY_LIMIT)
.clamp(1, MAX_POOL_RUNTIME_WINDOW_METRIC_KEY_LIMIT)
}
fn bounded_runtime_window_metric_key_ids(key_ids: &[String], limit: usize) -> &[String] {
let end = key_ids.len().min(limit.max(1));
&key_ids[..end]
}
pub(crate) async fn read_admin_provider_pool_cooldown_counts(
runtime: &RuntimeState,
provider_ids: &[String],
@@ -54,8 +73,21 @@ pub(crate) async fn read_admin_provider_pool_runtime_state(
) -> AdminProviderPoolRuntimeState {
let mut state = AdminProviderPoolRuntimeState::default();
let cooldown_keys = pool_cooldown_keys(provider_id, key_ids);
let cost_keys = pool_cost_keys(provider_id, key_ids);
let latency_keys = pool_latency_keys(provider_id, key_ids);
let metric_key_limit = pool_runtime_window_metric_key_limit();
let metric_key_ids = bounded_runtime_window_metric_key_ids(key_ids, metric_key_limit);
if metric_key_ids.len() < key_ids.len() {
info!(
event_name = "admin_pool_runtime_window_metrics_truncated",
log_type = "event",
provider_id,
total_key_count = key_ids.len(),
scanned_key_count = metric_key_ids.len(),
metric_key_limit,
"gateway limited admin pool runtime cost/latency window reads"
);
}
let cost_keys = pool_cost_keys(provider_id, metric_key_ids);
let latency_keys = pool_latency_keys(provider_id, metric_key_ids);
let sticky_sessions_enabled = pool_config.sticky_session_ttl_seconds > 0
&& admin_provider_pool_cache_affinity_enabled(pool_config);
@@ -179,7 +211,7 @@ pub(crate) async fn read_admin_provider_pool_runtime_state(
.map(|cost_key| runtime.score_range_by_min(cost_key, cost_window_start)),
)
.await;
for (key_id, members) in key_ids.iter().zip(cost_results) {
for (key_id, members) in metric_key_ids.iter().zip(cost_results) {
let total = members
.unwrap_or_default()
.iter()
@@ -197,7 +229,7 @@ pub(crate) async fn read_admin_provider_pool_runtime_state(
.map(|latency_key| runtime.score_range_by_min(latency_key, latency_window_start)),
)
.await;
for (key_id, members) in key_ids.iter().zip(latency_results) {
for (key_id, members) in metric_key_ids.iter().zip(latency_results) {
let samples = members
.unwrap_or_default()
.iter()
@@ -265,3 +297,30 @@ pub(crate) async fn read_admin_provider_pool_key_cooldown_reason(
.kv_get(&pool_cooldown_key(provider_id, key_id))
.await
}
#[cfg(test)]
mod tests {
use super::bounded_runtime_window_metric_key_ids;
#[test]
fn runtime_window_metric_key_ids_are_bounded() {
let key_ids = vec![
"key-1".to_string(),
"key-2".to_string(),
"key-3".to_string(),
];
let bounded = bounded_runtime_window_metric_key_ids(&key_ids, 2);
assert_eq!(bounded, &key_ids[..2]);
}
#[test]
fn runtime_window_metric_key_ids_keep_at_least_one_key() {
let key_ids = vec!["key-1".to_string(), "key-2".to_string()];
let bounded = bounded_runtime_window_metric_key_ids(&key_ids, 0);
assert_eq!(bounded, &key_ids[..1]);
}
}
@@ -13,6 +13,7 @@ use aether_data_contracts::repository::provider_catalog::{
StoredProviderCatalogEndpoint, StoredProviderCatalogKey,
};
use aether_data_contracts::repository::usage::StoredProviderApiKeyWindowUsageSummary;
use aether_scheduler_core::provider_key_circuit_payload_is_active_open_at;
use serde_json::json;
use std::collections::BTreeMap;
use std::time::{SystemTime, UNIX_EPOCH};
@@ -933,19 +934,14 @@ fn admin_pool_health_score(key: &StoredProviderCatalogKey) -> f64 {
}
}
fn admin_pool_circuit_breaker_open(key: &StoredProviderCatalogKey) -> bool {
fn admin_pool_circuit_breaker_open(key: &StoredProviderCatalogKey, now_unix_secs: u64) -> bool {
key.circuit_breaker_by_format
.as_ref()
.and_then(serde_json::Value::as_object)
.map(|formats| {
formats
.values()
.filter_map(serde_json::Value::as_object)
.any(|item| {
item.get("open")
.and_then(serde_json::Value::as_bool)
.unwrap_or(false)
})
.any(|item| provider_key_circuit_payload_is_active_open_at(item, now_unix_secs))
})
.unwrap_or(false)
}
@@ -1044,7 +1040,7 @@ pub(super) fn build_admin_pool_key_payload(
.as_ref()
.and_then(|_| runtime.cooldown_ttl_by_key.get(&key.id).copied());
let health_score = admin_pool_health_score(key);
let circuit_breaker_open = admin_pool_circuit_breaker_open(key);
let circuit_breaker_open = admin_pool_circuit_breaker_open(key, now_unix_secs);
let auth_semantics = provider_key_auth_semantics(key, provider_type);
let account_quota_exhausted = pool_config
.as_ref()
@@ -99,6 +99,7 @@ static PROVIDER_QUERY_POOL_LOAD_BALANCE_SEQUENCE: AtomicU64 = AtomicU64::new(0);
struct ProviderQueryKeyFetchResult {
models: Vec<Value>,
error: Option<String>,
warning: Option<String>,
from_cache: bool,
has_success: bool,
}
@@ -288,6 +289,7 @@ fn provider_query_codex_preset_fallback(
Some(ProviderQueryKeyFetchResult {
models: aggregate_models_for_cache(&models),
error: None,
warning: None,
from_cache: false,
has_success: true,
})
@@ -427,6 +429,7 @@ async fn provider_query_fetch_models_for_key(
return Ok(ProviderQueryKeyFetchResult {
models,
error: None,
warning: None,
from_cache: true,
has_success: true,
});
@@ -444,6 +447,7 @@ async fn provider_query_fetch_models_for_key(
return Ok(ProviderQueryKeyFetchResult {
models,
error: None,
warning: None,
from_cache: false,
has_success: true,
});
@@ -451,6 +455,7 @@ async fn provider_query_fetch_models_for_key(
return Ok(ProviderQueryKeyFetchResult {
models: Vec::new(),
error: Some(ADMIN_PROVIDER_QUERY_NO_ACTIVE_ENDPOINT_DETAIL.to_string()),
warning: None,
from_cache: false,
has_success: false,
});
@@ -477,6 +482,7 @@ async fn provider_query_fetch_models_for_key(
return Ok(ProviderQueryKeyFetchResult {
models: Vec::new(),
error: Some(all_errors.join("; ")),
warning: None,
from_cache: false,
has_success: false,
});
@@ -492,6 +498,7 @@ async fn provider_query_fetch_models_for_key(
return Ok(ProviderQueryKeyFetchResult {
models: Vec::new(),
error: Some(all_errors.join("; ")),
warning: None,
from_cache: false,
has_success: false,
});
@@ -528,18 +535,25 @@ async fn provider_query_fetch_models_for_key(
}
}
let mut error = if all_errors.is_empty() {
None
} else {
let has_models = !unique_models.is_empty();
let mut error = if !has_models && !all_errors.is_empty() {
Some(all_errors.join("; "))
} else {
None
};
if unique_models.is_empty() && error.is_none() {
let warning = if has_models && !all_errors.is_empty() {
Some(all_errors.join("; "))
} else {
None
};
if !has_models && error.is_none() {
error = Some(ADMIN_PROVIDER_QUERY_NO_MODELS_FROM_ENDPOINT_DETAIL.to_string());
}
Ok(ProviderQueryKeyFetchResult {
models: provider_query_filter_models_for_key(provider, key, unique_models),
error,
warning,
from_cache: false,
has_success: outcome.has_success,
})
@@ -600,6 +614,7 @@ pub(crate) async fn build_admin_provider_query_models_response(
"data": {
"models": models,
"error": result.error,
"warning": result.warning,
"from_cache": result.from_cache,
},
"provider": provider_query_provider_payload(&provider),
@@ -632,6 +647,7 @@ pub(crate) async fn build_admin_provider_query_models_response(
"data": {
"models": models,
"error": serde_json::Value::Null,
"warning": serde_json::Value::Null,
"from_cache": true,
"keys_total": active_key_count,
"keys_cached": active_key_count,
@@ -655,6 +671,7 @@ pub(crate) async fn build_admin_provider_query_models_response(
let mut all_models = Vec::new();
let mut all_errors = Vec::new();
let mut all_warnings = Vec::new();
let mut cache_hit_count = 0usize;
let mut fetch_count = 0usize;
for key in &ordered_keys {
@@ -669,6 +686,13 @@ pub(crate) async fn build_admin_provider_query_models_response(
error
));
}
if let Some(warning) = result.warning {
all_warnings.push(format!(
"Key {}: {}",
provider_query_key_display_name(key),
warning
));
}
if result.from_cache {
cache_hit_count += 1;
} else {
@@ -694,10 +718,17 @@ pub(crate) async fn build_admin_provider_query_models_response(
provider_query_write_provider_cached_models(state, &provider.id, &models).await;
}
let success = !models.is_empty();
let mut error = if all_errors.is_empty() {
None
let mut all_issues = all_errors;
all_issues.extend(all_warnings);
let mut error = if !success && !all_issues.is_empty() {
Some(all_issues.join("; "))
} else {
Some(all_errors.join("; "))
None
};
let warning = if success && !all_issues.is_empty() {
Some(all_issues.join("; "))
} else {
None
};
if !success && error.is_none() {
error = Some(ADMIN_PROVIDER_QUERY_NO_MODELS_FROM_KEY_DETAIL.to_string());
@@ -709,6 +740,7 @@ pub(crate) async fn build_admin_provider_query_models_response(
"data": {
"models": models,
"error": error,
"warning": warning,
"from_cache": fetch_count == 0 && cache_hit_count > 0,
"keys_total": active_key_count,
"keys_cached": cache_hit_count,
@@ -1,5 +1,5 @@
use super::super::payload::{
provider_query_extract_api_key_id, provider_query_extract_force_refresh,
provider_query_extract_api_key_ids, provider_query_extract_force_refresh,
provider_query_extract_model, provider_query_extract_provider_id,
provider_query_extract_request_id,
};
@@ -68,6 +68,7 @@ use aether_model_fetch::{
aggregate_models_for_cache, fetch_models_from_transports, json_string_list,
preset_models_for_provider, selected_models_fetch_endpoints,
};
use aether_scheduler_core::provider_key_circuit_payload_is_active_open_at;
use axum::{
body::{to_bytes, Body},
http::{self, HeaderMap, HeaderName, HeaderValue},
@@ -122,6 +123,7 @@ const ADMIN_PROVIDER_QUERY_NO_ACTIVE_TEST_CANDIDATE_DETAIL: &str =
"No active endpoint or API key found";
const ADMIN_PROVIDER_QUERY_INVALID_MAPPED_MODEL_DETAIL: &str =
"mapped_model_name is not valid for the selected model and endpoint";
const PROVIDER_QUERY_KEY_MODEL_NOT_ALLOWED_SKIP_REASON: &str = "key_model_not_allowed";
const ANTIGRAVITY_PROVIDER_CACHE_KEY_PREFIX: &str = "upstream_models_provider:";
const DEFAULT_PROVIDER_QUERY_TEST_MESSAGE: &str = "Hello! This is a test message.";
static PROVIDER_QUERY_POOL_LOAD_BALANCE_SEQUENCE: AtomicU64 = AtomicU64::new(0);
@@ -859,7 +861,7 @@ async fn provider_query_select_preferred_non_kiro_endpoint(
provider: &StoredProviderCatalogProvider,
endpoints: &[StoredProviderCatalogEndpoint],
keys: &[StoredProviderCatalogKey],
selected_key_id: Option<&str>,
selected_key_ids: Option<&BTreeSet<String>>,
) -> Option<StoredProviderCatalogEndpoint> {
for priority in 0..=2 {
for endpoint in endpoints.iter().filter(|endpoint| endpoint.is_active) {
@@ -872,7 +874,7 @@ async fn provider_query_select_preferred_non_kiro_endpoint(
}
for key in keys {
if !key.is_active
|| selected_key_id.is_some_and(|value| value != key.id.as_str())
|| !provider_query_selected_key_ids_allow_key(selected_key_ids, &key.id)
|| !provider_query_key_supports_endpoint(
key,
&provider.provider_type,
@@ -904,7 +906,7 @@ async fn provider_query_select_preferred_non_kiro_endpoint(
endpoint.is_active
&& keys.iter().any(|key| {
key.is_active
&& selected_key_id.is_none_or(|value| value == key.id.as_str())
&& provider_query_selected_key_ids_allow_key(selected_key_ids, &key.id)
&& provider_query_key_supports_endpoint(
key,
&provider.provider_type,
@@ -916,10 +918,56 @@ async fn provider_query_select_preferred_non_kiro_endpoint(
.cloned()
}
fn provider_query_selected_key_ids_allow_key(
selected_key_ids: Option<&BTreeSet<String>>,
key_id: &str,
) -> bool {
selected_key_ids.is_none_or(|ids| ids.contains(key_id))
}
fn provider_query_selected_key_ids_all_exist(
selected_key_ids: &BTreeSet<String>,
keys: &[StoredProviderCatalogKey],
) -> bool {
selected_key_ids
.iter()
.all(|id| keys.iter().any(|key| key.id == *id))
}
fn provider_query_model_name_matches(left: &str, right: &str) -> bool {
let left = left.trim();
let right = right.trim();
!left.is_empty() && !right.is_empty() && left.eq_ignore_ascii_case(right)
}
fn provider_query_key_allows_effective_test_model(
key: &StoredProviderCatalogKey,
requested_model: &str,
effective_model: &str,
) -> bool {
let allowed_models = json_string_list(key.allowed_models.as_ref());
if key.allowed_models.is_none() || allowed_models.is_empty() {
return true;
}
let requested_base_model = crate::ai_serving::model_directive_base_model(requested_model);
allowed_models
.iter()
.map(String::as_str)
.any(|allowed_model| {
provider_query_model_name_matches(allowed_model, requested_model)
|| provider_query_model_name_matches(allowed_model, effective_model)
|| requested_base_model.as_deref().is_some_and(|base_model| {
provider_query_model_name_matches(allowed_model, base_model)
})
})
}
fn provider_query_test_key_sort_key(
provider_type: &str,
key: &StoredProviderCatalogKey,
endpoint_api_format: &str,
now_unix_secs: u64,
) -> (u8, u8, i32, u64, i32) {
let quota_exhausted =
admin_provider_pool_pure::admin_pool_key_account_quota_exhausted(key, provider_type);
@@ -928,10 +976,7 @@ fn provider_query_test_key_sort_key(
.as_ref()
.and_then(Value::as_object)
.and_then(|value| value.get(endpoint_api_format))
.and_then(Value::as_object)
.and_then(|value| value.get("open"))
.and_then(Value::as_bool)
.unwrap_or(false);
.is_some_and(|value| provider_key_circuit_payload_is_active_open_at(value, now_unix_secs));
let health_score = key
.health_by_format
.as_ref()
@@ -1263,7 +1308,7 @@ async fn provider_query_build_kiro_test_candidates(
ADMIN_PROVIDER_QUERY_NO_ACTIVE_API_KEY_DETAIL,
)
})?;
let selected_key_id = provider_query_extract_api_key_id(payload);
let selected_key_ids = provider_query_extract_api_key_ids(payload);
let requested_endpoint_id = provider_query_extract_endpoint_id(payload);
let requested_api_format = provider_query_extract_api_format(payload);
let endpoint = if requested_endpoint_id.is_none()
@@ -1275,7 +1320,7 @@ async fn provider_query_build_kiro_test_candidates(
provider,
&endpoints,
&all_keys,
selected_key_id.as_deref(),
selected_key_ids.as_ref(),
)
.await
.ok_or_else(|| {
@@ -1308,22 +1353,11 @@ async fn provider_query_build_kiro_test_candidates(
}
};
if let Some(api_key_id) = selected_key_id.as_deref() {
let Some(key) = all_keys.iter().find(|key| key.id == api_key_id) else {
if let Some(selected_key_ids) = selected_key_ids.as_ref() {
if !provider_query_selected_key_ids_all_exist(selected_key_ids, &all_keys) {
return Err(build_admin_provider_query_not_found_response(
ADMIN_PROVIDER_QUERY_API_KEY_NOT_FOUND_DETAIL,
));
};
if !key.is_active
|| !provider_query_key_supports_endpoint(
key,
&provider.provider_type,
&endpoint.api_format,
)
{
return Err(build_admin_provider_query_not_found_response(
ADMIN_PROVIDER_QUERY_NO_ACTIVE_TEST_CANDIDATE_DETAIL,
));
}
}
@@ -1378,20 +1412,43 @@ async fn provider_query_build_kiro_test_candidates(
.unwrap_or(requested_model.clone())
};
let mut keys = all_keys
let now_unix_secs = current_unix_ms() / 1000;
let mut keys = Vec::new();
let mut model_skipped_candidates = Vec::new();
for key in all_keys
.into_iter()
.filter(|key| key.is_active)
.filter(|key| {
selected_key_id
.as_deref()
.is_none_or(|value| value == key.id.as_str())
})
.filter(|key| provider_query_selected_key_ids_allow_key(selected_key_ids.as_ref(), &key.id))
.filter(|key| {
provider_query_key_supports_endpoint(key, &provider.provider_type, &endpoint.api_format)
})
.collect::<Vec<_>>();
{
if provider_query_key_allows_effective_test_model(&key, &requested_model, &effective_model)
{
keys.push(key);
} else {
model_skipped_candidates.push(ProviderQueryTestCandidate {
endpoint: endpoint.clone(),
key,
effective_model: effective_model.clone(),
scheduler_skip_reason: Some(
PROVIDER_QUERY_KEY_MODEL_NOT_ALLOWED_SKIP_REASON.to_string(),
),
});
}
}
let candidates = if test_mode.eq_ignore_ascii_case("pool") {
model_skipped_candidates.sort_by_key(|candidate| {
provider_query_test_key_sort_key(
provider.provider_type.as_str(),
&candidate.key,
&endpoint.api_format,
now_unix_secs,
)
});
let scheduled_candidates = if test_mode.eq_ignore_ascii_case("pool") {
if let Some(pool_config) =
admin_provider_pool_config_from_config_value(provider.config.as_ref())
{
@@ -1411,6 +1468,7 @@ async fn provider_query_build_kiro_test_candidates(
provider.provider_type.as_str(),
key,
&endpoint.api_format,
now_unix_secs,
)
});
keys.into_iter()
@@ -1428,6 +1486,7 @@ async fn provider_query_build_kiro_test_candidates(
provider.provider_type.as_str(),
key,
&endpoint.api_format,
now_unix_secs,
)
});
keys.into_iter()
@@ -1439,6 +1498,8 @@ async fn provider_query_build_kiro_test_candidates(
})
.collect::<Vec<_>>()
};
let mut candidates = model_skipped_candidates;
candidates.extend(scheduled_candidates);
if candidates.is_empty() {
return Err(build_admin_provider_query_not_found_response(
@@ -64,6 +64,68 @@ fn sample_openai_image_transport(provider_type: &str) -> AdminGatewayProviderTra
}
}
fn sample_catalog_key_with_allowed_models(
allowed_models: Option<serde_json::Value>,
) -> aether_data_contracts::repository::provider_catalog::StoredProviderCatalogKey {
let mut key =
aether_data_contracts::repository::provider_catalog::StoredProviderCatalogKey::new(
"key-1".to_string(),
"provider-1".to_string(),
"key".to_string(),
"api_key".to_string(),
None,
true,
)
.expect("sample provider key should build");
key.allowed_models = allowed_models;
key
}
#[test]
fn provider_query_model_test_allows_keys_without_model_restrictions() {
let unrestricted = sample_catalog_key_with_allowed_models(None);
let empty = sample_catalog_key_with_allowed_models(Some(json!([])));
assert!(provider_query_key_allows_effective_test_model(
&unrestricted,
"model-b",
"model-b-upstream",
));
assert!(provider_query_key_allows_effective_test_model(
&empty,
"model-b",
"model-b-upstream",
));
}
#[test]
fn provider_query_model_test_filters_key_disallowed_for_requested_model() {
let key = sample_catalog_key_with_allowed_models(Some(json!(["model-a"])));
assert!(!provider_query_key_allows_effective_test_model(
&key,
"model-b",
"model-b-upstream",
));
}
#[test]
fn provider_query_model_test_allows_key_for_requested_or_mapped_model() {
let requested_allowed = sample_catalog_key_with_allowed_models(Some(json!(["model-b"])));
let mapped_allowed = sample_catalog_key_with_allowed_models(Some(json!(["MODEL-B-UPSTREAM"])));
assert!(provider_query_key_allows_effective_test_model(
&requested_allowed,
"model-b",
"model-b-upstream",
));
assert!(provider_query_key_allows_effective_test_model(
&mapped_allowed,
"model-b",
"model-b-upstream",
));
}
#[test]
fn provider_query_test_request_body_preserves_custom_model() {
let payload = json!({
@@ -232,6 +294,27 @@ fn provider_query_request_body_model_uses_non_empty_string_only() {
);
}
#[test]
fn provider_query_model_test_extracts_multiple_selected_key_ids() {
let payload = json!({
"api_key_ids": [" key-b ", "", "key-a", "key-b"],
"api_key_id": "key-c"
});
let ids = provider_query_extract_api_key_ids(&payload)
.expect("non-empty key selection should be extracted")
.into_iter()
.collect::<Vec<_>>();
assert_eq!(ids, vec!["key-a", "key-b", "key-c"]);
}
#[test]
fn provider_query_model_test_empty_selected_key_ids_keep_default_selection() {
assert!(provider_query_extract_api_key_ids(&json!({})).is_none());
assert!(provider_query_extract_api_key_ids(&json!({ "api_key_ids": [] })).is_none());
}
#[test]
fn provider_query_standard_test_resolves_codex_responses_upstream_streaming() {
assert!(provider_query_resolve_standard_test_upstream_is_stream(
@@ -1,6 +1,7 @@
use axum::body::Bytes;
use axum::response::{IntoResponse, Response};
use serde_json::json;
use std::collections::BTreeSet;
pub(crate) fn parse_admin_provider_query_body(
request_body: Option<&Bytes>,
@@ -36,6 +37,47 @@ pub(crate) fn provider_query_extract_api_key_id(payload: &serde_json::Value) ->
.map(ToOwned::to_owned)
}
fn provider_query_insert_api_key_id(ids: &mut BTreeSet<String>, value: &str) {
let value = value.trim();
if !value.is_empty() {
ids.insert(value.to_string());
}
}
pub(crate) fn provider_query_extract_api_key_ids(
payload: &serde_json::Value,
) -> Option<BTreeSet<String>> {
let mut ids = BTreeSet::new();
if let Some(value) = payload
.get("api_key_ids")
.or_else(|| payload.get("provider_key_ids"))
.or_else(|| payload.get("key_ids"))
{
match value {
serde_json::Value::Array(items) => {
for item in items {
if let Some(value) = item.as_str() {
provider_query_insert_api_key_id(&mut ids, value);
}
}
}
serde_json::Value::String(value) => {
for item in value.split(',') {
provider_query_insert_api_key_id(&mut ids, item);
}
}
_ => {}
}
}
if let Some(api_key_id) = provider_query_extract_api_key_id(payload) {
ids.insert(api_key_id);
}
(!ids.is_empty()).then_some(ids)
}
pub(crate) fn provider_query_extract_force_refresh(payload: &serde_json::Value) -> bool {
payload
.get("force_refresh")
@@ -662,10 +662,5 @@ fn admin_provider_oauth_decode_response_bytes(
}
fn admin_provider_oauth_gateway_error_message(error: GatewayError) -> String {
match error {
GatewayError::UpstreamUnavailable { message, .. }
| GatewayError::ControlUnavailable { message, .. }
| GatewayError::Client { message, .. }
| GatewayError::Internal(message) => message,
}
error.into_message()
}
@@ -1,4 +1,7 @@
use super::*;
use aether_data_contracts::repository::provider_catalog::{
StoredProviderCatalogKey, StoredProviderCatalogProvider,
};
use axum::{
body::Body,
http,
@@ -193,8 +196,6 @@ impl<'a> AdminAppState<'a> {
&self,
provider_id: &str,
) -> Result<Response<Body>, GatewayError> {
use aether_admin::provider::pool as admin_provider_pool_pure;
let Some(provider) = self
.read_provider_catalog_providers_by_ids(std::slice::from_ref(&provider_id.to_string()))
.await?
@@ -208,6 +209,30 @@ impl<'a> AdminAppState<'a> {
.into_response());
};
let affected = self
.cleanup_known_banned_provider_catalog_keys(&provider)
.await?;
if affected == 0 {
return Ok(Json(
aether_admin::provider::pool::build_admin_pool_cleanup_empty_payload(
"未发现可清理的异常账号",
),
)
.into_response());
}
Ok(
Json(aether_admin::provider::pool::build_admin_pool_cleanup_result_payload(affected))
.into_response(),
)
}
pub(crate) async fn cleanup_known_banned_provider_catalog_keys(
&self,
provider: &StoredProviderCatalogProvider,
) -> Result<usize, GatewayError> {
use aether_admin::provider::pool as admin_provider_pool_pure;
let banned_keys = self
.list_provider_catalog_keys_by_provider_ids(std::slice::from_ref(&provider.id))
.await?
@@ -215,12 +240,7 @@ impl<'a> AdminAppState<'a> {
.filter(admin_provider_pool_pure::admin_pool_key_is_known_banned)
.collect::<Vec<_>>();
if banned_keys.is_empty() {
return Ok(Json(
admin_provider_pool_pure::build_admin_pool_cleanup_empty_payload(
"未发现可清理的异常账号",
),
)
.into_response());
return Ok(0);
}
let deleted_key_ids = banned_keys
@@ -243,10 +263,42 @@ impl<'a> AdminAppState<'a> {
self.cleanup_deleted_provider_catalog_refs(&provider.id, &[], &deleted_key_ids)
.await?;
Ok(
Json(admin_provider_pool_pure::build_admin_pool_cleanup_result_payload(affected))
.into_response(),
)
Ok(affected)
}
pub(crate) async fn cleanup_provider_catalog_key_if_current<F>(
&self,
provider: &StoredProviderCatalogProvider,
key_id: &str,
should_delete: F,
) -> Result<bool, GatewayError>
where
F: FnOnce(&StoredProviderCatalogKey) -> bool,
{
let key_ids = [key_id.to_string()];
let Some(key) = self
.read_provider_catalog_keys_by_ids(&key_ids)
.await?
.into_iter()
.next()
else {
return Ok(false);
};
if key.provider_id != provider.id || !should_delete(&key) {
return Ok(false);
}
self.clear_admin_provider_pool_cooldown(&provider.id, &key.id)
.await;
self.reset_admin_provider_pool_cost(&provider.id, &key.id)
.await;
let deleted = self.delete_provider_catalog_key(&key.id).await?;
if deleted {
let deleted_key_ids = [key.id.clone()];
self.cleanup_deleted_provider_catalog_refs(&provider.id, &[], &deleted_key_ids)
.await?;
}
Ok(deleted)
}
pub(crate) async fn build_admin_pool_batch_action_response(
@@ -40,6 +40,7 @@ impl<'a> AdminAppState<'a> {
.map(|model| AdminSystemConfigGlobalModel {
name: model.name.clone(),
display_name: model.display_name.clone(),
usage_count: Some(model.usage_count),
default_price_per_request: model.default_price_per_request,
default_tiered_pricing: model.default_tiered_pricing.clone(),
supported_capabilities: model.supported_capabilities.as_ref().and_then(|value| {
@@ -169,6 +170,13 @@ impl<'a> AdminAppState<'a> {
) -> Result<serde_json::Value, GatewayError> {
let users = self.list_non_admin_export_users().await?;
let user_ids = users.iter().map(|user| user.id.clone()).collect::<Vec<_>>();
let user_usage_totals = self
.app
.summarize_usage_totals_by_user_ids(&user_ids)
.await?
.into_iter()
.map(|totals| (totals.user_id.clone(), totals))
.collect::<BTreeMap<_, _>>();
let user_wallets = self.list_wallet_snapshots_by_user_ids(&user_ids).await?;
let user_api_keys = self
.list_auth_api_key_export_records_by_user_ids(&user_ids)
@@ -185,6 +193,7 @@ impl<'a> AdminAppState<'a> {
let standalone_wallets = self
.list_wallet_snapshots_by_api_key_ids(&standalone_api_key_ids)
.await?;
let usage_aggregates = self.export_admin_system_usage_aggregates().await?;
let wallets_by_user_id = user_wallets
.into_iter()
@@ -260,8 +269,10 @@ impl<'a> AdminAppState<'a> {
self.build_admin_system_users_export_api_key_payload(key, None, true)
})
.collect::<Vec<_>>();
let usage_totals = user_usage_totals.get(&user.id);
json!({
"id": user.id.clone(),
"email": user.email.clone(),
"email_verified": user.email_verified,
"username": user.username.clone(),
@@ -284,6 +295,12 @@ impl<'a> AdminAppState<'a> {
.unwrap_or(false),
"wallet": wallet_payload,
"is_active": user.is_active,
"request_count": usage_totals
.map(|totals| totals.request_count)
.unwrap_or(0),
"total_tokens": usage_totals
.map(|totals| totals.total_tokens)
.unwrap_or(0),
"api_keys": api_keys_payload,
})
})
@@ -306,6 +323,7 @@ impl<'a> AdminAppState<'a> {
"user_groups": user_groups_data,
"users": users_data,
"standalone_keys": standalone_keys_data,
"usage_aggregates": usage_aggregates,
}))
}
@@ -330,6 +348,7 @@ impl<'a> AdminAppState<'a> {
include_is_standalone: bool,
) -> serde_json::Value {
let mut payload = serde_json::Map::from_iter([
("api_key_id".to_string(), json!(key.api_key_id.clone())),
("key_hash".to_string(), json!(key.key_hash.clone())),
("name".to_string(), json!(key.name.clone())),
(
File diff suppressed because it is too large Load Diff
@@ -70,6 +70,26 @@ impl<'a> AdminAppState<'a> {
self.app.purge_admin_system_data(target).await
}
pub(crate) async fn export_admin_system_usage_aggregates(
&self,
) -> Result<aether_data::repository::system::AdminSystemUsageAggregateSnapshot, GatewayError>
{
self.app.export_admin_system_usage_aggregates().await
}
pub(crate) async fn import_admin_system_usage_aggregates(
&self,
snapshot: &aether_data::repository::system::AdminSystemUsageAggregateSnapshot,
user_id_map: &std::collections::BTreeMap<String, String>,
api_key_id_map: &std::collections::BTreeMap<String, String>,
mode: aether_data::repository::system::AdminSystemUsageAggregateImportMode,
) -> Result<aether_data::repository::system::AdminSystemUsageAggregateImportSummary, GatewayError>
{
self.app
.import_admin_system_usage_aggregates(snapshot, user_id_map, api_key_id_map, mode)
.await
}
pub(crate) async fn run_admin_system_cleanup_once(
&self,
) -> Result<crate::maintenance::AdminSystemCleanupSummary, GatewayError> {
@@ -706,6 +706,19 @@ impl<'a> AdminAppState<'a> {
.await
}
pub(crate) async fn set_api_key_usage_totals(
&self,
api_key_id: &str,
total_requests: u64,
total_tokens: u64,
total_cost_usd: f64,
) -> Result<Option<aether_data::repository::auth::StoredAuthApiKeyExportRecord>, GatewayError>
{
self.app
.set_api_key_usage_totals(api_key_id, total_requests, total_tokens, total_cost_usd)
.await
}
pub(crate) async fn delete_user_api_key(
&self,
user_id: &str,
@@ -13,7 +13,8 @@ pub(crate) use crate::handlers::shared::{
effective_catalog_encryption_key, encrypt_catalog_secret_with_fallbacks, json_string_list,
masked_catalog_api_key, normalize_json_array, normalize_json_object, normalize_string_list,
parse_catalog_auth_config_json, provider_catalog_key_supports_format,
provider_key_health_summary, provider_key_status_snapshot_payload, query_param_bool,
query_param_optional_bool, query_param_value, take_secret_prefix, take_secret_suffix,
unix_secs_to_rfc3339, OFFICIAL_EXTERNAL_MODEL_PROVIDERS,
provider_key_health_summary, provider_key_health_summary_at,
provider_key_status_snapshot_payload, query_param_bool, query_param_optional_bool,
query_param_value, take_secret_prefix, take_secret_suffix, unix_secs_to_rfc3339,
OFFICIAL_EXTERNAL_MODEL_PROVIDERS,
};
@@ -13,10 +13,16 @@ use crate::handlers::admin::system::shared::paths::{
};
use crate::handlers::admin::system::shared::settings::{
apply_admin_system_settings_update, build_admin_api_formats_payload,
build_admin_system_check_update_payload_from_release, build_admin_system_settings_payload,
build_admin_system_stats_payload, current_aether_version, fetch_latest_admin_system_release,
build_admin_system_check_update_payload_from_release, build_admin_system_releases_list_payload,
build_admin_system_settings_payload, build_admin_system_stats_payload, current_aether_version,
fetch_admin_system_releases, fetch_latest_admin_system_release, resolve_update_target,
};
use crate::handlers::admin::system::shared::smtp::build_admin_smtp_test_payload;
use crate::handlers::admin::system::shared::update::{
build_admin_system_update_capability_payload, current_self_update_blocker,
prepare_admin_system_update_task, read_update_history, read_update_task_status,
self_update_supported, start_admin_system_rollback_task, start_admin_system_update_task,
};
use crate::important_notification::build_important_notification_test_payload;
use crate::maintenance::{ManualUsageCleanupMode, ManualUsageCleanupOptions};
use crate::GatewayError;
@@ -58,7 +64,8 @@ pub(super) async fn maybe_build_local_admin_core_system_response(
&& request_method == http::Method::GET
&& request_path == "/api/admin/system/check-update"
{
let (latest_release, error) = fetch_latest_admin_system_release().await;
let force = query_flag(request_context.query_string(), "force");
let (latest_release, error) = fetch_latest_admin_system_release(force).await;
return Ok(Some(
Json(build_admin_system_check_update_payload_from_release(
latest_release,
@@ -68,6 +75,133 @@ pub(super) async fn maybe_build_local_admin_core_system_response(
));
}
if decision.route_kind.as_deref() == Some("releases")
&& request_method == http::Method::GET
&& request_path == "/api/admin/system/releases"
{
let force = query_flag(request_context.query_string(), "force");
let (releases, error) = fetch_admin_system_releases(force).await;
return Ok(Some(
Json(build_admin_system_releases_list_payload(releases, error)).into_response(),
));
}
if decision.route_kind.as_deref() == Some("update_capability")
&& request_method == http::Method::GET
&& request_path == "/api/admin/system/update-capability"
{
return Ok(Some(
Json(build_admin_system_update_capability_payload()).into_response(),
));
}
if decision.route_kind.as_deref() == Some("prepare_update")
&& request_method == http::Method::POST
&& request_path == "/api/admin/system/prepare-update"
{
if !self_update_supported() {
return Ok(Some(
(
http::StatusCode::PRECONDITION_REQUIRED,
Json(json!({ "detail": current_self_update_blocker() })),
)
.into_response(),
));
}
let target_version = request_body
.filter(|b| !b.is_empty())
.and_then(|body| serde_json::from_slice::<serde_json::Value>(body).ok())
.and_then(|v| v.get("version").and_then(|v| v.as_str().map(String::from)));
let (version, tarball_url, sha256sums_url) =
match resolve_update_target(target_version).await {
Ok(result) => result,
Err((status, payload)) => {
return Ok(Some((status, Json(payload)).into_response()));
}
};
return Ok(Some(
match prepare_admin_system_update_task(version, tarball_url, sha256sums_url).await? {
Ok(payload) => attach_admin_audit_response(
Json(payload).into_response(),
"admin_system_update_prepared",
"prepare_system_update",
"system_update",
"global",
),
Err((status, payload)) => (status, Json(payload)).into_response(),
},
));
}
if decision.route_kind.as_deref() == Some("apply_update")
&& request_method == http::Method::POST
&& request_path == "/api/admin/system/apply-update"
{
let version = request_body
.filter(|b| !b.is_empty())
.and_then(|body| serde_json::from_slice::<serde_json::Value>(body).ok())
.and_then(|v| v.get("version").and_then(|v| v.as_str().map(String::from)));
return Ok(Some(
match start_admin_system_update_task(version).await? {
Ok(payload) => attach_admin_audit_response(
Json(payload).into_response(),
"admin_system_update_started",
"apply_system_update",
"system_update",
"global",
),
Err((status, payload)) => (status, Json(payload)).into_response(),
},
));
}
if decision.route_kind.as_deref() == Some("rollback")
&& request_method == http::Method::POST
&& request_path == "/api/admin/system/rollback"
{
return Ok(Some(match start_admin_system_rollback_task().await? {
Ok(payload) => attach_admin_audit_response(
Json(payload).into_response(),
"admin_system_rollback_started",
"rollback_system_update",
"system_rollback",
"global",
),
Err((status, payload)) => (status, Json(payload)).into_response(),
}));
}
if decision.route_kind.as_deref() == Some("update_status")
&& request_method == http::Method::GET
&& request_path == "/api/admin/system/update-status"
{
let status = read_update_task_status();
return Ok(Some(
Json(json!({
"phase": status.phase,
"error": status.error,
"output": status.output,
"progress_label": status.progress_label,
"downloaded_bytes": status.downloaded_bytes,
"total_bytes": status.total_bytes,
"progress_percent": status.progress_percent,
}))
.into_response(),
));
}
if decision.route_kind.as_deref() == Some("update_history")
&& request_method == http::Method::GET
&& request_path == "/api/admin/system/update-history"
{
let entries = read_update_history();
return Ok(Some(Json(json!({ "entries": entries })).into_response()));
}
if decision.route_kind.as_deref() == Some("aws_regions")
&& request_method == http::Method::GET
&& request_path == "/api/admin/system/aws-regions"
@@ -233,6 +367,39 @@ pub(super) async fn maybe_build_local_admin_core_system_response(
));
}
if decision.route_kind.as_deref() == Some("s3_backup_run")
&& request_method == http::Method::POST
&& request_path == "/api/admin/system/backups/s3/run"
{
return Ok(Some(
match crate::backup::task::start_s3_backup_task(
state.cloned_app(),
"manual",
decision
.admin_principal
.as_ref()
.map(|principal| principal.user_id.as_str()),
)
.await
{
Ok(task) => attach_admin_audit_response(
Json(json!({
"message": "S3 备份任务已提交",
"task": task,
}))
.into_response(),
"admin_system_s3_backup_task_started",
"run_s3_backup",
"s3_backup",
"global",
),
Err(error) => {
(error.status(), Json(json!({ "detail": error.detail() }))).into_response()
}
},
));
}
if decision.route_kind.as_deref() == Some("smtp_test")
&& request_method == http::Method::POST
&& request_path == "/api/admin/system/smtp/test"
@@ -942,6 +1109,15 @@ fn query_param(query_string: Option<&str>, name: &str) -> Option<String> {
.find_map(|(key, value)| (key == name && !value.is_empty()).then(|| value.into_owned()))
}
fn query_flag(query_string: Option<&str>, name: &str) -> bool {
query_param(query_string, name).is_some_and(|value| {
matches!(
value.trim().to_ascii_lowercase().as_str(),
"1" | "true" | "yes" | "on"
)
})
}
fn parse_older_than_days_query(query_string: Option<&str>) -> Result<Option<u32>, Response<Body>> {
let Some(value) = query_param(query_string, "older_than_days") else {
return Ok(None);
@@ -63,6 +63,10 @@ struct ProxyNodeRegisterRequest {
proxy_version: Option<String>,
#[serde(default)]
tunnel_mode: Option<bool>,
#[serde(default)]
tunnel_security: Option<String>,
#[serde(default)]
tunnel_encryption_key: Option<String>,
}
#[derive(Debug, Deserialize)]
@@ -349,9 +353,21 @@ pub(crate) async fn maybe_build_local_admin_proxy_nodes_response(
Ok(mutation) => mutation,
Err(response) => return Ok(Some(response)),
};
let tunnel_encryption_key = mutation
.proxy_metadata
.as_ref()
.and_then(|metadata| metadata.pointer("/tunnel_security/encryption_key"))
.and_then(|value| value.as_str())
.map(str::to_string);
let Some(node) = state.register_proxy_node(&mutation).await? else {
return Ok(Some(build_admin_proxy_nodes_data_unavailable_response()));
};
if let Some(key) = tunnel_encryption_key {
state
.app()
.tunnel
.register_secure_tunnel_key(node.id.clone(), key);
}
return Ok(Some(
Json(json!({
"node_id": node.id,
@@ -1358,6 +1374,9 @@ fn build_tunnel_probe_relay_envelope(
method: "GET".to_string(),
url: probe_url.trim().to_string(),
headers: std::collections::HashMap::new(),
stream: false,
request_timeout_ms: None,
stream_first_byte_timeout_ms: None,
timeout: timeout_secs,
follow_redirects: Some(false),
http1_only: false,
@@ -1420,12 +1439,43 @@ fn validate_register_request(
}
validate_optional_object(input.hardware_info.as_ref(), "hardware_info")?;
validate_optional_object(input.proxy_metadata.as_ref(), "proxy_metadata")?;
let tunnel_security =
normalize_optional_string(input.tunnel_security.as_deref(), "tunnel_security", 64)?;
let tunnel_encryption_key = normalize_optional_string(
input.tunnel_encryption_key.as_deref(),
"tunnel_encryption_key",
128,
)?;
let registered_by = request_context
.decision()
.and_then(|decision| decision.admin_principal.as_ref())
.map(|principal| principal.user_id.clone());
let mut proxy_metadata = input.proxy_metadata;
if tunnel_security.as_deref()
== Some(aether_contracts::tunnel_security::TUNNEL_SECURITY_NON_TLS_REQUIRED)
{
let key = tunnel_encryption_key.as_deref().ok_or_else(|| {
bad_request_response(
"tunnel_encryption_key is required when tunnel_security=non_tls_required",
)
})?;
aether_contracts::tunnel_security::decode_psk(key)
.map_err(|err| bad_request_response(err.to_string()))?;
let mut metadata = proxy_metadata
.and_then(|value| value.as_object().cloned())
.unwrap_or_default();
metadata.insert(
"tunnel_security".to_string(),
json!({
"mode": aether_contracts::tunnel_security::TUNNEL_SECURITY_NON_TLS_REQUIRED,
"encryption_key": key,
}),
);
proxy_metadata = Some(Value::Object(metadata));
}
Ok(
aether_data::repository::proxy_nodes::ProxyNodeRegistrationMutation {
name,
@@ -1438,7 +1488,7 @@ fn validate_register_request(
avg_latency_ms: input.avg_latency_ms,
hardware_info: input.hardware_info,
estimated_max_concurrency: input.estimated_max_concurrency,
proxy_metadata: input.proxy_metadata,
proxy_metadata,
proxy_version: normalize_optional_string(
input.proxy_version.as_deref(),
"proxy_version",
@@ -4,3 +4,5 @@ pub(crate) mod modules;
pub(crate) mod paths;
pub(crate) mod settings;
pub(crate) mod smtp;
pub(crate) mod update;
pub(crate) mod update_client;
@@ -1,3 +1,4 @@
use crate::backup::config::S3BackupConfig;
use crate::bark_push::bark_push_configured;
use crate::handlers::admin::request::AdminAppState;
use crate::handlers::shared::{module_available_from_env, system_config_bool};
@@ -112,7 +113,7 @@ pub(crate) const ADMIN_MODULE_DEFINITIONS: &[AdminModuleDefinition] = &[
AdminModuleDefinition {
name: "model_directives",
display_name: "模型后缀参数",
description: "允许通过模型名后缀覆盖推理参数",
description: "允许通过模型名后缀覆盖推理参数或服务层级",
category: "integration",
env_key: "MODEL_DIRECTIVES_AVAILABLE",
default_available: true,
@@ -121,6 +122,18 @@ pub(crate) const ADMIN_MODULE_DEFINITIONS: &[AdminModuleDefinition] = &[
admin_menu_group: None,
admin_menu_order: 59,
},
AdminModuleDefinition {
name: "s3_backup",
display_name: "S3 备份",
description: "将配置、用户或完整数据定期备份到 S3-compatible 对象存储",
category: "integration",
env_key: "S3_BACKUP_AVAILABLE",
default_available: true,
admin_route: Some("/admin/modules/s3-backup"),
admin_menu_icon: Some("CloudUpload"),
admin_menu_group: None,
admin_menu_order: 60,
},
AdminModuleDefinition {
name: "gemini_files",
display_name: "文件缓存",
@@ -183,6 +196,7 @@ pub(crate) struct AdminModuleRuntimeState {
important_notification_configured: bool,
server_chan_push_configured: bool,
bark_push_configured: bool,
s3_backup_configured: bool,
}
pub(crate) fn admin_module_by_name(name: &str) -> Option<&'static AdminModuleDefinition> {
@@ -209,6 +223,8 @@ pub(crate) fn admin_module_enabled_config_key(module: &AdminModuleDefinition) ->
ENABLE_MODEL_DIRECTIVES_CONFIG_KEY.to_string()
} else if module.name == "important_notification" {
IMPORTANT_NOTIFICATION_ENABLED_KEY.to_string()
} else if module.name == "s3_backup" {
crate::backup::S3_BACKUP_ENABLED_KEY.to_string()
} else {
format!("module.{}.enabled", module.name)
}
@@ -272,6 +288,7 @@ pub(crate) async fn build_admin_module_runtime_state(
let notification_configured = important_notification_configured(state.app()).await?;
let server_chan_configured = server_chan_push_configured(state.app()).await?;
let bark_configured = bark_push_configured(state.app()).await?;
let backup_configured = s3_backup_configured(state.app()).await;
Ok(AdminModuleRuntimeState {
oauth_providers,
@@ -280,21 +297,36 @@ pub(crate) async fn build_admin_module_runtime_state(
important_notification_configured: notification_configured,
server_chan_push_configured: server_chan_configured,
bark_push_configured: bark_configured,
s3_backup_configured: backup_configured,
})
}
async fn s3_backup_configured(app: &crate::AppState) -> bool {
let Ok(mut values) = crate::backup::task::load_s3_backup_config_values(app).await else {
return false;
};
values.insert(
crate::backup::S3_BACKUP_ENABLED_KEY.to_string(),
json!(true),
);
S3BackupConfig::from_json_map(&values).is_ok()
}
pub(crate) fn build_admin_module_validation_result(
module: &AdminModuleDefinition,
runtime: &AdminModuleRuntimeState,
) -> (bool, Option<String>) {
admin_system_kernel::build_admin_module_validation_result(
module.name,
&runtime.oauth_providers,
runtime.ldap_config.as_ref(),
runtime.gemini_files_has_capable_key,
runtime.important_notification_configured,
runtime.server_chan_push_configured,
runtime.bark_push_configured,
admin_system_kernel::AdminModuleValidationInput {
module_name: module.name,
oauth_providers: &runtime.oauth_providers,
ldap_config: runtime.ldap_config.as_ref(),
gemini_files_has_capable_key: runtime.gemini_files_has_capable_key,
important_notification_configured: runtime.important_notification_configured,
server_chan_push_configured: runtime.server_chan_push_configured,
bark_push_configured: runtime.bark_push_configured,
s3_backup_configured: runtime.s3_backup_configured,
},
)
}
@@ -1,11 +1,18 @@
use crate::handlers::admin::request::AdminAppState;
use crate::handlers::admin::shared::build_admin_usage_counter_health_payload;
use crate::handlers::admin::system::shared::update::{
current_self_update_blocker, self_update_supported,
};
use crate::handlers::admin::system::shared::update_client::{
build_direct_update_http_client, build_update_http_client, has_explicit_update_proxy_env,
update_github_token_from_env,
};
use crate::handlers::shared::{system_config_bool, system_config_string};
use crate::GatewayError;
use aether_admin::system::{
build_admin_api_formats_payload as build_admin_api_formats_payload_pure,
build_admin_system_check_update_payload as build_admin_system_check_update_payload_pure,
build_admin_system_check_update_payload_with_release,
build_admin_system_check_update_payload_with_release, build_admin_system_releases_payload,
build_admin_system_settings_payload as build_admin_system_settings_payload_pure,
build_admin_system_settings_updated_payload,
build_admin_system_stats_payload as build_admin_system_stats_payload_pure,
@@ -15,12 +22,18 @@ use axum::body::Bytes;
use axum::http;
#[cfg(not(test))]
use serde::Deserialize;
use serde_json::json;
#[cfg(not(test))]
use serde_json::{json, Value};
use std::time::Duration;
const AETHER_RELEASES_API_URL: &str =
"https://api.github.com/repos/fawney19/Aether/releases?per_page=20";
const AETHER_RELEASE_TAG_URL_BASE: &str = "https://github.com/fawney19/Aether/releases/tag";
const SOURCE_BUILD_UPDATE_BLOCKER: &str = "当前为源码构建,请使用 git pull 后重新编译。";
const SOURCE_BUILD_RELEASE_BLOCKER: &str = "当前为源码构建,请手动切换到对应标签后重新编译。";
/// Minimum interval between actual GitHub API requests. Within this window
/// the cached result is reused.
const RELEASE_CACHE_TTL: Duration = Duration::from_secs(1200);
pub(crate) fn current_aether_version() -> String {
option_env!("AETHER_BUILD_VERSION")
@@ -37,56 +50,424 @@ pub(crate) fn build_admin_system_check_update_payload_from_release(
latest_release: Option<AdminSystemUpdateRelease>,
error: Option<String>,
) -> serde_json::Value {
build_admin_system_check_update_payload_with_release(
let mut payload = build_admin_system_check_update_payload_with_release(
current_aether_version(),
latest_release,
error,
)
);
apply_self_update_check_update_override(&mut payload, self_update_supported());
payload
}
pub(crate) fn build_admin_system_releases_list_payload(
releases: Vec<AdminSystemUpdateRelease>,
error: Option<String>,
) -> serde_json::Value {
let mut payload =
build_admin_system_releases_payload(current_aether_version(), releases, error);
apply_self_update_releases_override(&mut payload, self_update_supported());
payload
}
fn current_build_is_release() -> bool {
option_env!("AETHER_BUILD_TYPE").unwrap_or("source") == "release"
}
fn apply_self_update_check_update_override(payload: &mut Value, supported: bool) {
apply_self_update_check_update_override_with_blocker(
payload,
supported,
current_self_update_blocker(),
);
}
fn apply_self_update_check_update_override_with_blocker(
payload: &mut Value,
supported: bool,
blocker: &str,
) {
if supported {
return;
}
if payload.get("has_update").and_then(Value::as_bool) != Some(true) {
return;
}
payload["updatable"] = json!(false);
payload["update_blocker"] = json!(blocker);
}
fn apply_self_update_releases_override(payload: &mut Value, supported: bool) {
apply_self_update_releases_override_with_blocker(
payload,
supported,
current_self_update_release_blocker(),
);
}
fn apply_self_update_releases_override_with_blocker(
payload: &mut Value,
supported: bool,
blocker: &str,
) {
if supported {
return;
}
let Some(releases) = payload.get_mut("releases").and_then(Value::as_array_mut) else {
return;
};
for release in releases {
if release.get("is_current").and_then(Value::as_bool) == Some(true) {
continue;
}
release["updatable"] = json!(false);
if release.get("update_blocker").is_none() || release["update_blocker"].is_null() {
release["update_blocker"] = json!(blocker);
}
}
}
fn current_self_update_release_blocker() -> &'static str {
if !current_build_is_release() {
return SOURCE_BUILD_RELEASE_BLOCKER;
}
if self_update_supported() {
""
} else {
current_self_update_blocker()
}
}
#[cfg(not(test))]
struct CachedReleases {
all: Vec<AdminSystemUpdateRelease>,
latest: Option<AdminSystemUpdateRelease>,
error: Option<String>,
fetched_at: std::time::Instant,
}
#[cfg(not(test))]
fn releases_cache() -> &'static std::sync::Mutex<Option<CachedReleases>> {
static CACHE: std::sync::OnceLock<std::sync::Mutex<Option<CachedReleases>>> =
std::sync::OnceLock::new();
CACHE.get_or_init(|| std::sync::Mutex::new(None))
}
#[cfg(not(test))]
async fn ensure_releases_cached(force: bool) {
{
if let Ok(guard) = releases_cache().lock() {
if let Some(cached) = guard.as_ref() {
if should_reuse_releases_cache(
force,
cached.error.is_some(),
cached.fetched_at.elapsed(),
) {
return;
}
}
}
}
let (all, latest, error) = match fetch_admin_system_releases_inner().await {
Ok(releases) => {
let latest = releases.first().cloned();
(releases, latest, None)
}
Err(err) => (Vec::new(), None, Some(err)),
};
if let Ok(mut guard) = releases_cache().lock() {
*guard = Some(CachedReleases {
all,
latest,
error,
fetched_at: std::time::Instant::now(),
});
}
}
fn should_reuse_releases_cache(force: bool, has_error: bool, age: Duration) -> bool {
!force && !has_error && age < RELEASE_CACHE_TTL
}
#[cfg(not(test))]
pub(crate) async fn fetch_latest_admin_system_release(
force: bool,
) -> (Option<AdminSystemUpdateRelease>, Option<String>) {
match fetch_latest_admin_system_release_inner().await {
Ok(release) => (release, None),
Err(err) => (None, Some(err)),
ensure_releases_cached(force).await;
if let Ok(guard) = releases_cache().lock() {
if let Some(cached) = guard.as_ref() {
return (cached.latest.clone(), cached.error.clone());
}
}
(None, None)
}
#[cfg(not(test))]
pub(crate) async fn fetch_admin_system_releases(
force: bool,
) -> (Vec<AdminSystemUpdateRelease>, Option<String>) {
ensure_releases_cached(force).await;
if let Ok(guard) = releases_cache().lock() {
if let Some(cached) = guard.as_ref() {
return (cached.all.clone(), cached.error.clone());
}
}
(Vec::new(), None)
}
#[cfg(test)]
pub(crate) async fn fetch_admin_system_releases(
_force: bool,
) -> (Vec<AdminSystemUpdateRelease>, Option<String>) {
(
Vec::new(),
Some("测试环境未请求 GitHub Releases".to_string()),
)
}
#[cfg(test)]
pub(crate) async fn fetch_latest_admin_system_release(
_force: bool,
) -> (Option<AdminSystemUpdateRelease>, Option<String>) {
(None, Some("测试环境未请求 GitHub Releases".to_string()))
}
#[cfg(not(test))]
async fn fetch_latest_admin_system_release_inner(
) -> Result<Option<AdminSystemUpdateRelease>, String> {
let releases: Vec<GitHubRelease> = reqwest::Client::builder()
.timeout(Duration::from_secs(8))
.build()
.map_err(|err| format!("创建更新检查客户端失败: {err}"))?
.get(AETHER_RELEASES_API_URL)
.header(reqwest::header::USER_AGENT, "Aether-Gateway update-check")
.header(reqwest::header::ACCEPT, "application/vnd.github+json")
.send()
.await
.map_err(|err| format!("请求 GitHub Releases 失败: {err}"))?
.error_for_status()
.map_err(|err| format!("GitHub Releases 返回错误: {err}"))?
.json()
.await
.map_err(|err| format!("解析 GitHub Releases 失败: {err}"))?;
pub(crate) async fn resolve_update_target(
version: Option<String>,
) -> Result<(String, String, Option<String>), (http::StatusCode, serde_json::Value)> {
let (releases, error) = fetch_admin_system_releases(false).await;
if releases.is_empty() {
return Err((
http::StatusCode::SERVICE_UNAVAILABLE,
json!({ "detail": error.unwrap_or_else(|| "无法获取版本信息".to_string()) }),
));
}
let release = match version {
Some(ref v) => releases.into_iter().find(|r| r.version == *v),
None => releases.into_iter().next(),
};
let release = release.ok_or_else(|| {
(
http::StatusCode::NOT_FOUND,
json!({ "detail": "未找到指定版本" }),
)
})?;
let tarball_url = release.tarball_url.ok_or_else(|| {
(
http::StatusCode::PRECONDITION_REQUIRED,
json!({ "detail": format!("版本 {} 没有适用于当前平台的安装包", release.version) }),
)
})?;
let sha256sums_url = release.sha256sums_url.ok_or_else(|| {
(
http::StatusCode::PRECONDITION_REQUIRED,
json!({ "detail": format!("版本 {} 缺少 SHA256SUMS 校验文件,已拒绝在线更新", release.version) }),
)
})?;
Ok((release.version, tarball_url, Some(sha256sums_url)))
}
#[cfg(test)]
pub(crate) async fn resolve_update_target(
_version: Option<String>,
) -> Result<(String, String, Option<String>), (http::StatusCode, serde_json::Value)> {
Err((
http::StatusCode::SERVICE_UNAVAILABLE,
json!({ "detail": "测试环境不支持更新" }),
))
}
#[cfg(not(test))]
async fn fetch_admin_system_releases_inner() -> Result<Vec<AdminSystemUpdateRelease>, String> {
let current_channel = update_channel_for_version(&current_aether_version());
let timeout = Duration::from_secs(8);
let github_token = update_github_token_from_env();
let client = build_update_http_client(timeout, "更新检查")?;
let releases = match fetch_github_releases_with_client(&client, github_token.as_deref()).await {
Ok(releases) => releases,
Err(err) if err.rate_limited && !has_explicit_update_proxy_env() => {
let direct_client = build_direct_update_http_client(timeout, "更新检查直连重试")?;
fetch_github_releases_with_client(&direct_client, github_token.as_deref())
.await
.map_err(|retry_err| retry_err.message)?
}
Err(err) => return Err(err.message),
};
Ok(releases
.into_iter()
.find(|release| !release.draft && release.tag_name.starts_with('v'))
.map(|release| AdminSystemUpdateRelease {
version: release.tag_name,
release_url: Some(release.html_url),
release_notes: release.body.filter(|body| !body.trim().is_empty()),
published_at: release.published_at,
}))
.filter(|release| should_include_release_for_channel(release, current_channel))
.map(|release| {
let (tarball_url, sha256sums_url) = select_release_tarball_urls(&release);
AdminSystemUpdateRelease {
version: release.tag_name.clone(),
release_url: Some(github_release_tag_url(&release.tag_name)),
release_notes: release.body.filter(|body| !body.trim().is_empty()),
published_at: release.published_at,
tarball_url,
sha256sums_url,
}
})
.collect())
}
#[derive(Debug)]
struct GitHubReleaseFetchError {
message: String,
rate_limited: bool,
}
#[cfg(not(test))]
async fn fetch_github_releases_with_client(
client: &reqwest::Client,
github_token: Option<&str>,
) -> Result<Vec<GitHubRelease>, GitHubReleaseFetchError> {
let mut request = client
.get(AETHER_RELEASES_API_URL)
.header(reqwest::header::USER_AGENT, "Aether-Gateway update-check")
.header(reqwest::header::ACCEPT, "application/vnd.github+json");
if let Some(token) = github_token {
request = request.bearer_auth(token);
}
let response = request
.send()
.await
.map_err(|err| GitHubReleaseFetchError {
message: format!("请求 GitHub Releases 失败: {err}"),
rate_limited: false,
})?;
let status = response.status();
if !status.is_success() {
let body = response.text().await.unwrap_or_default();
return Err(github_release_response_error(status, &body));
}
response
.json()
.await
.map_err(|err| GitHubReleaseFetchError {
message: format!("解析 GitHub Releases 失败: {err}"),
rate_limited: false,
})
}
fn github_release_response_error(
status: reqwest::StatusCode,
response_body: &str,
) -> GitHubReleaseFetchError {
if is_github_rate_limit_error(status, response_body) {
return GitHubReleaseFetchError {
message: "GitHub Releases API 已触发限流;当前共享代理出口的匿名额度已用尽,更新检查将自动尝试直连。若仍失败,请配置 AETHER_UPDATE_GITHUB_TOKEN / GITHUB_TOKEN / GH_TOKEN,或为 GitHub 更新检查单独设置可用代理。".to_string(),
rate_limited: true,
};
}
let detail = parse_github_error_message(response_body);
let message = if detail.is_empty() {
format!("GitHub Releases 返回错误: HTTP {status}")
} else {
format!("GitHub Releases 返回错误: HTTP {status}; {detail}")
};
GitHubReleaseFetchError {
message,
rate_limited: false,
}
}
fn is_github_rate_limit_error(status: reqwest::StatusCode, response_body: &str) -> bool {
status == reqwest::StatusCode::FORBIDDEN
&& response_body
.to_ascii_lowercase()
.contains("rate limit exceeded")
}
fn parse_github_error_message(response_body: &str) -> String {
serde_json::from_str::<serde_json::Value>(response_body)
.ok()
.and_then(|value| {
value
.get("message")
.and_then(|message| message.as_str())
.map(str::trim)
.filter(|message| !message.is_empty())
.map(ToString::to_string)
})
.unwrap_or_default()
}
#[cfg(not(test))]
fn should_include_release_for_channel(
release: &GitHubRelease,
current_channel: UpdateChannel,
) -> bool {
if release.draft || !release.tag_name.starts_with('v') {
return false;
}
if !release.prerelease {
return true;
}
current_channel.allows_prerelease(&release.tag_name)
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum UpdateChannel {
Stable,
Rc,
Beta,
OtherPrerelease,
}
impl UpdateChannel {
fn allows_prerelease(self, release_version: &str) -> bool {
match self {
Self::Stable => false,
Self::Rc => update_channel_for_version(release_version) == Self::Rc,
Self::Beta => update_channel_for_version(release_version) == Self::Beta,
Self::OtherPrerelease => {
matches!(
update_channel_for_version(release_version),
Self::Rc | Self::Beta | Self::OtherPrerelease
)
}
}
}
}
fn update_channel_for_version(version: &str) -> UpdateChannel {
let normalized = version
.trim()
.strip_prefix('v')
.or_else(|| version.trim().strip_prefix('V'))
.unwrap_or(version.trim());
let Some((_, prerelease)) = normalized.split_once('-') else {
return UpdateChannel::Stable;
};
let prerelease = prerelease.to_ascii_lowercase();
if prerelease.starts_with("rc") {
UpdateChannel::Rc
} else if prerelease.starts_with("beta") {
UpdateChannel::Beta
} else {
UpdateChannel::OtherPrerelease
}
}
#[cfg(not(test))]
#[derive(Debug, Deserialize)]
struct GitHubReleaseAsset {
name: String,
browser_download_url: String,
}
#[cfg(not(test))]
@@ -100,6 +481,43 @@ struct GitHubRelease {
published_at: Option<String>,
#[serde(default)]
draft: bool,
#[serde(default)]
prerelease: bool,
#[serde(default)]
assets: Vec<GitHubReleaseAsset>,
}
fn github_release_tag_url(tag_name: &str) -> String {
format!("{AETHER_RELEASE_TAG_URL_BASE}/{tag_name}")
}
#[cfg(not(test))]
fn select_release_tarball_urls(release: &GitHubRelease) -> (Option<String>, Option<String>) {
let platform = if cfg!(target_os = "macos") {
"macos"
} else {
"linux"
};
let arch = if cfg!(target_arch = "aarch64") {
"arm64"
} else {
"amd64"
};
let expected_name = format!("aether-{}-{}-{}.tar.gz", release.tag_name, platform, arch);
let tarball_url = release
.assets
.iter()
.find(|a| a.name == expected_name)
.map(|a| a.browser_download_url.clone());
let sha256sums_url = release
.assets
.iter()
.find(|a| a.name == "SHA256SUMS")
.map(|a| a.browser_download_url.clone());
(tarball_url, sha256sums_url)
}
pub(crate) async fn build_admin_system_stats_payload(
@@ -249,3 +667,145 @@ pub(crate) async fn apply_admin_system_settings_update(
pub(crate) fn build_admin_api_formats_payload() -> serde_json::Value {
build_admin_api_formats_payload_pure()
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn update_channel_detects_stable_rc_beta_and_other_prerelease() {
assert_eq!(update_channel_for_version("v1.2.3"), UpdateChannel::Stable);
assert_eq!(update_channel_for_version("1.2.3-rc1"), UpdateChannel::Rc);
assert_eq!(
update_channel_for_version("1.2.3-beta.2"),
UpdateChannel::Beta
);
assert_eq!(
update_channel_for_version("1.2.3-alpha.1"),
UpdateChannel::OtherPrerelease
);
}
#[test]
fn stable_channel_does_not_allow_prereleases() {
assert!(!UpdateChannel::Stable.allows_prerelease("v1.2.3-rc1"));
}
#[test]
fn prerelease_channels_only_follow_matching_channel() {
assert!(UpdateChannel::Rc.allows_prerelease("v1.2.3-rc2"));
assert!(!UpdateChannel::Rc.allows_prerelease("v1.2.3-beta.1"));
assert!(UpdateChannel::Beta.allows_prerelease("v1.2.3-beta.2"));
assert!(!UpdateChannel::Beta.allows_prerelease("v1.2.3-rc1"));
assert!(UpdateChannel::OtherPrerelease.allows_prerelease("v1.2.3-alpha.1"));
assert!(UpdateChannel::OtherPrerelease.allows_prerelease("v1.2.3-rc1"));
}
#[test]
fn github_release_tag_url_points_to_explicit_tag_page() {
assert_eq!(
github_release_tag_url("v0.7.3"),
"https://github.com/fawney19/Aether/releases/tag/v0.7.3"
);
}
#[test]
fn successful_release_cache_is_reused_within_ttl() {
assert!(should_reuse_releases_cache(
false,
false,
Duration::from_secs(60)
));
}
#[test]
fn failed_release_cache_is_not_reused() {
assert!(!should_reuse_releases_cache(
false,
true,
Duration::from_secs(60)
));
}
#[test]
fn force_refresh_bypasses_release_cache() {
assert!(!should_reuse_releases_cache(
true,
false,
Duration::from_secs(60)
));
}
#[test]
fn github_rate_limit_response_is_marked_retryable() {
let err = github_release_response_error(
reqwest::StatusCode::FORBIDDEN,
r#"{"message":"API rate limit exceeded for 1.2.3.4."}"#,
);
assert!(err.rate_limited);
assert!(err.message.contains("GitHub Releases API 已触发限流"));
}
#[test]
fn non_rate_limit_github_response_keeps_http_status_message() {
let err = github_release_response_error(
reqwest::StatusCode::FORBIDDEN,
r#"{"message":"Resource not accessible"}"#,
);
assert!(!err.rate_limited);
assert!(err.message.contains("HTTP 403 Forbidden"));
assert!(err.message.contains("Resource not accessible"));
}
#[test]
fn self_update_check_update_override_marks_latest_release_non_updatable() {
let mut payload = json!({
"has_update": true,
"updatable": true,
"update_blocker": serde_json::Value::Null
});
apply_self_update_check_update_override_with_blocker(
&mut payload,
false,
SOURCE_BUILD_UPDATE_BLOCKER,
);
assert_eq!(payload["updatable"], false);
assert_eq!(payload["update_blocker"], SOURCE_BUILD_UPDATE_BLOCKER);
}
#[test]
fn self_update_releases_override_marks_non_current_entries_non_updatable() {
let mut payload = json!({
"releases": [
{
"version": "v0.7.3",
"is_current": false,
"updatable": true,
"update_blocker": serde_json::Value::Null
},
{
"version": "v0.7.2",
"is_current": true,
"updatable": true,
"update_blocker": "当前版本"
}
]
});
apply_self_update_releases_override_with_blocker(
&mut payload,
false,
SOURCE_BUILD_RELEASE_BLOCKER,
);
assert_eq!(payload["releases"][0]["updatable"], false);
assert_eq!(
payload["releases"][0]["update_blocker"],
SOURCE_BUILD_RELEASE_BLOCKER
);
assert_eq!(payload["releases"][1]["updatable"], true);
assert_eq!(payload["releases"][1]["update_blocker"], "当前版本");
}
}
File diff suppressed because it is too large Load Diff
@@ -0,0 +1,68 @@
use std::time::Duration;
const EXPLICIT_UPDATE_PROXY_ENV_KEYS: &[&str] = &["AETHER_UPDATE_PROXY_URL", "UPDATE_PROXY_URL"];
const UPDATE_PROXY_ENV_KEYS: &[&str] = &[
"AETHER_UPDATE_PROXY_URL",
"UPDATE_PROXY_URL",
"HTTPS_PROXY",
"https_proxy",
"ALL_PROXY",
"all_proxy",
"HTTP_PROXY",
"http_proxy",
];
const UPDATE_GITHUB_TOKEN_ENV_KEYS: &[&str] =
&["AETHER_UPDATE_GITHUB_TOKEN", "GITHUB_TOKEN", "GH_TOKEN"];
pub(crate) fn build_update_http_client(
timeout: Duration,
label: &str,
) -> Result<reqwest::Client, String> {
let mut builder = base_update_http_client_builder(timeout);
if let Some(proxy_url) = update_proxy_url_from_env() {
let proxy = reqwest::Proxy::all(proxy_url)
.map_err(|_| format!("创建{label}代理失败,请检查更新代理环境变量"))?
.no_proxy(reqwest::NoProxy::from_env());
builder = builder.proxy(proxy);
}
builder
.build()
.map_err(|err| format!("创建{label}客户端失败: {err}"))
}
pub(crate) fn build_direct_update_http_client(
timeout: Duration,
label: &str,
) -> Result<reqwest::Client, String> {
base_update_http_client_builder(timeout)
.no_proxy()
.build()
.map_err(|err| format!("创建{label}客户端失败: {err}"))
}
pub(crate) fn has_explicit_update_proxy_env() -> bool {
read_nonempty_env_value(EXPLICIT_UPDATE_PROXY_ENV_KEYS).is_some()
}
fn base_update_http_client_builder(timeout: Duration) -> reqwest::ClientBuilder {
reqwest::Client::builder().timeout(timeout)
}
fn update_proxy_url_from_env() -> Option<String> {
read_nonempty_env_value(UPDATE_PROXY_ENV_KEYS)
}
pub(crate) fn update_github_token_from_env() -> Option<String> {
read_nonempty_env_value(UPDATE_GITHUB_TOKEN_ENV_KEYS)
}
fn read_nonempty_env_value(keys: &[&str]) -> Option<String> {
keys.iter().find_map(|key| {
std::env::var(key)
.ok()
.map(|value| value.trim().to_string())
.filter(|value| !value.is_empty())
})
}
@@ -417,6 +417,7 @@ async fn resolve_admin_user_selection(
group_id: filters
.as_ref()
.and_then(|filters| filters.group_id.clone()),
..Default::default()
})
.await
.map_err(|_| "用户数据不可用".to_string())?
@@ -37,6 +37,13 @@ pub(in super::super) async fn build_admin_list_users_response(
.map(|value| value.trim().to_string())
.filter(|value| !value.is_empty());
let sort_by = query_param_value(request_context.query_string(), "sort_by")
.and_then(|value| aether_data::repository::users::UserExportSortBy::parse(&value))
.unwrap_or_default();
let sort_order = query_param_value(request_context.query_string(), "sort_order")
.and_then(|value| aether_data::repository::users::UserExportSortOrder::parse(&value))
.unwrap_or_default();
let query = aether_data::repository::users::UserExportListQuery {
skip,
limit,
@@ -44,6 +51,8 @@ pub(in super::super) async fn build_admin_list_users_response(
is_active,
search,
group_id,
sort_by,
sort_order,
};
let (paged_rows_result, total_result) = tokio::join!(
state.list_export_users_page(&query),
@@ -152,12 +152,7 @@ pub(in super::super) async fn build_admin_update_user_response(
};
let effective_role = role.as_deref().unwrap_or(existing_user.role.as_str());
let group_ids = if field_presence.contains("group_ids") {
let requested_group_ids = normalize_admin_user_group_ids(payload.group_ids);
Some(
state
.include_default_user_group_ids_for_role(&requested_group_ids, effective_role)
.await?,
)
Some(normalize_admin_user_group_ids(payload.group_ids))
} else if role.is_some() {
let requested_group_ids = state
.list_user_groups_for_user(&user_id)