Files
sub2api/backend/internal/service/cn_provider_quota_service.go
T
李建琦 6d655c9903
Release / update-version (push) Has been cancelled
Release / build-frontend (push) Has been cancelled
Release / release (push) Has been cancelled
Release / sync-version-file (push) Has been cancelled
CI / shell (push) Canceled after 0s
CI / test (push) Canceled after 0s
CI / frontend (push) Canceled after 0s
CI / golangci-lint (push) Canceled after 0s
Security Scan / backend-security (push) Canceled after 0s
Security Scan / frontend-security (push) Canceled after 0s
Sub2API v1.0 - AI API 网关(二开初始版本,基于上游 Wei-Shaw/sub2api)
2026-08-21 18:30:13 +08:00

583 lines
19 KiB
Go
Raw Blame History

This file contains ambiguous Unicode characters
This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.
package service
import (
"context"
"encoding/json"
"fmt"
"io"
"log/slog"
"net/http"
"sort"
"strconv"
"strings"
"time"
"github.com/Wei-Shaw/sub2api/internal/config"
infraerrors "github.com/Wei-Shaw/sub2api/internal/pkg/errors"
"github.com/tidwall/gjson"
"golang.org/x/sync/singleflight"
)
// 国产供应商 Coding Plan 滚动窗口额度探测服务(Kimi For Coding / 智谱 GLM Coding Plan)。
//
// 与 grok_quota_service 不同:CN 供应商走数据面 API Key(无 OAuth token provider),
// 额度端点为只读 GET,解析 5h + weekly 两档滚动窗口并落 account.Extra 快照,
// 供账号调度阈值评估(account_scheduling_threshold_eval.go)做主动停调。
//
// 解析逻辑对齐 cc-switchfarion1231/cc-switchservices/coding_plan.rs 的
// query_kimi / query_zhipu,包括智谱 unit 字段优先分类与 reset 兜底启发式。
const (
cnQuotaUpstreamTimeout = 15 * time.Second
cnQuotaMaxBodyBytes = 256 * 1024
// Extra 快照键后缀(加 provider 前缀,如 kimi_5h_used_percent)。
cnExtraSuffix5hUsed = "5h_used_percent"
cnExtraSuffix5hReset = "5h_reset_at"
cnExtraSuffixWeeklyUsed = "weekly_used_percent"
cnExtraSuffixWeeklyReset = "weekly_reset_at"
cnExtraSuffixUsageUpdated = "usage_updated_at"
)
// cnExtraKey 拼接 provider 维度的 extra 键。
func cnExtraKey(provider, suffix string) string { return provider + "_" + suffix }
// CNQuotaTier 表示一个滚动用量窗口档位(5h / weekly)。
type CNQuotaTier struct {
Window string `json:"window"` // "5h" | "weekly"
UsedPercent float64 `json:"used_percent"` // 已用百分比(0-100+,不做裁剪)
ResetAt string `json:"reset_at,omitempty"` // RFC3339,空表示无重置时间
}
// CNProviderQuotaProbeResult 是 Coding Plan 额度探测的返回结构(管理端 + UI 消费)。
type CNProviderQuotaProbeResult struct {
Provider string `json:"provider"`
Source string `json:"source"`
Success bool `json:"success"`
CredentialValid bool `json:"credential_valid"` // false = 401/403 鉴权失败
Tiers []CNQuotaTier `json:"tiers,omitempty"`
PlanLevel string `json:"plan_level,omitempty"` // 智谱套餐等级
StatusCode int `json:"status_code,omitempty"`
FetchedAt int64 `json:"fetched_at"`
Persisted bool `json:"persisted"`
Error string `json:"error,omitempty"`
}
// CNProviderQuotaService 探测 Kimi / Zhipu Coding Plan 的滚动窗口用量。
type CNProviderQuotaService struct {
accountRepo AccountRepository
proxyRepo ProxyRepository
httpUpstream HTTPUpstream
cfg *config.Config
flight singleflight.Group
}
// NewCNProviderQuotaService 构造 Coding Plan 额度探测服务。
func NewCNProviderQuotaService(
accountRepo AccountRepository,
proxyRepo ProxyRepository,
httpUpstream HTTPUpstream,
cfg *config.Config,
) *CNProviderQuotaService {
return &CNProviderQuotaService{
accountRepo: accountRepo,
proxyRepo: proxyRepo,
httpUpstream: httpUpstream,
cfg: cfg,
}
}
// QueryUsage 探测指定账号的 Coding Plan 滚动窗口用量并落 Extra 快照。
// 同一账号的并发探测会被 singleflight 合并。
func (s *CNProviderQuotaService) QueryUsage(ctx context.Context, accountID int64) (*CNProviderQuotaProbeResult, error) {
account, err := s.loadCodingPlanAccount(ctx, accountID)
if err != nil {
return nil, err
}
return s.QueryUsageForAccount(ctx, account)
}
// QueryUsageForAccount 探测已加载账号(配额监控 fetcher 复用,避免二次 GetByID)。
// singleflight key 与 QueryUsage 相同,按账号 ID 与 admin 侧并发探测合并。
func (s *CNProviderQuotaService) QueryUsageForAccount(ctx context.Context, account *Account) (*CNProviderQuotaProbeResult, error) {
if s == nil || s.accountRepo == nil || s.httpUpstream == nil {
return nil, infraerrors.New(http.StatusInternalServerError, "CN_QUOTA_NOT_CONFIGURED", "cn provider quota service is not configured")
}
if err := validateCodingPlanAccount(account); err != nil {
return nil, err
}
key := "cn_quota:" + strconv.FormatInt(account.ID, 10)
resultCh := s.flight.DoChan(key, func() (any, error) {
probeCtx, cancel := context.WithTimeout(context.Background(), cnQuotaUpstreamTimeout+5*time.Second)
defer cancel()
return s.queryUsageForAccount(probeCtx, account)
})
select {
case <-ctx.Done():
return nil, ctx.Err()
case flightResult := <-resultCh:
if flightResult.Err != nil {
return nil, flightResult.Err
}
result, ok := flightResult.Val.(*CNProviderQuotaProbeResult)
if !ok || result == nil {
return nil, infraerrors.New(http.StatusInternalServerError, "CN_QUOTA_PROBE_RESULT_INVALID", "invalid cn provider quota probe result")
}
cloned := *result
return &cloned, nil
}
}
func (s *CNProviderQuotaService) queryUsageForAccount(ctx context.Context, account *Account) (*CNProviderQuotaProbeResult, error) {
provider := account.GetCodingPlanProvider()
if provider != PlatformKimi && provider != PlatformZhipu {
return nil, infraerrors.New(http.StatusBadRequest, "CN_QUOTA_NOT_CODING_PLAN", "account is not a kimi/zhipu coding plan account")
}
apiKey := strings.TrimSpace(account.GetCNAPIKey())
if apiKey == "" {
return nil, infraerrors.New(http.StatusBadRequest, "CN_QUOTA_NO_APIKEY", "account api_key is empty")
}
baseURL := account.GetOpenAIBaseURL()
var (
targetURL string
authHeader string
)
switch provider {
case PlatformKimi:
targetURL = kimiQuotaURL(baseURL)
authHeader = "Bearer " + apiKey
case PlatformZhipu:
targetURL = zhipuQuotaURL(baseURL)
authHeader = apiKey // 智谱额度端点鉴权不加 Bearer 前缀
}
// 探测发起前过出站 URL 安全策略(与网关转发/Grok 探测同一套校验):
// 端点多由账号 base_url 衍生,不得把 API key 发往策略外主机。
validatedURL, err := cnValidateProbeURL(s.cfg, targetURL)
if err != nil {
return nil, infraerrors.New(http.StatusForbidden, "CN_QUOTA_URL_REJECTED", err.Error())
}
targetURL = validatedURL
proxyURL := s.resolveProxyURL(ctx, account)
callCtx, cancel := context.WithTimeout(ctx, cnQuotaUpstreamTimeout)
defer cancel()
req, err := http.NewRequestWithContext(callCtx, http.MethodGet, targetURL, nil)
if err != nil {
return nil, infraerrors.Newf(http.StatusInternalServerError, "CN_QUOTA_REQUEST_BUILD_FAILED", "build request: %v", err)
}
req.Header.Set("Authorization", authHeader)
req.Header.Set("Accept", "application/json")
if provider == PlatformZhipu {
req.Header.Set("Content-Type", "application/json")
req.Header.Set("Accept-Language", "en-US,en")
}
// 探测与真实转发保持同一套账号级请求头覆写,避免探测通过但转发失败。
account.ApplyHeaderOverrides(req.Header)
resp, err := s.httpUpstream.Do(req, proxyURL, account.ID, maxInt(account.Concurrency, 1))
if err != nil {
return nil, infraerrors.Newf(http.StatusBadGateway, "CN_QUOTA_REQUEST_FAILED", "upstream request failed: %v", err)
}
defer func() { _ = resp.Body.Close() }()
bodyBytes, _ := io.ReadAll(io.LimitReader(resp.Body, cnQuotaMaxBodyBytes))
now := time.Now().UTC()
result := &CNProviderQuotaProbeResult{
Provider: provider,
Source: "coding_plan",
FetchedAt: now.Unix(),
StatusCode: resp.StatusCode,
}
if resp.StatusCode == http.StatusUnauthorized || resp.StatusCode == http.StatusForbidden {
// 鉴权失败:不落快照(不覆盖之前的有效值),仅返回失败结果供前端提示。
result.Error = fmt.Sprintf("Authentication failed (HTTP %d)", resp.StatusCode)
return result, nil
}
if resp.StatusCode < 200 || resp.StatusCode >= 300 {
result.Error = fmt.Sprintf("API error (HTTP %d): %s", resp.StatusCode, truncate(strings.TrimSpace(string(bodyBytes)), 240))
return result, nil
}
// 智谱业务级错误(HTTP 2xx 但 success=false)。
if provider == PlatformZhipu {
if success := gjson.GetBytes(bodyBytes, "success"); success.Exists() && !success.Bool() {
msg := strings.TrimSpace(gjson.GetBytes(bodyBytes, "msg").String())
if msg == "" {
msg = "unknown zhipu quota error"
}
result.Error = "API error: " + msg
return result, nil
}
}
var tiers []CNQuotaTier
switch provider {
case PlatformKimi:
tiers = parseKimiUsageTiers(bodyBytes)
case PlatformZhipu:
tiers = parseZhipuTokenTiers(gjson.GetBytes(bodyBytes, "data"))
result.PlanLevel = strings.TrimSpace(gjson.GetBytes(bodyBytes, "data.level").String())
}
result.Tiers = tiers
result.Success = true
result.CredentialValid = true
updates := cnQuotaExtraUpdates(provider, tiers, now)
if err := s.accountRepo.UpdateExtra(ctx, account.ID, updates); err != nil {
slog.Warn("cn_quota_persist_failed", "account_id", account.ID, "provider", provider, "error", err)
} else {
result.Persisted = true
}
return result, nil
}
func (s *CNProviderQuotaService) loadCodingPlanAccount(ctx context.Context, accountID int64) (*Account, error) {
account, err := s.accountRepo.GetByID(ctx, accountID)
if err != nil {
return nil, infraerrors.Newf(http.StatusNotFound, "CN_QUOTA_ACCOUNT_NOT_FOUND", "account not found: %v", err)
}
if err := validateCodingPlanAccount(account); err != nil {
return nil, err
}
return account, nil
}
// validateCodingPlanAccount 加载后的非 DB 校验(ForAccount 入口同样复用,
// 保证直传 account 也不绕过平台/模式检查)。
func validateCodingPlanAccount(account *Account) error {
if account == nil {
return infraerrors.New(http.StatusNotFound, "CN_QUOTA_ACCOUNT_NOT_FOUND", "account not found")
}
if !account.IsCNProvider() {
return infraerrors.New(http.StatusBadRequest, "CN_QUOTA_INVALID_PLATFORM", "account is not a CN provider account")
}
if !account.IsCodingPlan() {
return infraerrors.New(http.StatusBadRequest, "CN_QUOTA_NOT_CODING_PLAN", "account is not a coding plan account")
}
return nil
}
func (s *CNProviderQuotaService) resolveProxyURL(ctx context.Context, account *Account) string {
if account == nil || account.ProxyID == nil {
return ""
}
if account.Proxy != nil {
return account.Proxy.URL()
}
if s != nil && s.proxyRepo != nil {
if proxy, err := s.proxyRepo.GetByID(ctx, *account.ProxyID); err == nil && proxy != nil {
account.Proxy = proxy
return proxy.URL()
}
}
return ""
}
// zhipuQuotaURL 根据 base_url 解析智谱额度端点(与数据面推理域名同主机)。
func zhipuQuotaURL(baseURL string) string {
return zhipuQuotaHost(baseURL) + "/api/monitor/usage/quota/limit"
}
// kimiQuotaURL 根据 base_url 解析 Kimi For Coding 额度端点。
// cc-switch query_kimi 固定探测 https://api.kimi.com/coding/v1/usages
// (实测 /coding/usages 无 /v1 → 404)。coding/v1CC 协议默认)与
// codingAnthropic 协议默认)两种 base 统一剥掉尾部后拼回 /v1/usages,
// 协议切换不影响额度探测端点。
func kimiQuotaURL(baseURL string) string {
base := strings.TrimSuffix(strings.TrimRight(baseURL, "/"), "/v1")
return base + "/v1/usages"
}
func zhipuQuotaHost(baseURL string) string {
switch u := strings.ToLower(baseURL); {
case strings.Contains(u, "bigmodel.cn"):
return "https://open.bigmodel.cn"
case strings.Contains(u, "z.ai"):
return "https://api.z.ai"
default:
// 国产优先:未知域名回落国内站(与前端 zhipu 预设一致)。
return "https://open.bigmodel.cn"
}
}
// parseKimiUsageTiers 解析 Kimi For Coding 的 /usages 响应。
//
// - limits[].detail.{limit,remaining,resetTime} → 5h 窗口(取首个 detail
// - usage.{limit,remaining,resetTime} → 周窗口
//
// utilization = (limit-remaining)/limit*100。
func parseKimiUsageTiers(body []byte) []CNQuotaTier {
var tiers []CNQuotaTier
if limits := gjson.GetBytes(body, "limits"); limits.IsArray() {
limits.ForEach(func(_, item gjson.Result) bool {
detail := item.Get("detail")
if !detail.Exists() {
return true
}
limit, _ := cnParseF64(detail.Get("limit").Value())
remaining, _ := cnParseF64(detail.Get("remaining").Value())
used := limit - remaining
if used < 0 {
used = 0
}
var util float64
if limit > 0 {
util = used / limit * 100
}
tiers = append(tiers, CNQuotaTier{
Window: "5h",
UsedPercent: util,
ResetAt: cnNormalizeResetTime(detail.Get("resetTime").Value()),
})
return false // 取首个 detail 作为 5h 窗口
})
}
if usage := gjson.GetBytes(body, "usage"); usage.Exists() {
limit, _ := cnParseF64(usage.Get("limit").Value())
remaining, _ := cnParseF64(usage.Get("remaining").Value())
used := limit - remaining
if used < 0 {
used = 0
}
var util float64
if limit > 0 {
util = used / limit * 100
}
tiers = append(tiers, CNQuotaTier{
Window: "weekly",
UsedPercent: util,
ResetAt: cnNormalizeResetTime(usage.Get("resetTime").Value()),
})
}
return tiers
}
// cnZhipuWindow 标识智谱 TOKENS_LIMIT 条目所属窗口。
type cnZhipuWindow int
const (
cnZhipuWindowUnknown cnZhipuWindow = iota
cnZhipuWindow5h
cnZhipuWindowWeekly
)
// classifyZhipuWindowUnit 按 unit 字段判定窗口类型(3=5h6=weekly)。
// unit 缺失或未识别时返回 Unknown,由调用方走 reset 时间启发式兜底。
func classifyZhipuWindowUnit(unit int64) cnZhipuWindow {
switch unit {
case 3:
return cnZhipuWindow5h
case 6:
return cnZhipuWindowWeekly
default:
return cnZhipuWindowUnknown
}
}
// parseZhipuTokenTiers 解析智谱额度响应 data.limits 为 5h + weekly 两档。
//
// 分类优先级(对齐 cc-switch parse_zhipu_token_tiersissue #3036):
// 1. 显式 unit 字段(3=5h / 6=weekly)——不能用 reset 排序代替,周期末尾
// 周窗口会比 5h 更早重置,时间排序必然标反。
// 2. unit 缺失/未识别:无 nextResetTime 的条目优先归 5h0% 状态下 5h 桶可能
// 没有 reset),其余按 reset 升序依次填入仍空缺的槽位。
//
// CREDIT_LIMIT(信用额度)与 TOKENS_LIMIT(token 窗口)度量不同:两者同时返回时
// 只让 TOKENS_LIMIT 参与 5h/weekly 槽位竞争,避免信用额度百分比污染阈值停调
// 快照;仅当无任何 TOKENS_LIMIT 条目时才降级用 CREDIT_LIMIT 展示。
// 老套餐只回 1 条 TOKENS_LIMIT,自然降级为仅 5h;新套餐回 2 条。
func parseZhipuTokenTiers(data gjson.Result) []CNQuotaTier {
type entry struct {
resetMs int64
hasReset bool
percentage float64
resetISO string
}
var (
fiveHour entry
fiveHourSet bool
weekly entry
weeklySet bool
unclassified []entry
)
classify := func(item gjson.Result, e entry) {
switch classifyZhipuWindowUnit(item.Get("unit").Int()) {
case cnZhipuWindow5h:
if !fiveHourSet {
fiveHour, fiveHourSet = e, true
} else {
unclassified = append(unclassified, e)
}
case cnZhipuWindowWeekly:
if !weeklySet {
weekly, weeklySet = e, true
} else {
unclassified = append(unclassified, e)
}
default:
unclassified = append(unclassified, e)
}
}
var creditFallback []entry
hasTokensLimit := false
data.Get("limits").ForEach(func(_, item gjson.Result) bool {
limitType := strings.ToUpper(strings.TrimSpace(item.Get("type").String()))
if limitType != "TOKENS_LIMIT" && limitType != "CREDIT_LIMIT" {
return true
}
percentage := 0.0
if p, ok := cnParseF64(item.Get("percentage").Value()); ok {
percentage = p
}
var (
resetMs int64
hasReset bool
resetISO string
)
if nr := item.Get("nextResetTime"); nr.Exists() {
switch nr.Type {
case gjson.Number:
resetMs = nr.Int()
hasReset = resetMs > 0
resetISO = cnMillisToRFC3339(resetMs)
case gjson.String:
resetISO = cnNormalizeResetTime(nr.String())
hasReset = resetISO != ""
}
}
e := entry{resetMs: resetMs, hasReset: hasReset, percentage: percentage, resetISO: resetISO}
if limitType == "TOKENS_LIMIT" {
hasTokensLimit = true
classify(item, e)
} else {
creditFallback = append(creditFallback, e)
}
return true
})
// 无任何 TOKENS_LIMIT 条目(部分套餐只报信用额度):降级用 CREDIT_LIMIT 展示。
if !hasTokensLimit {
unclassified = append(unclassified, creditFallback...)
}
// 无 reset 的条目排前,再按 reset 升序,依次填入仍空缺的槽位。
sort.SliceStable(unclassified, func(i, j int) bool {
if unclassified[i].hasReset != unclassified[j].hasReset {
return !unclassified[i].hasReset
}
return unclassified[i].resetMs < unclassified[j].resetMs
})
for _, e := range unclassified {
switch {
case !fiveHourSet:
fiveHour, fiveHourSet = e, true
case !weeklySet:
weekly, weeklySet = e, true
}
}
var tiers []CNQuotaTier
if fiveHourSet {
tiers = append(tiers, CNQuotaTier{Window: "5h", UsedPercent: fiveHour.percentage, ResetAt: fiveHour.resetISO})
}
if weeklySet {
tiers = append(tiers, CNQuotaTier{Window: "weekly", UsedPercent: weekly.percentage, ResetAt: weekly.resetISO})
}
return tiers
}
// cnQuotaExtraUpdates 根据 tier 列表构造 provider 维度的 Extra 快照更新。
func cnQuotaExtraUpdates(provider string, tiers []CNQuotaTier, now time.Time) map[string]any {
updates := map[string]any{
cnExtraKey(provider, cnExtraSuffixUsageUpdated): now.Format(time.RFC3339),
}
for _, t := range tiers {
switch t.Window {
case "5h":
updates[cnExtraKey(provider, cnExtraSuffix5hUsed)] = t.UsedPercent
if t.ResetAt != "" {
updates[cnExtraKey(provider, cnExtraSuffix5hReset)] = t.ResetAt
}
case "weekly":
updates[cnExtraKey(provider, cnExtraSuffixWeeklyUsed)] = t.UsedPercent
if t.ResetAt != "" {
updates[cnExtraKey(provider, cnExtraSuffixWeeklyReset)] = t.ResetAt
}
}
}
return updates
}
// cnParseF64 把 JSON 数值或字符串解析为 float64(兼容 "100" 与 100)。
func cnParseF64(raw any) (float64, bool) {
switch v := raw.(type) {
case float64:
return v, true
case float32:
return float64(v), true
case int:
return float64(v), true
case int64:
return float64(v), true
case json.Number:
f, err := v.Float64()
return f, err == nil
case string:
f, err := strconv.ParseFloat(strings.TrimSpace(v), 64)
return f, err == nil
default:
return 0, false
}
}
// cnNormalizeResetTime 把上游重置时间(ISO8601 字符串 / 秒级 / 毫秒级数字)归一化为
// RFC3339 字符串;无法识别或非正时间戳返回空串。
func cnNormalizeResetTime(raw any) string {
switch v := raw.(type) {
case string:
s := strings.TrimSpace(v)
if s == "" {
return ""
}
if ts, err := parseSchedulingTime(s); err == nil {
return ts.UTC().Format(time.RFC3339)
}
return ""
case float64:
return cnMillisToRFC3339(int64(v))
case int:
return cnMillisToRFC3339(int64(v))
case int64:
return cnMillisToRFC3339(v)
case json.Number:
if n, err := v.Int64(); err == nil {
return cnMillisToRFC3339(n)
}
return ""
default:
return ""
}
}
// cnMillisToRFC3339 把秒级(<1e12)或毫秒级时间戳转为 RFC3339 字符串;非正返回空串。
func cnMillisToRFC3339(n int64) string {
if n <= 0 {
return ""
}
var ms int64
if n < 1_000_000_000_000 {
ms = n * 1000
} else {
ms = n
}
return time.UnixMilli(ms).UTC().Format(time.RFC3339)
}