fn json_u64(value: &serde_json::Value) -> Option { value .as_u64() .or_else(|| value.as_i64().and_then(|raw| u64::try_from(raw).ok())) } fn admin_provider_pool_lru_enabled( raw_pool_advanced: &serde_json::Map, ) -> bool { if let Some(explicit) = raw_pool_advanced .get("lru_enabled") .and_then(serde_json::Value::as_bool) { return explicit; } let Some(presets) = raw_pool_advanced .get("scheduling_presets") .and_then(serde_json::Value::as_array) else { return false; }; let Some(first) = presets.first() else { return false; }; if first.is_string() { return raw_pool_advanced .get("lru_enabled") .and_then(serde_json::Value::as_bool) .unwrap_or(true); } presets .iter() .filter_map(serde_json::Value::as_object) .any(|item| { item.get("preset") .and_then(serde_json::Value::as_str) .is_some_and(|preset| preset.eq_ignore_ascii_case("lru")) && item .get("enabled") .and_then(serde_json::Value::as_bool) .unwrap_or(true) }) } fn admin_provider_pool_config( provider: &StoredProviderCatalogProvider, ) -> Option { let raw_pool_advanced = provider .config .as_ref() .and_then(serde_json::Value::as_object) .and_then(|config| config.get("pool_advanced"))?; let Some(pool_advanced) = raw_pool_advanced.as_object() else { return Some(AdminProviderPoolConfig { lru_enabled: false, cost_window_seconds: 18_000, cost_limit_per_key_tokens: None, }); }; Some(AdminProviderPoolConfig { lru_enabled: admin_provider_pool_lru_enabled(pool_advanced), cost_window_seconds: pool_advanced .get("cost_window_seconds") .and_then(json_u64) .filter(|value| *value > 0) .unwrap_or(18_000), cost_limit_per_key_tokens: pool_advanced .get("cost_limit_per_key_tokens") .and_then(json_u64), }) } fn pool_sticky_pattern(keyspace: &RedisKeyspace, provider_id: &str) -> String { keyspace.key(&format!("ap:{provider_id}:sticky:*")) } fn pool_lru_key(keyspace: &RedisKeyspace, provider_id: &str) -> String { keyspace.key(&format!("ap:{provider_id}:lru")) } fn pool_cooldown_key(keyspace: &RedisKeyspace, provider_id: &str, key_id: &str) -> String { keyspace.key(&format!("ap:{provider_id}:cooldown:{key_id}")) } fn pool_cooldown_index_key(keyspace: &RedisKeyspace, provider_id: &str) -> String { keyspace.key(&format!("ap:{provider_id}:cooldown_idx")) } fn pool_cost_key(keyspace: &RedisKeyspace, provider_id: &str, key_id: &str) -> String { keyspace.key(&format!("ap:{provider_id}:cost:{key_id}")) } fn parse_pool_cost_member(member: &str) -> u64 { member .rsplit_once(':') .and_then(|(_, suffix)| suffix.parse::().ok()) .unwrap_or(0) } async fn scan_redis_keys( connection: &mut redis::aio::MultiplexedConnection, pattern: &str, ) -> Result, GatewayError> { let mut cursor = 0u64; let mut keys = Vec::new(); loop { let (next_cursor, batch): (u64, Vec) = redis::cmd("SCAN") .arg(cursor) .arg("MATCH") .arg(pattern) .arg("COUNT") .arg(ADMIN_PROVIDER_POOL_SCAN_BATCH) .query_async(connection) .await .map_err(|err| GatewayError::Internal(err.to_string()))?; keys.extend(batch); if next_cursor == 0 { break; } cursor = next_cursor; } Ok(keys) } async fn read_admin_provider_pool_runtime_state( runner: &RedisKvRunner, provider_id: &str, key_ids: &[String], pool_config: AdminProviderPoolConfig, ) -> AdminProviderPoolRuntimeState { let mut runtime = AdminProviderPoolRuntimeState::default(); let Ok(mut connection) = runner.client().get_multiplexed_async_connection().await else { warn!("gateway admin provider pool: failed to connect redis for provider {provider_id}"); return runtime; }; let keyspace = runner.keyspace().clone(); let sticky_keys = match scan_redis_keys( &mut connection, &pool_sticky_pattern(&keyspace, provider_id), ) .await { Ok(keys) => keys, Err(err) => { warn!( "gateway admin provider pool: failed to scan sticky keys for provider {provider_id}: {:?}", err ); Vec::new() } }; runtime.total_sticky_sessions = sticky_keys.len(); if !sticky_keys.is_empty() { for chunk in sticky_keys.chunks(ADMIN_PROVIDER_POOL_SCAN_BATCH as usize) { let values = redis::cmd("MGET") .arg(chunk) .query_async::>>(&mut connection) .await; let Ok(values) = values else { warn!("gateway admin provider pool: failed to read sticky bindings for provider {provider_id}"); break; }; for bound_key_id in values.into_iter().flatten() { *runtime .sticky_sessions_by_key .entry(bound_key_id) .or_insert(0) += 1; } } } for key_id in key_ids { let cooldown_key = pool_cooldown_key(&keyspace, provider_id, key_id); let reason = redis::cmd("GET") .arg(&cooldown_key) .query_async::>(&mut connection) .await; if let Ok(Some(reason)) = reason { runtime .cooldown_reason_by_key .insert(key_id.clone(), reason); let ttl = redis::cmd("TTL") .arg(&cooldown_key) .query_async::(&mut connection) .await .ok() .and_then(|value| u64::try_from(value).ok()); if let Some(ttl_seconds) = ttl.filter(|value| *value > 0) { runtime .cooldown_ttl_by_key .insert(key_id.clone(), ttl_seconds); } } let window_start = SystemTime::now() .duration_since(UNIX_EPOCH) .unwrap_or_default() .as_secs() .saturating_sub(pool_config.cost_window_seconds); let members = redis::cmd("ZRANGEBYSCORE") .arg(pool_cost_key(&keyspace, provider_id, key_id)) .arg(window_start) .arg("+inf") .query_async::>(&mut connection) .await .unwrap_or_default(); let total = members .iter() .map(|member| parse_pool_cost_member(member)) .sum::(); runtime .cost_window_usage_by_key .insert(key_id.clone(), total); } if pool_config.lru_enabled && !key_ids.is_empty() { let mut command = redis::cmd("ZMSCORE"); command.arg(pool_lru_key(&keyspace, provider_id)); for key_id in key_ids { command.arg(key_id); } if let Ok(scores) = command .query_async::>>(&mut connection) .await { for (key_id, score) in key_ids.iter().zip(scores.into_iter()) { if let Some(score) = score { runtime.lru_score_by_key.insert(key_id.clone(), score); } } } } runtime } async fn read_admin_provider_pool_cooldown_count( runner: &RedisKvRunner, provider_id: &str, ) -> usize { let Ok(mut connection) = runner.client().get_multiplexed_async_connection().await else { warn!("gateway admin provider pool: failed to connect redis for provider {provider_id}"); return 0; }; let keyspace = runner.keyspace().clone(); redis::cmd("SCARD") .arg(pool_cooldown_index_key(&keyspace, provider_id)) .query_async::(&mut connection) .await .map(|value| value as usize) .unwrap_or(0) } async fn read_admin_provider_pool_cooldown_key_ids( runner: &RedisKvRunner, provider_id: &str, ) -> Vec { let Ok(mut connection) = runner.client().get_multiplexed_async_connection().await else { warn!("gateway admin provider pool: failed to connect redis for provider {provider_id}"); return Vec::new(); }; let keyspace = runner.keyspace().clone(); redis::cmd("SMEMBERS") .arg(pool_cooldown_index_key(&keyspace, provider_id)) .query_async::>(&mut connection) .await .unwrap_or_default() } async fn build_admin_provider_pool_status_payload( state: &AppState, provider_id: &str, ) -> Option { if !state.has_provider_catalog_data_reader() { return None; } let provider = state .read_provider_catalog_providers_by_ids(&[provider_id.to_string()]) .await .ok() .and_then(|mut providers| providers.drain(..).next())?; let Some(pool_config) = admin_provider_pool_config(&provider) else { return Some(json!({ "provider_id": provider.id, "provider_name": provider.name, "pool_enabled": false, "total_keys": 0, "total_sticky_sessions": 0, "keys": [], })); }; let keys = state .list_provider_catalog_keys_by_provider_ids(std::slice::from_ref(&provider.id)) .await .ok() .unwrap_or_default(); let key_ids = keys.iter().map(|key| key.id.clone()).collect::>(); let runtime = match state.redis_kv_runner() { Some(runner) => { read_admin_provider_pool_runtime_state(&runner, &provider.id, &key_ids, pool_config) .await } None => AdminProviderPoolRuntimeState::default(), }; let key_payloads = keys .into_iter() .map(|key| { let cooldown_reason = runtime.cooldown_reason_by_key.get(&key.id).cloned(); json!({ "key_id": key.id, "key_name": key.name, "is_active": key.is_active, "cooldown_reason": cooldown_reason, "cooldown_ttl_seconds": cooldown_reason .as_ref() .and_then(|_| runtime.cooldown_ttl_by_key.get(&key.id).copied()), "cost_window_usage": runtime.cost_window_usage_by_key.get(&key.id).copied().unwrap_or(0), "cost_limit": pool_config.cost_limit_per_key_tokens, "sticky_sessions": runtime.sticky_sessions_by_key.get(&key.id).copied().unwrap_or(0), "lru_score": runtime.lru_score_by_key.get(&key.id).copied(), }) }) .collect::>(); Some(json!({ "provider_id": provider.id, "provider_name": provider.name, "pool_enabled": true, "total_keys": key_payloads.len(), "total_sticky_sessions": runtime.total_sticky_sessions, "keys": key_payloads, })) } async fn clear_admin_provider_pool_cooldown(state: &AppState, provider_id: &str, key_id: &str) { let Some(runner) = state.redis_kv_runner() else { return; }; let Ok(mut connection) = runner.client().get_multiplexed_async_connection().await else { warn!("gateway admin provider pool: failed to connect redis to clear cooldown for key {key_id}"); return; }; let keyspace = runner.keyspace().clone(); let _: Result<(), _> = redis::pipe() .cmd("DEL") .arg(pool_cooldown_key(&keyspace, provider_id, key_id)) .ignore() .cmd("SREM") .arg(pool_cooldown_index_key(&keyspace, provider_id)) .arg(key_id) .ignore() .query_async(&mut connection) .await; } async fn reset_admin_provider_pool_cost(state: &AppState, provider_id: &str, key_id: &str) { let Some(runner) = state.redis_kv_runner() else { return; }; let _ = runner.del(&format!("ap:{provider_id}:cost:{key_id}")).await; }