feat: add adaptive pool metrics and self-check

This commit is contained in:
fawney19
2026-05-15 01:46:24 +08:00
parent bf511f9f8c
commit 54a8312e46
34 changed files with 4521 additions and 300 deletions

View File

@@ -2,14 +2,15 @@ pub(crate) use crate::handlers::admin::{
admin_provider_ops_local_action_response, admin_provider_pool_config, admin_provider_ops_local_action_response, admin_provider_pool_config,
build_internal_control_error_response, create_provider_oauth_catalog_key, build_internal_control_error_response, create_provider_oauth_catalog_key,
find_duplicate_provider_oauth_key, maybe_build_local_admin_pool_response, find_duplicate_provider_oauth_key, maybe_build_local_admin_pool_response,
maybe_build_local_admin_response, provider_oauth_maintenance_endpoint_for_provider, maybe_build_local_admin_response, persist_provider_quota_refresh_state,
provider_oauth_runtime_endpoint_for_provider, provider_quota_refresh_endpoint_for_provider, provider_oauth_maintenance_endpoint_for_provider, provider_oauth_runtime_endpoint_for_provider,
provider_type_supports_quota_refresh, reconcile_admin_fixed_provider_template_endpoints, provider_quota_refresh_endpoint_for_provider, provider_type_supports_quota_refresh,
reconcile_admin_fixed_provider_template_endpoints,
refresh_provider_oauth_account_state_after_update, refresh_provider_pool_quota_locally, refresh_provider_oauth_account_state_after_update, refresh_provider_pool_quota_locally,
update_existing_provider_oauth_catalog_key, AdminAppState, update_existing_provider_oauth_catalog_key, AdminAppState,
AdminGatewayProviderTransportSnapshot, AdminLocalOAuthRefreshError, AdminRequestContext, AdminGatewayProviderTransportSnapshot, AdminLocalOAuthRefreshError, AdminRequestContext,
AdminRouteRequest, AdminRouteResponse, AdminRouteResult, AdminStatsTimeRange, AdminRouteRequest, AdminRouteResponse, AdminRouteResult, AdminStatsTimeRange,
AdminStatsUsageFilter, AdminStatsUsageFilter, OAUTH_ACCOUNT_BLOCK_PREFIX, OAUTH_REQUEST_FAILED_PREFIX,
}; };
use crate::handlers::admin::{ use crate::handlers::admin::{

View File

@@ -35,9 +35,11 @@ use crate::handlers::shared::provider_pool::{
AdminProviderPoolRuntimeState, AdminProviderPoolRuntimeState,
}; };
use crate::handlers::shared::{parse_catalog_auth_config_json, provider_key_health_summary}; use crate::handlers::shared::{parse_catalog_auth_config_json, provider_key_health_summary};
use crate::maintenance::spawn_pool_quota_probe_replenish_for_request;
use crate::orchestration::LocalExecutionCandidateMetadata; use crate::orchestration::LocalExecutionCandidateMetadata;
static LOAD_BALANCE_SEQUENCE: AtomicU64 = AtomicU64::new(0); static LOAD_BALANCE_SEQUENCE: AtomicU64 = AtomicU64::new(0);
const POOL_ACTIVE_PROBE_SEALED_SKIP_REASON: &str = "pool_active_probe_sealed";
type PoolCatalogKeyContext = PoolMemberSignals; type PoolCatalogKeyContext = PoolMemberSignals;
@@ -108,6 +110,7 @@ async fn schedule_pool_page_candidates(
let key_context_by_id = read_pool_catalog_key_contexts_by_id(state, &candidates).await; let key_context_by_id = read_pool_catalog_key_contexts_by_id(state, &candidates).await;
let mut runtime_by_provider = BTreeMap::new(); let mut runtime_by_provider = BTreeMap::new();
let mut burst_provider_ids = BTreeSet::<String>::new();
for (provider_id, (pool_config, key_ids)) in provider_runtime_requirements { for (provider_id, (pool_config, key_ids)) in provider_runtime_requirements {
let key_ids = key_ids.into_iter().collect::<Vec<_>>(); let key_ids = key_ids.into_iter().collect::<Vec<_>>();
let runtime = if key_ids.is_empty() { let runtime = if key_ids.is_empty() {
@@ -122,14 +125,29 @@ async fn schedule_pool_page_candidates(
) )
.await .await
}; };
if should_trigger_active_probe_burst_for_request(&pool_config, &runtime) {
burst_provider_ids.insert(provider_id.clone());
}
runtime_by_provider.insert(provider_id, runtime); runtime_by_provider.insert(provider_id, runtime);
} }
apply_local_execution_pool_scheduler_with_runtime_map( let (scheduled, skipped) = apply_local_execution_pool_scheduler_with_runtime_map(
candidates, candidates,
&runtime_by_provider, &runtime_by_provider,
&key_context_by_id, &key_context_by_id,
) );
for skipped_candidate in &skipped {
if skipped_candidate.skip_reason == POOL_ACTIVE_PROBE_SEALED_SKIP_REASON {
burst_provider_ids.insert(skipped_candidate.candidate.provider_id.clone());
}
}
for provider_id in burst_provider_ids {
let _ = spawn_pool_quota_probe_replenish_for_request(state.app().clone(), provider_id);
}
(scheduled, skipped)
} }
async fn expand_pool_group_candidate( async fn expand_pool_group_candidate(
@@ -847,48 +865,69 @@ fn apply_local_execution_pool_scheduler_with_runtime_map(
Vec<EligibleLocalExecutionCandidate>, Vec<EligibleLocalExecutionCandidate>,
Vec<SkippedLocalExecutionCandidate>, Vec<SkippedLocalExecutionCandidate>,
) { ) {
let runtime_by_provider = runtime_by_provider let scheduler_runtime_by_provider = runtime_by_provider
.iter() .iter()
.map(|(provider_id, runtime)| (provider_id.clone(), pool_runtime_state(runtime))) .map(|(provider_id, runtime)| (provider_id.clone(), pool_runtime_state(runtime)))
.collect::<BTreeMap<_, _>>(); .collect::<BTreeMap<_, _>>();
let inputs = candidates let mut inputs = Vec::new();
.into_iter() let mut skipped_candidates = Vec::new();
.map(|candidate| { for candidate in candidates {
let key_context = key_context_by_id let key_context = key_context_by_id
.get(&candidate.candidate.key_id) .get(&candidate.candidate.key_id)
.cloned() .cloned()
.unwrap_or_default(); .unwrap_or_default();
PoolCandidateInput { let admin_pool_config = pool_config_for_candidate(&candidate);
facts: pool_candidate_facts(&candidate),
pool_config: pool_config_for_candidate(&candidate).map(|config| { if let Some(config) = admin_pool_config.as_ref() {
pool_scheduling_config( if should_enforce_active_probe_sealed_pool(config) {
config, let active_member_ids = runtime_by_provider
candidate.transport.provider.provider_type.as_str(), .get(&candidate.candidate.provider_id)
) .map(|runtime| &runtime.active_probe_member_ids);
}), if !active_member_ids
key_context, .is_some_and(|members| members.contains(&candidate.candidate.key_id))
candidate, {
skipped_candidates.push(SkippedLocalExecutionCandidate {
candidate: candidate.candidate.clone(),
skip_reason: POOL_ACTIVE_PROBE_SEALED_SKIP_REASON,
transport: Some(candidate.transport.clone()),
ranking: candidate.ranking.clone(),
extra_data: None,
});
continue;
}
} }
}) }
.collect::<Vec<_>>();
let outcome = run_pool_scheduler(inputs, &runtime_by_provider, pool_sort_seed().as_str()); let pool_config = admin_pool_config.map(|config| {
pool_scheduling_config(config, candidate.transport.provider.provider_type.as_str())
});
inputs.push(PoolCandidateInput {
facts: pool_candidate_facts(&candidate),
pool_config,
key_context,
candidate,
});
}
let outcome = run_pool_scheduler(
inputs,
&scheduler_runtime_by_provider,
pool_sort_seed().as_str(),
);
let candidates = outcome let candidates = outcome
.candidates .candidates
.into_iter() .into_iter()
.map(|scheduled| apply_pool_orchestration(scheduled.candidate, scheduled.orchestration)) .map(|scheduled| apply_pool_orchestration(scheduled.candidate, scheduled.orchestration))
.collect::<Vec<_>>(); .collect::<Vec<_>>();
let skipped_candidates = outcome skipped_candidates.extend(outcome.skipped_candidates.into_iter().map(|skipped| {
.skipped_candidates SkippedLocalExecutionCandidate {
.into_iter()
.map(|skipped| SkippedLocalExecutionCandidate {
candidate: skipped.candidate.candidate, candidate: skipped.candidate.candidate,
skip_reason: skipped.skip_reason, skip_reason: skipped.skip_reason,
transport: Some(skipped.candidate.transport), transport: Some(skipped.candidate.transport),
ranking: skipped.candidate.ranking, ranking: skipped.candidate.ranking,
extra_data: None, extra_data: None,
}) }
.collect::<Vec<_>>(); }));
(candidates, skipped_candidates) (candidates, skipped_candidates)
} }
@@ -899,6 +938,24 @@ fn pool_config_for_candidate(
admin_provider_pool_config_from_config_value(candidate.transport.provider.config.as_ref()) admin_provider_pool_config_from_config_value(candidate.transport.provider.config.as_ref())
} }
fn should_enforce_active_probe_sealed_pool(pool_config: &AdminProviderPoolConfig) -> bool {
pool_config.probing_enabled
}
fn should_trigger_active_probe_burst_for_request(
pool_config: &AdminProviderPoolConfig,
runtime: &AdminProviderPoolRuntimeState,
) -> bool {
if !should_enforce_active_probe_sealed_pool(pool_config) {
return false;
}
if runtime.provider_burst_pending {
return false;
}
let active_count = runtime.active_probe_member_ids.len();
runtime.provider_desired_hot > 0 && active_count < runtime.provider_desired_hot
}
fn pool_key_candidate_order_for_group( fn pool_key_candidate_order_for_group(
group: &EligibleLocalExecutionCandidate, group: &EligibleLocalExecutionCandidate,
) -> StoredPoolKeyCandidateOrder { ) -> StoredPoolKeyCandidateOrder {
@@ -1012,7 +1069,8 @@ fn apply_pool_orchestration(
mod tests { mod tests {
use super::{ use super::{
apply_local_execution_pool_scheduler_with_runtime_map, build_pool_catalog_key_context, apply_local_execution_pool_scheduler_with_runtime_map, build_pool_catalog_key_context,
pool_config_for_candidate, PoolCatalogKeyContext, PoolKeyCursor, pool_config_for_candidate, should_trigger_active_probe_burst_for_request,
PoolCatalogKeyContext, PoolKeyCursor, POOL_ACTIVE_PROBE_SEALED_SKIP_REASON,
}; };
use crate::ai_serving::{ use crate::ai_serving::{
apply_local_runtime_candidate_terminal_reason, EligibleLocalExecutionCandidate, apply_local_runtime_candidate_terminal_reason, EligibleLocalExecutionCandidate,
@@ -1040,7 +1098,7 @@ mod tests {
}; };
use aether_scheduler_core::SchedulerMinimalCandidateSelectionCandidate; use aether_scheduler_core::SchedulerMinimalCandidateSelectionCandidate;
use serde_json::json; use serde_json::json;
use std::collections::{BTreeMap, VecDeque}; use std::collections::{BTreeMap, BTreeSet, VecDeque};
use std::sync::Arc; use std::sync::Arc;
#[test] #[test]
@@ -1387,6 +1445,168 @@ mod tests {
); );
} }
#[test]
fn pool_scheduler_uses_only_active_probe_members_when_active_probe_enabled() {
let provider_config = Some(json!({
"pool_advanced": {
"probing_enabled": true
}
}));
let key_active = sample_eligible_candidate(
"provider-pool",
"endpoint-1",
"key-active",
10,
provider_config.clone(),
);
let key_sealed = sample_eligible_candidate(
"provider-pool",
"endpoint-1",
"key-sealed",
10,
provider_config,
);
let runtime_by_provider = BTreeMap::from([(
"provider-pool".to_string(),
AdminProviderPoolRuntimeState {
active_probe_member_ids: BTreeSet::from(["key-active".to_string()]),
..AdminProviderPoolRuntimeState::default()
},
)]);
let (scheduled, skipped) = apply_local_execution_pool_scheduler_with_runtime_map(
vec![key_sealed, key_active],
&runtime_by_provider,
&BTreeMap::new(),
);
assert_eq!(
scheduled
.iter()
.map(|item| item.candidate.key_id.as_str())
.collect::<Vec<_>>(),
vec!["key-active"]
);
assert_eq!(
skipped
.iter()
.map(|item| (item.candidate.key_id.as_str(), item.skip_reason))
.collect::<Vec<_>>(),
vec![("key-sealed", POOL_ACTIVE_PROBE_SEALED_SKIP_REASON)]
);
}
#[test]
fn pool_scheduler_keeps_pool_out_keys_sealed_when_active_probe_pool_is_empty() {
let provider_config = Some(json!({
"pool_advanced": {
"probing_enabled": true
}
}));
let key_a = sample_eligible_candidate(
"provider-pool",
"endpoint-1",
"key-a",
10,
provider_config.clone(),
);
let key_b =
sample_eligible_candidate("provider-pool", "endpoint-1", "key-b", 10, provider_config);
let runtime_by_provider = BTreeMap::from([(
"provider-pool".to_string(),
AdminProviderPoolRuntimeState::default(),
)]);
let (scheduled, skipped) = apply_local_execution_pool_scheduler_with_runtime_map(
vec![key_a, key_b],
&runtime_by_provider,
&BTreeMap::new(),
);
assert!(scheduled.is_empty());
assert_eq!(
skipped
.iter()
.map(|item| (item.candidate.key_id.as_str(), item.skip_reason))
.collect::<Vec<_>>(),
vec![
("key-a", POOL_ACTIVE_PROBE_SEALED_SKIP_REASON),
("key-b", POOL_ACTIVE_PROBE_SEALED_SKIP_REASON),
]
);
}
#[test]
fn pool_scheduler_triggers_burst_when_auto_hot_target_has_gap() {
let provider_config = Some(json!({
"pool_advanced": {
"probing_enabled": true
}
}));
let candidate =
sample_eligible_candidate("provider-pool", "endpoint-1", "key-a", 10, provider_config);
let pool_config = pool_config_for_candidate(&candidate).expect("pool config should parse");
let runtime = AdminProviderPoolRuntimeState {
active_probe_member_ids: BTreeSet::from(["key-a".to_string()]),
provider_desired_hot: 3,
..AdminProviderPoolRuntimeState::default()
};
assert!(should_trigger_active_probe_burst_for_request(
&pool_config,
&runtime
));
}
#[test]
fn pool_scheduler_ignores_legacy_threshold_fields_for_burst_target() {
let provider_config = Some(json!({
"pool_advanced": {
"probing_enabled": true,
"probing_target_percent": 60,
"probing_target_count": 10
}
}));
let candidate =
sample_eligible_candidate("provider-pool", "endpoint-1", "key-a", 10, provider_config);
let pool_config = pool_config_for_candidate(&candidate).expect("pool config should parse");
let runtime = AdminProviderPoolRuntimeState {
active_probe_member_ids: BTreeSet::from(["key-a".to_string(), "key-b".to_string()]),
provider_desired_hot: 2,
..AdminProviderPoolRuntimeState::default()
};
assert!(!should_trigger_active_probe_burst_for_request(
&pool_config,
&runtime
));
}
#[test]
fn pool_scheduler_skips_burst_when_active_probe_target_is_met() {
let provider_config = Some(json!({
"pool_advanced": {
"probing_enabled": true,
"probing_target_count": 1
}
}));
let candidate =
sample_eligible_candidate("provider-pool", "endpoint-1", "key-a", 10, provider_config);
let pool_config = pool_config_for_candidate(&candidate).expect("pool config should parse");
let runtime = AdminProviderPoolRuntimeState {
active_probe_member_ids: BTreeSet::from(["key-a".to_string()]),
provider_desired_hot: 1,
..AdminProviderPoolRuntimeState::default()
};
assert!(!should_trigger_active_probe_burst_for_request(
&pool_config,
&runtime
));
}
#[test] #[test]
fn pool_scheduler_applies_distribution_mode_before_strategy_presets() { fn pool_scheduler_applies_distribution_mode_before_strategy_presets() {
let key_a = sample_eligible_candidate( let key_a = sample_eligible_candidate(

View File

@@ -89,6 +89,9 @@ use crate::orchestration::{
LocalExecutionEffect, LocalExecutionEffectContext, LocalHealthFailureEffect, LocalExecutionEffect, LocalExecutionEffectContext, LocalHealthFailureEffect,
LocalHealthSuccessEffect, LocalOAuthInvalidationEffect, LocalPoolErrorEffect, LocalHealthSuccessEffect, LocalOAuthInvalidationEffect, LocalPoolErrorEffect,
}; };
use crate::provider_pool_demand::{
acquire_provider_pool_in_flight_guard, ProviderPoolInFlightGuard,
};
use crate::request_candidate_runtime::{ use crate::request_candidate_runtime::{
ensure_execution_request_candidate_slot, record_local_request_candidate_status, ensure_execution_request_candidate_slot, record_local_request_candidate_status,
record_local_request_candidate_status_snapshot, snapshot_local_request_candidate_status, record_local_request_candidate_status_snapshot, snapshot_local_request_candidate_status,
@@ -514,6 +517,14 @@ pub(crate) async fn execute_execution_runtime_stream(
.and_then(|context| context.candidate_index) .and_then(|context| context.candidate_index)
.map(|value| value.to_string()) .map(|value| value.to_string())
.unwrap_or_else(|| "-".to_string()); .unwrap_or_else(|| "-".to_string());
let mut provider_pool_in_flight_guard = acquire_provider_pool_in_flight_guard(
state.runtime_state.clone(),
&plan.provider_id,
plan.request_id.as_str(),
plan.candidate_id.as_deref(),
key_id.as_str(),
)
.await;
match maybe_execute_kiro_web_search_stream(state, &plan, report_context.as_ref()).await { match maybe_execute_kiro_web_search_stream(state, &plan, report_context.as_ref()).await {
Ok(Some(kiro_web_search)) => { Ok(Some(kiro_web_search)) => {
return execute_stream_from_frame_stream( return execute_stream_from_frame_stream(
@@ -527,6 +538,7 @@ pub(crate) async fn execute_execution_runtime_stream(
candidate_started_unix_secs, candidate_started_unix_secs,
stream_started_at, stream_started_at,
kiro_web_search.frame_stream, kiro_web_search.frame_stream,
provider_pool_in_flight_guard.take(),
) )
.await; .await;
} }
@@ -578,6 +590,7 @@ pub(crate) async fn execute_execution_runtime_stream(
candidate_started_unix_secs, candidate_started_unix_secs,
stream_started_at, stream_started_at,
chatgpt_web_image.frame_stream, chatgpt_web_image.frame_stream,
provider_pool_in_flight_guard.take(),
) )
.await; .await;
} }
@@ -673,6 +686,7 @@ pub(crate) async fn execute_execution_runtime_stream(
candidate_started_unix_secs, candidate_started_unix_secs,
stream_started_at, stream_started_at,
frame_stream, frame_stream,
provider_pool_in_flight_guard.take(),
) )
.await; .await;
} }
@@ -737,6 +751,7 @@ pub(crate) async fn execute_execution_runtime_stream(
candidate_started_unix_secs, candidate_started_unix_secs,
stream_started_at, stream_started_at,
frame_stream, frame_stream,
provider_pool_in_flight_guard.take(),
) )
.await; .await;
} }
@@ -822,6 +837,7 @@ pub(crate) async fn execute_execution_runtime_stream(
candidate_started_unix_secs, candidate_started_unix_secs,
stream_started_at, stream_started_at,
frame_stream, frame_stream,
provider_pool_in_flight_guard.take(),
) )
.await; .await;
} }
@@ -1089,6 +1105,7 @@ async fn execute_stream_from_frame_stream(
candidate_started_unix_secs: u64, candidate_started_unix_secs: u64,
stream_started_at: Instant, stream_started_at: Instant,
frame_stream: BoxStream<'static, Result<Bytes, IoError>>, frame_stream: BoxStream<'static, Result<Bytes, IoError>>,
in_flight_guard: Option<ProviderPoolInFlightGuard>,
) -> Result<Option<Response<Body>>, GatewayError> { ) -> Result<Option<Response<Body>>, GatewayError> {
let request_id = plan.request_id.as_str(); let request_id = plan.request_id.as_str();
let request_id_for_log = short_request_id(request_id); let request_id_for_log = short_request_id(request_id);
@@ -1923,7 +1940,9 @@ async fn execute_stream_from_frame_stream(
}; };
let plan_kind_for_report = plan_kind.to_string(); let plan_kind_for_report = plan_kind.to_string();
let stream_started_at_for_report = stream_started_at; let stream_started_at_for_report = stream_started_at;
let provider_pool_in_flight_guard_for_report = in_flight_guard;
tokio::spawn(async move { tokio::spawn(async move {
let _provider_pool_in_flight_guard = provider_pool_in_flight_guard_for_report;
let mut provider_buffered_body = Vec::new(); let mut provider_buffered_body = Vec::new();
let mut buffered_body = Vec::new(); let mut buffered_body = Vec::new();
let mut provider_body_truncated = false; let mut provider_body_truncated = false;
@@ -3136,6 +3155,7 @@ mod tests {
crate::clock::current_unix_ms(), crate::clock::current_unix_ms(),
Instant::now(), Instant::now(),
frame_stream, frame_stream,
None,
) )
.await .await
.expect("execution should succeed") .expect("execution should succeed")

View File

@@ -59,6 +59,7 @@ use crate::orchestration::{
LocalExecutionEffect, LocalExecutionEffectContext, LocalHealthFailureEffect, LocalExecutionEffect, LocalExecutionEffectContext, LocalHealthFailureEffect,
LocalHealthSuccessEffect, LocalOAuthInvalidationEffect, LocalPoolErrorEffect, LocalHealthSuccessEffect, LocalOAuthInvalidationEffect, LocalPoolErrorEffect,
}; };
use crate::provider_pool_demand::acquire_provider_pool_in_flight_guard;
use crate::request_candidate_runtime::{ use crate::request_candidate_runtime::{
ensure_execution_request_candidate_slot, record_local_request_candidate_extra_data, ensure_execution_request_candidate_slot, record_local_request_candidate_extra_data,
record_local_request_candidate_status, record_local_request_candidate_status,
@@ -1066,6 +1067,14 @@ async fn execute_execution_runtime_sync_impl(
}, },
) )
.await; .await;
let _provider_pool_in_flight_guard = acquire_provider_pool_in_flight_guard(
state.runtime_state.clone(),
&plan.provider_id,
plan_request_id.as_str(),
plan_candidate_id.as_deref(),
key_id.as_str(),
)
.await;
#[cfg(not(test))] #[cfg(not(test))]
let mut result = { let mut result = {
match maybe_execute_chatgpt_web_image_sync(state, &plan, report_context.as_ref()).await { match maybe_execute_chatgpt_web_image_sync(state, &plan, report_context.as_ref()).await {

View File

@@ -26,8 +26,10 @@ pub(crate) use self::provider::oauth::provisioning::{
create_provider_oauth_catalog_key, update_existing_provider_oauth_catalog_key, create_provider_oauth_catalog_key, update_existing_provider_oauth_catalog_key,
}; };
pub(crate) use self::provider::oauth::quota::dispatch::refresh_provider_pool_quota_locally; pub(crate) use self::provider::oauth::quota::dispatch::refresh_provider_pool_quota_locally;
pub(crate) use self::provider::oauth::quota::shared::provider_quota_refresh_endpoint_for_provider; pub(crate) use self::provider::oauth::quota::shared::{
pub(crate) use self::provider::oauth::quota::shared::provider_type_supports_quota_refresh; persist_provider_quota_refresh_state, provider_quota_refresh_endpoint_for_provider,
provider_type_supports_quota_refresh,
};
pub(crate) use self::provider::oauth::runtime::{ pub(crate) use self::provider::oauth::runtime::{
provider_oauth_maintenance_endpoint_for_provider, provider_oauth_runtime_endpoint_for_provider, provider_oauth_maintenance_endpoint_for_provider, provider_oauth_runtime_endpoint_for_provider,
refresh_provider_oauth_account_state_after_update, refresh_provider_oauth_account_state_after_update,
@@ -35,6 +37,9 @@ pub(crate) use self::provider::oauth::runtime::{
pub(crate) use self::provider::ops::providers::actions::admin_provider_ops_local_action_response; pub(crate) use self::provider::ops::providers::actions::admin_provider_ops_local_action_response;
pub(crate) use self::provider::pool::config::admin_provider_pool_config; pub(crate) use self::provider::pool::config::admin_provider_pool_config;
pub(crate) use self::provider::pool_admin::maybe_build_local_admin_pool_response; pub(crate) use self::provider::pool_admin::maybe_build_local_admin_pool_response;
pub(crate) use self::provider::shared::payloads::{
OAUTH_ACCOUNT_BLOCK_PREFIX, OAUTH_REQUEST_FAILED_PREFIX,
};
pub(crate) use self::provider::write::provider::reconcile_admin_fixed_provider_template_endpoints; pub(crate) use self::provider::write::provider::reconcile_admin_fixed_provider_template_endpoints;
pub(crate) use self::provider::{ pub(crate) use self::provider::{
maybe_build_local_admin_provider_oauth_response, maybe_build_local_admin_providers_response, maybe_build_local_admin_provider_oauth_response, maybe_build_local_admin_providers_response,

View File

@@ -37,6 +37,38 @@ fn json_f64(value: &Value) -> Option<f64> {
}) })
} }
fn parse_pool_probe_target_percent(pool_advanced: &Map<String, Value>) -> Option<f64> {
pool_advanced
.get("probing_target_percent")
.or_else(|| pool_advanced.get("probing_active_target_percent"))
.or_else(|| pool_advanced.get("active_probe_target_percent"))
.and_then(json_f64)
.filter(|value| value.is_finite() && *value > 0.0)
.map(|value| value.clamp(0.0, 100.0))
}
fn parse_pool_probe_target_count(pool_advanced: &Map<String, Value>) -> Option<u64> {
pool_advanced
.get("probing_target_count")
.or_else(|| pool_advanced.get("probing_active_target_count"))
.or_else(|| pool_advanced.get("active_probe_target_count"))
.and_then(json_u64)
.filter(|value| *value > 0)
.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 { fn pool_score_weight(object: &Map<String, Value>, names: &[&str], current: f64) -> f64 {
names names
.iter() .iter()
@@ -388,7 +420,14 @@ pub(crate) fn admin_provider_pool_config_from_config_value(
health_policy_enabled: true, health_policy_enabled: true,
probing_enabled: false, probing_enabled: false,
probing_interval_minutes: 10, probing_interval_minutes: 10,
probing_target_percent: None,
probing_target_count: None,
probe_concurrency: 4, probe_concurrency: 4,
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_top_n: 128,
score_fallback_scan_limit: 1024, score_fallback_scan_limit: 1024,
score_rules: PoolMemberScoreRules::default(), score_rules: PoolMemberScoreRules::default(),
@@ -456,12 +495,39 @@ pub(crate) fn admin_provider_pool_config_from_config_value(
.filter(|value| *value > 0) .filter(|value| *value > 0)
.map(|value| value.min(1440)) .map(|value| value.min(1440))
.unwrap_or(10), .unwrap_or(10),
probing_target_percent: parse_pool_probe_target_percent(pool_advanced),
probing_target_count: parse_pool_probe_target_count(pool_advanced),
probe_concurrency: pool_advanced probe_concurrency: pool_advanced
.get("probe_concurrency") .get("probe_concurrency")
.and_then(json_u64) .and_then(json_u64)
.filter(|value| *value > 0) .filter(|value| *value > 0)
.map(|value| value.min(64)) .map(|value| value.min(64))
.unwrap_or(4), .unwrap_or(4),
account_self_check_enabled: pool_advanced
.get("account_self_check_enabled")
.or_else(|| pool_advanced.get("self_check_enabled"))
.and_then(Value::as_bool)
.unwrap_or(false),
account_self_check_interval_minutes: pool_advanced
.get("account_self_check_interval_minutes")
.or_else(|| pool_advanced.get("self_check_interval_minutes"))
.and_then(json_u64)
.filter(|value| *value > 0)
.map(|value| value.min(1440))
.unwrap_or(60),
account_self_check_concurrency: pool_advanced
.get("account_self_check_concurrency")
.or_else(|| pool_advanced.get("self_check_concurrency"))
.and_then(json_u64)
.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 score_top_n: pool_advanced
.get("score_top_n") .get("score_top_n")
.and_then(json_u64) .and_then(json_u64)
@@ -547,7 +613,18 @@ mod tests {
"health_policy_enabled": false, "health_policy_enabled": false,
"probing_enabled": true, "probing_enabled": true,
"probing_interval_minutes": 20, "probing_interval_minutes": 20,
"probing_target_percent": 25,
"probing_target_count": 3,
"probe_concurrency": 6, "probe_concurrency": 6,
"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_top_n": 256,
"score_fallback_scan_limit": 2048, "score_fallback_scan_limit": 2048,
"score_rules": { "score_rules": {
@@ -584,7 +661,21 @@ mod tests {
assert!(!config.health_policy_enabled); assert!(!config.health_policy_enabled);
assert!(config.probing_enabled); assert!(config.probing_enabled);
assert_eq!(config.probing_interval_minutes, 20); assert_eq!(config.probing_interval_minutes, 20);
assert_eq!(config.probing_target_percent, Some(25.0));
assert_eq!(config.probing_target_count, Some(3));
assert_eq!(config.probe_concurrency, 6); assert_eq!(config.probe_concurrency, 6);
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_top_n, 256);
assert_eq!(config.score_fallback_scan_limit, 2048); assert_eq!(config.score_fallback_scan_limit, 2048);
assert_eq!(config.score_rules.weights.manual_priority, 0.4); assert_eq!(config.score_rules.weights.manual_priority, 0.4);

View File

@@ -5,10 +5,15 @@ use super::keys::{
}; };
use crate::handlers::admin::provider::pool::config::admin_provider_pool_cache_affinity_enabled; use crate::handlers::admin::provider::pool::config::admin_provider_pool_cache_affinity_enabled;
use crate::handlers::admin::provider::shared::support::{ use crate::handlers::admin::provider::shared::support::{
AdminProviderPoolConfig, AdminProviderPoolRuntimeState, admin_provider_pool_quota_probe_active_members_key, AdminProviderPoolConfig,
AdminProviderPoolRuntimeState,
};
use crate::maintenance::PoolQuotaProbeWorkerConfig;
use crate::provider_pool_demand::{
provider_pool_burst_pending, read_provider_pool_demand_snapshot,
}; };
use aether_runtime_state::{DataLayerError, RuntimeState}; use aether_runtime_state::{DataLayerError, RuntimeState};
use std::collections::BTreeMap; use std::collections::{BTreeMap, BTreeSet};
use std::time::{SystemTime, UNIX_EPOCH}; use std::time::{SystemTime, UNIX_EPOCH};
use tracing::warn; use tracing::warn;
@@ -19,6 +24,10 @@ fn current_unix_secs() -> u64 {
.as_secs() .as_secs()
} }
fn should_load_active_probe_members(pool_config: &AdminProviderPoolConfig) -> bool {
pool_config.probing_enabled
}
pub(crate) async fn read_admin_provider_pool_cooldown_counts( pub(crate) async fn read_admin_provider_pool_cooldown_counts(
runtime: &RuntimeState, runtime: &RuntimeState,
provider_ids: &[String], provider_ids: &[String],
@@ -102,6 +111,40 @@ pub(crate) async fn read_admin_provider_pool_runtime_state(
} }
} }
if should_load_active_probe_members(pool_config) {
state.active_probe_member_ids = runtime
.set_members(&admin_provider_pool_quota_probe_active_members_key(
provider_id,
))
.await
.map(|values| {
values
.into_iter()
.map(|value| value.trim().to_string())
.filter(|value| !value.is_empty())
.collect::<BTreeSet<_>>()
})
.unwrap_or_default();
}
let probe_config = PoolQuotaProbeWorkerConfig::from_env();
let demand_snapshot = read_provider_pool_demand_snapshot(
runtime,
provider_id,
key_ids.len(),
probe_config.max_keys_per_provider,
)
.await;
state.provider_in_flight = demand_snapshot.in_flight;
state.provider_ema_in_flight = demand_snapshot.ema_in_flight;
state.provider_desired_hot = if pool_config.probing_enabled {
demand_snapshot.desired_hot
} else {
0
};
state.provider_burst_pending =
pool_config.probing_enabled && provider_pool_burst_pending(runtime, provider_id).await;
if !cooldown_keys.is_empty() { if !cooldown_keys.is_empty() {
let cooldown_reasons = runtime let cooldown_reasons = runtime
.kv_get_many(&cooldown_keys) .kv_get_many(&cooldown_keys)

View File

@@ -23,6 +23,11 @@ pub(crate) async fn build_admin_provider_pool_status_payload(
"pool_enabled": false, "pool_enabled": false,
"total_keys": 0, "total_keys": 0,
"total_sticky_sessions": 0, "total_sticky_sessions": 0,
"provider_hot_count": 0,
"provider_desired_hot": 0,
"provider_in_flight": 0,
"provider_ema_in_flight": 0.0,
"provider_burst_pending": false,
"keys": [], "keys": [],
})); }));
}; };
@@ -68,6 +73,11 @@ pub(crate) async fn build_admin_provider_pool_status_payload(
"pool_enabled": true, "pool_enabled": true,
"total_keys": key_payloads.len(), "total_keys": key_payloads.len(),
"total_sticky_sessions": runtime.total_sticky_sessions, "total_sticky_sessions": runtime.total_sticky_sessions,
"provider_hot_count": runtime.active_probe_member_ids.len(),
"provider_desired_hot": runtime.provider_desired_hot,
"provider_in_flight": runtime.provider_in_flight,
"provider_ema_in_flight": runtime.provider_ema_in_flight,
"provider_burst_pending": runtime.provider_burst_pending,
"keys": key_payloads, "keys": key_payloads,
})) }))
} }

View File

@@ -615,7 +615,14 @@ mod tests {
health_policy_enabled: true, health_policy_enabled: true,
probing_enabled: false, probing_enabled: false,
probing_interval_minutes: 10, probing_interval_minutes: 10,
probing_target_percent: None,
probing_target_count: None,
probe_concurrency: 4, probe_concurrency: 4,
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_top_n: 128,
score_fallback_scan_limit: 1024, score_fallback_scan_limit: 1024,
score_rules: aether_pool_core::PoolMemberScoreRules::default(), score_rules: aether_pool_core::PoolMemberScoreRules::default(),

View File

@@ -5,7 +5,12 @@ use super::{
read_admin_provider_pool_cooldown_counts, read_admin_provider_pool_cooldown_counts,
ADMIN_POOL_PROVIDER_CATALOG_READER_UNAVAILABLE_DETAIL, ADMIN_POOL_PROVIDER_CATALOG_READER_UNAVAILABLE_DETAIL,
}; };
use crate::handlers::admin::provider::shared::support::admin_provider_pool_quota_probe_active_members_key;
use crate::handlers::admin::request::AdminAppState; use crate::handlers::admin::request::AdminAppState;
use crate::maintenance::PoolQuotaProbeWorkerConfig;
use crate::provider_pool_demand::{
provider_pool_burst_pending, read_provider_pool_demand_snapshot,
};
use crate::GatewayError; use crate::GatewayError;
use aether_admin::provider::pool as admin_provider_pool_pure; use aether_admin::provider::pool as admin_provider_pool_pure;
use axum::{ use axum::{
@@ -14,6 +19,7 @@ use axum::{
response::{IntoResponse, Response}, response::{IntoResponse, Response},
Json, Json,
}; };
use serde_json::{json, Value};
pub(super) async fn build_admin_pool_overview_response( pub(super) async fn build_admin_pool_overview_response(
state: &AdminAppState<'_>, state: &AdminAppState<'_>,
@@ -60,17 +66,77 @@ pub(super) async fn build_admin_pool_overview_response(
.map(|item| (item.provider_id.clone(), item)) .map(|item| (item.provider_id.clone(), item))
.collect::<BTreeMap<_, _>>(); .collect::<BTreeMap<_, _>>();
let probe_config = PoolQuotaProbeWorkerConfig::from_env();
let mut runtime_metrics_by_provider = BTreeMap::new();
for (provider, pool_config) in &pool_enabled_providers {
let active_keys = key_stats_by_provider
.get(&provider.id)
.map(|item| item.active_keys as usize)
.unwrap_or(0);
let hot_count = if pool_config.probing_enabled {
state
.runtime_state()
.set_len(&admin_provider_pool_quota_probe_active_members_key(
&provider.id,
))
.await
.unwrap_or(0)
} else {
0
};
let demand_snapshot = read_provider_pool_demand_snapshot(
state.runtime_state(),
&provider.id,
active_keys,
probe_config.max_keys_per_provider,
)
.await;
let burst_pending = pool_config.probing_enabled
&& provider_pool_burst_pending(state.runtime_state(), &provider.id).await;
runtime_metrics_by_provider.insert(
provider.id.clone(),
json!({
"provider_hot_count": hot_count,
"provider_desired_hot": if pool_config.probing_enabled {
demand_snapshot.desired_hot
} else {
0
},
"provider_in_flight": demand_snapshot.in_flight,
"provider_ema_in_flight": demand_snapshot.ema_in_flight,
"provider_burst_pending": burst_pending,
}),
);
}
let providers = pool_enabled_providers let providers = pool_enabled_providers
.into_iter() .into_iter()
.map(|(provider, _)| provider) .map(|(provider, _)| provider)
.collect::<Vec<_>>(); .collect::<Vec<_>>();
Ok( let mut payload = admin_provider_pool_pure::build_admin_pool_overview_payload(
Json(admin_provider_pool_pure::build_admin_pool_overview_payload( &providers,
&providers, &key_stats_by_provider,
&key_stats_by_provider, &cooldown_counts_by_provider,
&cooldown_counts_by_provider, );
)) if let Some(items) = payload.get_mut("items").and_then(Value::as_array_mut) {
.into_response(), for item in items {
) let Some(provider_id) = item.get("provider_id").and_then(Value::as_str) else {
continue;
};
let Some(metrics) = runtime_metrics_by_provider.get(provider_id) else {
continue;
};
let Some(item_object) = item.as_object_mut() else {
continue;
};
if let Some(metrics_object) = metrics.as_object() {
for (key, value) in metrics_object {
item_object.insert(key.clone(), value.clone());
}
}
}
}
Ok(Json(payload).into_response())
} }

View File

@@ -2,15 +2,21 @@ use crate::handlers::admin::request::AdminAppState;
use crate::LocalProviderDeleteTaskState; use crate::LocalProviderDeleteTaskState;
use aether_pool_core::PoolMemberScoreRules; use aether_pool_core::PoolMemberScoreRules;
use serde_json::json; use serde_json::json;
use std::collections::BTreeMap; use std::collections::{BTreeMap, BTreeSet};
pub(crate) const ADMIN_PROVIDER_MAPPING_PREVIEW_MAX_KEYS: usize = 200; pub(crate) const ADMIN_PROVIDER_MAPPING_PREVIEW_MAX_KEYS: usize = 200;
pub(crate) const ADMIN_PROVIDER_MAPPING_PREVIEW_MAX_MODELS: usize = 500; pub(crate) const ADMIN_PROVIDER_MAPPING_PREVIEW_MAX_MODELS: usize = 500;
pub(crate) const ADMIN_PROVIDER_MAPPING_PREVIEW_FETCH_LIMIT: usize = 10_000; pub(crate) const ADMIN_PROVIDER_MAPPING_PREVIEW_FETCH_LIMIT: usize = 10_000;
pub(crate) const ADMIN_PROVIDER_POOL_SCAN_BATCH: u64 = 200; pub(crate) const ADMIN_PROVIDER_POOL_SCAN_BATCH: u64 = 200;
pub(crate) const ADMIN_PROVIDER_POOL_QUOTA_PROBE_ACTIVE_SET_PREFIX: &str =
"ap:quota_probe:active_members";
pub(crate) const ADMIN_PROVIDER_OAUTH_DATA_UNAVAILABLE_DETAIL: &str = pub(crate) const ADMIN_PROVIDER_OAUTH_DATA_UNAVAILABLE_DETAIL: &str =
"Admin provider OAuth data unavailable"; "Admin provider OAuth data unavailable";
pub(crate) fn admin_provider_pool_quota_probe_active_members_key(provider_id: &str) -> String {
format!("{ADMIN_PROVIDER_POOL_QUOTA_PROBE_ACTIVE_SET_PREFIX}:{provider_id}")
}
#[derive(Debug, Clone, PartialEq, Eq)] #[derive(Debug, Clone, PartialEq, Eq)]
pub(crate) struct AdminProviderPoolSchedulingPreset { pub(crate) struct AdminProviderPoolSchedulingPreset {
pub(crate) preset: String, pub(crate) preset: String,
@@ -40,7 +46,14 @@ pub(crate) struct AdminProviderPoolConfig {
pub(crate) health_policy_enabled: bool, pub(crate) health_policy_enabled: bool,
pub(crate) probing_enabled: bool, pub(crate) probing_enabled: bool,
pub(crate) probing_interval_minutes: u64, pub(crate) probing_interval_minutes: u64,
pub(crate) probing_target_percent: Option<f64>,
pub(crate) probing_target_count: Option<u64>,
pub(crate) probe_concurrency: u64, pub(crate) probe_concurrency: u64,
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_top_n: u64,
pub(crate) score_fallback_scan_limit: u64, pub(crate) score_fallback_scan_limit: u64,
pub(crate) score_rules: PoolMemberScoreRules, pub(crate) score_rules: PoolMemberScoreRules,
@@ -54,6 +67,11 @@ pub(crate) struct AdminProviderPoolRuntimeState {
pub(crate) total_sticky_sessions: usize, pub(crate) total_sticky_sessions: usize,
pub(crate) sticky_sessions_by_key: BTreeMap<String, usize>, pub(crate) sticky_sessions_by_key: BTreeMap<String, usize>,
pub(crate) sticky_bound_key_id: Option<String>, pub(crate) sticky_bound_key_id: Option<String>,
pub(crate) active_probe_member_ids: BTreeSet<String>,
pub(crate) provider_in_flight: usize,
pub(crate) provider_ema_in_flight: f64,
pub(crate) provider_desired_hot: usize,
pub(crate) provider_burst_pending: bool,
pub(crate) cooldown_reason_by_key: BTreeMap<String, String>, pub(crate) cooldown_reason_by_key: BTreeMap<String, String>,
pub(crate) cooldown_ttl_by_key: BTreeMap<String, u64>, pub(crate) cooldown_ttl_by_key: BTreeMap<String, u64>,
pub(crate) cost_window_usage_by_key: BTreeMap<String, u64>, pub(crate) cost_window_usage_by_key: BTreeMap<String, u64>,

View File

@@ -8,6 +8,7 @@ pub(crate) use super::super::admin::provider::pool::runtime::{
release_admin_provider_pool_key_lease, release_admin_provider_pool_key_lease,
}; };
pub(crate) use super::super::admin::provider::shared::support::{ pub(crate) use super::super::admin::provider::shared::support::{
AdminProviderPoolConfig, AdminProviderPoolRuntimeState, AdminProviderPoolSchedulingPreset, admin_provider_pool_quota_probe_active_members_key, AdminProviderPoolConfig,
AdminProviderPoolRuntimeState, AdminProviderPoolSchedulingPreset,
AdminProviderPoolUnschedulableRule, ADMIN_PROVIDER_POOL_SCAN_BATCH, AdminProviderPoolUnschedulableRule, ADMIN_PROVIDER_POOL_SCAN_BATCH,
}; };

View File

@@ -53,6 +53,7 @@ mod oauth;
mod orchestration; mod orchestration;
mod privacy; mod privacy;
mod provider_key_auth; mod provider_key_auth;
mod provider_pool_demand;
pub(crate) use aether_provider_transport as provider_transport; pub(crate) use aether_provider_transport as provider_transport;
mod rate_limit; mod rate_limit;
mod request_candidate_runtime; mod request_candidate_runtime;

View File

@@ -5,25 +5,27 @@ mod tests;
pub(crate) use runtime::{ pub(crate) use runtime::{
cancel_proxy_upgrade_rollout, clear_proxy_upgrade_rollout_conflicts, cancel_proxy_upgrade_rollout, clear_proxy_upgrade_rollout_conflicts,
ensure_provider_key_pool_scores_for_keys, inspect_proxy_upgrade_rollout, ensure_provider_key_pool_scores_for_keys, inspect_proxy_upgrade_rollout,
list_admin_cleanup_run_records, perform_oauth_token_refresh_once, list_admin_cleanup_run_records, perform_account_self_check_once,
perform_pool_quota_probe_once, perform_provider_checkin_once, preview_manual_usage_cleanup, perform_oauth_token_refresh_once, perform_pool_quota_probe_once, perform_provider_checkin_once,
rebuild_admin_stats_once, record_completed_cleanup_run, record_proxy_upgrade_traffic_success, pool_quota_probe_target_count, preview_manual_usage_cleanup, rebuild_admin_stats_once,
record_completed_cleanup_run, record_proxy_upgrade_traffic_success,
restore_proxy_upgrade_rollout_skipped_nodes, retry_proxy_upgrade_rollout_node, restore_proxy_upgrade_rollout_skipped_nodes, retry_proxy_upgrade_rollout_node,
run_admin_system_cleanup_once, run_manual_usage_cleanup_once, skip_proxy_upgrade_rollout_node, run_admin_system_cleanup_once, run_manual_usage_cleanup_once, skip_proxy_upgrade_rollout_node,
spawn_audit_cleanup_worker, spawn_db_maintenance_worker, spawn_account_self_check_worker, spawn_audit_cleanup_worker, spawn_db_maintenance_worker,
spawn_gemini_file_mapping_cleanup_worker, spawn_oauth_token_refresh_worker, spawn_gemini_file_mapping_cleanup_worker, spawn_oauth_token_refresh_worker,
spawn_pending_cleanup_worker, spawn_pool_monitor_worker, spawn_pool_quota_probe_worker, spawn_pending_cleanup_worker, spawn_pool_monitor_worker,
spawn_pool_quota_probe_replenish_for_request, spawn_pool_quota_probe_worker,
spawn_pool_score_rebuild_worker, spawn_provider_checkin_worker, spawn_pool_score_rebuild_worker, spawn_provider_checkin_worker,
spawn_proxy_node_metrics_cleanup_worker, spawn_proxy_node_stale_cleanup_worker, spawn_proxy_node_metrics_cleanup_worker, spawn_proxy_node_stale_cleanup_worker,
spawn_proxy_upgrade_rollout_worker, spawn_request_candidate_cleanup_worker, spawn_proxy_upgrade_rollout_worker, spawn_request_candidate_cleanup_worker,
spawn_stats_aggregation_worker, spawn_stats_hourly_aggregation_worker, spawn_stats_aggregation_worker, spawn_stats_hourly_aggregation_worker,
spawn_usage_cleanup_worker, spawn_wallet_daily_usage_aggregation_worker, spawn_usage_cleanup_worker, spawn_wallet_daily_usage_aggregation_worker,
start_admin_request_body_cleanup_task, start_admin_system_purge_task, start_admin_request_body_cleanup_task, start_admin_system_purge_task,
start_proxy_upgrade_rollout, AdminCleanupRunRecord, AdminCleanupTaskKind, start_proxy_upgrade_rollout, AccountSelfCheckRunSummary, AdminCleanupRunRecord,
AdminStatsRebuildSummary, AdminSystemCleanupSummary, ManualUsageCleanupError, AdminCleanupTaskKind, AdminStatsRebuildSummary, AdminSystemCleanupSummary,
OAuthTokenRefreshRunSummary, PoolQuotaProbeRunSummary, ProviderCheckinRunSummary, ManualUsageCleanupError, OAuthTokenRefreshRunSummary, PoolQuotaProbeRunSummary,
ProxyUpgradeRolloutCancelSummary, ProxyUpgradeRolloutConflictClearSummary, PoolQuotaProbeWorkerConfig, ProviderCheckinRunSummary, ProxyUpgradeRolloutCancelSummary,
ProxyUpgradeRolloutNodeActionSummary, ProxyUpgradeRolloutProbeConfig, ProxyUpgradeRolloutConflictClearSummary, ProxyUpgradeRolloutNodeActionSummary,
ProxyUpgradeRolloutSkippedRestoreSummary, ProxyUpgradeRolloutStatus, ProxyUpgradeRolloutProbeConfig, ProxyUpgradeRolloutSkippedRestoreSummary,
ProxyUpgradeRolloutTrackedNodeState, ProxyUpgradeRolloutStatus, ProxyUpgradeRolloutTrackedNodeState,
}; };

View File

@@ -6,6 +6,8 @@ use crate::admin_api::admin_provider_ops_local_action_response;
use crate::data::GatewayDataState; use crate::data::GatewayDataState;
use crate::{AppState, GatewayError}; use crate::{AppState, GatewayError};
#[path = "runtime/account_self_check.rs"]
mod account_self_check;
#[path = "runtime/audit_cleanup.rs"] #[path = "runtime/audit_cleanup.rs"]
mod audit_cleanup; mod audit_cleanup;
#[path = "runtime/cleanup_runs.rs"] #[path = "runtime/cleanup_runs.rs"]
@@ -49,6 +51,11 @@ mod usage_cleanup;
mod wallet_daily_usage; mod wallet_daily_usage;
#[path = "runtime/workers.rs"] #[path = "runtime/workers.rs"]
mod workers; mod workers;
pub(crate) use account_self_check::{
perform_account_self_check_once, perform_account_self_check_once_with_config,
select_account_self_check_key_ids, spawn_account_self_check_worker, AccountSelfCheckRunSummary,
AccountSelfCheckWorkerConfig,
};
pub(crate) use aether_data_contracts::repository::usage::{ pub(crate) use aether_data_contracts::repository::usage::{
UsageCleanupSummary, UsageCleanupWindow, UsageCleanupSummary, UsageCleanupWindow,
}; };
@@ -65,9 +72,10 @@ pub(crate) use oauth_token_refresh::{
}; };
use pending_cleanup::*; use pending_cleanup::*;
pub(crate) use pool_quota_probe::{ pub(crate) use pool_quota_probe::{
perform_pool_quota_probe_once, perform_pool_quota_probe_once_with_config, perform_pool_quota_probe_once, perform_pool_quota_probe_once_for_provider_with_config,
select_pool_quota_probe_key_ids, spawn_pool_quota_probe_worker, PoolQuotaProbeRunSummary, perform_pool_quota_probe_once_with_config, pool_quota_probe_target_count,
PoolQuotaProbeWorkerConfig, select_pool_quota_probe_key_ids, spawn_pool_quota_probe_replenish_for_request,
spawn_pool_quota_probe_worker, PoolQuotaProbeRunSummary, PoolQuotaProbeWorkerConfig,
}; };
pub(crate) use pool_score_rebuild::{ pub(crate) use pool_score_rebuild::{
ensure_provider_key_pool_scores_for_keys, perform_pool_score_rebuild_once, ensure_provider_key_pool_scores_for_keys, perform_pool_score_rebuild_once,

File diff suppressed because it is too large Load Diff

File diff suppressed because it is too large Load Diff

View File

@@ -20,15 +20,15 @@ use super::{
next_stats_hourly_aggregation_run_after, pending_cleanup_batch_size, next_stats_hourly_aggregation_run_after, pending_cleanup_batch_size,
pending_cleanup_timeout_minutes, plan_pending_cleanup_batch, provider_checkin_schedule, pending_cleanup_timeout_minutes, plan_pending_cleanup_batch, provider_checkin_schedule,
proxy_node_metrics_cleanup_settings, record_proxy_upgrade_traffic_success, proxy_node_metrics_cleanup_settings, record_proxy_upgrade_traffic_success,
run_db_maintenance_with, run_proxy_upgrade_rollout_once, spawn_audit_cleanup_worker, run_db_maintenance_with, run_proxy_upgrade_rollout_once, spawn_account_self_check_worker,
spawn_db_maintenance_worker, spawn_oauth_token_refresh_worker, spawn_pending_cleanup_worker, spawn_audit_cleanup_worker, spawn_db_maintenance_worker, spawn_oauth_token_refresh_worker,
spawn_pool_monitor_worker, spawn_pool_quota_probe_worker, spawn_provider_checkin_worker, spawn_pending_cleanup_worker, spawn_pool_monitor_worker, spawn_pool_quota_probe_worker,
spawn_proxy_node_stale_cleanup_worker, spawn_proxy_upgrade_rollout_worker, spawn_provider_checkin_worker, spawn_proxy_node_stale_cleanup_worker,
spawn_stats_aggregation_worker, spawn_stats_hourly_aggregation_worker, spawn_proxy_upgrade_rollout_worker, spawn_stats_aggregation_worker,
spawn_usage_cleanup_worker, spawn_wallet_daily_usage_aggregation_worker, spawn_stats_hourly_aggregation_worker, spawn_usage_cleanup_worker,
start_proxy_upgrade_rollout, stats_aggregation_target_day, spawn_wallet_daily_usage_aggregation_worker, start_proxy_upgrade_rollout,
stats_hourly_aggregation_target_hour, summarize_database_pool, usage_cleanup_settings, stats_aggregation_target_day, stats_hourly_aggregation_target_hour, summarize_database_pool,
usage_cleanup_window, usage_cleanup_window_with_override, usage_cleanup_settings, usage_cleanup_window, usage_cleanup_window_with_override,
wallet_daily_usage_aggregation_target, AppState, DbMaintenanceRunSummary, wallet_daily_usage_aggregation_target, AppState, DbMaintenanceRunSummary,
FailedPendingUsageRow, GatewayDataState, ProxyNodeMetricsCleanupSettings, FailedPendingUsageRow, GatewayDataState, ProxyNodeMetricsCleanupSettings,
ProxyUpgradeRolloutProbeConfig, StalePendingUsageRow, UsageCleanupSettings, USAGE_CLEANUP_HOUR, ProxyUpgradeRolloutProbeConfig, StalePendingUsageRow, UsageCleanupSettings, USAGE_CLEANUP_HOUR,
@@ -98,6 +98,15 @@ async fn spawn_pool_quota_probe_worker_skips_when_provider_catalog_unavailable()
assert!(spawn_pool_quota_probe_worker(state).is_none()); assert!(spawn_pool_quota_probe_worker(state).is_none());
} }
#[tokio::test]
async fn spawn_account_self_check_worker_skips_when_provider_catalog_unavailable() {
let state = AppState::new()
.expect("gateway state should build")
.with_data_state_for_tests(GatewayDataState::disabled());
assert!(spawn_account_self_check_worker(state).is_none());
}
fn sample_connected_proxy_node( fn sample_connected_proxy_node(
node_id: &str, node_id: &str,
heartbeat_interval: i32, heartbeat_interval: i32,

View File

@@ -0,0 +1,416 @@
use std::sync::{
atomic::{AtomicBool, Ordering},
Arc,
};
use std::time::{Duration, SystemTime, UNIX_EPOCH};
use aether_runtime_state::RuntimeState;
use serde::{Deserialize, Serialize};
use tokio::task::JoinHandle;
use tracing::debug;
use uuid::Uuid;
const PROVIDER_POOL_IN_FLIGHT_TOKENS_PREFIX: &str = "ap:provider_pool:in_flight";
const PROVIDER_POOL_DEMAND_SNAPSHOT_PREFIX: &str = "ap:provider_pool:demand";
const PROVIDER_POOL_BURST_PENDING_PREFIX: &str = "ap:quota_probe:burst_pending";
const PROVIDER_POOL_IN_FLIGHT_TOKEN_TTL_MS: u64 = 120_000;
const PROVIDER_POOL_IN_FLIGHT_RENEW_MS: u64 = 30_000;
const PROVIDER_POOL_DEMAND_SNAPSHOT_TTL_SECONDS: u64 = 6 * 60 * 60;
const PROVIDER_POOL_DEMAND_ALPHA: f64 = 0.2;
const PROVIDER_POOL_DEMAND_HEADROOM: f64 = 1.2;
const PROVIDER_POOL_DEMAND_FLOOR: usize = 2;
#[derive(Debug, Clone, Copy, PartialEq, Serialize, Deserialize)]
pub(crate) struct ProviderPoolDemandSnapshot {
pub(crate) in_flight: usize,
pub(crate) ema_in_flight: f64,
pub(crate) desired_hot: usize,
pub(crate) sampled_at_unix_ms: u64,
}
pub(crate) struct ProviderPoolInFlightGuard {
runtime: Arc<RuntimeState>,
tokens_key: String,
token: String,
stop_renewal: Arc<AtomicBool>,
renew_handle: Option<JoinHandle<()>>,
released: bool,
}
impl ProviderPoolInFlightGuard {
pub(crate) async fn release(mut self) {
self.release_inner().await;
}
async fn release_inner(&mut self) {
if self.released {
return;
}
self.released = true;
self.stop_renewal.store(true, Ordering::Release);
if let Some(handle) = self.renew_handle.take() {
handle.abort();
}
if let Err(err) = self
.runtime
.score_remove(&self.tokens_key, &self.token)
.await
{
debug!(
error = ?err,
"gateway provider pool demand: failed to release in-flight token"
);
}
}
}
impl Drop for ProviderPoolInFlightGuard {
fn drop(&mut self) {
if self.released {
return;
}
self.released = true;
self.stop_renewal.store(true, Ordering::Release);
if let Some(handle) = self.renew_handle.take() {
handle.abort();
}
let runtime = self.runtime.clone();
let tokens_key = self.tokens_key.clone();
let token = self.token.clone();
if let Ok(handle) = tokio::runtime::Handle::try_current() {
handle.spawn(async move {
if let Err(err) = runtime.score_remove(&tokens_key, &token).await {
debug!(
error = ?err,
"gateway provider pool demand: failed to release dropped in-flight token"
);
}
});
}
}
}
fn current_unix_ms() -> u64 {
let millis = SystemTime::now()
.duration_since(UNIX_EPOCH)
.unwrap_or_default()
.as_millis();
u64::try_from(millis).unwrap_or(u64::MAX)
}
fn in_flight_tokens_key(provider_id: &str) -> String {
format!("{PROVIDER_POOL_IN_FLIGHT_TOKENS_PREFIX}:{provider_id}")
}
fn demand_snapshot_key(provider_id: &str) -> String {
format!("{PROVIDER_POOL_DEMAND_SNAPSHOT_PREFIX}:{provider_id}")
}
pub(crate) fn provider_pool_burst_pending_key(provider_id: &str) -> String {
format!("{PROVIDER_POOL_BURST_PENDING_PREFIX}:{provider_id}")
}
fn build_in_flight_token(request_id: &str, candidate_id: Option<&str>, key_id: &str) -> String {
let request_id = request_id.trim();
let candidate_id = candidate_id
.map(str::trim)
.filter(|value| !value.is_empty())
.unwrap_or("-");
let key_id = key_id.trim();
format!(
"{}:{}:{}:{}",
current_unix_ms(),
request_id,
candidate_id,
if key_id.is_empty() { "-" } else { key_id },
) + &format!(":{}", Uuid::new_v4())
}
fn token_expiry_score(now_ms: u64) -> f64 {
now_ms.saturating_add(PROVIDER_POOL_IN_FLIGHT_TOKEN_TTL_MS) as f64
}
fn spawn_in_flight_renewal(
runtime: Arc<RuntimeState>,
tokens_key: String,
token: String,
stop: Arc<AtomicBool>,
) -> JoinHandle<()> {
tokio::spawn(async move {
let mut interval =
tokio::time::interval(Duration::from_millis(PROVIDER_POOL_IN_FLIGHT_RENEW_MS));
interval.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Delay);
loop {
interval.tick().await;
if stop.load(Ordering::Acquire) {
break;
}
if let Err(err) = runtime
.score_set(&tokens_key, &token, token_expiry_score(current_unix_ms()))
.await
{
debug!(
error = ?err,
"gateway provider pool demand: failed to renew in-flight token"
);
}
}
})
}
pub(crate) async fn acquire_provider_pool_in_flight_guard(
runtime: Arc<RuntimeState>,
provider_id: &str,
request_id: &str,
candidate_id: Option<&str>,
key_id: &str,
) -> Option<ProviderPoolInFlightGuard> {
let provider_id = provider_id.trim();
if provider_id.is_empty() {
return None;
}
let tokens_key = in_flight_tokens_key(provider_id);
let token = build_in_flight_token(request_id, candidate_id, key_id);
if let Err(err) = runtime
.score_set(&tokens_key, &token, token_expiry_score(current_unix_ms()))
.await
{
debug!(
provider_id,
error = ?err,
"gateway provider pool demand: failed to acquire in-flight token"
);
return None;
}
let stop_renewal = Arc::new(AtomicBool::new(false));
let renew_handle = spawn_in_flight_renewal(
runtime.clone(),
tokens_key.clone(),
token.clone(),
stop_renewal.clone(),
);
Some(ProviderPoolInFlightGuard {
runtime,
tokens_key,
token,
stop_renewal,
renew_handle: Some(renew_handle),
released: false,
})
}
pub(crate) async fn provider_pool_live_in_flight_count(
runtime: &RuntimeState,
provider_id: &str,
) -> usize {
let provider_id = provider_id.trim();
if provider_id.is_empty() {
return 0;
}
let key = in_flight_tokens_key(provider_id);
let now_ms = current_unix_ms() as f64;
if let Err(err) = runtime.score_remove_by_score(&key, now_ms).await {
debug!(
provider_id,
error = ?err,
"gateway provider pool demand: failed to prune expired in-flight tokens"
);
}
runtime.score_len(&key).await.unwrap_or(0)
}
pub(crate) async fn provider_pool_burst_pending(runtime: &RuntimeState, provider_id: &str) -> bool {
let provider_id = provider_id.trim();
if provider_id.is_empty() {
return false;
}
runtime
.kv_exists(&provider_pool_burst_pending_key(provider_id))
.await
.unwrap_or(false)
}
pub(crate) fn provider_pool_desired_hot(
in_flight: usize,
ema_in_flight: f64,
total_active_keys: usize,
max_keys_per_provider: usize,
) -> usize {
let cap = total_active_keys.min(max_keys_per_provider);
if cap == 0 {
return 0;
}
let signal = ema_in_flight
.max(in_flight as f64)
.max(0.0)
.min(usize::MAX as f64);
let desired = (signal * PROVIDER_POOL_DEMAND_HEADROOM).ceil() as usize;
desired.max(PROVIDER_POOL_DEMAND_FLOOR.min(cap)).min(cap)
}
fn parse_stored_demand_snapshot(raw: Option<String>) -> Option<ProviderPoolDemandSnapshot> {
let mut snapshot: ProviderPoolDemandSnapshot = serde_json::from_str(&raw?).ok()?;
if !snapshot.ema_in_flight.is_finite() || snapshot.ema_in_flight < 0.0 {
snapshot.ema_in_flight = 0.0;
}
Some(snapshot)
}
async fn stored_demand_snapshot(
runtime: &RuntimeState,
provider_id: &str,
) -> Option<ProviderPoolDemandSnapshot> {
runtime
.kv_get(&demand_snapshot_key(provider_id))
.await
.ok()
.and_then(parse_stored_demand_snapshot)
}
pub(crate) async fn read_provider_pool_demand_snapshot(
runtime: &RuntimeState,
provider_id: &str,
total_active_keys: usize,
max_keys_per_provider: usize,
) -> ProviderPoolDemandSnapshot {
let in_flight = provider_pool_live_in_flight_count(runtime, provider_id).await;
let stored = stored_demand_snapshot(runtime, provider_id).await;
let ema_in_flight = stored
.as_ref()
.map(|snapshot| snapshot.ema_in_flight)
.unwrap_or(in_flight as f64);
ProviderPoolDemandSnapshot {
in_flight,
ema_in_flight,
desired_hot: provider_pool_desired_hot(
in_flight,
ema_in_flight,
total_active_keys,
max_keys_per_provider,
),
sampled_at_unix_ms: stored
.map(|snapshot| snapshot.sampled_at_unix_ms)
.unwrap_or(0),
}
}
pub(crate) async fn sample_provider_pool_demand(
runtime: &RuntimeState,
provider_id: &str,
total_active_keys: usize,
max_keys_per_provider: usize,
) -> ProviderPoolDemandSnapshot {
let in_flight = provider_pool_live_in_flight_count(runtime, provider_id).await;
let previous = stored_demand_snapshot(runtime, provider_id).await;
let previous_ema = previous
.as_ref()
.map(|snapshot| snapshot.ema_in_flight)
.unwrap_or(in_flight as f64);
let ema_in_flight = if previous.is_some() {
previous_ema.mul_add(
1.0 - PROVIDER_POOL_DEMAND_ALPHA,
in_flight as f64 * PROVIDER_POOL_DEMAND_ALPHA,
)
} else {
in_flight as f64
}
.max(0.0);
let sampled_at_unix_ms = current_unix_ms();
let snapshot = ProviderPoolDemandSnapshot {
in_flight,
ema_in_flight,
desired_hot: provider_pool_desired_hot(
in_flight,
ema_in_flight,
total_active_keys,
max_keys_per_provider,
),
sampled_at_unix_ms,
};
if let Ok(serialized) = serde_json::to_string(&snapshot) {
if let Err(err) = runtime
.kv_set(
&demand_snapshot_key(provider_id),
serialized,
Some(Duration::from_secs(
PROVIDER_POOL_DEMAND_SNAPSHOT_TTL_SECONDS,
)),
)
.await
{
debug!(
provider_id,
error = ?err,
"gateway provider pool demand: failed to persist demand snapshot"
);
}
}
snapshot
}
#[cfg(test)]
mod tests {
use super::*;
use aether_runtime_state::{MemoryRuntimeStateConfig, RuntimeState};
#[tokio::test]
async fn in_flight_guard_tracks_and_releases_provider_tokens() {
let runtime = Arc::new(RuntimeState::memory(MemoryRuntimeStateConfig::default()));
let guard = acquire_provider_pool_in_flight_guard(
runtime.clone(),
"provider-1",
"request-1",
Some("candidate-1"),
"key-1",
)
.await
.expect("guard should be acquired");
assert_eq!(
provider_pool_live_in_flight_count(runtime.as_ref(), "provider-1").await,
1
);
guard.release().await;
assert_eq!(
provider_pool_live_in_flight_count(runtime.as_ref(), "provider-1").await,
0
);
}
#[tokio::test]
async fn demand_snapshot_uses_instant_in_flight_for_fast_rise_and_ema_for_fall() {
let runtime = RuntimeState::memory(MemoryRuntimeStateConfig::default());
for idx in 0..10 {
let guard = acquire_provider_pool_in_flight_guard(
Arc::new(runtime.clone()),
"provider-1",
"request-1",
Some(&format!("candidate-{idx}")),
"key-1",
)
.await
.expect("guard");
std::mem::forget(guard);
}
let high = sample_provider_pool_demand(&runtime, "provider-1", 100, 50).await;
assert_eq!(high.in_flight, 10);
assert_eq!(high.desired_hot, 12);
let key = in_flight_tokens_key("provider-1");
let _ = runtime.score_remove_by_score(&key, f64::INFINITY).await;
let low = sample_provider_pool_demand(&runtime, "provider-1", 100, 50).await;
assert_eq!(low.in_flight, 0);
assert!(low.ema_in_flight > 0.0);
assert!(low.desired_hot >= PROVIDER_POOL_DEMAND_FLOOR);
assert!(low.desired_hot < high.desired_hot);
}
}

View File

@@ -40,6 +40,7 @@ use super::super::router::RequestAdmissionError;
use super::super::{control::GatewayControlDecision, error::GatewayError}; use super::super::{control::GatewayControlDecision, error::GatewayError};
use super::super::{provider_transport, usage}; use super::super::{provider_transport, usage};
use crate::maintenance::spawn_account_self_check_worker;
use crate::maintenance::spawn_audit_cleanup_worker; use crate::maintenance::spawn_audit_cleanup_worker;
use crate::maintenance::spawn_db_maintenance_worker; use crate::maintenance::spawn_db_maintenance_worker;
use crate::maintenance::spawn_gemini_file_mapping_cleanup_worker; use crate::maintenance::spawn_gemini_file_mapping_cleanup_worker;
@@ -1172,6 +1173,10 @@ impl AppState {
crate::task_runtime::TASK_KEY_POOL_QUOTA_PROBE, crate::task_runtime::TASK_KEY_POOL_QUOTA_PROBE,
spawn_pool_quota_probe_worker(self.clone()), spawn_pool_quota_probe_worker(self.clone()),
); );
supervise_worker(
crate::task_runtime::TASK_KEY_ACCOUNT_SELF_CHECK,
spawn_account_self_check_worker(self.clone()),
);
supervise_worker( supervise_worker(
crate::task_runtime::TASK_KEY_POOL_SCORE_REBUILD, crate::task_runtime::TASK_KEY_POOL_SCORE_REBUILD,
spawn_pool_score_rebuild_worker(self.clone()), spawn_pool_score_rebuild_worker(self.clone()),

View File

@@ -22,6 +22,7 @@ pub(crate) const TASK_KEY_VIDEO_TASK_POLLER: &str = "video.task.poller";
pub(crate) const TASK_KEY_MODEL_FETCH_WORKER: &str = "model.fetch.worker"; pub(crate) const TASK_KEY_MODEL_FETCH_WORKER: &str = "model.fetch.worker";
pub(crate) const TASK_KEY_PROVIDER_QUOTA_RESET: &str = "provider.quota.reset.worker"; pub(crate) const TASK_KEY_PROVIDER_QUOTA_RESET: &str = "provider.quota.reset.worker";
pub(crate) const TASK_KEY_POOL_QUOTA_PROBE: &str = "pool.quota.probe.worker"; pub(crate) const TASK_KEY_POOL_QUOTA_PROBE: &str = "pool.quota.probe.worker";
pub(crate) const TASK_KEY_ACCOUNT_SELF_CHECK: &str = "account.self_check.worker";
pub(crate) const TASK_KEY_POOL_SCORE_REBUILD: &str = "pool.score.rebuild.worker"; pub(crate) const TASK_KEY_POOL_SCORE_REBUILD: &str = "pool.score.rebuild.worker";
pub(crate) const TASK_KEY_POOL_MONITOR: &str = "pool.monitor.worker"; pub(crate) const TASK_KEY_POOL_MONITOR: &str = "pool.monitor.worker";
pub(crate) const TASK_KEY_AUDIT_CLEANUP: &str = "maintenance.audit.cleanup"; pub(crate) const TASK_KEY_AUDIT_CLEANUP: &str = "maintenance.audit.cleanup";
@@ -102,6 +103,14 @@ const TASK_DEFINITIONS: &[TaskDefinition] = &[
true, true,
RETRY_ONCE, RETRY_ONCE,
), ),
TaskDefinition::new(
TASK_KEY_ACCOUNT_SELF_CHECK,
TaskKind::Scheduled,
"interval",
true,
true,
RETRY_ONCE,
),
TaskDefinition::new( TaskDefinition::new(
TASK_KEY_POOL_SCORE_REBUILD, TASK_KEY_POOL_SCORE_REBUILD,
TaskKind::Scheduled, TaskKind::Scheduled,

View File

@@ -236,6 +236,11 @@ async fn gateway_handles_admin_pool_overview_locally_with_trusted_admin_principa
assert_eq!(items[0]["active_keys"], 1); assert_eq!(items[0]["active_keys"], 1);
assert_eq!(items[0]["cooldown_count"], 0); assert_eq!(items[0]["cooldown_count"], 0);
assert_eq!(items[0]["pool_enabled"], true); assert_eq!(items[0]["pool_enabled"], true);
assert_eq!(items[0]["provider_hot_count"], 0);
assert_eq!(items[0]["provider_desired_hot"], 0);
assert_eq!(items[0]["provider_in_flight"], 0);
assert_eq!(items[0]["provider_ema_in_flight"], 0.0);
assert_eq!(items[0]["provider_burst_pending"], false);
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0); assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
gateway_handle.abort(); gateway_handle.abort();

View File

@@ -1839,6 +1839,11 @@ async fn gateway_handles_admin_provider_pool_status_locally_with_trusted_admin_p
assert_eq!(payload["pool_enabled"], true); assert_eq!(payload["pool_enabled"], true);
assert_eq!(payload["total_keys"], 1); assert_eq!(payload["total_keys"], 1);
assert_eq!(payload["total_sticky_sessions"], 0); assert_eq!(payload["total_sticky_sessions"], 0);
assert_eq!(payload["provider_hot_count"], 0);
assert_eq!(payload["provider_desired_hot"], 0);
assert_eq!(payload["provider_in_flight"], 0);
assert_eq!(payload["provider_ema_in_flight"], 0.0);
assert_eq!(payload["provider_burst_pending"], false);
let keys = payload["keys"].as_array().expect("keys should be an array"); let keys = payload["keys"].as_array().expect("keys should be an array");
assert_eq!(keys.len(), 1); assert_eq!(keys.len(), 1);
assert_eq!(keys[0]["key_id"], "key-openai-pool"); assert_eq!(keys[0]["key_id"], "key-openai-pool");

View File

@@ -807,6 +807,22 @@ impl RuntimeState {
} }
} }
pub async fn score_remove(&self, key: &str, member: &str) -> Result<bool, DataLayerError> {
match self.backend.as_ref() {
RuntimeStateBackend::Memory(memory) => Ok(memory.score_remove(key, member).await),
RuntimeStateBackend::Redis(redis) => {
let key = redis.keyspace.key(key);
let removed = redis_query_i64(redis, "runtime score remove", {
let mut command = redis_cmd("ZREM");
command.arg(&key).arg(member);
command
})
.await?;
Ok(removed > 0)
}
}
}
pub async fn score_remove_by_rank( pub async fn score_remove_by_rank(
&self, &self,
key: &str, key: &str,

View File

@@ -349,6 +349,14 @@ impl MemoryRuntimeBackend {
before.saturating_sub(set.len()) before.saturating_sub(set.len())
} }
pub(crate) async fn score_remove(&self, key: &str, member: &str) -> bool {
self.scores
.lock()
.await
.get_mut(key)
.is_some_and(|set| set.remove(member).is_some())
}
pub(crate) async fn score_len(&self, key: &str) -> usize { pub(crate) async fn score_len(&self, key: &str) -> usize {
self.scores.lock().await.get(key).map_or(0, BTreeMap::len) self.scores.lock().await.get(key).map_or(0, BTreeMap::len)
} }

View File

@@ -28,6 +28,11 @@ export interface PoolStatusResponse {
pool_enabled: boolean pool_enabled: boolean
total_keys: number total_keys: number
total_sticky_sessions: number total_sticky_sessions: number
provider_hot_count: number
provider_desired_hot: number
provider_in_flight: number
provider_ema_in_flight: number
provider_burst_pending: boolean
keys: PoolKeyStatus[] keys: PoolKeyStatus[]
} }
@@ -77,6 +82,11 @@ export interface PoolOverviewItem {
active_keys: number active_keys: number
cooldown_count: number cooldown_count: number
pool_enabled: boolean pool_enabled: boolean
provider_hot_count?: number
provider_desired_hot?: number
provider_in_flight?: number
provider_ema_in_flight?: number
provider_burst_pending?: boolean
} }
export interface PoolOverviewResponse { export interface PoolOverviewResponse {

View File

@@ -584,6 +584,15 @@ export interface PoolAdvancedConfig {
score_rules?: PoolScoreRules | null score_rules?: PoolScoreRules | null
probing_enabled?: boolean probing_enabled?: boolean
probing_interval_minutes?: number | null probing_interval_minutes?: number | null
// deprecated: retained only for backward-compatible reads
probing_target_percent?: number | null
// deprecated: retained only for backward-compatible reads
probing_target_count?: number | null
account_self_check_enabled?: boolean
account_self_check_interval_minutes?: number | null
account_self_check_concurrency?: number | null
account_self_check_method?: 'quota_refresh' | 'custom_request' | string
account_self_check_request?: Record<string, unknown> | null
auto_remove_banned_keys?: boolean auto_remove_banned_keys?: boolean
} }

View File

@@ -89,6 +89,75 @@
</div> </div>
</div> </div>
<div
v-if="form.account_self_check_enabled"
class="space-y-3 rounded-xl border border-dashed border-primary/25 bg-primary/5 p-4"
>
<div class="grid gap-3 sm:grid-cols-3">
<div class="space-y-1.5">
<Label>
自检间隔
<span class="text-xs text-muted-foreground">(分钟)</span>
</Label>
<Input
:model-value="form.account_self_check_interval_minutes ?? ''"
type="number"
min="1"
max="1440"
placeholder="60"
@update:model-value="(v) => form.account_self_check_interval_minutes = parseNum(v)"
/>
</div>
<div class="space-y-1.5">
<Label>
自检并发
</Label>
<Input
:model-value="form.account_self_check_concurrency ?? ''"
type="number"
min="1"
max="64"
placeholder="4"
@update:model-value="(v) => form.account_self_check_concurrency = parseNum(v)"
/>
</div>
<div class="space-y-1.5">
<Label>
自检方式
</Label>
<div class="flex w-fit gap-0.5 rounded-md bg-muted/40 p-0.5">
<button
v-for="opt in accountSelfCheckMethodOptions"
:key="opt.value"
type="button"
class="rounded px-2.5 py-1 text-xs font-medium transition-all"
:class="[
form.account_self_check_method === opt.value
? 'bg-primary text-primary-foreground shadow-sm'
: 'text-muted-foreground hover:bg-background/50 hover:text-foreground'
]"
@click="form.account_self_check_method = opt.value"
>
{{ opt.label }}
</button>
</div>
</div>
</div>
<div
v-if="form.account_self_check_method === 'custom_request'"
class="space-y-1.5"
>
<Label>请求配置</Label>
<Textarea
v-model="form.account_self_check_request_text"
class="min-h-[160px] font-mono text-xs leading-5"
spellcheck="false"
placeholder='{"path":"/v1/models","success_status_codes":[200],"blocked_status_codes":[401,403]}'
/>
</div>
</div>
<div <div
class="grid gap-3 sm:grid-cols-2" class="grid gap-3 sm:grid-cols-2"
:class="cooldownFieldLayout.desktopColumnsClass" :class="cooldownFieldLayout.desktopColumnsClass"
@@ -301,7 +370,7 @@
</span> </span>
</div> </div>
<p class="text-xs leading-5 text-muted-foreground"> <p class="text-xs leading-5 text-muted-foreground">
调整主动探测、健康、额度、延迟和使用成本进入号池候选排序时的权重。 调整探测结果、健康、额度、延迟和使用成本进入号池候选排序时的权重。
</p> </p>
</div> </div>
@@ -606,7 +675,7 @@
<script setup lang="ts"> <script setup lang="ts">
import { computed, ref, watch } from 'vue' import { computed, ref, watch } from 'vue'
import { CircleHelp } from 'lucide-vue-next' import { CircleHelp } from 'lucide-vue-next'
import { Dialog, Button, Input, Label, Switch, Tooltip, TooltipContent, TooltipProvider, TooltipTrigger } from '@/components/ui' import { Dialog, Button, Input, Label, Switch, Textarea, Tooltip, TooltipContent, TooltipProvider, TooltipTrigger } from '@/components/ui'
import { useToast } from '@/composables/useToast' import { useToast } from '@/composables/useToast'
import { parseApiError } from '@/utils/errorParser' import { parseApiError } from '@/utils/errorParser'
import { updateProvider } from '@/api/endpoints' import { updateProvider } from '@/api/endpoints'
@@ -647,6 +716,12 @@ const healthToggleCards = buildPoolHealthToggleCards()
const cooldownFieldLayout = buildPoolCooldownFieldLayout() const cooldownFieldLayout = buildPoolCooldownFieldLayout()
const costFieldLayout = buildPoolCostFieldLayout() const costFieldLayout = buildPoolCostFieldLayout()
const secondarySectionLayout = buildPoolSecondarySectionLayout() const secondarySectionLayout = buildPoolSecondarySectionLayout()
const accountSelfCheckMethodOptions = [
{ value: 'quota_refresh', label: '刷新额度' },
{ value: 'custom_request', label: '自定义请求' },
] as const
type AccountSelfCheckMethod = typeof accountSelfCheckMethodOptions[number]['value']
const form = ref({ const form = ref({
global_priority: null as number | null | undefined, global_priority: null as number | null | undefined,
@@ -674,6 +749,11 @@ const form = ref({
probe_failure_cooldown_threshold: null as number | null | undefined, probe_failure_cooldown_threshold: null as number | null | undefined,
probing_enabled: false, probing_enabled: false,
probing_interval_minutes: null as number | null | undefined, probing_interval_minutes: null as number | null | undefined,
account_self_check_enabled: false,
account_self_check_interval_minutes: null as number | null | undefined,
account_self_check_concurrency: null as number | null | undefined,
account_self_check_method: 'quota_refresh' as AccountSelfCheckMethod,
account_self_check_request_text: '',
auto_remove_banned_keys: false, auto_remove_banned_keys: false,
skip_exhausted_accounts: false, skip_exhausted_accounts: false,
}) })
@@ -704,12 +784,45 @@ function parseNum(v: string | number): number | undefined {
return Number.isNaN(n) ? undefined : n return Number.isNaN(n) ? undefined : n
} }
function normalizeAccountSelfCheckMethod(value: unknown): AccountSelfCheckMethod {
return value === 'custom_request' ? 'custom_request' : 'quota_refresh'
}
function formatJsonForTextarea(value: unknown): string {
if (!value || typeof value !== 'object' || Array.isArray(value)) {
return ''
}
return JSON.stringify(value, null, 2)
}
function parseJsonObjectText(text: string): Record<string, unknown> | undefined | null {
const trimmed = text.trim()
if (!trimmed) return undefined
let parsed: unknown
try {
parsed = JSON.parse(trimmed)
} catch {
showError('账号自检请求 JSON 格式不正确')
return null
}
if (!parsed || typeof parsed !== 'object' || Array.isArray(parsed)) {
showError('账号自检请求必须是 JSON 对象')
return null
}
return parsed as Record<string, unknown>
}
function getHealthToggleValue(key: PoolHealthToggleKey): boolean { function getHealthToggleValue(key: PoolHealthToggleKey): boolean {
switch (key) { switch (key) {
case 'health_policy_enabled': case 'health_policy_enabled':
return form.value.health_policy_enabled return form.value.health_policy_enabled
case 'probing_enabled': case 'probing_enabled':
return form.value.probing_enabled return form.value.probing_enabled
case 'account_self_check_enabled':
return form.value.account_self_check_enabled
case 'auto_remove_banned_keys': case 'auto_remove_banned_keys':
return form.value.auto_remove_banned_keys return form.value.auto_remove_banned_keys
case 'skip_exhausted_accounts': case 'skip_exhausted_accounts':
@@ -725,6 +838,9 @@ function updateHealthToggleValue(key: PoolHealthToggleKey, value: boolean): void
case 'probing_enabled': case 'probing_enabled':
form.value.probing_enabled = value form.value.probing_enabled = value
return return
case 'account_self_check_enabled':
form.value.account_self_check_enabled = value
return
case 'auto_remove_banned_keys': case 'auto_remove_banned_keys':
form.value.auto_remove_banned_keys = value form.value.auto_remove_banned_keys = value
return return
@@ -765,6 +881,11 @@ watch(() => props.modelValue, (open) => {
probe_failure_cooldown_threshold: scoreRules?.probe_failure_cooldown_threshold ?? null, probe_failure_cooldown_threshold: scoreRules?.probe_failure_cooldown_threshold ?? null,
probing_enabled: cfg?.probing_enabled ?? false, probing_enabled: cfg?.probing_enabled ?? false,
probing_interval_minutes: cfg?.probing_interval_minutes ?? null, probing_interval_minutes: cfg?.probing_interval_minutes ?? null,
account_self_check_enabled: cfg?.account_self_check_enabled ?? false,
account_self_check_interval_minutes: cfg?.account_self_check_interval_minutes ?? null,
account_self_check_concurrency: cfg?.account_self_check_concurrency ?? null,
account_self_check_method: normalizeAccountSelfCheckMethod(cfg?.account_self_check_method),
account_self_check_request_text: formatJsonForTextarea(cfg?.account_self_check_request),
auto_remove_banned_keys: cfg?.auto_remove_banned_keys ?? false, auto_remove_banned_keys: cfg?.auto_remove_banned_keys ?? false,
skip_exhausted_accounts: cfg?.skip_exhausted_accounts ?? false, skip_exhausted_accounts: cfg?.skip_exhausted_accounts ?? false,
} }
@@ -784,6 +905,12 @@ watch(() => props.modelValue, (open) => {
async function handleSave() { async function handleSave() {
loading.value = true loading.value = true
try { try {
const accountSelfCheckRequest = form.value.account_self_check_enabled
&& form.value.account_self_check_method === 'custom_request'
? parseJsonObjectText(form.value.account_self_check_request_text)
: undefined
if (accountSelfCheckRequest === null) return
const scoreRules = { const scoreRules = {
...(props.currentConfig?.score_rules ?? {}), ...(props.currentConfig?.score_rules ?? {}),
weights: { weights: {
@@ -801,9 +928,20 @@ async function handleSave() {
request_failure_penalty: form.value.request_failure_penalty ?? undefined, request_failure_penalty: form.value.request_failure_penalty ?? undefined,
probe_failure_cooldown_threshold: form.value.probe_failure_cooldown_threshold ?? undefined, probe_failure_cooldown_threshold: form.value.probe_failure_cooldown_threshold ?? undefined,
} }
const existingPoolAdvanced: Record<string, unknown> = { ...(props.currentConfig ?? {}) }
for (const key of [
'probing_target_percent',
'probing_target_count',
'probing_active_target_percent',
'probing_active_target_count',
'active_probe_target_percent',
'active_probe_target_count',
]) {
delete existingPoolAdvanced[key]
}
// 合并已有配置(保留 scheduling_presets 等不在此对话框编辑的字段) // 合并已有配置(保留 scheduling_presets 等不在此对话框编辑的字段)
const poolAdvanced: Record<string, unknown> = { const poolAdvanced: Record<string, unknown> = {
...(props.currentConfig ?? {}), ...existingPoolAdvanced,
global_priority: form.value.global_priority ?? undefined, global_priority: form.value.global_priority ?? undefined,
sticky_session_ttl_seconds: form.value.sticky_session_ttl_seconds ?? undefined, sticky_session_ttl_seconds: form.value.sticky_session_ttl_seconds ?? undefined,
cost_window_seconds: form.value.cost_window_seconds ?? undefined, cost_window_seconds: form.value.cost_window_seconds ?? undefined,
@@ -821,6 +959,20 @@ async function handleSave() {
probing_interval_minutes: form.value.probing_enabled probing_interval_minutes: form.value.probing_enabled
? (form.value.probing_interval_minutes ?? undefined) ? (form.value.probing_interval_minutes ?? undefined)
: undefined, : undefined,
account_self_check_enabled: form.value.account_self_check_enabled,
account_self_check_interval_minutes: form.value.account_self_check_enabled
? (form.value.account_self_check_interval_minutes ?? undefined)
: undefined,
account_self_check_concurrency: form.value.account_self_check_enabled
? (form.value.account_self_check_concurrency ?? undefined)
: undefined,
account_self_check_method: form.value.account_self_check_enabled
? form.value.account_self_check_method
: undefined,
account_self_check_request: form.value.account_self_check_enabled
&& form.value.account_self_check_method === 'custom_request'
? accountSelfCheckRequest
: undefined,
auto_remove_banned_keys: form.value.auto_remove_banned_keys, auto_remove_banned_keys: form.value.auto_remove_banned_keys,
skip_exhausted_accounts: form.value.skip_exhausted_accounts, skip_exhausted_accounts: form.value.skip_exhausted_accounts,
} }

View File

@@ -0,0 +1,311 @@
<template>
<Dialog
:model-value="modelValue"
:no-padding="true"
size="3xl"
@update:model-value="emit('update:modelValue', $event)"
>
<template #header>
<div class="border-b border-border px-4 py-4 sm:px-6">
<div class="flex items-start gap-3">
<div class="flex-1 min-w-0">
<h3 class="text-lg font-semibold text-foreground leading-tight">
自适应热池指标
</h3>
<p class="text-xs text-muted-foreground">
{{ providerName || '当前 Provider' }}
</p>
</div>
<Button
variant="ghost"
size="icon"
class="h-8 w-8 shrink-0"
title="关闭"
@click="emit('update:modelValue', false)"
>
<X class="h-4 w-4" />
</Button>
</div>
</div>
</template>
<div class="max-h-[calc(100dvh-13rem)] space-y-4 overflow-y-auto overscroll-contain pr-1 sm:max-h-[min(72vh,42rem)] sm:pr-2">
<div class="grid gap-3 sm:grid-cols-2 xl:grid-cols-4">
<div
v-for="item in summaryCards"
:key="item.label"
class="rounded-lg border border-border/60 bg-card/70 px-3 py-3"
>
<div class="text-xs text-muted-foreground">
{{ item.label }}
</div>
<div class="mt-2 text-xl font-semibold tabular-nums">
{{ item.value }}
</div>
<div class="mt-1 text-[11px] text-muted-foreground">
{{ item.hint }}
</div>
</div>
</div>
<section class="rounded-lg border border-border/60 bg-card/70 p-4">
<div class="flex flex-col gap-2 sm:flex-row sm:items-center sm:justify-between">
<div>
<h3 class="text-sm font-semibold">
最近采样
</h3>
<p class="text-xs text-muted-foreground">
{{ sampleWindowText }}
</p>
</div>
<div class="flex flex-wrap items-center gap-2 text-[11px]">
<span
v-for="item in legendItems"
:key="item.label"
class="inline-flex items-center gap-1.5 text-muted-foreground"
>
<span
class="h-2 w-2 rounded-full"
:class="item.dotClass"
/>
{{ item.label }}
</span>
</div>
</div>
<div
v-if="samples.length === 0"
class="mt-4 flex h-56 items-center justify-center rounded-lg border border-dashed border-border/70 bg-muted/20 text-sm text-muted-foreground"
>
暂无采样
</div>
<div
v-else
class="mt-4 h-56 rounded-lg border border-border/50 bg-background/60 p-3"
>
<svg
class="h-full w-full overflow-visible"
:viewBox="`0 0 ${chartWidth} ${chartHeight}`"
preserveAspectRatio="none"
role="img"
aria-label="自适应热池趋势"
>
<line
v-for="tick in yTicks"
:key="tick.y"
x1="0"
:x2="chartWidth"
:y1="tick.y"
:y2="tick.y"
class="stroke-border/70"
stroke-width="1"
/>
<polyline
v-if="desiredHotLine"
:points="desiredHotLine"
fill="none"
stroke="rgb(99 102 241)"
stroke-width="2.4"
stroke-linecap="round"
stroke-linejoin="round"
vector-effect="non-scaling-stroke"
/>
<polyline
v-if="hotLine"
:points="hotLine"
fill="none"
stroke="rgb(16 185 129)"
stroke-width="2.2"
stroke-linecap="round"
stroke-linejoin="round"
vector-effect="non-scaling-stroke"
/>
<polyline
v-if="inFlightLine"
:points="inFlightLine"
fill="none"
stroke="rgb(245 158 11)"
stroke-width="2.2"
stroke-linecap="round"
stroke-linejoin="round"
vector-effect="non-scaling-stroke"
/>
<polyline
v-if="emaLine"
:points="emaLine"
fill="none"
stroke="rgb(14 165 233)"
stroke-width="2.2"
stroke-linecap="round"
stroke-linejoin="round"
vector-effect="non-scaling-stroke"
/>
<circle
v-for="burst in burstPoints"
:key="`${burst.x}-${burst.y}`"
:cx="burst.x"
:cy="burst.y"
r="3"
fill="rgb(239 68 68)"
vector-effect="non-scaling-stroke"
/>
</svg>
</div>
<div class="mt-3 flex items-center justify-between gap-3 text-[11px] text-muted-foreground">
<span>{{ firstSampleTime }}</span>
<span>峰值 {{ maxChartValueText }}</span>
<span>{{ lastSampleTime }}</span>
</div>
</section>
</div>
</Dialog>
</template>
<script setup lang="ts">
import { computed } from 'vue'
import { Button, Dialog } from '@/components/ui'
import { X } from 'lucide-vue-next'
export interface PoolDemandMetricSample {
providerId: string
sampledAt: number
hotCount: number
desiredHot: number
inFlight: number
emaInFlight: number
burstPending: boolean
}
const props = defineProps<{
modelValue: boolean
providerName?: string | null
samples: PoolDemandMetricSample[]
}>()
const emit = defineEmits<{
'update:modelValue': [value: boolean]
}>()
const chartWidth = 640
const chartHeight = 220
const yTickRatios = [0, 0.25, 0.5, 0.75, 1]
const latest = computed(() => props.samples.at(-1) ?? null)
const maxChartValue = computed(() => {
const maxValue = props.samples.reduce((max, sample) => {
return Math.max(
max,
sample.hotCount,
sample.desiredHot,
sample.inFlight,
sample.emaInFlight,
)
}, 1)
return Math.max(1, Math.ceil(maxValue))
})
const yTicks = computed(() => {
return yTickRatios.map(ratio => ({
y: chartHeight * ratio,
}))
})
function formatMetric(value: number, fractionDigits = 0): string {
if (!Number.isFinite(value) || value <= 0) {
return fractionDigits > 0 ? '0.0' : '0'
}
return value.toFixed(fractionDigits)
}
function formatSampleTime(timestamp: number | undefined): string {
if (!timestamp) return '--:--:--'
const date = new Date(timestamp)
const pad = (value: number) => String(value).padStart(2, '0')
return `${pad(date.getHours())}:${pad(date.getMinutes())}:${pad(date.getSeconds())}`
}
function buildLine(valueOf: (sample: PoolDemandMetricSample) => number): string {
if (props.samples.length === 0) return ''
const maxValue = maxChartValue.value
const widthStep = props.samples.length > 1
? chartWidth / (props.samples.length - 1)
: 0
return props.samples
.map((sample, index) => {
const x = props.samples.length > 1 ? index * widthStep : chartWidth / 2
const normalized = Math.max(0, Math.min(valueOf(sample), maxValue))
const y = chartHeight - ((normalized / maxValue) * chartHeight)
return `${x.toFixed(2)},${y.toFixed(2)}`
})
.join(' ')
}
const hotLine = computed(() => buildLine(sample => sample.hotCount))
const desiredHotLine = computed(() => buildLine(sample => sample.desiredHot))
const inFlightLine = computed(() => buildLine(sample => sample.inFlight))
const emaLine = computed(() => buildLine(sample => sample.emaInFlight))
const burstPoints = computed(() => {
if (props.samples.length === 0) return []
const widthStep = props.samples.length > 1
? chartWidth / (props.samples.length - 1)
: 0
const maxValue = maxChartValue.value
return props.samples
.map((sample, index) => {
if (!sample.burstPending) return null
const x = props.samples.length > 1 ? index * widthStep : chartWidth / 2
const normalized = Math.max(0, Math.min(sample.desiredHot, maxValue))
const y = chartHeight - ((normalized / maxValue) * chartHeight)
return { x, y }
})
.filter((point): point is { x: number, y: number } => point !== null)
})
const summaryCards = computed(() => {
const sample = latest.value
return [
{
label: '热池',
value: sample ? `${sample.hotCount} / ${sample.desiredHot}` : '0 / 0',
hint: '当前 / 目标',
},
{
label: 'in-flight',
value: sample ? formatMetric(sample.inFlight) : '0',
hint: '正在执行',
},
{
label: 'EMA',
value: sample ? formatMetric(sample.emaInFlight, 1) : '0.0',
hint: '平滑热度',
},
{
label: 'Burst',
value: sample?.burstPending ? '补热中' : '空闲',
hint: '异步补位',
},
]
})
const legendItems = [
{ label: '目标', dotClass: 'bg-indigo-500' },
{ label: '热池', dotClass: 'bg-emerald-500' },
{ label: 'in-flight', dotClass: 'bg-amber-500' },
{ label: 'EMA', dotClass: 'bg-sky-500' },
{ label: 'Burst', dotClass: 'bg-red-500' },
]
const sampleWindowText = computed(() => {
const count = props.samples.length
if (count === 0) return '等待下一次采样'
return `最近 ${count} 个采样点`
})
const firstSampleTime = computed(() => formatSampleTime(props.samples[0]?.sampledAt))
const lastSampleTime = computed(() => formatSampleTime(latest.value?.sampledAt))
const maxChartValueText = computed(() => formatMetric(maxChartValue.value))
</script>

View File

@@ -12,6 +12,7 @@ describe('poolAdvancedDialog', () => {
expect(buildPoolHealthToggleCards().map(item => item.key)).toEqual([ expect(buildPoolHealthToggleCards().map(item => item.key)).toEqual([
'health_policy_enabled', 'health_policy_enabled',
'probing_enabled', 'probing_enabled',
'account_self_check_enabled',
'auto_remove_banned_keys', 'auto_remove_banned_keys',
'skip_exhausted_accounts', 'skip_exhausted_accounts',
]) ])
@@ -27,7 +28,12 @@ describe('poolAdvancedDialog', () => {
{ {
key: 'probing_enabled', key: 'probing_enabled',
label: '主动探测', label: '主动探测',
description: '按固定间隔刷新 Key 的状态与额度,减少号池状态滞后。', description: '自动维护热池,缺口时异步补位。',
},
{
key: 'account_self_check_enabled',
label: '账号自检',
description: '定时确认封号状态,默认刷新额度,也可使用自定义请求。',
}, },
{ {
key: 'auto_remove_banned_keys', key: 'auto_remove_banned_keys',

View File

@@ -1,6 +1,7 @@
export type PoolHealthToggleKey = export type PoolHealthToggleKey =
| 'health_policy_enabled' | 'health_policy_enabled'
| 'probing_enabled' | 'probing_enabled'
| 'account_self_check_enabled'
| 'auto_remove_banned_keys' | 'auto_remove_banned_keys'
| 'skip_exhausted_accounts' | 'skip_exhausted_accounts'
@@ -34,7 +35,12 @@ export function buildPoolHealthToggleCards(): PoolHealthToggleCard[] {
{ {
key: 'probing_enabled', key: 'probing_enabled',
label: '主动探测', label: '主动探测',
description: '按固定间隔刷新 Key 的状态与额度,减少号池状态滞后。', description: '自动维护热池,缺口时异步补位。',
},
{
key: 'account_self_check_enabled',
label: '账号自检',
description: '定时确认封号状态,默认刷新额度,也可使用自定义请求。',
}, },
{ {
key: 'auto_remove_banned_keys', key: 'auto_remove_banned_keys',

View File

@@ -3,7 +3,7 @@
<!-- Header --> <!-- Header -->
<div class="p-4 border-b border-border/60"> <div class="p-4 border-b border-border/60">
<div class="flex items-center justify-between"> <div class="flex items-center justify-between">
<div class="flex items-center gap-2"> <div class="flex flex-wrap items-center gap-2">
<h3 class="text-sm font-semibold"> <h3 class="text-sm font-semibold">
号池状态 号池状态
</h3> </h3>
@@ -21,6 +21,34 @@
> >
{{ poolStatus.total_sticky_sessions }} 个粘性会话 {{ poolStatus.total_sticky_sessions }} 个粘性会话
</Badge> </Badge>
<Badge
v-if="poolStatus && poolStatus.provider_desired_hot > 0"
variant="outline"
class="text-xs"
>
热池 {{ poolStatus.provider_hot_count }} / {{ poolStatus.provider_desired_hot }}
</Badge>
<Badge
v-if="poolStatus && poolStatus.provider_in_flight > 0"
variant="outline"
class="text-xs"
>
in-flight {{ poolStatus.provider_in_flight }}
</Badge>
<Badge
v-if="poolStatus && poolStatus.provider_desired_hot > 0"
variant="outline"
class="text-xs"
>
EMA {{ formatEmaHeat(poolStatus.provider_ema_in_flight) }}
</Badge>
<Badge
v-if="poolStatus?.provider_burst_pending"
variant="secondary"
class="text-xs"
>
补热中
</Badge>
</div> </div>
<RefreshButton <RefreshButton
:loading="refreshing" :loading="refreshing"
@@ -265,6 +293,11 @@ function formatTokens(tokens: number): string {
return String(tokens) return String(tokens)
} }
function formatEmaHeat(value: number): string {
if (!Number.isFinite(value) || value <= 0) return '0.0'
return value.toFixed(1)
}
function getCostBarColor(usage: number, limit: number): string { function getCostBarColor(usage: number, limit: number): string {
const ratio = usage / limit const ratio = usage / limit
if (ratio >= 0.9) return 'bg-red-500' if (ratio >= 0.9) return 'bg-red-500'

View File

@@ -146,6 +146,21 @@
<Plug class="w-3.5 h-3.5" /> <Plug class="w-3.5 h-3.5" />
</Button> </Button>
</div> </div>
<div
v-if="showAdaptiveHotPoolMetricsButton"
class="min-w-0 flex-1 flex justify-center"
>
<Button
variant="ghost"
size="icon"
class="h-8 w-8 shrink-0"
data-testid="pool-demand-metrics-button"
title="查看自适应热池指标"
@click="showDemandMetricsDialog = true"
>
<Activity class="w-3.5 h-3.5" />
</Button>
</div>
<div class="min-w-0 flex-1 flex justify-center"> <div class="min-w-0 flex-1 flex justify-center">
<Button <Button
variant="ghost" variant="ghost"
@@ -293,6 +308,17 @@
> >
<Plug class="w-3.5 h-3.5" /> <Plug class="w-3.5 h-3.5" />
</Button> </Button>
<Button
v-if="showAdaptiveHotPoolMetricsButton"
variant="ghost"
size="icon"
class="h-8 w-8"
data-testid="pool-demand-metrics-button"
title="查看自适应热池指标"
@click="showDemandMetricsDialog = true"
>
<Activity class="w-3.5 h-3.5" />
</Button>
<Button <Button
v-if="selectedProviderId" v-if="selectedProviderId"
variant="ghost" variant="ghost"
@@ -1393,6 +1419,11 @@
:current-claude-config="selectedProviderClaudeConfig" :current-claude-config="selectedProviderClaudeConfig"
@saved="handleSchedulingSaved" @saved="handleSchedulingSaved"
/> />
<PoolDemandMetricsDialog
v-model="showDemandMetricsDialog"
:provider-name="selectedProviderOverview?.provider_name"
:samples="providerDemandMetricSamples"
/>
<ProviderFormDialog <ProviderFormDialog
v-model="providerEditDialogOpen" v-model="providerEditDialogOpen"
:provider="providerToEdit" :provider="providerToEdit"
@@ -1451,6 +1482,7 @@ import {
Upload, Upload,
ChevronDown, ChevronDown,
RefreshCw, RefreshCw,
Activity,
Power, Power,
Database, Database,
KeyRound, KeyRound,
@@ -1533,6 +1565,7 @@ import { getProvider, getProviderEndpoints, updateProvider } from '@/api/endpoin
import { useProxyNodesStore } from '@/stores/proxy-nodes' import { useProxyNodesStore } from '@/stores/proxy-nodes'
import PoolSchedulingDialog from '@/features/pool/components/PoolSchedulingDialog.vue' import PoolSchedulingDialog from '@/features/pool/components/PoolSchedulingDialog.vue'
import PoolAdvancedDialog from '@/features/pool/components/PoolAdvancedDialog.vue' import PoolAdvancedDialog from '@/features/pool/components/PoolAdvancedDialog.vue'
import PoolDemandMetricsDialog from '@/features/pool/components/PoolDemandMetricsDialog.vue'
import PoolAccountBatchDialog from '@/features/pool/components/PoolAccountBatchDialog.vue' import PoolAccountBatchDialog from '@/features/pool/components/PoolAccountBatchDialog.vue'
import ProviderProxyPopover from '@/features/pool/components/ProviderProxyPopover.vue' import ProviderProxyPopover from '@/features/pool/components/ProviderProxyPopover.vue'
import KeyAllowedModelsEditDialog from '@/features/providers/components/KeyAllowedModelsEditDialog.vue' import KeyAllowedModelsEditDialog from '@/features/providers/components/KeyAllowedModelsEditDialog.vue'
@@ -1621,11 +1654,28 @@ let selectProviderRequestId = 0
let providerDataRequestId = 0 let providerDataRequestId = 0
let keysRequestId = 0 let keysRequestId = 0
let keysSearchDebounceTimer: number | null = null let keysSearchDebounceTimer: number | null = null
let demandMetricsPollingTimer: number | null = null
let demandMetricsRequestId = 0
let suppressFiltersWatch = false let suppressFiltersWatch = false
let hasHydratedInitialProviderSelection = false let hasHydratedInitialProviderSelection = false
const POOL_OVERVIEW_CACHE_TTL_MS = 10 * 1000 const POOL_OVERVIEW_CACHE_TTL_MS = 10 * 1000
const POOL_KEYS_CACHE_TTL_MS = 10 * 1000 const POOL_KEYS_CACHE_TTL_MS = 10 * 1000
const POOL_SCHEDULING_PRESETS_CACHE_TTL_MS = 5 * 60 * 1000 const POOL_SCHEDULING_PRESETS_CACHE_TTL_MS = 5 * 60 * 1000
const POOL_DEMAND_METRICS_SAMPLES_LIMIT = 120
const POOL_DEMAND_METRICS_POLL_INTERVAL_MS = 10 * 1000
interface PoolDemandMetricSample {
providerId: string
sampledAt: number
hotCount: number
desiredHot: number
inFlight: number
emaInFlight: number
burstPending: boolean
}
const showDemandMetricsDialog = ref(false)
const providerDemandMetricSamples = ref<PoolDemandMetricSample[]>([])
const poolKeyStatusFilterOptions: Array<{ value: PoolManagementViewState['status'], label: string }> = [ const poolKeyStatusFilterOptions: Array<{ value: PoolManagementViewState['status'], label: string }> = [
{ value: 'all', label: '全部状态' }, { value: 'all', label: '全部状态' },
{ value: 'active', label: '可调度' }, { value: 'active', label: '可调度' },
@@ -1651,9 +1701,11 @@ const poolScoreProbeStatusOptions = [
{ value: 'in_progress', label: '探测中' }, { value: 'in_progress', label: '探测中' },
] ]
async function loadOverview(options: { cacheTtlMs?: number } = {}) { async function loadOverview(options: { cacheTtlMs?: number, silent?: boolean } = {}) {
const requestId = ++overviewRequestId const requestId = ++overviewRequestId
overviewLoading.value = true if (!options.silent) {
overviewLoading.value = true
}
try { try {
const res = await getPoolOverview({ cacheTtlMs: options.cacheTtlMs ?? 0 }) const res = await getPoolOverview({ cacheTtlMs: options.cacheTtlMs ?? 0 })
if (requestId !== overviewRequestId) return if (requestId !== overviewRequestId) return
@@ -1711,9 +1763,11 @@ async function loadOverview(options: { cacheTtlMs?: number } = {}) {
} }
} catch (err) { } catch (err) {
if (requestId !== overviewRequestId) return if (requestId !== overviewRequestId) return
showError(parseApiError(err)) if (!options.silent) {
showError(parseApiError(err))
}
} finally { } finally {
if (requestId === overviewRequestId) { if (requestId === overviewRequestId && !options.silent) {
overviewLoading.value = false overviewLoading.value = false
} }
} }
@@ -1795,6 +1849,97 @@ const selectedProviderOverview = computed<PoolOverviewItem | null>(() => {
return poolProviders.value.find(item => item.provider_id === selectedId) || null return poolProviders.value.find(item => item.provider_id === selectedId) || null
}) })
const showAdaptiveHotPoolMetricsButton = computed(() => {
if (!selectedProviderId.value) return false
return selectedProviderConfig.value?.probing_enabled === true
})
function normalizeDemandMetricNumber(value: unknown): number {
const normalized = Number(value ?? 0)
if (!Number.isFinite(normalized) || normalized <= 0) return 0
return normalized
}
function buildDemandMetricSample(overview: PoolOverviewItem): PoolDemandMetricSample {
return {
providerId: overview.provider_id,
sampledAt: Date.now(),
hotCount: Math.floor(normalizeDemandMetricNumber(overview.provider_hot_count)),
desiredHot: Math.floor(normalizeDemandMetricNumber(overview.provider_desired_hot)),
inFlight: Math.floor(normalizeDemandMetricNumber(overview.provider_in_flight)),
emaInFlight: normalizeDemandMetricNumber(overview.provider_ema_in_flight),
burstPending: overview.provider_burst_pending === true,
}
}
function appendDemandMetricSample(overview: PoolOverviewItem | null): void {
if (!overview || !showDemandMetricsDialog.value || !showAdaptiveHotPoolMetricsButton.value) return
const nextSample = buildDemandMetricSample(overview)
const existing = providerDemandMetricSamples.value.filter(
sample => sample.providerId === overview.provider_id,
)
const lastSample = existing.at(-1)
if (
lastSample
&& nextSample.sampledAt - lastSample.sampledAt < 1000
&& lastSample.hotCount === nextSample.hotCount
&& lastSample.desiredHot === nextSample.desiredHot
&& lastSample.inFlight === nextSample.inFlight
&& lastSample.emaInFlight === nextSample.emaInFlight
&& lastSample.burstPending === nextSample.burstPending
) {
providerDemandMetricSamples.value = existing
return
}
providerDemandMetricSamples.value = [...existing, nextSample]
.slice(-POOL_DEMAND_METRICS_SAMPLES_LIMIT)
}
function stopDemandMetricsPolling(): void {
if (demandMetricsPollingTimer !== null) {
window.clearInterval(demandMetricsPollingTimer)
demandMetricsPollingTimer = null
}
}
async function refreshDemandMetricsOverview(): Promise<void> {
const providerId = selectedProviderId.value
if (!showDemandMetricsDialog.value || !showAdaptiveHotPoolMetricsButton.value || !providerId) {
return
}
const requestId = ++demandMetricsRequestId
try {
const res = await getPoolOverview({ cacheTtlMs: 0 })
if (
requestId !== demandMetricsRequestId
|| !showDemandMetricsDialog.value
|| selectedProviderId.value !== providerId
) {
return
}
const allProviders = Array.isArray(res.items) ? res.items : []
const enabledProviders = allProviders.filter(item => item.pool_enabled)
poolProviders.value = enabledProviders
appendDemandMetricSample(
enabledProviders.find(item => item.provider_id === providerId) || null,
)
} catch {
// 指标弹窗只做尽力刷新,失败不打断主流程。
}
}
function startDemandMetricsPolling(): void {
stopDemandMetricsPolling()
appendDemandMetricSample(selectedProviderOverview.value)
void refreshDemandMetricsOverview()
demandMetricsPollingTimer = window.setInterval(() => {
if (!showDemandMetricsDialog.value || !showAdaptiveHotPoolMetricsButton.value) return
if (document.visibilityState === 'hidden') return
void refreshDemandMetricsOverview()
}, POOL_DEMAND_METRICS_POLL_INTERVAL_MS)
}
const poolSchedulingLabel = computed(() => { const poolSchedulingLabel = computed(() => {
if (!selectedProviderConfig.value && selectedProviderOverview.value?.pool_enabled === false) { if (!selectedProviderConfig.value && selectedProviderOverview.value?.pool_enabled === false) {
return '未启用' return '未启用'
@@ -1864,11 +2009,63 @@ const selectedProviderStatusText = computed(() => {
return '' return ''
}) })
function formatDemandEma(value: number | undefined): string {
const normalized = Number(value ?? 0)
if (!Number.isFinite(normalized) || normalized <= 0) return '0.0'
return normalized.toFixed(1)
}
const selectedProviderDemandMetaText = computed(() => {
const overview = selectedProviderOverview.value
if (!overview) return ''
const segments: string[] = []
const desiredHot = Number(overview.provider_desired_hot ?? 0)
const hotCount = Number(overview.provider_hot_count ?? 0)
const inFlight = Number(overview.provider_in_flight ?? 0)
if (Number.isFinite(desiredHot) && desiredHot > 0) {
segments.push(`热池 ${hotCount} / ${desiredHot}`)
segments.push(`EMA ${formatDemandEma(overview.provider_ema_in_flight)}`)
}
if (Number.isFinite(inFlight) && inFlight > 0) {
segments.push(`in-flight ${inFlight}`)
}
if (overview.provider_burst_pending) {
segments.push('补热中')
}
return segments.join(' | ')
})
const poolHeaderMetaText = computed(() => { const poolHeaderMetaText = computed(() => {
const providerType = selectedProviderType.value return [
const status = selectedProviderStatusText.value selectedProviderType.value,
if (providerType && status) return `${providerType} | ${status}` selectedProviderStatusText.value,
return providerType || status || '' selectedProviderDemandMetaText.value,
].filter(Boolean).join(' | ')
})
watch(showDemandMetricsDialog, (open) => {
if (open) {
startDemandMetricsPolling()
} else {
stopDemandMetricsPolling()
}
})
watch(selectedProviderId, () => {
providerDemandMetricSamples.value = []
if (showDemandMetricsDialog.value) {
appendDemandMetricSample(selectedProviderOverview.value)
}
})
watch(selectedProviderOverview, (overview) => {
appendDemandMetricSample(overview)
})
watch(showAdaptiveHotPoolMetricsButton, (enabled) => {
if (!enabled && showDemandMetricsDialog.value) {
showDemandMetricsDialog.value = false
}
}) })
const showAccountQuotaColumn = computed(() => { const showAccountQuotaColumn = computed(() => {
@@ -3955,6 +4152,7 @@ onMounted(() => {
}) })
onBeforeUnmount(() => { onBeforeUnmount(() => {
stopDemandMetricsPolling()
if (keysSearchDebounceTimer !== null) { if (keysSearchDebounceTimer !== null) {
clearTimeout(keysSearchDebounceTimer) clearTimeout(keysSearchDebounceTimer)
keysSearchDebounceTimer = null keysSearchDebounceTimer = null

View File

@@ -125,6 +125,7 @@ vi.mock('lucide-vue-next', async () => {
Upload: Icon, Upload: Icon,
ChevronDown: Icon, ChevronDown: Icon,
RefreshCw: Icon, RefreshCw: Icon,
Activity: Icon,
Power: Icon, Power: Icon,
Database: Icon, Database: Icon,
KeyRound: Icon, KeyRound: Icon,
@@ -140,6 +141,8 @@ vi.mock('lucide-vue-next', async () => {
Settings2: Icon, Settings2: Icon,
SlidersHorizontal: Icon, SlidersHorizontal: Icon,
CircleHelp: Icon, CircleHelp: Icon,
Edit: Icon,
Plug: Icon,
} }
}) })
@@ -312,6 +315,17 @@ vi.mock('@/features/pool/components/PoolAdvancedDialog.vue', async () => {
}), }),
} }
}) })
vi.mock('@/features/pool/components/PoolDemandMetricsDialog.vue', async () => {
const { defineComponent } = await import('vue')
return {
default: defineComponent({
name: 'PoolDemandMetricsDialogStub',
setup() {
return () => null
},
}),
}
})
vi.mock('@/features/pool/components/PoolAccountBatchDialog.vue', async () => { vi.mock('@/features/pool/components/PoolAccountBatchDialog.vue', async () => {
const { defineComponent } = await import('vue') const { defineComponent } = await import('vue')
return { return {
@@ -334,6 +348,28 @@ vi.mock('@/features/pool/components/ProviderProxyPopover.vue', async () => {
}), }),
} }
}) })
vi.mock('@/features/providers/components/EndpointFormDialog.vue', async () => {
const { defineComponent } = await import('vue')
return {
default: defineComponent({
name: 'EndpointFormDialogStub',
setup() {
return () => null
},
}),
}
})
vi.mock('@/features/providers/components/ProviderFormDialog.vue', async () => {
const { defineComponent } = await import('vue')
return {
default: defineComponent({
name: 'ProviderFormDialogStub',
setup() {
return () => null
},
}),
}
})
vi.mock('@/features/providers/components/KeyAllowedModelsEditDialog.vue', async () => { vi.mock('@/features/providers/components/KeyAllowedModelsEditDialog.vue', async () => {
const { defineComponent } = await import('vue') const { defineComponent } = await import('vue')
return { return {
@@ -404,7 +440,7 @@ function createOverview(providerType: string): PoolOverviewItem {
} }
} }
function createProvider(providerType: string) { function createProvider(providerType: string, overrides: Record<string, unknown> = {}) {
return { return {
id: `${providerType}-provider`, id: `${providerType}-provider`,
name: `${providerType} Provider`, name: `${providerType} Provider`,
@@ -414,6 +450,7 @@ function createProvider(providerType: string) {
proxy: null, proxy: null,
pool_advanced: null, pool_advanced: null,
claude_code_advanced: null, claude_code_advanced: null,
...overrides,
} }
} }
@@ -772,4 +809,37 @@ describe('PoolManagement Codex cycle stats mode', () => {
expect(root.textContent).toContain('3.5K') expect(root.textContent).toContain('3.5K')
expect(root.textContent).toContain('$1.25') expect(root.textContent).toContain('$1.25')
}) })
it('shows adaptive hot pool metrics entry only when probing is enabled', async () => {
endpointMocks.getPoolOverview.mockResolvedValue({
items: [{ ...createOverview('codex'), provider_desired_hot: 4, provider_in_flight: 2, provider_ema_in_flight: 1.8 }],
})
endpointMocks.listPoolKeys.mockResolvedValue(createKeyPage(createPoolKey('codex')))
endpointMocks.getProvider.mockResolvedValue(createProvider('codex', {
pool_advanced: {
probing_enabled: true,
},
}))
const enabledRoot = mountPoolManagement()
await settle()
expect(enabledRoot.querySelectorAll('[data-testid="pool-demand-metrics-button"]').length).toBeGreaterThan(0)
for (const { app, root } of mountedApps.splice(0)) {
app.unmount()
root.remove()
}
endpointMocks.getProvider.mockResolvedValue(createProvider('codex', {
pool_advanced: {
probing_enabled: false,
},
}))
const disabledRoot = mountPoolManagement()
await settle()
expect(disabledRoot.querySelector('[data-testid="pool-demand-metrics-button"]')).toBeNull()
})
}) })