chore: 提交所有工作进度 — SaaS 后端增强、Admin UI、桌面端集成
包含大量 SaaS 平台改进、Admin 管理后台更新、桌面端集成完善、 文档同步、测试文件重构等内容。为 QA 测试准备干净工作树。
This commit is contained in:
iven
320
crates/zclaw-saas/src/relay/key_pool.rs
Normal file
320
crates/zclaw-saas/src/relay/key_pool.rs
Normal file
@@ -0,0 +1,320 @@
|
||||
//! Provider Key Pool 服务
|
||||
//!
|
||||
//! 管理 provider 的多个 API Key，实现智能轮转绕过限额。
|
||||
|
||||
use sqlx::PgPool;
|
||||
use crate::error::{SaasError, SaasResult};
|
||||
use crate::crypto;
|
||||
|
||||
/// 解密 key_value (如果已加密)，否则原样返回
|
||||
fn decrypt_key_value(encrypted: &str, enc_key: &[u8; 32]) -> SaasResult<String> {
|
||||
if crypto::is_encrypted(encrypted) {
|
||||
crypto::decrypt_value(encrypted, enc_key)
|
||||
.map_err(|e| SaasError::Internal(e.to_string()))
|
||||
} else {
|
||||
// 兼容旧的明文格式
|
||||
Ok(encrypted.to_string())
|
||||
}
|
||||
}
|
||||
|
||||
/// Key Pool 中的可用 Key
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct PoolKey {
|
||||
pub id: String,
|
||||
pub key_value: String,
|
||||
pub priority: i32,
|
||||
pub max_rpm: Option<i64>,
|
||||
pub max_tpm: Option<i64>,
|
||||
pub quota_reset_interval: Option<String>,
|
||||
}
|
||||
|
||||
/// Key 选择结果
|
||||
pub struct KeySelection {
|
||||
pub key: PoolKey,
|
||||
pub key_id: String,
|
||||
}
|
||||
|
||||
/// 从 provider 的 Key Pool 中选择最佳可用 Key
|
||||
pub async fn select_best_key(db: &PgPool, provider_id: &str, enc_key: &[u8; 32]) -> SaasResult<KeySelection> {
|
||||
let now = chrono::Utc::now().to_rfc3339();
|
||||
let current_minute = chrono::Utc::now().format("%Y-%m-%dT%H:%M").to_string();
|
||||
|
||||
// 获取所有活跃 Key
|
||||
let rows: Vec<(String, String, i32, Option<i64>, Option<i64>, Option<String>)> =
|
||||
sqlx::query_as(
|
||||
"SELECT id, key_value, priority, max_rpm, max_tpm, quota_reset_interval
|
||||
FROM provider_keys
|
||||
WHERE provider_id = $1 AND is_active = TRUE AND (cooldown_until IS NULL OR cooldown_until <= $2)
|
||||
ORDER BY priority ASC"
|
||||
).bind(provider_id).bind(&now).fetch_all(db).await?;
|
||||
|
||||
if rows.is_empty() {
|
||||
// 检查是否有冷却中的 Key，返回预计等待时间
|
||||
let cooldown_row: Option<(String,)> = sqlx::query_as(
|
||||
"SELECT cooldown_until FROM provider_keys
|
||||
WHERE provider_id = $1 AND is_active = TRUE AND cooldown_until IS NOT NULL AND cooldown_until > $2
|
||||
ORDER BY cooldown_until ASC
|
||||
LIMIT 1"
|
||||
).bind(provider_id).bind(&now).fetch_optional(db).await?;
|
||||
|
||||
if let Some((earliest,)) = cooldown_row {
|
||||
// 尝试解析时间差
|
||||
let wait_secs = parse_cooldown_remaining(&earliest, &now);
|
||||
return Err(SaasError::RateLimited(
|
||||
format!("所有 Key 均在冷却中，预计 {} 秒后可用", wait_secs)
|
||||
));
|
||||
}
|
||||
|
||||
// 检查 provider 级别的单 Key
|
||||
let provider_key: Option<String> = sqlx::query_scalar(
|
||||
"SELECT api_key FROM providers WHERE id = $1"
|
||||
).bind(provider_id).fetch_optional(db).await?.flatten();
|
||||
|
||||
if let Some(key) = provider_key {
|
||||
let decrypted = decrypt_key_value(&key, enc_key)?;
|
||||
return Ok(KeySelection {
|
||||
key: PoolKey {
|
||||
id: "provider-fallback".to_string(),
|
||||
key_value: decrypted,
|
||||
priority: 0,
|
||||
max_rpm: None,
|
||||
max_tpm: None,
|
||||
quota_reset_interval: None,
|
||||
},
|
||||
key_id: "provider-fallback".to_string(),
|
||||
});
|
||||
}
|
||||
|
||||
return Err(SaasError::NotFound(format!("Provider {} 没有可用的 API Key", provider_id)));
|
||||
}
|
||||
|
||||
// 检查滑动窗口使用量
|
||||
for (id, key_value, priority, max_rpm, max_tpm, quota_reset_interval) in rows {
|
||||
// 检查 RPM 限额
|
||||
if let Some(rpm_limit) = max_rpm {
|
||||
if rpm_limit > 0 {
|
||||
let window: Option<(i64,)> = sqlx::query_as(
|
||||
"SELECT COALESCE(SUM(request_count), 0) FROM key_usage_window
|
||||
WHERE key_id = $1 AND window_minute = $2"
|
||||
).bind(&id).bind(¤t_minute).fetch_optional(db).await?;
|
||||
|
||||
if let Some((count,)) = window {
|
||||
if count >= rpm_limit {
|
||||
tracing::debug!("Key {} hit RPM limit ({}/{})", id, count, rpm_limit);
|
||||
continue;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// 检查 TPM 限额
|
||||
if let Some(tpm_limit) = max_tpm {
|
||||
if tpm_limit > 0 {
|
||||
let window: Option<(i64,)> = sqlx::query_as(
|
||||
"SELECT COALESCE(SUM(token_count), 0) FROM key_usage_window
|
||||
WHERE key_id = $1 AND window_minute = $2"
|
||||
).bind(&id).bind(¤t_minute).fetch_optional(db).await?;
|
||||
|
||||
if let Some((tokens,)) = window {
|
||||
if tokens >= tpm_limit {
|
||||
tracing::debug!("Key {} hit TPM limit ({}/{})", id, tokens, tpm_limit);
|
||||
continue;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// 此 Key 可用 — 解密 key_value
|
||||
let decrypted_kv = decrypt_key_value(&key_value, enc_key)?;
|
||||
return Ok(KeySelection {
|
||||
key: PoolKey {
|
||||
id: id.clone(),
|
||||
key_value: decrypted_kv,
|
||||
priority,
|
||||
max_rpm,
|
||||
max_tpm,
|
||||
quota_reset_interval,
|
||||
},
|
||||
key_id: id,
|
||||
});
|
||||
}
|
||||
|
||||
// 所有 Key 都超限，回退到 provider 单 Key
|
||||
let provider_key: Option<String> = sqlx::query_scalar(
|
||||
"SELECT api_key FROM providers WHERE id = $1"
|
||||
).bind(provider_id).fetch_optional(db).await?.flatten();
|
||||
|
||||
if let Some(key) = provider_key {
|
||||
let decrypted = decrypt_key_value(&key, enc_key)?;
|
||||
return Ok(KeySelection {
|
||||
key: PoolKey {
|
||||
id: "provider-fallback".to_string(),
|
||||
key_value: decrypted,
|
||||
priority: 0,
|
||||
max_rpm: None,
|
||||
max_tpm: None,
|
||||
quota_reset_interval: None,
|
||||
},
|
||||
key_id: "provider-fallback".to_string(),
|
||||
});
|
||||
}
|
||||
|
||||
Err(SaasError::RateLimited(
|
||||
format!("Provider {} 所有 Key 均已达限额", provider_id)
|
||||
))
|
||||
}
|
||||
|
||||
/// 记录 Key 使用量（滑动窗口）
|
||||
pub async fn record_key_usage(
|
||||
db: &PgPool,
|
||||
key_id: &str,
|
||||
tokens: Option<i64>,
|
||||
) -> SaasResult<()> {
|
||||
let current_minute = chrono::Utc::now().format("%Y-%m-%dT%H:%M").to_string();
|
||||
|
||||
sqlx::query(
|
||||
"INSERT INTO key_usage_window (key_id, window_minute, request_count, token_count)
|
||||
VALUES ($1, $2, 1, $3)
|
||||
ON CONFLICT (key_id, window_minute) DO UPDATE
|
||||
SET request_count = key_usage_window.request_count + 1,
|
||||
token_count = key_usage_window.token_count + $3"
|
||||
)
|
||||
.bind(key_id).bind(¤t_minute).bind(tokens.unwrap_or(0))
|
||||
.execute(db).await?;
|
||||
|
||||
// 更新 Key 的累计统计
|
||||
sqlx::query(
|
||||
"UPDATE provider_keys SET total_requests = total_requests + 1, total_tokens = total_tokens + COALESCE($1, 0), updated_at = $2
|
||||
WHERE id = $3"
|
||||
)
|
||||
.bind(tokens).bind(&chrono::Utc::now().to_rfc3339()).bind(key_id)
|
||||
.execute(db).await?;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// 标记 Key 收到 429，设置冷却期
|
||||
pub async fn mark_key_429(
|
||||
db: &PgPool,
|
||||
key_id: &str,
|
||||
retry_after_seconds: Option<u64>,
|
||||
) -> SaasResult<()> {
|
||||
let cooldown = if let Some(secs) = retry_after_seconds {
|
||||
(chrono::Utc::now() + chrono::Duration::seconds(secs as i64)).to_rfc3339()
|
||||
} else {
|
||||
// 默认 5 分钟冷却
|
||||
(chrono::Utc::now() + chrono::Duration::minutes(5)).to_rfc3339()
|
||||
};
|
||||
|
||||
let now = chrono::Utc::now().to_rfc3339();
|
||||
|
||||
sqlx::query(
|
||||
"UPDATE provider_keys SET last_429_at = $1, cooldown_until = $2, updated_at = $3
|
||||
WHERE id = $4"
|
||||
)
|
||||
.bind(&now).bind(&cooldown).bind(&now).bind(key_id)
|
||||
.execute(db).await?;
|
||||
|
||||
tracing::warn!(
|
||||
"Key {} 收到 429，冷却至 {}",
|
||||
key_id,
|
||||
cooldown
|
||||
);
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// 获取 provider 的所有 Key（管理用）
|
||||
pub async fn list_provider_keys(
|
||||
db: &PgPool,
|
||||
provider_id: &str,
|
||||
) -> SaasResult<Vec<serde_json::Value>> {
|
||||
let rows: Vec<(String, String, String, i32, Option<i64>, Option<i64>, Option<String>, bool, Option<String>, Option<String>, i64, i64, String, String)> =
|
||||
sqlx::query_as(
|
||||
"SELECT id, provider_id, key_label, priority, max_rpm, max_tpm, quota_reset_interval, is_active,
|
||||
last_429_at, cooldown_until, total_requests, total_tokens, created_at, updated_at
|
||||
FROM provider_keys WHERE provider_id = $1 ORDER BY priority ASC"
|
||||
).bind(provider_id).fetch_all(db).await?;
|
||||
|
||||
Ok(rows.into_iter().map(|r| {
|
||||
serde_json::json!({
|
||||
"id": r.0,
|
||||
"provider_id": r.1,
|
||||
"key_label": r.2,
|
||||
"priority": r.3,
|
||||
"max_rpm": r.4,
|
||||
"max_tpm": r.5,
|
||||
"quota_reset_interval": r.6,
|
||||
"is_active": r.7,
|
||||
"last_429_at": r.8,
|
||||
"cooldown_until": r.9,
|
||||
"total_requests": r.10,
|
||||
"total_tokens": r.11,
|
||||
"created_at": r.12,
|
||||
"updated_at": r.13,
|
||||
})
|
||||
}).collect())
|
||||
}
|
||||
|
||||
/// 添加 Key 到 Pool
|
||||
pub async fn add_provider_key(
|
||||
db: &PgPool,
|
||||
provider_id: &str,
|
||||
key_label: &str,
|
||||
key_value: &str,
|
||||
priority: i32,
|
||||
max_rpm: Option<i64>,
|
||||
max_tpm: Option<i64>,
|
||||
quota_reset_interval: Option<&str>,
|
||||
) -> SaasResult<String> {
|
||||
let id = uuid::Uuid::new_v4().to_string();
|
||||
let now = chrono::Utc::now().to_rfc3339();
|
||||
|
||||
sqlx::query(
|
||||
"INSERT INTO provider_keys (id, provider_id, key_label, key_value, priority, max_rpm, max_tpm, quota_reset_interval, is_active, total_requests, total_tokens, created_at, updated_at)
|
||||
VALUES ($1, $2, $3, $4, $5, $6, $7, $8, TRUE, 0, 0, $9, $9)"
|
||||
)
|
||||
.bind(&id).bind(provider_id).bind(key_label).bind(key_value)
|
||||
.bind(priority).bind(max_rpm).bind(max_tpm).bind(quota_reset_interval).bind(&now)
|
||||
.execute(db).await?;
|
||||
|
||||
tracing::info!("Added key '{}' to provider {}", key_label, provider_id);
|
||||
Ok(id)
|
||||
}
|
||||
|
||||
/// 切换 Key 活跃状态
|
||||
pub async fn toggle_key_active(
|
||||
db: &PgPool,
|
||||
key_id: &str,
|
||||
active: bool,
|
||||
) -> SaasResult<()> {
|
||||
let now = chrono::Utc::now().to_rfc3339();
|
||||
sqlx::query(
|
||||
"UPDATE provider_keys SET is_active = $1, updated_at = $2 WHERE id = $3"
|
||||
).bind(active).bind(&now).bind(key_id).execute(db).await?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// 删除 Key
|
||||
pub async fn delete_provider_key(
|
||||
db: &PgPool,
|
||||
key_id: &str,
|
||||
) -> SaasResult<()> {
|
||||
sqlx::query("DELETE FROM provider_keys WHERE id = $1")
|
||||
.bind(key_id).execute(db).await?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// 解析冷却剩余时间（秒）
|
||||
fn parse_cooldown_remaining(cooldown_until: &str, now: &str) -> i64 {
|
||||
let cooldown = chrono::DateTime::parse_from_rfc3339(cooldown_until);
|
||||
let current = chrono::DateTime::parse_from_rfc3339(now);
|
||||
|
||||
match (cooldown, current) {
|
||||
(Ok(c), Ok(n)) => {
|
||||
let diff = c.signed_duration_since(n);
|
||||
diff.num_seconds().max(0)
|
||||
}
|
||||
_ => 300, // 默认 5 分钟
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user