diff --git a/crates/tw-api/src/lib.rs b/crates/tw-api/src/lib.rs index 5392a914..03386352 100644 --- a/crates/tw-api/src/lib.rs +++ b/crates/tw-api/src/lib.rs @@ -597,7 +597,15 @@ pub const MSG_CODES: &str = include_str!("../msg-codes.txt"); /// 同一版起**记录说得出请求带没带 DeepSeek Harness 的会话日志**：`RequestStarted` 和 /// [`HistoryRow`] 多了 `session_log_bytes`（`dsh_session_log` 的字节数，没带的没有）。 /// 照 22 写的界面看不到它。 -pub const CONTROL_API_VERSION: u32 = 23; +/// +/// **24 起 GLM 的额度不再数 MCP 工具的调用**：老套餐额度接口里 `TIME_LIMIT` 那一项数的是 +/// Z.ai 自家 MCP 工具（search-prime、web-reader、zread）的调用次数，这些调用不经过网关， +/// core 不再读它 —— 窗口的词表里没有 `monthly` 了，模型请求上的 1310 只标 `weekly`， +/// [`QuotaWindow::credits`] 只有积分制套餐（`CREDIT_LIMIT`）的窗口才有。同一版起 +/// **额度撤下也有事件**：GLM 的 key 被判定没有套餐、core 清掉这一家的额度时，报一条 +/// `windows` 为空的 [`Event::QuotaSeen`]，界面当场收起那一格，不用等下一次读 `/quota`。 +/// 照 23 写的界面不知道空的 `QuotaSeen` 是撤下，还给一个再也不会来的 `monthly` 留着位置。 +pub const CONTROL_API_VERSION: u32 = 24; #[derive(Debug, Clone, Serialize, Deserialize)] #[cfg_attr(feature = "ts", derive(ts_rs::TS))] @@ -1084,10 +1092,14 @@ pub enum Event { status: Option, at_ms: u64, }, - /// 上游在响应头里报了订阅额度。 + /// 一家上游的订阅额度：**这一家此刻的全部窗口**，收到就整个换掉，不和之前的合并。 + /// + /// 大多来自响应头，**零成本**：不发额外请求，顺着真实流量白捡。按量付费的账号没有 + /// 这些头，那时这个事件根本不会出现 —— 而不是报一个「用了 0%」。GLM Coding Plan 和 + /// ChatGPT 账号的额度也会去问账号的接口。 /// - /// **零成本**：不发额外请求，顺着真实流量白捡。按量付费的账号没有 - /// 这些头，那时这个事件根本不会出现 —— 而不是报一个「用了 0%」。 + /// **`windows` 是空的：这一家的额度撤下了**（GLM 的 key 被判定没有套餐）。之前报过的 + /// 窗口都不再作数，`/quota` 里也没有它了 —— 那一格收起来，不是「用了 0%」。 QuotaSeen { id: u64, provider: String, @@ -1256,8 +1268,7 @@ pub struct RoutingView { #[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] #[cfg_attr(feature = "ts", derive(ts_rs::TS))] pub struct QuotaWindow { - /// `5h` / `7d`（Anthropic）/ `weekly`（Codex、GLM）/ `monthly`（GLM 老套餐每月的 - /// MCP 调用次数） + /// `5h` / `7d`（Anthropic）/ `weekly`（Codex、GLM） pub window: String, pub used_percent: f64, /// 什么时候重置，Unix 毫秒。**是时刻，不是「还有多少秒」**：上游报的秒数 @@ -1267,12 +1278,15 @@ pub struct QuotaWindow { pub resets_at_ms: Option, #[serde(default, skip_serializing_if = "Option::is_none")] pub status: Option, - /// 积分制套餐（GLM Coding Plan）这个窗口的积分。别的套餐没有 + /// 这个窗口的积分。**只有 GLM Coding Plan 积分制套餐的窗口有**（额度接口里的 + /// `CREDIT_LIMIT`）；按 token 算的老套餐、Anthropic、Codex 的窗口都没有 #[serde(default, skip_serializing_if = "Option::is_none")] pub credits: Option, } -/// 一个额度窗口的积分，三个数都是上游给的原数。 +/// 积分制套餐一个额度窗口的积分：GLM Coding Plan 额度接口里 `CREDIT_LIMIT` 那一项的 +/// 总额（`usage`）、已用（`currentValue`）、剩余（`remaining`）。**只有这一种**，别的 +/// 窗口不带它。三个数都是上游给的原数。 /// /// **剩余不是总额减已用算出来的**：上游给的三个数不一定对得上（实测总额 2000、已用 23、 /// 剩余 1976），界面要显示剩余就显示它说的剩余。 diff --git a/crates/tw-control/src/lib.rs b/crates/tw-control/src/lib.rs index 46db48ed..9d563330 100644 --- a/crates/tw-control/src/lib.rs +++ b/crates/tw-control/src/lib.rs @@ -986,7 +986,7 @@ async fn request_detail( /// 金额还是百分比，两种人格共用同一块地方。 /// /// GLM Coding Plan 的额度不在响应头里：**界面来要时去问**（60 秒内合成一次），问得慢的 -/// 不等，问完由 `QuotaSeen` 补上 +/// 不等，问完由 `QuotaSeen` 补上；判定没有套餐的，由一条窗口为空的 `QuotaSeen` 撤下 async fn quota(State(s): State) -> Json> { s.gateway.refresh_glm_quotas(GLM_QUOTA_WAIT).await; let mut out: Vec = s diff --git a/crates/tw-gateway/src/glm.rs b/crates/tw-gateway/src/glm.rs index 6baa79c7..2699f891 100644 --- a/crates/tw-gateway/src/glm.rs +++ b/crates/tw-gateway/src/glm.rs @@ -135,21 +135,30 @@ const CLOCK_SLACK_MS: u64 = MINUTE_MS; /// **按 `unit` / `number` 分，不按 `nextResetTime` 的先后猜**：周期末尾每周窗口会比 /// 5 小时窗口先重置，按先后猜就把两个标反了。认不出来的一项不要 —— 没有名字的百分比 /// 放上界面只会被当成另一个窗口。 +/// +/// **只认 token（老套餐）和积分（积分制）两种。**老套餐还有一项 `TIME_LIMIT`，数的是 +/// Z.ai 自家 MCP 工具（search-prime、web-reader、zread）的调用次数。这些调用不经过网关， +/// 这一项整个不要：不报、不显示，用满了也不算用完 fn window_of(item: &Value) -> Option<(&'static str, u64)> { let kind = item["type"].as_str().unwrap_or_default(); - // 5 小时和每周是 token（老套餐）或积分（积分制）；每月是老套餐的 MCP 调用次数 - let usage = - kind.eq_ignore_ascii_case("TOKENS_LIMIT") || kind.eq_ignore_ascii_case("CREDIT_LIMIT"); - let calls = kind.eq_ignore_ascii_case("TIME_LIMIT"); + if !(kind.eq_ignore_ascii_case("TOKENS_LIMIT") || kind.eq_ignore_ascii_case("CREDIT_LIMIT")) { + return None; + } match (item["unit"].as_i64(), item["number"].as_i64()) { - (Some(3), Some(5)) if usage => Some(("5h", 5 * HOUR_MS)), + (Some(3), Some(5)) => Some(("5h", 5 * HOUR_MS)), // 每周窗口的 `number` 见过 7 也见过 1，只认 `unit` - (Some(6), _) if usage => Some(("weekly", 7 * DAY_MS)), - (Some(5), Some(1)) if calls => Some(("monthly", 31 * DAY_MS)), + (Some(6), _) => Some(("weekly", 7 * DAY_MS)), _ => None, } } +/// 积分制套餐的一项（`CREDIT_LIMIT`）。**只有它带的是积分** +fn is_credit(item: &Value) -> bool { + item["type"] + .as_str() + .is_some_and(|k| k.eq_ignore_ascii_case("CREDIT_LIMIT")) +} + fn windows(items: &[Value], now_ms: u64) -> Vec { let mut out: Vec = Vec::new(); for item in items { @@ -164,12 +173,15 @@ fn windows(items: &[Value], now_ms: u64) -> Vec { continue; }; let used_percent = used_percent.clamp(0.0, 100.0); + // **积分只来自积分制套餐。**老套餐的 token 窗口就算也带着这三个数，数的也是 + // token，当成积分显示就错了 let credits = match ( + is_credit(item), item["usage"].as_f64(), item["currentValue"].as_f64(), item["remaining"].as_f64(), ) { - (Some(total), Some(used), Some(remaining)) => Some(tw_api::QuotaCredits { + (true, Some(total), Some(used), Some(remaining)) => Some(tw_api::QuotaCredits { total, used, remaining, @@ -211,9 +223,13 @@ pub struct Exhausted { /// 读 429 的 body：是额度用完的话，是哪个窗口、什么时候重置。 /// /// 业务码在 Anthropic 端点上是 `error.type`（字符串），在 OpenAI 端点上是 `error.code`。 -/// **算用完的只有这几个**：1308（5 小时）、1310（每周或每月）、1316–1321（团队超额）。 +/// **算用完的只有这几个**：1308（5 小时）、1310（每周）、1316–1321（团队超额）。 /// 套餐到期（1309）、公平使用限制（1313）、临时限流（1302、1305）都不是额度用完， /// 记成用完会让界面报一个重置之后也好不了的「用完了」。 +/// +/// **1310 只标每周。**它本来说的是「每周或每月」，可每月的那个数的是 MCP 工具的调用次数， +/// 那些调用不经过网关（见 `window_of`）：模型请求上的 1310 说的只能是每周额度，消息 +/// 怎么写都一样 pub fn exhausted(body: &[u8]) -> Option { let v: Value = serde_json::from_slice(body).ok()?; let err = v.get("error")?; @@ -225,10 +241,11 @@ pub fn exhausted(body: &[u8]) -> Option { return None; } let message = err["message"].as_str().unwrap_or_default(); - let window = window_in(message).or(match code { - 1308 => Some("5h"), - _ => None, - }); + let window = match code { + 1310 => Some("weekly"), + 1308 => window_in(message).or(Some("5h")), + _ => window_in(message), + }; Some(Exhausted { code, window, @@ -244,8 +261,6 @@ fn window_in(message: &str) -> Option<&'static str> { Some("5h") } else if m.contains("week") || m.contains("周") { Some("weekly") - } else if m.contains("month") || m.contains("月") { - Some("monthly") } else { None } @@ -276,17 +291,13 @@ fn reset_in(message: &str) -> Option { .ok() } -/// 业务码和消息都没说是哪个窗口时，按已知的额度挑：候选里用得最多的那个。 -/// 一个都不知道就是每周 —— 1310 说的是「每周或每月」，新套餐只有每周 -pub fn guess_window(code: i64, known: &Quota) -> String { - let candidates: &[&str] = match code { - 1310 => &["weekly", "monthly"], - _ => &["5h", "weekly", "monthly"], - }; +/// 业务码和消息都没说是哪个窗口时（团队超额的 1316–1321），按已知的额度挑：5 小时和 +/// 每周里用得最多的那个。一个都不知道就是每周 +pub fn guess_window(known: &Quota) -> String { known .windows .iter() - .filter(|w| candidates.contains(&w.window.as_str())) + .filter(|w| matches!(w.window.as_str(), "5h" | "weekly")) .max_by(|a, b| a.used_percent.total_cmp(&b.used_percent)) .map(|w| w.window.clone()) .unwrap_or_else(|| "weekly".to_string()) @@ -512,10 +523,10 @@ mod tests { assert_eq!(week.credits.map(|c| c.remaining), Some(9731.0)); } - /// 老套餐 V1：一个 5 小时的 token 窗口，加每月的 MCP 调用次数。 - /// 每月窗口放在前面，它先重置 —— **不按重置先后猜** + /// 老套餐 V1：一个 5 小时的 token 窗口，加 MCP 工具的调用次数（`TIME_LIMIT`）。 + /// **调用次数不要**：那些调用不经过网关 #[test] - fn an_old_v1_plan_has_five_hours_and_monthly_calls() { + fn an_old_v1_plan_has_only_the_five_hour_window() { let q = quota(&format!( r#"{{"code":200,"success":true,"data":{{"level":"pro","limits":[ {{"type":"TIME_LIMIT","unit":5,"number":1,"usage":1000,"currentValue":40,"remaining":960,"percentage":4,"nextResetTime":{m}}}, @@ -524,32 +535,43 @@ mod tests { m = NOW + 10 * MINUTE_MS, h = NOW + 3 * HOUR_MS, )); - assert_eq!(q.windows.len(), 2); - assert_eq!(window(&q, "5h").used_percent, 37.0); - assert_eq!(window(&q, "5h").credits, None); - let month = window(&q, "monthly"); - assert_eq!(month.used_percent, 4.0); - assert_eq!(month.resets_at_ms, Some(NOW + 10 * MINUTE_MS)); + assert_eq!(q.windows.len(), 1, "{q:?}"); + let five = window(&q, "5h"); + assert_eq!(five.used_percent, 37.0); + assert_eq!(five.resets_at_ms, Some(NOW + 3 * HOUR_MS)); + assert_eq!(five.credits, None); } - /// 老套餐 V2：多了每周的 token 窗口，`number` 是 7 + /// 老套餐 V2：多了每周的 token 窗口，`number` 是 7。**`TIME_LIMIT` 照样不要**， + /// 调用次数用满了也不要：那不是模型请求的额度 #[test] - fn an_old_v2_plan_adds_a_weekly_window() { + fn an_old_v2_plan_adds_a_weekly_window_and_ignores_the_mcp_calls() { let q = quota(&format!( r#"{{"code":0,"data":{{"limits":[ {{"type":"TOKENS_LIMIT","unit":6,"number":7,"percentage":81,"nextResetTime":{w}}}, {{"type":"TOKENS_LIMIT","unit":3,"number":5,"percentage":100,"nextResetTime":{h}}}, - {{"type":"TIME_LIMIT","unit":5,"number":1,"percentage":0}} + {{"type":"TIME_LIMIT","unit":5,"number":1,"usage":1000,"currentValue":1000,"remaining":0,"percentage":100,"nextResetTime":{m}}} ]}}}}"#, w = NOW + 2 * HOUR_MS, h = NOW + 4 * HOUR_MS, + m = NOW + 10 * DAY_MS, )); let names: Vec<_> = q.windows.iter().map(|w| w.window.as_str()).collect(); - assert_eq!(names, ["weekly", "5h", "monthly"]); - assert_eq!(window(&q, "weekly").used_percent, 81.0); + assert_eq!(names, ["weekly", "5h"]); + let week = window(&q, "weekly"); + assert_eq!(week.used_percent, 81.0); + assert_eq!(week.status, None); // 用满就是被拒 assert_eq!(window(&q, "5h").status.as_deref(), Some("rejected")); - assert_eq!(window(&q, "monthly").resets_at_ms, None); + } + + /// 积分只来自积分制套餐：token 窗口带着同样三个数，那也是 token + #[test] + fn only_a_credit_window_has_credits() { + let q = quota( + r#"{"success":true,"data":{"limits":[{"type":"TOKENS_LIMIT","unit":3,"number":5,"usage":40000000,"currentValue":1000000,"remaining":39000000,"percentage":2}]}}"#, + ); + assert_eq!(window(&q, "5h").credits, None); } #[test] @@ -622,15 +644,20 @@ mod tests { let q = quota(&format!( r#"{{"success":true,"data":{{"limits":[ {{"type":"TOKENS_LIMIT","unit":3,"number":5,"percentage":20,"nextResetTime":{late}}}, - {{"type":"TOKENS_LIMIT","unit":6,"number":1,"percentage":20,"nextResetTime":{past}}}, - {{"type":"TIME_LIMIT","unit":5,"number":1,"percentage":20,"nextResetTime":"明天"}} + {{"type":"TOKENS_LIMIT","unit":6,"number":1,"percentage":20,"nextResetTime":{past}}} ]}}}}"#, // 5 小时窗口却在 6 小时后重置 late = NOW + 6 * HOUR_MS, // 已经过去了 past = NOW - MINUTE_MS, )); + assert_eq!(q.windows.len(), 2, "{q:?}"); assert!(q.windows.iter().all(|w| w.resets_at_ms.is_none()), "{q:?}"); + // 不是时刻的也不要 + let q = quota( + r#"{"success":true,"data":{"limits":[{"type":"CREDIT_LIMIT","unit":6,"number":1,"percentage":20,"nextResetTime":"明天"}]}}"#, + ); + assert_eq!(window(&q, "weekly").resets_at_ms, None); // 差一点的钟不该让一个真实的时刻被丢掉 let q = quota(&format!( r#"{{"success":true,"data":{{"limits":[{{"type":"TOKENS_LIMIT","unit":3,"number":5,"percentage":20,"nextResetTime":{t}}}]}}}}"#, @@ -703,19 +730,43 @@ mod tests { assert!(e.resets_at_ms.is_some()); } + /// **1310 只标每周**：每月的那个上限数的是 MCP 工具的调用，不经过网关。消息里 + /// 说的是每月、是 5 小时，都还是每周 #[test] - fn monthly_and_team_limits_count_too() { - let e = exhausted( - br#"{"type":"error","error":{"type":"1310","message":"Monthly limit exhausted."}}"#, - ) - .unwrap(); - assert_eq!(e.window, Some("monthly")); - assert_eq!(e.resets_at_ms, None, "消息里没有时刻就没有"); - for code in 1316..=1321 { + fn a_1310_is_the_weekly_window_whatever_the_message_says() { + for message in [ + "Weekly/Monthly Limit Exhausted.", + "Monthly limit exhausted.", + "已达到每月使用上限。", + "Usage limit reached for 5 hour.", + "", + ] { let body = - format!(r#"{{"error":{{"code":"{code}","message":"Team quota exceeded"}}}}"#); - let e = exhausted(body.as_bytes()).unwrap_or_else(|| panic!("{code}")); - assert_eq!(e.window, None, "{code}：说不清是哪个窗口"); + format!(r#"{{"type":"error","error":{{"type":"1310","message":"{message}"}}}}"#); + let e = exhausted(body.as_bytes()).unwrap_or_else(|| panic!("{message}")); + assert_eq!(e.code, 1310); + assert_eq!(e.window, Some("weekly"), "{message}"); + assert_eq!(e.resets_at_ms, None, "消息里没有时刻就没有"); + } + } + + /// 团队超额：消息说了是 5 小时还是每周就照着标，说不清的交给 `guess_window`。 + /// **说每月的也是说不清**：没有每月这个窗口 + #[test] + fn team_limits_count_too() { + for code in 1316..=1321 { + let e = |message: &str| { + let body = format!(r#"{{"error":{{"code":"{code}","message":"{message}"}}}}"#); + exhausted(body.as_bytes()).unwrap_or_else(|| panic!("{code}")) + }; + assert_eq!( + e("Team quota exceeded").window, + None, + "{code}：说不清是哪个窗口" + ); + assert_eq!(e("Team limit reached for 5 hour").window, Some("5h")); + assert_eq!(e("团队每周额度已用完").window, Some("weekly")); + assert_eq!(e("Team monthly limit exhausted").window, None, "{code}"); } } @@ -741,12 +792,15 @@ mod tests { #[test] fn an_unnamed_window_is_the_tightest_known_candidate() { let known = Quota { - windows: vec![w("5h", 100.0), w("weekly", 40.0), w("monthly", 90.0)], + windows: vec![w("5h", 100.0), w("weekly", 40.0)], + }; + assert_eq!(guess_window(&known), "5h"); + let known = Quota { + windows: vec![w("5h", 10.0), w("weekly", 90.0)], }; - // 1310 说的是每周或每月，不会是 5 小时 - assert_eq!(guess_window(1310, &known), "monthly"); - assert_eq!(guess_window(1318, &known), "5h"); - assert_eq!(guess_window(1310, &Quota::default()), "weekly"); + assert_eq!(guess_window(&known), "weekly"); + // 一个都不知道就是每周 + assert_eq!(guess_window(&Quota::default()), "weekly"); } fn w(name: &str, used: f64) -> Window { diff --git a/crates/tw-gateway/src/quota.rs b/crates/tw-gateway/src/quota.rs index e90c4501..7729c44b 100644 --- a/crates/tw-gateway/src/quota.rs +++ b/crates/tw-gateway/src/quota.rs @@ -25,7 +25,7 @@ use serde::{Deserialize, Serialize}; /// 一个额度窗口的状态。**每个字段都直接来自上游（响应头或额度接口），没有一个是推算的。** #[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] pub struct Window { - /// 哪个窗口：`5h` / `7d`（Anthropic）/ `weekly`（Codex、GLM）/ `monthly`（GLM） + /// 哪个窗口：`5h` / `7d`（Anthropic）/ `weekly`（Codex、GLM） pub window: String, /// 用了百分之多少。0–100 pub used_percent: f64, @@ -39,7 +39,7 @@ pub struct Window { /// `allowed` / `allowed_warning` / `rejected` #[serde(default, skip_serializing_if = "Option::is_none")] pub status: Option, - /// 积分制套餐的积分（GLM Coding Plan） + /// 积分制套餐的积分（GLM Coding Plan 的 `CREDIT_LIMIT`）。别的窗口都没有 #[serde(default, skip_serializing_if = "Option::is_none")] pub credits: Option, } diff --git a/crates/tw-gateway/src/state/glm.rs b/crates/tw-gateway/src/state/glm.rs index 0d670ed1..5bf1fbb6 100644 --- a/crates/tw-gateway/src/state/glm.rs +++ b/crates/tw-gateway/src/state/glm.rs @@ -78,7 +78,7 @@ impl AppState { }; match &answer { Answer::Quota(q) => state.record_quota(state.bus.next_id(), &p.name, q.clone()), - // 没有套餐就没有额度数据：之前记着的也不再作数 + // 没有套餐就没有额度数据：之前记着的也不再作数，界面上的那一格马上收起 Answer::NoPlan => state.forget_quota(&p.name), Answer::Rejected => { tracing::warn!(provider = %p.name, "the GLM quota endpoint rejected the key") @@ -144,7 +144,7 @@ impl AppState { let window = hit .window .map(str::to_string) - .unwrap_or_else(|| glm::guess_window(hit.code, "a)); + .unwrap_or_else(|| glm::guess_window("a)); let known = quota .windows .iter() @@ -265,6 +265,79 @@ mod tests { assert!(w.rejected()); } + /// 1310 只标每周：消息里写着每月，已知额度里 5 小时用得更多，都一样 + #[tokio::test] + async fn a_1310_marks_only_the_weekly_window() { + let s = state(); + let window = |name: &str, used_percent: f64| crate::quota::Window { + window: name.into(), + used_percent, + resets_at_ms: None, + status: None, + credits: None, + }; + s.record_quota( + 1, + "glm", + crate::quota::Quota { + windows: vec![window("5h", 99.0), window("weekly", 40.0)], + }, + ); + let mut rx = s.bus.subscribe(); + s.note_glm_exhausted( + 2, + "glm", + br#"{"type":"error","error":{"type":"1310","message":"Monthly Limit Exhausted. Your limit will reset at 2099-10-06 10:00:00"}}"#, + ); + let q = &s.quotas()["glm"]; + let names: Vec<_> = q.windows.iter().map(|w| w.window.as_str()).collect(); + assert_eq!(names, ["5h", "weekly"]); + let five = &q.windows[0]; + assert_eq!(five.used_percent, 99.0); + assert!(!five.rejected()); + let week = &q.windows[1]; + assert_eq!(week.used_percent, 100.0); + assert!(week.rejected()); + + let mut exhausted = Vec::new(); + while let Ok(e) = rx.try_recv() { + if let tw_api::Event::QuotaExhausted { window, .. } = e { + exhausted.push(window); + } + } + assert_eq!(exhausted, ["weekly"]); + } + + /// 判定没有套餐、记着的额度清掉了：**报一条窗口为空的 `QuotaSeen`**，界面当场收起。 + /// 本来就没记着的，没有什么可收的，不报 + #[tokio::test] + async fn forgetting_a_quota_tells_the_ui_with_an_empty_quota_seen() { + let s = state(); + s.note_glm_exhausted(1, "glm", FIVE_HOURS); + let mut rx = s.bus.subscribe(); + s.forget_quota("glm"); + assert!(s.quotas().is_empty()); + let seen: Vec<_> = std::iter::from_fn(|| rx.try_recv().ok()) + .filter_map(|e| match e { + tw_api::Event::QuotaSeen { + provider, windows, .. + } => Some((provider, windows)), + _ => None, + }) + .collect(); + assert_eq!(seen, [("glm".to_string(), vec![])]); + + s.forget_quota("glm"); + assert!(rx.try_recv().is_err(), "没记着的不报"); + + // 「用完」也不再作数：之后再用完，照样再报一次 + s.note_glm_exhausted(2, "glm", FIVE_HOURS); + let exhausted = std::iter::from_fn(|| rx.try_recv().ok()) + .filter(|e| matches!(e, tw_api::Event::QuotaExhausted { .. })) + .count(); + assert_eq!(exhausted, 1); + } + #[tokio::test] async fn a_rate_limit_that_is_not_the_quota_changes_nothing() { let s = state(); diff --git a/crates/tw-gateway/src/state/upstream.rs b/crates/tw-gateway/src/state/upstream.rs index b706b907..f0b06d88 100644 --- a/crates/tw-gateway/src/state/upstream.rs +++ b/crates/tw-gateway/src/state/upstream.rs @@ -266,14 +266,27 @@ impl AppState { } } - /// 这一家没有额度数据了（比如 key 没有开通套餐）：记着的额度和「用完」都不再作数 + /// 这一家没有额度数据了（比如 key 没有开通套餐）：记着的额度和「用完」都不再作数。 + /// + /// 记着的话，**报一条窗口为空的 `QuotaSeen`**：界面照着这个事件换上这一家的额度， + /// 空的就是没有了，当场收起，不用等下一次读 `/quota`。本来就没记着的不报 pub(crate) fn forget_quota(&self, provider: &str) { - if let Ok(mut g) = self.quotas.lock() { - g.remove(provider); - } + let had = self + .quotas + .lock() + .map(|mut g| g.remove(provider).is_some()) + .unwrap_or(false); if let Ok(mut g) = self.exhausted.lock() { g.retain(|(p, _)| p != provider); } + if had { + self.bus.emit(tw_api::Event::QuotaSeen { + id: self.bus.next_id(), + provider: provider.to_string(), + windows: Vec::new(), + at_ms: now_ms(), + }); + } } /// 每个上游最近一次报的订阅额度。 diff --git a/crates/tw-gateway/tests/glm_quota.rs b/crates/tw-gateway/tests/glm_quota.rs index 94ead7ca..26cb1669 100644 --- a/crates/tw-gateway/tests/glm_quota.rs +++ b/crates/tw-gateway/tests/glm_quota.rs @@ -135,6 +135,41 @@ async fn a_key_without_a_plan_has_no_quota_and_is_not_asked_again_soon() { assert_eq!(g.asked.lock().unwrap().len(), 1); } +/// 判定没有套餐、清掉了记着的额度：**当场报一条窗口为空的 `QuotaSeen`**，界面不用等 +/// 下一次读 `/quota` 才收起那一格。这里是 key 换成了一把没开通套餐的 +#[tokio::test] +async fn a_key_found_to_have_no_plan_takes_its_quota_down_at_once() { + let g = Arc::new(Glm::default()); + *g.quota.lock().unwrap() = credit_plan(now_ms()); + let state = state_for(start_glm(g.clone()).await).await; + state.refresh_glm_quotas(Duration::from_secs(5)).await; + assert_eq!(state.quotas()["glm"].windows.len(), 2); + + *g.quota.lock().unwrap() = + r#"{"code":500,"msg":"当前用户不存在coding plan","success":false}"#.into(); + let mut cfg = (*state.config()).clone(); + cfg.providers[0].key = Some("another-fake-glm-key".into()); + state.reload(cfg).unwrap(); + let mut rx = state.bus.subscribe(); + state.refresh_glm_quotas(Duration::from_secs(5)).await; + + assert!(state.quotas().is_empty()); + let windows = loop { + match tokio::time::timeout(Duration::from_secs(3), rx.recv()).await { + Ok(Ok(tw_api::Event::QuotaSeen { + provider, windows, .. + })) if provider == "glm" => break windows, + Ok(Ok(_)) => continue, + other => panic!("没等到撤下额度的事件：{other:?}"), + } + }; + assert!(windows.is_empty(), "{windows:?}"); + assert_eq!( + *g.asked.lock().unwrap(), + ["fake-glm-key", "another-fake-glm-key"] + ); +} + #[tokio::test] async fn a_used_up_quota_in_a_429_is_reported_and_the_quota_is_asked() { let g = Arc::new(Glm::default());