|
| 1 | +//! AI 余额查询（REQ-139，v0.8.0 M1 使能层）。 |
| 2 | +//! |
| 3 | +//! @ai-context: GET {base_url}/user/balance（Bearer 密钥）→ SiliconFlow 返回 |
| 4 | +//! {total_balance, grants_balance, topped_up_balance, currency}。 |
| 5 | +//! 超时+重试沿用适配器模式（ai_text_filter 先例：429/5xx/传输 |
| 6 | +//! 错误指数退避，4xx 非瞬态不重试）。 |
| 7 | +//! @ai-context: 解析容错（防御性编程）：字段缺失/非法数值按 0 处理并记录 |
| 8 | +//! 缺失项（不 panic、不整单丢弃——余额展示尽力而为）； |
| 9 | +//! currency 缺失默认 "CNY"。 |
| 10 | +//! @ai-context: 本模块只做余额数据与查询，不 gate 授权（查询为读操作， |
| 11 | +//! 不涉及内容上传——门控见 ai_settings.rs content_gate）。 |
| 12 | +
|
| 13 | +use serde::{Deserialize, Serialize}; |
| 14 | + |
| 15 | +/// 余额快照（SiliconFlow /v1/user/balance 响应；字段缺失容错解析）。 |
| 16 | +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] |
| 17 | +#[serde(rename_all = "camelCase")] |
| 18 | +pub struct AiBalance { |
| 19 | + pub total_balance: f64, |
| 20 | + pub grants_balance: f64, |
| 21 | + pub topped_up_balance: f64, |
| 22 | + pub currency: String, |
| 23 | +} |
| 24 | + |
| 25 | +/// 解析余额响应（纯函数；容错：字段缺失/非法 → 0，currency 缺失 → "CNY"）。 |
| 26 | +/// |
| 27 | +/// @ai-context: 返回 Err 仅当整体非 JSON 对象（无法尽力而为）； |
| 28 | +/// 字段缺失按 0 处理——SiliconFlow 分项字段可能因账户类型缺失。 |
| 29 | +pub fn parse_balance(body: &str) -> Result<AiBalance, String> { |
| 30 | + let v: serde_json::Value = |
| 31 | + serde_json::from_str(body).map_err(|e| format!("余额响应 JSON 解析失败: {}", e))?; |
| 32 | + let obj = v |
| 33 | + .as_object() |
| 34 | + .ok_or_else(|| "余额响应不是 JSON 对象".to_string())?; |
| 35 | + let num = |k: &str| -> f64 { |
| 36 | + obj.get(k).and_then(|x| x.as_f64()).or_else(|| { |
| 37 | + obj.get(k).and_then(|x| x.as_str()).and_then(|s| s.parse().ok()) |
| 38 | + }).unwrap_or(0.0) |
| 39 | + }; |
| 40 | + let currency = obj |
| 41 | + .get("currency") |
| 42 | + .and_then(|c| c.as_str()) |
| 43 | + .filter(|s| !s.is_empty()) |
| 44 | + .unwrap_or("CNY") |
| 45 | + .to_string(); |
| 46 | + Ok(AiBalance { |
| 47 | + total_balance: num("total_balance"), |
| 48 | + grants_balance: num("grants_balance"), |
| 49 | + topped_up_balance: num("topped_up_balance"), |
| 50 | + currency, |
| 51 | + }) |
| 52 | +} |
| 53 | + |
| 54 | +/// 低余额提醒（纯函数：total < 阈值 → Some 提示；负数余额也触发）。 |
| 55 | +pub fn low_balance_warning(balance: &AiBalance, threshold: f64) -> Option<String> { |
| 56 | + if balance.total_balance < threshold { |
| 57 | + Some(format!( |
| 58 | + "余额不足：当前 {:.2} {}（低于提醒阈值 {:.2}），AI 生成可能中断，请及时充值", |
| 59 | + balance.total_balance, balance.currency, threshold |
| 60 | + )) |
| 61 | + } else { |
| 62 | + None |
| 63 | + } |
| 64 | +} |
| 65 | + |
| 66 | +/// 余额查询适配器（阻塞调用——command 层 spawn_blocking 包裹）。 |
| 67 | +#[derive(Debug, Clone)] |
| 68 | +pub struct AiBalanceAdapter { |
| 69 | + pub base_url: String, |
| 70 | + pub api_key: String, |
| 71 | + pub timeout_secs: u64, |
| 72 | + pub max_retries: u32, |
| 73 | +} |
| 74 | + |
| 75 | +impl AiBalanceAdapter { |
| 76 | + /// 查询余额（HTTP 429/5xx/传输错误指数退避重试；4xx 非瞬态不重试）。 |
| 77 | + /// |
| 78 | + /// @ai-context: 401/403 → 明确"密钥无效"错误（REQ-138 一键连通性验证 |
| 79 | + /// 复用本函数——错误密钥验证失败有明确提示）。 |
| 80 | + pub fn fetch(&self) -> Result<AiBalance, String> { |
| 81 | + if self.api_key.trim().is_empty() { |
| 82 | + return Err("未配置 API 密钥（设置页保存密钥或配置环境变量 SILICONFLOW_API_KEY）".to_string()); |
| 83 | + } |
| 84 | + let url = format!( |
| 85 | + "{}/user/balance", |
| 86 | + self.base_url.trim_end_matches('/').trim_end_matches("/v1") |
| 87 | + ); |
| 88 | + let agent = ureq::AgentBuilder::new() |
| 89 | + .timeout(std::time::Duration::from_secs(self.timeout_secs.max(5))) |
| 90 | + .build(); |
| 91 | + let mut last_err = "未发起请求".to_string(); |
| 92 | + for attempt in 0..=self.max_retries { |
| 93 | + if attempt > 0 { |
| 94 | + let backoff = (500u64 << attempt.min(4)).min(8000); |
| 95 | + std::thread::sleep(std::time::Duration::from_millis(backoff)); |
| 96 | + } |
| 97 | + match agent |
| 98 | + .get(&url) |
| 99 | + .set("Authorization", &format!("Bearer {}", self.api_key.trim())) |
| 100 | + .call() |
| 101 | + { |
| 102 | + Ok(resp) => { |
| 103 | + let body = resp |
| 104 | + .into_string() |
| 105 | + .map_err(|e| format!("读取余额响应失败: {}", e))?; |
| 106 | + return parse_balance(&body); |
| 107 | + } |
| 108 | + Err(ureq::Error::Status(401 | 403, _)) => { |
| 109 | + return Err("API 密钥无效或无权限（请检查设置页密钥）".to_string()); |
| 110 | + } |
| 111 | + Err(ureq::Error::Status(429, _)) => { |
| 112 | + last_err = "请求过频（HTTP 429）".to_string(); |
| 113 | + if attempt == self.max_retries { |
| 114 | + break; |
| 115 | + } |
| 116 | + } |
| 117 | + Err(ureq::Error::Status(code, _)) if code >= 500 => { |
| 118 | + last_err = format!("服务端错误 HTTP {}", code); |
| 119 | + if attempt == self.max_retries { |
| 120 | + break; |
| 121 | + } |
| 122 | + } |
| 123 | + Err(ureq::Error::Status(code, _)) => { |
| 124 | + return Err(format!("余额查询被拒绝 HTTP {}（不重试——4xx 非瞬态）", code)); |
| 125 | + } |
| 126 | + Err(e) => { |
| 127 | + last_err = format!("传输错误: {}", e); |
| 128 | + if attempt == self.max_retries { |
| 129 | + break; |
| 130 | + } |
| 131 | + } |
| 132 | + } |
| 133 | + } |
| 134 | + Err(format!("重试 {} 次后仍失败: {}", self.max_retries, last_err)) |
| 135 | + } |
| 136 | +} |
| 137 | + |
| 138 | +/// 单测独立文件（保持本文件 ≤300 行，AGENTS.md §3）。 |
| 139 | +#[cfg(test)] |
| 140 | +#[path = "ai_balance_tests.rs"] |
| 141 | +mod tests; |
0 commit comments