Sub2API v1.0 - AI API 网关（二开初始版本，基于上游 Wei-Shaw/sub2api）

backend/internal/service/cn_provider_balance_service.go

@@ -0,0 +1,295 @@
+package service
+
+import (
+	"context"
+	"fmt"
+	"io"
+	"log/slog"
+	"net/http"
+	"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"
+)
+
+// 国产供应商 payg（按量付费）账号余额探测服务。
+//
+// 仅覆盖有公开余额端点的供应商：
+//   - Kimi/Moonshot：GET https://api.moonshot.cn/v1/users/me/balance (Bearer) → data.available_balance
+//   - DeepSeek：     GET https://api.deepseek.com/user/balance (Bearer) → balance_infos[].total_balance + is_available
+//
+// 智谱（zhipu）无公开余额端点（OpenAPI 规格验证），仅靠响应式 429/402（见
+// ratelimit_cn_providers.go）。解析逻辑对齐 cc-switch services/balance.rs::query_deepseek。
+const (
+	cnBalanceUpstreamTimeout = 15 * time.Second
+	cnBalanceMaxBodyBytes    = 256 * 1024
+
+	// Extra 余额快照键后缀（加 provider 前缀）。
+	cnBalanceExtraSuffixBalance   = "balance"
+	cnBalanceExtraSuffixCurrency  = "balance_currency"
+	cnBalanceExtraSuffixAvailable = "balance_available" // deepseek is_available 健康标记
+	cnBalanceExtraSuffixUpdated   = "balance_updated_at"
+	cnBalanceExtraSuffixBalances  = "balances" // 多币种明细（deepseek USD+CNY）
+)
+
+// CNProviderBalanceEntry 是单一币种的余额明细。
+type CNProviderBalanceEntry struct {
+	Currency string  `json:"currency"`
+	Balance  float64 `json:"balance"`
+}
+
+// CNProviderBalanceResult 是余额探测的返回结构（管理端 + UI 消费）。
+type CNProviderBalanceResult struct {
+	Provider string `json:"provider"`
+	Success  bool   `json:"success"`
+	// Balance/Currency 为主币种（balance_infos 首条，兼容单币种消费方）；
+	// 完整明细见 Balances（deepseek 双币种账号含 CNY + USD 两条）。
+	Balance    float64                  `json:"balance"`
+	Currency   string                   `json:"currency,omitempty"`
+	Balances   []CNProviderBalanceEntry `json:"balances,omitempty"`
+	Available  bool                     `json:"available"` // 健康标记（deepseek is_available；kimi 无此概念恒 true）
+	StatusCode int                      `json:"status_code,omitempty"`
+	FetchedAt  int64                    `json:"fetched_at"`
+	Persisted  bool                     `json:"persisted"`
+	Error      string                   `json:"error,omitempty"`
+}
+
+// CNProviderBalanceService 探测 Kimi / DeepSeek payg 账号的账户余额。
+type CNProviderBalanceService struct {
+	accountRepo  AccountRepository
+	proxyRepo    ProxyRepository
+	httpUpstream HTTPUpstream
+	cfg          *config.Config
+	flight       singleflight.Group
+}
+
+// NewCNProviderBalanceService 构造余额探测服务。
+func NewCNProviderBalanceService(
+	accountRepo AccountRepository,
+	proxyRepo ProxyRepository,
+	httpUpstream HTTPUpstream,
+	cfg *config.Config,
+) *CNProviderBalanceService {
+	return &CNProviderBalanceService{
+		accountRepo:  accountRepo,
+		proxyRepo:    proxyRepo,
+		httpUpstream: httpUpstream,
+		cfg:          cfg,
+	}
+}
+
+// QueryBalance 探测指定 payg 账号的余额并落 Extra 快照。
+func (s *CNProviderBalanceService) QueryBalance(ctx context.Context, accountID int64) (*CNProviderBalanceResult, error) {
+	account, err := s.loadPayGAccount(ctx, accountID)
+	if err != nil {
+		return nil, err
+	}
+	return s.QueryBalanceForAccount(ctx, account)
+}
+
+// QueryBalanceForAccount 探测已加载账号（配额监控 fetcher / 周期余额检测复用，
+// 避免二次 GetByID）。singleflight key 与 QueryBalance 相同，按账号 ID 合并。
+func (s *CNProviderBalanceService) QueryBalanceForAccount(ctx context.Context, account *Account) (*CNProviderBalanceResult, error) {
+	if s == nil || s.accountRepo == nil || s.httpUpstream == nil {
+		return nil, infraerrors.New(http.StatusInternalServerError, "CN_BALANCE_NOT_CONFIGURED", "cn provider balance service is not configured")
+	}
+	if err := validatePayGAccount(account); err != nil {
+		return nil, err
+	}
+	key := "cn_balance:" + strconv.FormatInt(account.ID, 10)
+	resultCh := s.flight.DoChan(key, func() (any, error) {
+		probeCtx, cancel := context.WithTimeout(context.Background(), cnBalanceUpstreamTimeout+5*time.Second)
+		defer cancel()
+		return s.queryBalanceForAccount(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.(*CNProviderBalanceResult)
+		if !ok || result == nil {
+			return nil, infraerrors.New(http.StatusInternalServerError, "CN_BALANCE_PROBE_RESULT_INVALID", "invalid cn provider balance probe result")
+		}
+		cloned := *result
+		return &cloned, nil
+	}
+}
+
+func (s *CNProviderBalanceService) queryBalanceForAccount(ctx context.Context, account *Account) (*CNProviderBalanceResult, error) {
+	provider := account.Platform
+	if provider != PlatformKimi && provider != PlatformDeepseek {
+		return nil, infraerrors.New(http.StatusBadRequest, "CN_BALANCE_NO_ENDPOINT", "account provider has no balance endpoint")
+	}
+
+	apiKey := strings.TrimSpace(account.GetCNAPIKey())
+	if apiKey == "" {
+		return nil, infraerrors.New(http.StatusBadRequest, "CN_BALANCE_NO_APIKEY", "account api_key is empty")
+	}
+
+	targetURL := cnBalanceURL(account)
+	// 探测发起前过出站 URL 安全策略（与网关转发/Grok 探测同一套校验）：
+	// DeepSeek 端点由账号 base_url 衍生，不得把 API key 发往策略外主机。
+	validatedURL, err := cnValidateProbeURL(s.cfg, targetURL)
+	if err != nil {
+		return nil, infraerrors.New(http.StatusForbidden, "CN_BALANCE_URL_REJECTED", err.Error())
+	}
+	targetURL = validatedURL
+	proxyURL := s.resolveProxyURL(ctx, account)
+	callCtx, cancel := context.WithTimeout(ctx, cnBalanceUpstreamTimeout)
+	defer cancel()
+	req, err := http.NewRequestWithContext(callCtx, http.MethodGet, targetURL, nil)
+	if err != nil {
+		return nil, infraerrors.Newf(http.StatusInternalServerError, "CN_BALANCE_REQUEST_BUILD_FAILED", "build request: %v", err)
+	}
+	req.Header.Set("Authorization", "Bearer "+apiKey)
+	req.Header.Set("Accept", "application/json")
+	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_BALANCE_REQUEST_FAILED", "upstream request failed: %v", err)
+	}
+	defer func() { _ = resp.Body.Close() }()
+	bodyBytes, _ := io.ReadAll(io.LimitReader(resp.Body, cnBalanceMaxBodyBytes))
+
+	now := time.Now().UTC()
+	result := &CNProviderBalanceResult{
+		Provider:   provider,
+		FetchedAt:  now.Unix(),
+		StatusCode: resp.StatusCode,
+		Available:  true,
+	}
+
+	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
+	}
+
+	var entries []CNProviderBalanceEntry
+	available := true
+	switch provider {
+	case PlatformKimi:
+		// Moonshot：code==0 成功；data.available_balance（number），单币种 CNY。
+		balance, _ := cnParseF64(gjson.GetBytes(bodyBytes, "data.available_balance").Value())
+		entries = append(entries, CNProviderBalanceEntry{Currency: "CNY", Balance: balance})
+	case PlatformDeepseek:
+		// is_available 缺省视为 true（健康）；显式存在时取其值。
+		if v := gjson.GetBytes(bodyBytes, "is_available"); v.Exists() {
+			available = v.Bool()
+		}
+		// balance_infos 逐条解析：双币种账号同时返回 CNY + USD（数组顺序即
+		// 主次序，首条为主币种）。
+		gjson.GetBytes(bodyBytes, "balance_infos").ForEach(func(_, info gjson.Result) bool {
+			currency := strings.ToUpper(strings.TrimSpace(info.Get("currency").String()))
+			balance, _ := cnParseF64(info.Get("total_balance").Value())
+			if currency == "" {
+				currency = "CNY"
+			}
+			entries = append(entries, CNProviderBalanceEntry{Currency: currency, Balance: balance})
+			return true
+		})
+		if len(entries) == 0 {
+			entries = append(entries, CNProviderBalanceEntry{Currency: "CNY"})
+		}
+	}
+	result.Balances = entries
+	result.Balance = entries[0].Balance
+	result.Currency = entries[0].Currency
+	result.Available = available
+	result.Success = true
+
+	balanceUpdates := make([]any, 0, len(entries))
+	for _, entry := range entries {
+		balanceUpdates = append(balanceUpdates, map[string]any{
+			"currency": entry.Currency,
+			"balance":  entry.Balance,
+		})
+	}
+	updates := map[string]any{
+		cnExtraKey(provider, cnBalanceExtraSuffixBalance):   result.Balance,
+		cnExtraKey(provider, cnBalanceExtraSuffixCurrency):  result.Currency,
+		cnExtraKey(provider, cnBalanceExtraSuffixAvailable): available,
+		cnExtraKey(provider, cnBalanceExtraSuffixUpdated):   now.Format(time.RFC3339),
+		cnExtraKey(provider, cnBalanceExtraSuffixBalances):  balanceUpdates,
+		// 余额探测成功即清除响应式 402/429 写下的 balance_low 标记。
+		cnExtraKey(provider, cnBalanceExtraSuffixLow): false,
+	}
+	if err := s.accountRepo.UpdateExtra(ctx, account.ID, updates); err != nil {
+		slog.Warn("cn_balance_persist_failed", "account_id", account.ID, "provider", provider, "error", err)
+	} else {
+		result.Persisted = true
+	}
+	return result, nil
+}
+
+// loadPayGAccount 加载 payg 模式的国产供应商账号（余额仅对 payg 有意义；coding 走额度）。
+func (s *CNProviderBalanceService) loadPayGAccount(ctx context.Context, accountID int64) (*Account, error) {
+	account, err := s.accountRepo.GetByID(ctx, accountID)
+	if err != nil {
+		return nil, infraerrors.Newf(http.StatusNotFound, "CN_BALANCE_ACCOUNT_NOT_FOUND", "account not found: %v", err)
+	}
+	if err := validatePayGAccount(account); err != nil {
+		return nil, err
+	}
+	return account, nil
+}
+
+// validatePayGAccount 加载后的非 DB 校验（ForAccount 入口同样复用，
+// 保证直传 account 也不绕过平台/模式检查）。
+func validatePayGAccount(account *Account) error {
+	if account == nil {
+		return infraerrors.New(http.StatusNotFound, "CN_BALANCE_ACCOUNT_NOT_FOUND", "account not found")
+	}
+	if !account.IsCNProvider() {
+		return infraerrors.New(http.StatusBadRequest, "CN_BALANCE_INVALID_PLATFORM", "account is not a CN provider account")
+	}
+	// coding 账号走额度探测，余额端点不适用。
+	if account.IsCodingPlan() {
+		return infraerrors.New(http.StatusBadRequest, "CN_BALANCE_CODING_PLAN", "coding plan account has no balance endpoint; use quota probe")
+	}
+	return nil
+}
+
+func (s *CNProviderBalanceService) 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 ""
+}
+
+// cnBalanceURL 解析账号的余额端点。
+//
+//   - Kimi：固定 https://api.moonshot.cn/v1/users/me/balance（与 base_url 无关，Moonshot 仅此一处）
+//   - DeepSeek：基于 base_url 拼接 /user/balance（支持自定义域名）
+func cnBalanceURL(account *Account) string {
+	switch account.Platform {
+	case PlatformKimi:
+		return "https://api.moonshot.cn/v1/users/me/balance"
+	case PlatformDeepseek:
+		// Anthropic 协议账号的凭证 base_url 指向 /anthropic 端点，余额探测需回退
+		// 到 OpenAI 格式 base（协议感知）再拼接 /user/balance。
+		return strings.TrimRight(account.GetOpenAIFormatBaseURL(), "/") + "/user/balance"
+	default:
+		return ""
+	}
+}