cn_provider_quota_service.go 23 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685
  1. package service
  2. import (
  3. "context"
  4. "encoding/json"
  5. "fmt"
  6. "io"
  7. "log/slog"
  8. "net/http"
  9. "sort"
  10. "strconv"
  11. "strings"
  12. "time"
  13. "github.com/Wei-Shaw/sub2api/internal/config"
  14. infraerrors "github.com/Wei-Shaw/sub2api/internal/pkg/errors"
  15. "github.com/tidwall/gjson"
  16. "golang.org/x/sync/singleflight"
  17. )
  18. // 国产供应商 Coding Plan 滚动窗口额度探测服务(Kimi For Coding / 智谱 GLM Coding Plan)。
  19. //
  20. // 与 grok_quota_service 不同:CN 供应商走数据面 API Key(无 OAuth token provider),
  21. // 额度端点为只读 GET,解析 5h + weekly 两档滚动窗口并落 account.Extra 快照,
  22. // 供账号调度阈值评估(account_scheduling_threshold_eval.go)做主动停调。
  23. //
  24. // 解析逻辑对齐 cc-switch(farion1231/cc-switch)services/coding_plan.rs 的
  25. // query_kimi / query_zhipu,包括智谱 unit 字段优先分类与 reset 兜底启发式。
  26. const (
  27. cnQuotaUpstreamTimeout = 15 * time.Second
  28. cnQuotaMaxBodyBytes = 256 * 1024
  29. // Extra 快照键后缀(加 provider 前缀,如 kimi_5h_used_percent)。
  30. cnExtraSuffix5hUsed = "5h_used_percent"
  31. cnExtraSuffix5hReset = "5h_reset_at"
  32. cnExtraSuffixWeeklyUsed = "weekly_used_percent"
  33. cnExtraSuffixWeeklyReset = "weekly_reset_at"
  34. cnExtraSuffixUsageUpdated = "usage_updated_at"
  35. )
  36. // cnExtraKey 拼接 provider 维度的 extra 键。
  37. func cnExtraKey(provider, suffix string) string { return provider + "_" + suffix }
  38. // CNQuotaTier 表示一个滚动用量窗口档位(5h / weekly)。
  39. type CNQuotaTier struct {
  40. Window string `json:"window"` // "5h" | "weekly"
  41. UsedPercent float64 `json:"used_percent"` // 已用百分比(0-100+,不做裁剪)
  42. ResetAt string `json:"reset_at,omitempty"` // RFC3339,空表示无重置时间
  43. }
  44. // CNProviderQuotaProbeResult 是 Coding Plan 额度探测的返回结构(管理端 + UI 消费)。
  45. type CNProviderQuotaProbeResult struct {
  46. Provider string `json:"provider"`
  47. Source string `json:"source"`
  48. Success bool `json:"success"`
  49. CredentialValid bool `json:"credential_valid"` // false = 401/403 鉴权失败
  50. Tiers []CNQuotaTier `json:"tiers,omitempty"`
  51. PlanLevel string `json:"plan_level,omitempty"` // 智谱套餐等级
  52. StatusCode int `json:"status_code,omitempty"`
  53. FetchedAt int64 `json:"fetched_at"`
  54. Persisted bool `json:"persisted"`
  55. Error string `json:"error,omitempty"`
  56. }
  57. // CNProviderQuotaService 探测 Kimi / Zhipu Coding Plan 的滚动窗口用量。
  58. type CNProviderQuotaService struct {
  59. accountRepo AccountRepository
  60. proxyRepo ProxyRepository
  61. httpUpstream HTTPUpstream
  62. cfg *config.Config
  63. flight singleflight.Group
  64. }
  65. // NewCNProviderQuotaService 构造 Coding Plan 额度探测服务。
  66. func NewCNProviderQuotaService(
  67. accountRepo AccountRepository,
  68. proxyRepo ProxyRepository,
  69. httpUpstream HTTPUpstream,
  70. cfg *config.Config,
  71. ) *CNProviderQuotaService {
  72. return &CNProviderQuotaService{
  73. accountRepo: accountRepo,
  74. proxyRepo: proxyRepo,
  75. httpUpstream: httpUpstream,
  76. cfg: cfg,
  77. }
  78. }
  79. // QueryUsage 探测指定账号的 Coding Plan 滚动窗口用量并落 Extra 快照。
  80. // 同一账号的并发探测会被 singleflight 合并。
  81. func (s *CNProviderQuotaService) QueryUsage(ctx context.Context, accountID int64) (*CNProviderQuotaProbeResult, error) {
  82. account, err := s.loadCodingPlanAccount(ctx, accountID)
  83. if err != nil {
  84. return nil, err
  85. }
  86. return s.QueryUsageForAccount(ctx, account)
  87. }
  88. // QueryUsageForAccount 探测已加载账号(配额监控 fetcher 复用,避免二次 GetByID)。
  89. // singleflight key 与 QueryUsage 相同,按账号 ID 与 admin 侧并发探测合并。
  90. func (s *CNProviderQuotaService) QueryUsageForAccount(ctx context.Context, account *Account) (*CNProviderQuotaProbeResult, error) {
  91. if s == nil || s.accountRepo == nil || s.httpUpstream == nil {
  92. return nil, infraerrors.New(http.StatusInternalServerError, "CN_QUOTA_NOT_CONFIGURED", "cn provider quota service is not configured")
  93. }
  94. if err := validateCodingPlanAccount(account); err != nil {
  95. return nil, err
  96. }
  97. key := "cn_quota:" + strconv.FormatInt(account.ID, 10)
  98. resultCh := s.flight.DoChan(key, func() (any, error) {
  99. probeCtx, cancel := context.WithTimeout(context.Background(), cnQuotaUpstreamTimeout+5*time.Second)
  100. defer cancel()
  101. return s.queryUsageForAccount(probeCtx, account)
  102. })
  103. select {
  104. case <-ctx.Done():
  105. return nil, ctx.Err()
  106. case flightResult := <-resultCh:
  107. if flightResult.Err != nil {
  108. return nil, flightResult.Err
  109. }
  110. result, ok := flightResult.Val.(*CNProviderQuotaProbeResult)
  111. if !ok || result == nil {
  112. return nil, infraerrors.New(http.StatusInternalServerError, "CN_QUOTA_PROBE_RESULT_INVALID", "invalid cn provider quota probe result")
  113. }
  114. cloned := *result
  115. return &cloned, nil
  116. }
  117. }
  118. func (s *CNProviderQuotaService) queryUsageForAccount(ctx context.Context, account *Account) (*CNProviderQuotaProbeResult, error) {
  119. provider := account.GetCodingPlanProvider()
  120. if provider != PlatformKimi && provider != PlatformZhipu && provider != PlatformMiniMax {
  121. return nil, infraerrors.New(http.StatusBadRequest, "CN_QUOTA_NOT_CODING_PLAN", "account is not a kimi/zhipu/minimax coding plan account")
  122. }
  123. apiKey := strings.TrimSpace(account.GetCNAPIKey())
  124. if apiKey == "" {
  125. return nil, infraerrors.New(http.StatusBadRequest, "CN_QUOTA_NO_APIKEY", "account api_key is empty")
  126. }
  127. baseURL := account.GetOpenAIBaseURL()
  128. var (
  129. targetURL string
  130. authHeader string
  131. zhipuOrg string
  132. )
  133. switch provider {
  134. case PlatformKimi:
  135. targetURL = kimiQuotaURL(baseURL)
  136. authHeader = "Bearer " + apiKey
  137. case PlatformZhipu:
  138. targetURL = zhipuQuotaURL(baseURL)
  139. authHeader = apiKey // 智谱额度端点鉴权不加 Bearer 前缀
  140. // 团队版 GLM Coding Plan:额度端点需 ?type=2 + 组织/项目请求头,
  141. // 否则官方 API 回「当前用户不存在coding plan」。组织 ID 存在即视为
  142. // 团队版(个人版凭据不含该字段,走原个人版查询路径)。
  143. zhipuOrg = strings.TrimSpace(account.GetCredential("zhipu_organization"))
  144. if zhipuOrg != "" {
  145. targetURL += "?type=2"
  146. }
  147. case PlatformMiniMax:
  148. targetURL = minimaxQuotaURL(baseURL)
  149. authHeader = "Bearer " + apiKey
  150. }
  151. // 探测发起前过出站 URL 安全策略(与网关转发/Grok 探测同一套校验):
  152. // 端点多由账号 base_url 衍生,不得把 API key 发往策略外主机。
  153. validatedURL, err := cnValidateProbeURL(s.cfg, targetURL)
  154. if err != nil {
  155. return nil, infraerrors.New(http.StatusForbidden, "CN_QUOTA_URL_REJECTED", err.Error())
  156. }
  157. targetURL = validatedURL
  158. proxyURL := s.resolveProxyURL(ctx, account)
  159. callCtx, cancel := context.WithTimeout(ctx, cnQuotaUpstreamTimeout)
  160. defer cancel()
  161. req, err := http.NewRequestWithContext(callCtx, http.MethodGet, targetURL, nil)
  162. if err != nil {
  163. return nil, infraerrors.Newf(http.StatusInternalServerError, "CN_QUOTA_REQUEST_BUILD_FAILED", "build request: %v", err)
  164. }
  165. req.Header.Set("Authorization", authHeader)
  166. req.Header.Set("Accept", "application/json")
  167. if provider == PlatformZhipu || provider == PlatformMiniMax {
  168. req.Header.Set("Content-Type", "application/json")
  169. req.Header.Set("Accept-Language", "en-US,en")
  170. if zhipuOrg != "" {
  171. req.Header.Set("bigmodel-organization", zhipuOrg)
  172. if project := strings.TrimSpace(account.GetCredential("zhipu_project")); project != "" {
  173. req.Header.Set("bigmodel-project", project)
  174. }
  175. }
  176. }
  177. // 探测与真实转发保持同一套账号级请求头覆写,避免探测通过但转发失败。
  178. account.ApplyHeaderOverrides(req.Header)
  179. resp, err := s.httpUpstream.Do(req, proxyURL, account.ID, maxInt(account.Concurrency, 1))
  180. if err != nil {
  181. return nil, infraerrors.Newf(http.StatusBadGateway, "CN_QUOTA_REQUEST_FAILED", "upstream request failed: %v", err)
  182. }
  183. defer func() { _ = resp.Body.Close() }()
  184. bodyBytes, _ := io.ReadAll(io.LimitReader(resp.Body, cnQuotaMaxBodyBytes))
  185. now := time.Now().UTC()
  186. result := &CNProviderQuotaProbeResult{
  187. Provider: provider,
  188. Source: "coding_plan",
  189. FetchedAt: now.Unix(),
  190. StatusCode: resp.StatusCode,
  191. }
  192. if resp.StatusCode == http.StatusUnauthorized || resp.StatusCode == http.StatusForbidden {
  193. // 鉴权失败:不落快照(不覆盖之前的有效值),仅返回失败结果供前端提示。
  194. result.Error = fmt.Sprintf("Authentication failed (HTTP %d)", resp.StatusCode)
  195. return result, nil
  196. }
  197. if resp.StatusCode < 200 || resp.StatusCode >= 300 {
  198. result.Error = fmt.Sprintf("API error (HTTP %d): %s", resp.StatusCode, truncate(strings.TrimSpace(string(bodyBytes)), 240))
  199. return result, nil
  200. }
  201. // 智谱业务级错误(HTTP 2xx 但 success=false)。
  202. if provider == PlatformZhipu {
  203. if success := gjson.GetBytes(bodyBytes, "success"); success.Exists() && !success.Bool() {
  204. msg := strings.TrimSpace(gjson.GetBytes(bodyBytes, "msg").String())
  205. if msg == "" {
  206. msg = "unknown zhipu quota error"
  207. }
  208. result.Error = "API error: " + msg
  209. return result, nil
  210. }
  211. }
  212. var tiers []CNQuotaTier
  213. switch provider {
  214. case PlatformKimi:
  215. tiers = parseKimiUsageTiers(bodyBytes)
  216. case PlatformZhipu:
  217. tiers = parseZhipuTokenTiers(gjson.GetBytes(bodyBytes, "data"))
  218. result.PlanLevel = strings.TrimSpace(gjson.GetBytes(bodyBytes, "data.level").String())
  219. case PlatformMiniMax:
  220. if status := gjson.GetBytes(bodyBytes, "base_resp.status_code"); status.Exists() && status.Int() != 0 {
  221. msg := strings.TrimSpace(gjson.GetBytes(bodyBytes, "base_resp.status_msg").String())
  222. if msg == "" {
  223. msg = "unknown minimax quota error"
  224. }
  225. result.Error = fmt.Sprintf("API error (%d): %s", status.Int(), msg)
  226. return result, nil
  227. }
  228. tiers = parseMiniMaxUsageTiers(bodyBytes)
  229. result.PlanLevel = strings.TrimSpace(gjson.GetBytes(bodyBytes, "current_subscribe_title").String())
  230. }
  231. result.Tiers = tiers
  232. result.Success = true
  233. result.CredentialValid = true
  234. updates := cnQuotaExtraUpdates(provider, tiers, now)
  235. if err := s.accountRepo.UpdateExtra(ctx, account.ID, updates); err != nil {
  236. slog.Warn("cn_quota_persist_failed", "account_id", account.ID, "provider", provider, "error", err)
  237. } else {
  238. result.Persisted = true
  239. }
  240. return result, nil
  241. }
  242. func (s *CNProviderQuotaService) loadCodingPlanAccount(ctx context.Context, accountID int64) (*Account, error) {
  243. account, err := s.accountRepo.GetByID(ctx, accountID)
  244. if err != nil {
  245. return nil, infraerrors.Newf(http.StatusNotFound, "CN_QUOTA_ACCOUNT_NOT_FOUND", "account not found: %v", err)
  246. }
  247. if err := validateCodingPlanAccount(account); err != nil {
  248. return nil, err
  249. }
  250. return account, nil
  251. }
  252. // validateCodingPlanAccount 加载后的非 DB 校验(ForAccount 入口同样复用,
  253. // 保证直传 account 也不绕过平台/模式检查)。
  254. func validateCodingPlanAccount(account *Account) error {
  255. if account == nil {
  256. return infraerrors.New(http.StatusNotFound, "CN_QUOTA_ACCOUNT_NOT_FOUND", "account not found")
  257. }
  258. if !account.IsCNProvider() {
  259. return infraerrors.New(http.StatusBadRequest, "CN_QUOTA_INVALID_PLATFORM", "account is not a CN provider account")
  260. }
  261. if !account.IsCodingPlan() {
  262. return infraerrors.New(http.StatusBadRequest, "CN_QUOTA_NOT_CODING_PLAN", "account is not a coding plan account")
  263. }
  264. return nil
  265. }
  266. func (s *CNProviderQuotaService) resolveProxyURL(ctx context.Context, account *Account) string {
  267. if account == nil || account.ProxyID == nil {
  268. return ""
  269. }
  270. if account.Proxy != nil {
  271. return account.Proxy.URL()
  272. }
  273. if s != nil && s.proxyRepo != nil {
  274. if proxy, err := s.proxyRepo.GetByID(ctx, *account.ProxyID); err == nil && proxy != nil {
  275. account.Proxy = proxy
  276. return proxy.URL()
  277. }
  278. }
  279. return ""
  280. }
  281. // zhipuQuotaURL 根据 base_url 解析智谱额度端点(与数据面推理域名同主机)。
  282. func zhipuQuotaURL(baseURL string) string {
  283. return zhipuQuotaHost(baseURL) + "/api/monitor/usage/quota/limit"
  284. }
  285. // kimiQuotaURL 根据 base_url 解析 Kimi For Coding 额度端点。
  286. // cc-switch query_kimi 固定探测 https://api.kimi.com/coding/v1/usages
  287. // (实测 /coding/usages 无 /v1 → 404)。coding/v1(CC 协议默认)与
  288. // coding(Anthropic 协议默认)两种 base 统一剥掉尾部后拼回 /v1/usages,
  289. // 协议切换不影响额度探测端点。
  290. func kimiQuotaURL(baseURL string) string {
  291. base := strings.TrimSuffix(strings.TrimRight(baseURL, "/"), "/v1")
  292. return base + "/v1/usages"
  293. }
  294. // minimaxQuotaURL 根据推理域名选择 Token Plan / Coding Plan 额度主机。
  295. // 官方 FAQ 写 www.minimax.io / www.minimaxi.com,实际以 Bearer Key 打 api.*。
  296. // 国际站 api.minimax.io;国内站 api.minimaxi.com(含 api.minimax.com 与自定义回落)。
  297. func minimaxQuotaURL(baseURL string) string {
  298. if strings.Contains(strings.ToLower(baseURL), "minimax.io") {
  299. return "https://api.minimax.io/v1/api/openplatform/coding_plan/remains"
  300. }
  301. return "https://api.minimaxi.com/v1/api/openplatform/coding_plan/remains"
  302. }
  303. func zhipuQuotaHost(baseURL string) string {
  304. switch u := strings.ToLower(baseURL); {
  305. case strings.Contains(u, "bigmodel.cn"):
  306. return "https://open.bigmodel.cn"
  307. case strings.Contains(u, "z.ai"):
  308. return "https://api.z.ai"
  309. default:
  310. // 国产优先:未知域名回落国内站(与前端 zhipu 预设一致)。
  311. return "https://open.bigmodel.cn"
  312. }
  313. }
  314. // parseKimiUsageTiers 解析 Kimi For Coding 的 /usages 响应。
  315. //
  316. // - limits[].detail.{limit,remaining,resetTime} → 5h 窗口(取首个 detail)
  317. // - usage.{limit,remaining,resetTime} → 周窗口
  318. //
  319. // utilization = (limit-remaining)/limit*100。
  320. func parseKimiUsageTiers(body []byte) []CNQuotaTier {
  321. var tiers []CNQuotaTier
  322. if limits := gjson.GetBytes(body, "limits"); limits.IsArray() {
  323. limits.ForEach(func(_, item gjson.Result) bool {
  324. detail := item.Get("detail")
  325. if !detail.Exists() {
  326. return true
  327. }
  328. limit, _ := cnParseF64(detail.Get("limit").Value())
  329. remaining, _ := cnParseF64(detail.Get("remaining").Value())
  330. used := limit - remaining
  331. if used < 0 {
  332. used = 0
  333. }
  334. var util float64
  335. if limit > 0 {
  336. util = used / limit * 100
  337. }
  338. tiers = append(tiers, CNQuotaTier{
  339. Window: "5h",
  340. UsedPercent: util,
  341. ResetAt: cnNormalizeResetTime(detail.Get("resetTime").Value()),
  342. })
  343. return false // 取首个 detail 作为 5h 窗口
  344. })
  345. }
  346. if usage := gjson.GetBytes(body, "usage"); usage.Exists() {
  347. limit, _ := cnParseF64(usage.Get("limit").Value())
  348. remaining, _ := cnParseF64(usage.Get("remaining").Value())
  349. used := limit - remaining
  350. if used < 0 {
  351. used = 0
  352. }
  353. var util float64
  354. if limit > 0 {
  355. util = used / limit * 100
  356. }
  357. tiers = append(tiers, CNQuotaTier{
  358. Window: "weekly",
  359. UsedPercent: util,
  360. ResetAt: cnNormalizeResetTime(usage.Get("resetTime").Value()),
  361. })
  362. }
  363. return tiers
  364. }
  365. // parseMiniMaxUsageTiers 解析 MiniMax Token Plan / Coding Plan remains 响应。
  366. //
  367. // 只取 model_name == "general"(编程套餐),跳过 video。字段是剩余百分比,
  368. // 展示已用 = 100 - remaining:
  369. // - 5h:current_interval_remaining_percent + end_time
  370. // - 周限额:仅 current_weekly_status == 1 时用 current_weekly_remaining_percent + weekly_end_time
  371. func parseMiniMaxUsageTiers(body []byte) []CNQuotaTier {
  372. remains := gjson.GetBytes(body, "model_remains")
  373. if !remains.IsArray() {
  374. return nil
  375. }
  376. var general gjson.Result
  377. remains.ForEach(func(_, item gjson.Result) bool {
  378. if strings.EqualFold(strings.TrimSpace(item.Get("model_name").String()), "general") {
  379. general = item
  380. return false
  381. }
  382. return true
  383. })
  384. if !general.Exists() {
  385. return nil
  386. }
  387. var tiers []CNQuotaTier
  388. if remaining, ok := cnParseF64(general.Get("current_interval_remaining_percent").Value()); ok {
  389. used := 100 - remaining
  390. if used < 0 {
  391. used = 0
  392. }
  393. tiers = append(tiers, CNQuotaTier{
  394. Window: "5h",
  395. UsedPercent: used,
  396. ResetAt: minimaxResetTime(general.Get("end_time")),
  397. })
  398. }
  399. if general.Get("current_weekly_status").Int() == 1 {
  400. if remaining, ok := cnParseF64(general.Get("current_weekly_remaining_percent").Value()); ok {
  401. used := 100 - remaining
  402. if used < 0 {
  403. used = 0
  404. }
  405. tiers = append(tiers, CNQuotaTier{
  406. Window: "weekly",
  407. UsedPercent: used,
  408. ResetAt: minimaxResetTime(general.Get("weekly_end_time")),
  409. })
  410. }
  411. }
  412. return tiers
  413. }
  414. func minimaxResetTime(v gjson.Result) string {
  415. if !v.Exists() {
  416. return ""
  417. }
  418. ms := v.Int()
  419. if ms <= 0 {
  420. return cnNormalizeResetTime(v.Value())
  421. }
  422. if ms < 1_000_000_000_000 {
  423. return time.Unix(ms, 0).UTC().Format(time.RFC3339)
  424. }
  425. return time.UnixMilli(ms).UTC().Format(time.RFC3339)
  426. }
  427. // cnZhipuWindow 标识智谱 TOKENS_LIMIT 条目所属窗口。
  428. type cnZhipuWindow int
  429. const (
  430. cnZhipuWindowUnknown cnZhipuWindow = iota
  431. cnZhipuWindow5h
  432. cnZhipuWindowWeekly
  433. )
  434. // classifyZhipuWindowUnit 按 unit 字段判定窗口类型(3=5h,6=weekly)。
  435. // unit 缺失或未识别时返回 Unknown,由调用方走 reset 时间启发式兜底。
  436. func classifyZhipuWindowUnit(unit int64) cnZhipuWindow {
  437. switch unit {
  438. case 3:
  439. return cnZhipuWindow5h
  440. case 6:
  441. return cnZhipuWindowWeekly
  442. default:
  443. return cnZhipuWindowUnknown
  444. }
  445. }
  446. // parseZhipuTokenTiers 解析智谱额度响应 data.limits 为 5h + weekly 两档。
  447. //
  448. // 分类优先级(对齐 cc-switch parse_zhipu_token_tiers,issue #3036):
  449. // 1. 显式 unit 字段(3=5h / 6=weekly)——不能用 reset 排序代替,周期末尾
  450. // 周窗口会比 5h 更早重置,时间排序必然标反。
  451. // 2. unit 缺失/未识别:无 nextResetTime 的条目优先归 5h(0% 状态下 5h 桶可能
  452. // 没有 reset),其余按 reset 升序依次填入仍空缺的槽位。
  453. //
  454. // CREDIT_LIMIT(信用额度)与 TOKENS_LIMIT(token 窗口)度量不同:两者同时返回时
  455. // 只让 TOKENS_LIMIT 参与 5h/weekly 槽位竞争,避免信用额度百分比污染阈值停调
  456. // 快照;仅当无任何 TOKENS_LIMIT 条目时才降级用 CREDIT_LIMIT 展示。
  457. // 老套餐只回 1 条 TOKENS_LIMIT,自然降级为仅 5h;新套餐回 2 条。
  458. func parseZhipuTokenTiers(data gjson.Result) []CNQuotaTier {
  459. type entry struct {
  460. resetMs int64
  461. hasReset bool
  462. percentage float64
  463. resetISO string
  464. }
  465. var (
  466. fiveHour entry
  467. fiveHourSet bool
  468. weekly entry
  469. weeklySet bool
  470. unclassified []entry
  471. )
  472. classify := func(item gjson.Result, e entry) {
  473. switch classifyZhipuWindowUnit(item.Get("unit").Int()) {
  474. case cnZhipuWindow5h:
  475. if !fiveHourSet {
  476. fiveHour, fiveHourSet = e, true
  477. } else {
  478. unclassified = append(unclassified, e)
  479. }
  480. case cnZhipuWindowWeekly:
  481. if !weeklySet {
  482. weekly, weeklySet = e, true
  483. } else {
  484. unclassified = append(unclassified, e)
  485. }
  486. default:
  487. unclassified = append(unclassified, e)
  488. }
  489. }
  490. var creditFallback []entry
  491. hasTokensLimit := false
  492. data.Get("limits").ForEach(func(_, item gjson.Result) bool {
  493. limitType := strings.ToUpper(strings.TrimSpace(item.Get("type").String()))
  494. if limitType != "TOKENS_LIMIT" && limitType != "CREDIT_LIMIT" {
  495. return true
  496. }
  497. percentage := 0.0
  498. if p, ok := cnParseF64(item.Get("percentage").Value()); ok {
  499. percentage = p
  500. }
  501. var (
  502. resetMs int64
  503. hasReset bool
  504. resetISO string
  505. )
  506. if nr := item.Get("nextResetTime"); nr.Exists() {
  507. switch nr.Type {
  508. case gjson.Number:
  509. resetMs = nr.Int()
  510. hasReset = resetMs > 0
  511. resetISO = cnMillisToRFC3339(resetMs)
  512. case gjson.String:
  513. resetISO = cnNormalizeResetTime(nr.String())
  514. hasReset = resetISO != ""
  515. }
  516. }
  517. e := entry{resetMs: resetMs, hasReset: hasReset, percentage: percentage, resetISO: resetISO}
  518. if limitType == "TOKENS_LIMIT" {
  519. hasTokensLimit = true
  520. classify(item, e)
  521. } else {
  522. creditFallback = append(creditFallback, e)
  523. }
  524. return true
  525. })
  526. // 无任何 TOKENS_LIMIT 条目(部分套餐只报信用额度):降级用 CREDIT_LIMIT 展示。
  527. if !hasTokensLimit {
  528. unclassified = append(unclassified, creditFallback...)
  529. }
  530. // 无 reset 的条目排前,再按 reset 升序,依次填入仍空缺的槽位。
  531. sort.SliceStable(unclassified, func(i, j int) bool {
  532. if unclassified[i].hasReset != unclassified[j].hasReset {
  533. return !unclassified[i].hasReset
  534. }
  535. return unclassified[i].resetMs < unclassified[j].resetMs
  536. })
  537. for _, e := range unclassified {
  538. switch {
  539. case !fiveHourSet:
  540. fiveHour, fiveHourSet = e, true
  541. case !weeklySet:
  542. weekly, weeklySet = e, true
  543. }
  544. }
  545. var tiers []CNQuotaTier
  546. if fiveHourSet {
  547. tiers = append(tiers, CNQuotaTier{Window: "5h", UsedPercent: fiveHour.percentage, ResetAt: fiveHour.resetISO})
  548. }
  549. if weeklySet {
  550. tiers = append(tiers, CNQuotaTier{Window: "weekly", UsedPercent: weekly.percentage, ResetAt: weekly.resetISO})
  551. }
  552. return tiers
  553. }
  554. // cnQuotaExtraUpdates 根据 tier 列表构造 provider 维度的 Extra 快照更新。
  555. func cnQuotaExtraUpdates(provider string, tiers []CNQuotaTier, now time.Time) map[string]any {
  556. updates := map[string]any{
  557. cnExtraKey(provider, cnExtraSuffixUsageUpdated): now.Format(time.RFC3339),
  558. }
  559. for _, t := range tiers {
  560. switch t.Window {
  561. case "5h":
  562. updates[cnExtraKey(provider, cnExtraSuffix5hUsed)] = t.UsedPercent
  563. if t.ResetAt != "" {
  564. updates[cnExtraKey(provider, cnExtraSuffix5hReset)] = t.ResetAt
  565. }
  566. case "weekly":
  567. updates[cnExtraKey(provider, cnExtraSuffixWeeklyUsed)] = t.UsedPercent
  568. if t.ResetAt != "" {
  569. updates[cnExtraKey(provider, cnExtraSuffixWeeklyReset)] = t.ResetAt
  570. }
  571. }
  572. }
  573. return updates
  574. }
  575. // cnParseF64 把 JSON 数值或字符串解析为 float64(兼容 "100" 与 100)。
  576. func cnParseF64(raw any) (float64, bool) {
  577. switch v := raw.(type) {
  578. case float64:
  579. return v, true
  580. case float32:
  581. return float64(v), true
  582. case int:
  583. return float64(v), true
  584. case int64:
  585. return float64(v), true
  586. case json.Number:
  587. f, err := v.Float64()
  588. return f, err == nil
  589. case string:
  590. f, err := strconv.ParseFloat(strings.TrimSpace(v), 64)
  591. return f, err == nil
  592. default:
  593. return 0, false
  594. }
  595. }
  596. // cnNormalizeResetTime 把上游重置时间(ISO8601 字符串 / 秒级 / 毫秒级数字)归一化为
  597. // RFC3339 字符串;无法识别或非正时间戳返回空串。
  598. func cnNormalizeResetTime(raw any) string {
  599. switch v := raw.(type) {
  600. case string:
  601. s := strings.TrimSpace(v)
  602. if s == "" {
  603. return ""
  604. }
  605. if ts, err := parseSchedulingTime(s); err == nil {
  606. return ts.UTC().Format(time.RFC3339)
  607. }
  608. return ""
  609. case float64:
  610. return cnMillisToRFC3339(int64(v))
  611. case int:
  612. return cnMillisToRFC3339(int64(v))
  613. case int64:
  614. return cnMillisToRFC3339(v)
  615. case json.Number:
  616. if n, err := v.Int64(); err == nil {
  617. return cnMillisToRFC3339(n)
  618. }
  619. return ""
  620. default:
  621. return ""
  622. }
  623. }
  624. // cnMillisToRFC3339 把秒级(<1e12)或毫秒级时间戳转为 RFC3339 字符串;非正返回空串。
  625. func cnMillisToRFC3339(n int64) string {
  626. if n <= 0 {
  627. return ""
  628. }
  629. var ms int64
  630. if n < 1_000_000_000_000 {
  631. ms = n * 1000
  632. } else {
  633. ms = n
  634. }
  635. return time.UnixMilli(ms).UTC().Format(time.RFC3339)
  636. }