mirror of
https://github.com/fawney19/Aether.git
synced 2026-09-02 01:10:23 +08:00
fix: simplify account self-check config
This commit is contained in:
@@ -3,9 +3,10 @@ pub(crate) use crate::handlers::admin::{
|
||||
build_internal_control_error_response, create_provider_oauth_catalog_key,
|
||||
find_duplicate_provider_oauth_key, maybe_build_local_admin_pool_response,
|
||||
maybe_build_local_admin_response, persist_provider_quota_refresh_state,
|
||||
provider_account_self_check_endpoint_for_provider,
|
||||
provider_oauth_maintenance_endpoint_for_provider, provider_oauth_runtime_endpoint_for_provider,
|
||||
provider_quota_refresh_endpoint_for_provider, provider_type_supports_quota_refresh,
|
||||
reconcile_admin_fixed_provider_template_endpoints,
|
||||
provider_quota_refresh_endpoint_for_provider, provider_type_supports_account_self_check,
|
||||
provider_type_supports_quota_refresh, reconcile_admin_fixed_provider_template_endpoints,
|
||||
refresh_provider_oauth_account_state_after_update, refresh_provider_pool_quota_locally,
|
||||
update_existing_provider_oauth_catalog_key, AdminAppState,
|
||||
AdminGatewayProviderTransportSnapshot, AdminLocalOAuthRefreshError, AdminRequestContext,
|
||||
|
||||
@@ -27,7 +27,8 @@ pub(crate) use self::provider::oauth::provisioning::{
|
||||
};
|
||||
pub(crate) use self::provider::oauth::quota::dispatch::refresh_provider_pool_quota_locally;
|
||||
pub(crate) use self::provider::oauth::quota::shared::{
|
||||
persist_provider_quota_refresh_state, provider_quota_refresh_endpoint_for_provider,
|
||||
persist_provider_quota_refresh_state, provider_account_self_check_endpoint_for_provider,
|
||||
provider_quota_refresh_endpoint_for_provider, provider_type_supports_account_self_check,
|
||||
provider_type_supports_quota_refresh,
|
||||
};
|
||||
pub(crate) use self::provider::oauth::runtime::{
|
||||
|
||||
@@ -76,6 +76,22 @@ pub(crate) fn provider_quota_refresh_endpoint_for_provider(
|
||||
)
|
||||
}
|
||||
|
||||
pub(crate) fn provider_type_supports_account_self_check(provider_type: &str) -> bool {
|
||||
ProviderPoolService::with_builtin_adapters().supports_account_self_check(provider_type)
|
||||
}
|
||||
|
||||
pub(crate) fn provider_account_self_check_endpoint_for_provider(
|
||||
provider_type: &str,
|
||||
endpoints: &[StoredProviderCatalogEndpoint],
|
||||
include_inactive: bool,
|
||||
) -> Option<StoredProviderCatalogEndpoint> {
|
||||
ProviderPoolService::with_builtin_adapters().account_self_check_endpoint_for_provider(
|
||||
provider_type,
|
||||
endpoints,
|
||||
include_inactive,
|
||||
)
|
||||
}
|
||||
|
||||
pub(crate) fn provider_quota_refresh_missing_endpoint_message(provider_type: &str) -> String {
|
||||
ProviderPoolService::with_builtin_adapters()
|
||||
.quota_refresh_missing_endpoint_message(provider_type)
|
||||
|
||||
@@ -57,18 +57,6 @@ fn parse_pool_probe_target_count(pool_advanced: &Map<String, Value>) -> Option<u
|
||||
.map(|value| value.min(100_000))
|
||||
}
|
||||
|
||||
fn parse_pool_account_self_check_method(pool_advanced: &Map<String, Value>) -> String {
|
||||
pool_advanced
|
||||
.get("account_self_check_method")
|
||||
.or_else(|| pool_advanced.get("self_check_method"))
|
||||
.and_then(Value::as_str)
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
.map(str::to_ascii_lowercase)
|
||||
.filter(|value| matches!(value.as_str(), "quota_refresh" | "custom_request"))
|
||||
.unwrap_or_else(|| "quota_refresh".to_string())
|
||||
}
|
||||
|
||||
fn pool_score_weight(object: &Map<String, Value>, names: &[&str], current: f64) -> f64 {
|
||||
names
|
||||
.iter()
|
||||
@@ -426,8 +414,6 @@ pub(crate) fn admin_provider_pool_config_from_config_value(
|
||||
account_self_check_enabled: false,
|
||||
account_self_check_interval_minutes: 60,
|
||||
account_self_check_concurrency: 4,
|
||||
account_self_check_method: "quota_refresh".to_string(),
|
||||
account_self_check_request: None,
|
||||
score_top_n: 128,
|
||||
score_fallback_scan_limit: 1024,
|
||||
score_rules: PoolMemberScoreRules::default(),
|
||||
@@ -522,12 +508,6 @@ pub(crate) fn admin_provider_pool_config_from_config_value(
|
||||
.filter(|value| *value > 0)
|
||||
.map(|value| value.min(64))
|
||||
.unwrap_or(4),
|
||||
account_self_check_method: parse_pool_account_self_check_method(pool_advanced),
|
||||
account_self_check_request: pool_advanced
|
||||
.get("account_self_check_request")
|
||||
.or_else(|| pool_advanced.get("self_check_request"))
|
||||
.filter(|value| value.is_object())
|
||||
.cloned(),
|
||||
score_top_n: pool_advanced
|
||||
.get("score_top_n")
|
||||
.and_then(json_u64)
|
||||
@@ -619,12 +599,6 @@ mod tests {
|
||||
"account_self_check_enabled": true,
|
||||
"account_self_check_interval_minutes": 90,
|
||||
"account_self_check_concurrency": 5,
|
||||
"account_self_check_method": "custom_request",
|
||||
"account_self_check_request": {
|
||||
"method": "GET",
|
||||
"path": "/v1/me",
|
||||
"success_status_codes": [200]
|
||||
},
|
||||
"score_top_n": 256,
|
||||
"score_fallback_scan_limit": 2048,
|
||||
"score_rules": {
|
||||
@@ -667,15 +641,6 @@ mod tests {
|
||||
assert!(config.account_self_check_enabled);
|
||||
assert_eq!(config.account_self_check_interval_minutes, 90);
|
||||
assert_eq!(config.account_self_check_concurrency, 5);
|
||||
assert_eq!(config.account_self_check_method, "custom_request");
|
||||
assert_eq!(
|
||||
config
|
||||
.account_self_check_request
|
||||
.as_ref()
|
||||
.and_then(|value| value.get("path"))
|
||||
.and_then(serde_json::Value::as_str),
|
||||
Some("/v1/me")
|
||||
);
|
||||
assert_eq!(config.score_top_n, 256);
|
||||
assert_eq!(config.score_fallback_scan_limit, 2048);
|
||||
assert_eq!(config.score_rules.weights.manual_priority, 0.4);
|
||||
|
||||
@@ -621,8 +621,6 @@ mod tests {
|
||||
account_self_check_enabled: false,
|
||||
account_self_check_interval_minutes: 60,
|
||||
account_self_check_concurrency: 4,
|
||||
account_self_check_method: "quota_refresh".to_string(),
|
||||
account_self_check_request: None,
|
||||
score_top_n: 128,
|
||||
score_fallback_scan_limit: 1024,
|
||||
score_rules: aether_pool_core::PoolMemberScoreRules::default(),
|
||||
|
||||
@@ -52,8 +52,6 @@ pub(crate) struct AdminProviderPoolConfig {
|
||||
pub(crate) account_self_check_enabled: bool,
|
||||
pub(crate) account_self_check_interval_minutes: u64,
|
||||
pub(crate) account_self_check_concurrency: u64,
|
||||
pub(crate) account_self_check_method: String,
|
||||
pub(crate) account_self_check_request: Option<serde_json::Value>,
|
||||
pub(crate) score_top_n: u64,
|
||||
pub(crate) score_fallback_scan_limit: u64,
|
||||
pub(crate) score_rules: PoolMemberScoreRules,
|
||||
|
||||
@@ -1,7 +1,6 @@
|
||||
use std::collections::{BTreeMap, BTreeSet};
|
||||
use std::collections::BTreeMap;
|
||||
use std::time::{Duration, SystemTime, UNIX_EPOCH};
|
||||
|
||||
use aether_contracts::{ExecutionPlan, ExecutionResult, RequestBody};
|
||||
use aether_data_contracts::repository::pool_scores::{
|
||||
PoolMemberHardState, PoolMemberIdentity, PoolMemberProbeAttempt, PoolMemberProbeResult,
|
||||
PoolMemberProbeStatus,
|
||||
@@ -10,17 +9,13 @@ use aether_data_contracts::repository::provider_catalog::{
|
||||
StoredProviderCatalogEndpoint, StoredProviderCatalogKey, StoredProviderCatalogProvider,
|
||||
};
|
||||
use aether_runtime_state::{RuntimeLockLease, RuntimeState};
|
||||
use base64::Engine as _;
|
||||
use futures_util::{stream, StreamExt};
|
||||
use serde_json::{json, Value};
|
||||
use tracing::{debug, info, warn};
|
||||
use uuid::Uuid;
|
||||
|
||||
use crate::admin_api::{
|
||||
admin_provider_pool_config, persist_provider_quota_refresh_state,
|
||||
provider_quota_refresh_endpoint_for_provider, provider_type_supports_quota_refresh,
|
||||
refresh_provider_pool_quota_locally, AdminAppState, AdminGatewayProviderTransportSnapshot,
|
||||
OAUTH_ACCOUNT_BLOCK_PREFIX, OAUTH_REQUEST_FAILED_PREFIX,
|
||||
admin_provider_pool_config, provider_account_self_check_endpoint_for_provider,
|
||||
provider_type_supports_account_self_check, refresh_provider_pool_quota_locally, AdminAppState,
|
||||
};
|
||||
use crate::{AppState, GatewayError};
|
||||
|
||||
@@ -94,37 +89,6 @@ impl AccountSelfCheckWorkerConfig {
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub(crate) struct AccountSelfCheckRequestConfig {
|
||||
pub(crate) method: String,
|
||||
pub(crate) url: Option<String>,
|
||||
pub(crate) path: Option<String>,
|
||||
pub(crate) headers: BTreeMap<String, String>,
|
||||
pub(crate) json_body: Option<Value>,
|
||||
pub(crate) body: Option<String>,
|
||||
pub(crate) body_bytes_b64: Option<String>,
|
||||
pub(crate) content_type: Option<String>,
|
||||
pub(crate) success_status_codes: BTreeSet<u16>,
|
||||
pub(crate) blocked_status_codes: BTreeSet<u16>,
|
||||
}
|
||||
|
||||
impl Default for AccountSelfCheckRequestConfig {
|
||||
fn default() -> Self {
|
||||
Self {
|
||||
method: "GET".to_string(),
|
||||
url: None,
|
||||
path: None,
|
||||
headers: BTreeMap::new(),
|
||||
json_body: None,
|
||||
body: None,
|
||||
body_bytes_b64: None,
|
||||
content_type: None,
|
||||
success_status_codes: BTreeSet::from([200]),
|
||||
blocked_status_codes: BTreeSet::from([401, 403, 423]),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
enum AccountSelfCheckOutcome {
|
||||
Success {
|
||||
status_code: Option<u16>,
|
||||
@@ -388,314 +352,6 @@ async fn select_keys_for_provider(
|
||||
result
|
||||
}
|
||||
|
||||
fn normalize_http_method(raw: Option<&Value>) -> String {
|
||||
let method = raw
|
||||
.and_then(Value::as_str)
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
.unwrap_or("GET")
|
||||
.to_ascii_uppercase();
|
||||
match method.as_str() {
|
||||
"GET" | "POST" | "PUT" | "PATCH" | "DELETE" | "HEAD" => method,
|
||||
_ => "GET".to_string(),
|
||||
}
|
||||
}
|
||||
|
||||
fn json_string(raw: Option<&Value>) -> Option<String> {
|
||||
raw.and_then(Value::as_str)
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
.map(ToOwned::to_owned)
|
||||
}
|
||||
|
||||
fn parse_status_codes(raw: Option<&Value>, fallback: &[u16]) -> BTreeSet<u16> {
|
||||
let mut out = BTreeSet::new();
|
||||
if let Some(array) = raw.and_then(Value::as_array) {
|
||||
for item in array {
|
||||
if let Some(value) = item
|
||||
.as_u64()
|
||||
.and_then(|value| u16::try_from(value).ok())
|
||||
.filter(|value| (100..=599).contains(value))
|
||||
{
|
||||
out.insert(value);
|
||||
}
|
||||
}
|
||||
}
|
||||
if out.is_empty() {
|
||||
out.extend(fallback.iter().copied());
|
||||
}
|
||||
out
|
||||
}
|
||||
|
||||
fn parse_headers(raw: Option<&Value>) -> BTreeMap<String, String> {
|
||||
let Some(object) = raw.and_then(Value::as_object) else {
|
||||
return BTreeMap::new();
|
||||
};
|
||||
object
|
||||
.iter()
|
||||
.filter_map(|(key, value)| {
|
||||
let key = key.trim().to_ascii_lowercase();
|
||||
let value = match value {
|
||||
Value::String(value) => value.trim().to_string(),
|
||||
_ => value.to_string(),
|
||||
};
|
||||
(!key.is_empty()).then_some((key, value))
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
|
||||
pub(crate) fn parse_account_self_check_request_config(
|
||||
raw: Option<&Value>,
|
||||
) -> AccountSelfCheckRequestConfig {
|
||||
let Some(object) = raw.and_then(Value::as_object) else {
|
||||
return AccountSelfCheckRequestConfig::default();
|
||||
};
|
||||
AccountSelfCheckRequestConfig {
|
||||
method: normalize_http_method(object.get("method")),
|
||||
url: json_string(object.get("url")),
|
||||
path: json_string(object.get("path")),
|
||||
headers: parse_headers(object.get("headers")),
|
||||
json_body: object
|
||||
.get("json_body")
|
||||
.or_else(|| object.get("json"))
|
||||
.or_else(|| object.get("body_json"))
|
||||
.cloned(),
|
||||
body: json_string(object.get("body")),
|
||||
body_bytes_b64: json_string(
|
||||
object
|
||||
.get("body_bytes_b64")
|
||||
.or_else(|| object.get("body_base64")),
|
||||
),
|
||||
content_type: json_string(object.get("content_type")),
|
||||
success_status_codes: parse_status_codes(object.get("success_status_codes"), &[200]),
|
||||
blocked_status_codes: parse_status_codes(
|
||||
object
|
||||
.get("blocked_status_codes")
|
||||
.or_else(|| object.get("banned_status_codes")),
|
||||
&[401, 403, 423],
|
||||
),
|
||||
}
|
||||
}
|
||||
|
||||
fn custom_check_url(
|
||||
state: &AdminAppState<'_>,
|
||||
transport: &AdminGatewayProviderTransportSnapshot,
|
||||
request_config: &AccountSelfCheckRequestConfig,
|
||||
) -> Option<String> {
|
||||
if let Some(url) = request_config.url.as_deref() {
|
||||
let parsed = url::Url::parse(url).ok()?;
|
||||
if !matches!(parsed.scheme(), "http" | "https") {
|
||||
return None;
|
||||
}
|
||||
return Some(url.to_string());
|
||||
}
|
||||
let path = request_config.path.as_deref()?;
|
||||
state.build_passthrough_path_url(&transport.endpoint.base_url, path, None, &["key"])
|
||||
}
|
||||
|
||||
async fn resolve_self_check_auth(
|
||||
state: &AdminAppState<'_>,
|
||||
transport: &AdminGatewayProviderTransportSnapshot,
|
||||
) -> Result<Option<(String, String)>, GatewayError> {
|
||||
if let Some(auth) = state.resolve_local_oauth_header_auth(transport).await? {
|
||||
return Ok(Some(auth));
|
||||
}
|
||||
if let Some(auth) = crate::provider_transport::auth::resolve_local_openai_bearer_auth(transport)
|
||||
{
|
||||
return Ok(Some(auth));
|
||||
}
|
||||
if let Some(auth) = crate::provider_transport::auth::resolve_local_standard_auth(transport) {
|
||||
return Ok(Some(auth));
|
||||
}
|
||||
Ok(state.resolve_local_gemini_auth(transport))
|
||||
}
|
||||
|
||||
fn custom_check_body(request_config: &AccountSelfCheckRequestConfig) -> RequestBody {
|
||||
if let Some(json_body) = request_config.json_body.clone() {
|
||||
return RequestBody::from_json(json_body);
|
||||
}
|
||||
if let Some(body_bytes_b64) = request_config.body_bytes_b64.as_ref() {
|
||||
return RequestBody {
|
||||
json_body: None,
|
||||
body_bytes_b64: Some(body_bytes_b64.clone()),
|
||||
body_ref: None,
|
||||
};
|
||||
}
|
||||
if let Some(body) = request_config.body.as_ref() {
|
||||
return RequestBody {
|
||||
json_body: None,
|
||||
body_bytes_b64: Some(base64::engine::general_purpose::STANDARD.encode(body.as_bytes())),
|
||||
body_ref: None,
|
||||
};
|
||||
}
|
||||
RequestBody {
|
||||
json_body: None,
|
||||
body_bytes_b64: None,
|
||||
body_ref: None,
|
||||
}
|
||||
}
|
||||
|
||||
fn build_custom_check_plan(
|
||||
state: &AdminAppState<'_>,
|
||||
transport: &AdminGatewayProviderTransportSnapshot,
|
||||
request_config: &AccountSelfCheckRequestConfig,
|
||||
url: String,
|
||||
auth: Option<(String, String)>,
|
||||
proxy: Option<aether_contracts::ProxySnapshot>,
|
||||
) -> ExecutionPlan {
|
||||
let mut headers = request_config.headers.clone();
|
||||
if let Some((auth_header, auth_value)) = auth {
|
||||
crate::provider_transport::ensure_upstream_auth_header(
|
||||
&mut headers,
|
||||
&auth_header,
|
||||
&auth_value,
|
||||
);
|
||||
}
|
||||
|
||||
let content_type = request_config.content_type.clone().or_else(|| {
|
||||
request_config
|
||||
.json_body
|
||||
.is_some()
|
||||
.then(|| "application/json".to_string())
|
||||
});
|
||||
if let Some(content_type) = content_type.as_ref() {
|
||||
headers
|
||||
.entry("content-type".to_string())
|
||||
.or_insert_with(|| content_type.clone());
|
||||
}
|
||||
|
||||
ExecutionPlan {
|
||||
request_id: format!("account-self-check-{}", Uuid::new_v4()),
|
||||
candidate_id: None,
|
||||
provider_name: Some(transport.provider.name.clone()),
|
||||
provider_id: transport.provider.id.clone(),
|
||||
endpoint_id: transport.endpoint.id.clone(),
|
||||
key_id: transport.key.id.clone(),
|
||||
method: request_config.method.clone(),
|
||||
url,
|
||||
headers,
|
||||
content_type,
|
||||
content_encoding: None,
|
||||
body: custom_check_body(request_config),
|
||||
stream: false,
|
||||
client_api_format: transport.endpoint.api_format.clone(),
|
||||
provider_api_format: transport.endpoint.api_format.clone(),
|
||||
model_name: None,
|
||||
proxy,
|
||||
transport_profile: state.resolve_transport_profile(transport),
|
||||
timeouts: state.resolve_transport_execution_timeouts(transport),
|
||||
}
|
||||
}
|
||||
|
||||
fn execution_error_message(result: &ExecutionResult) -> Option<String> {
|
||||
if let Some(body_json) = result
|
||||
.body
|
||||
.as_ref()
|
||||
.and_then(|body| body.json_body.as_ref())
|
||||
.and_then(Value::as_object)
|
||||
{
|
||||
if let Some(error) = body_json.get("error") {
|
||||
if let Some(message) = error
|
||||
.get("message")
|
||||
.or_else(|| error.get("error_description"))
|
||||
.and_then(Value::as_str)
|
||||
{
|
||||
let trimmed = message.trim();
|
||||
if !trimmed.is_empty() {
|
||||
return Some(trimmed.to_string());
|
||||
}
|
||||
}
|
||||
if let Some(text) = error
|
||||
.as_str()
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
{
|
||||
return Some(text.to_string());
|
||||
}
|
||||
}
|
||||
if let Some(message) = body_json
|
||||
.get("message")
|
||||
.or_else(|| body_json.get("error_description"))
|
||||
.and_then(Value::as_str)
|
||||
{
|
||||
let trimmed = message.trim();
|
||||
if !trimmed.is_empty() {
|
||||
return Some(trimmed.to_string());
|
||||
}
|
||||
}
|
||||
}
|
||||
result
|
||||
.error
|
||||
.as_ref()
|
||||
.map(|error| error.message.trim().to_string())
|
||||
.filter(|value| !value.is_empty())
|
||||
}
|
||||
|
||||
fn custom_result_to_outcome(
|
||||
result: ExecutionResult,
|
||||
request_config: &AccountSelfCheckRequestConfig,
|
||||
) -> AccountSelfCheckOutcome {
|
||||
let message = execution_error_message(&result);
|
||||
if request_config
|
||||
.success_status_codes
|
||||
.contains(&result.status_code)
|
||||
{
|
||||
return AccountSelfCheckOutcome::Success {
|
||||
status_code: Some(result.status_code),
|
||||
message,
|
||||
};
|
||||
}
|
||||
if request_config
|
||||
.blocked_status_codes
|
||||
.contains(&result.status_code)
|
||||
{
|
||||
let detail = message.unwrap_or_else(|| format!("HTTP {}", result.status_code));
|
||||
return AccountSelfCheckOutcome::Blocked {
|
||||
status_code: Some(result.status_code),
|
||||
message: detail,
|
||||
};
|
||||
}
|
||||
let detail = message.unwrap_or_else(|| format!("HTTP {}", result.status_code));
|
||||
AccountSelfCheckOutcome::Failed {
|
||||
status_code: Some(result.status_code),
|
||||
message: detail,
|
||||
}
|
||||
}
|
||||
|
||||
async fn perform_custom_request_check(
|
||||
state: &AdminAppState<'_>,
|
||||
provider: &StoredProviderCatalogProvider,
|
||||
endpoint: &StoredProviderCatalogEndpoint,
|
||||
key: &StoredProviderCatalogKey,
|
||||
request_config: &AccountSelfCheckRequestConfig,
|
||||
) -> Result<AccountSelfCheckOutcome, GatewayError> {
|
||||
let Some(transport) = state
|
||||
.read_provider_transport_snapshot(&provider.id, &endpoint.id, &key.id)
|
||||
.await?
|
||||
else {
|
||||
return Ok(AccountSelfCheckOutcome::Skipped {
|
||||
message: "Provider transport snapshot unavailable".to_string(),
|
||||
});
|
||||
};
|
||||
let Some(url) = custom_check_url(state, &transport, request_config) else {
|
||||
return Ok(AccountSelfCheckOutcome::Skipped {
|
||||
message: "account_self_check_request missing valid url/path".to_string(),
|
||||
});
|
||||
};
|
||||
let auth = resolve_self_check_auth(state, &transport).await?;
|
||||
let proxy = state
|
||||
.resolve_transport_proxy_snapshot_with_tunnel_affinity(&transport)
|
||||
.await;
|
||||
let plan = build_custom_check_plan(state, &transport, request_config, url, auth, proxy);
|
||||
match state.execute_execution_runtime_sync_plan(None, &plan).await {
|
||||
Ok(result) => Ok(custom_result_to_outcome(result, request_config)),
|
||||
Err(err) => Ok(AccountSelfCheckOutcome::Failed {
|
||||
status_code: None,
|
||||
message: gateway_error_message(err),
|
||||
}),
|
||||
}
|
||||
}
|
||||
|
||||
fn quota_payload_result_for_key(key_id: &str, payload: Option<Value>) -> AccountSelfCheckOutcome {
|
||||
let Some(payload) = payload else {
|
||||
return AccountSelfCheckOutcome::Failed {
|
||||
@@ -787,97 +443,6 @@ async fn perform_quota_refresh_check(
|
||||
Ok(quota_payload_result_for_key(&key_id, payload))
|
||||
}
|
||||
|
||||
fn account_block_reason(message: &str) -> String {
|
||||
let detail = message.trim();
|
||||
if detail.starts_with(OAUTH_ACCOUNT_BLOCK_PREFIX) {
|
||||
detail.to_string()
|
||||
} else {
|
||||
format!("{OAUTH_ACCOUNT_BLOCK_PREFIX}{detail}")
|
||||
}
|
||||
}
|
||||
|
||||
fn request_failure_reason(message: &str) -> String {
|
||||
let detail = message.trim();
|
||||
if detail.starts_with(OAUTH_REQUEST_FAILED_PREFIX) {
|
||||
detail.to_string()
|
||||
} else {
|
||||
format!("{OAUTH_REQUEST_FAILED_PREFIX}{detail}")
|
||||
}
|
||||
}
|
||||
|
||||
fn success_invalid_state(key: &StoredProviderCatalogKey) -> (Option<u64>, Option<String>) {
|
||||
let current_reason = key
|
||||
.oauth_invalid_reason
|
||||
.as_deref()
|
||||
.map(str::trim)
|
||||
.unwrap_or_default();
|
||||
if current_reason.starts_with("[REFRESH_FAILED] ") {
|
||||
return (
|
||||
key.oauth_invalid_at_unix_secs,
|
||||
(!current_reason.is_empty()).then_some(current_reason.to_string()),
|
||||
);
|
||||
}
|
||||
(None, None)
|
||||
}
|
||||
|
||||
async fn persist_custom_check_outcome(
|
||||
state: &AdminAppState<'_>,
|
||||
key: &StoredProviderCatalogKey,
|
||||
outcome: &AccountSelfCheckOutcome,
|
||||
now_ts: u64,
|
||||
) -> Result<bool, GatewayError> {
|
||||
match outcome {
|
||||
AccountSelfCheckOutcome::Success { .. } => {
|
||||
let (invalid_at, invalid_reason) = success_invalid_state(key);
|
||||
persist_provider_quota_refresh_state(
|
||||
state,
|
||||
&key.id,
|
||||
None,
|
||||
invalid_at,
|
||||
invalid_reason,
|
||||
None,
|
||||
)
|
||||
.await
|
||||
}
|
||||
AccountSelfCheckOutcome::Blocked { message, .. } => {
|
||||
persist_provider_quota_refresh_state(
|
||||
state,
|
||||
&key.id,
|
||||
None,
|
||||
Some(now_ts),
|
||||
Some(account_block_reason(message)),
|
||||
None,
|
||||
)
|
||||
.await
|
||||
}
|
||||
AccountSelfCheckOutcome::Failed { message, .. } => {
|
||||
persist_provider_quota_refresh_state(
|
||||
state,
|
||||
&key.id,
|
||||
None,
|
||||
Some(now_ts),
|
||||
Some(request_failure_reason(message)),
|
||||
None,
|
||||
)
|
||||
.await
|
||||
}
|
||||
AccountSelfCheckOutcome::Skipped { .. } => Ok(false),
|
||||
}
|
||||
}
|
||||
|
||||
async fn persist_self_check_outcome(
|
||||
state: &AdminAppState<'_>,
|
||||
method: &str,
|
||||
key: &StoredProviderCatalogKey,
|
||||
outcome: &AccountSelfCheckOutcome,
|
||||
now_ts: u64,
|
||||
) -> Result<bool, GatewayError> {
|
||||
if method == "custom_request" {
|
||||
return persist_custom_check_outcome(state, key, outcome, now_ts).await;
|
||||
}
|
||||
Ok(!matches!(outcome, AccountSelfCheckOutcome::Skipped { .. }))
|
||||
}
|
||||
|
||||
async fn record_score_probe_in_progress_for_key(
|
||||
state: &AppState,
|
||||
provider_id: &str,
|
||||
@@ -982,14 +547,7 @@ fn endpoint_for_self_check(
|
||||
provider_type: &str,
|
||||
endpoints: &[StoredProviderCatalogEndpoint],
|
||||
) -> Option<StoredProviderCatalogEndpoint> {
|
||||
provider_quota_refresh_endpoint_for_provider(provider_type, endpoints, true)
|
||||
.or_else(|| {
|
||||
endpoints
|
||||
.iter()
|
||||
.find(|endpoint| endpoint.is_active)
|
||||
.cloned()
|
||||
})
|
||||
.or_else(|| endpoints.first().cloned())
|
||||
provider_account_self_check_endpoint_for_provider(provider_type, endpoints, true)
|
||||
}
|
||||
|
||||
fn gateway_error_message(err: GatewayError) -> String {
|
||||
@@ -1079,9 +637,7 @@ pub(crate) async fn perform_account_self_check_once_with_config(
|
||||
summary.providers_skipped = summary.providers_skipped.saturating_add(1);
|
||||
continue;
|
||||
};
|
||||
if pool_config.account_self_check_method == "quota_refresh"
|
||||
&& !provider_type_supports_quota_refresh(&provider_type)
|
||||
{
|
||||
if !provider_type_supports_account_self_check(&provider_type) {
|
||||
summary.providers_skipped = summary.providers_skipped.saturating_add(1);
|
||||
continue;
|
||||
}
|
||||
@@ -1112,10 +668,6 @@ pub(crate) async fn perform_account_self_check_once_with_config(
|
||||
record_score_probe_in_progress_for_key(state, &provider.id, &key.id, now_ts).await;
|
||||
}
|
||||
|
||||
let method = pool_config.account_self_check_method.clone();
|
||||
let request_config = parse_account_self_check_request_config(
|
||||
pool_config.account_self_check_request.as_ref(),
|
||||
);
|
||||
let provider_short_id = provider.id.chars().take(8).collect::<String>();
|
||||
let concurrency = (pool_config.account_self_check_concurrency as usize)
|
||||
.clamp(1, 64)
|
||||
@@ -1126,29 +678,16 @@ pub(crate) async fn perform_account_self_check_once_with_config(
|
||||
let provider = &provider;
|
||||
let endpoint = &endpoint;
|
||||
let provider_type = provider_type.as_str();
|
||||
let method = method.as_str();
|
||||
let request_config = &request_config;
|
||||
async move {
|
||||
let key_for_check = key.clone();
|
||||
let result = if method == "custom_request" {
|
||||
perform_custom_request_check(
|
||||
admin_state,
|
||||
provider,
|
||||
endpoint,
|
||||
&key_for_check,
|
||||
request_config,
|
||||
)
|
||||
.await
|
||||
} else {
|
||||
perform_quota_refresh_check(
|
||||
admin_state,
|
||||
provider,
|
||||
endpoint,
|
||||
provider_type,
|
||||
key_for_check,
|
||||
)
|
||||
.await
|
||||
};
|
||||
let result = perform_quota_refresh_check(
|
||||
admin_state,
|
||||
provider,
|
||||
endpoint,
|
||||
provider_type,
|
||||
key_for_check,
|
||||
)
|
||||
.await;
|
||||
(key, result)
|
||||
}
|
||||
}))
|
||||
@@ -1164,15 +703,6 @@ pub(crate) async fn perform_account_self_check_once_with_config(
|
||||
message: gateway_error_message(err),
|
||||
},
|
||||
};
|
||||
let persisted =
|
||||
persist_self_check_outcome(&admin_state, &method, &key, &outcome, now_ts).await?;
|
||||
if !persisted && !matches!(outcome, AccountSelfCheckOutcome::Skipped { .. }) {
|
||||
warn!(
|
||||
provider_id = %provider.id,
|
||||
key_id = %key.id,
|
||||
"gateway account self-check: key state was not updated"
|
||||
);
|
||||
}
|
||||
record_score_probe_result_for_key(state, &provider.id, &key.id, now_ts, &outcome).await;
|
||||
update_summary_from_outcome(&mut summary, &outcome);
|
||||
}
|
||||
@@ -1222,12 +752,8 @@ pub(crate) fn spawn_account_self_check_worker(
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::{
|
||||
parse_account_self_check_request_config, select_account_self_check_key_ids,
|
||||
AccountSelfCheckRequestConfig,
|
||||
};
|
||||
use serde_json::json;
|
||||
use std::collections::{BTreeMap, BTreeSet};
|
||||
use super::select_account_self_check_key_ids;
|
||||
use std::collections::BTreeMap;
|
||||
|
||||
#[test]
|
||||
fn selects_never_and_stale_self_check_keys_first() {
|
||||
@@ -1242,37 +768,4 @@ mod tests {
|
||||
|
||||
assert_eq!(selected, vec!["never".to_string(), "stale".to_string()]);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn parses_custom_self_check_request_config() {
|
||||
let parsed = parse_account_self_check_request_config(Some(&json!({
|
||||
"method": "post",
|
||||
"path": "/v1/me?trace=1",
|
||||
"headers": {"X-Test": "yes"},
|
||||
"json_body": {"ping": true},
|
||||
"success_status_codes": [200, 204],
|
||||
"blocked_status_codes": [401, 403, 423, 451]
|
||||
})));
|
||||
|
||||
assert_eq!(parsed.method, "POST");
|
||||
assert_eq!(parsed.path.as_deref(), Some("/v1/me?trace=1"));
|
||||
assert_eq!(
|
||||
parsed.headers.get("x-test").map(String::as_str),
|
||||
Some("yes")
|
||||
);
|
||||
assert_eq!(parsed.json_body, Some(json!({"ping": true})));
|
||||
assert_eq!(parsed.success_status_codes, BTreeSet::from([200, 204]));
|
||||
assert_eq!(
|
||||
parsed.blocked_status_codes,
|
||||
BTreeSet::from([401, 403, 423, 451])
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn defaults_custom_self_check_request_config() {
|
||||
assert_eq!(
|
||||
parse_account_self_check_request_config(None),
|
||||
AccountSelfCheckRequestConfig::default()
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user