273 lines
8.4 KiB
Go
273 lines
8.4 KiB
Go
|
package service
|
|||
|
|
|
|||
|
|
import (
|
|||
|
|
"context"
|
|||
|
|
"encoding/json"
|
|||
|
|
"errors"
|
|||
|
|
"fmt"
|
|||
|
|
"log/slog"
|
|||
|
|
"regexp"
|
|||
|
|
"strings"
|
|||
|
|
"time"
|
|||
|
|
|
|||
|
|
"github.com/Wei-Shaw/sub2api/internal/pkg/antigravity"
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
const (
|
|||
|
|
forbiddenTypeValidation = "validation"
|
|||
|
|
forbiddenTypeViolation = "violation"
|
|||
|
|
forbiddenTypeForbidden = "forbidden"
|
|||
|
|
|
|||
|
|
// 机器可读的错误码
|
|||
|
|
errorCodeForbidden = "forbidden"
|
|||
|
|
errorCodeUnauthenticated = "unauthenticated"
|
|||
|
|
errorCodeRateLimited = "rate_limited"
|
|||
|
|
errorCodeNetworkError = "network_error"
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
// AntigravityQuotaFetcher 从 Antigravity API 获取额度
|
|||
|
|
type AntigravityQuotaFetcher struct {
|
|||
|
|
proxyRepo ProxyRepository
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
// NewAntigravityQuotaFetcher 创建 AntigravityQuotaFetcher
|
|||
|
|
func NewAntigravityQuotaFetcher(proxyRepo ProxyRepository) *AntigravityQuotaFetcher {
|
|||
|
|
return &AntigravityQuotaFetcher{proxyRepo: proxyRepo}
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
// CanFetch 检查是否可以获取此账户的额度
|
|||
|
|
func (f *AntigravityQuotaFetcher) CanFetch(account *Account) bool {
|
|||
|
|
if account.Platform != PlatformAntigravity {
|
|||
|
|
return false
|
|||
|
|
}
|
|||
|
|
accessToken := account.GetCredential("access_token")
|
|||
|
|
return accessToken != ""
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
// FetchQuota 获取 Antigravity 账户额度信息
|
|||
|
|
func (f *AntigravityQuotaFetcher) FetchQuota(ctx context.Context, account *Account, proxyURL string) (*QuotaResult, error) {
|
|||
|
|
accessToken := account.GetCredential("access_token")
|
|||
|
|
projectID := account.GetCredential("project_id")
|
|||
|
|
|
|||
|
|
client, err := antigravity.NewClient(proxyURL)
|
|||
|
|
if err != nil {
|
|||
|
|
return nil, fmt.Errorf("create antigravity client failed: %w", err)
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
// 调用 API 获取配额
|
|||
|
|
modelsResp, modelsRaw, err := client.FetchAvailableModels(ctx, accessToken, projectID)
|
|||
|
|
if err != nil {
|
|||
|
|
// 403 Forbidden: 不报错，返回 is_forbidden 标记
|
|||
|
|
var forbiddenErr *antigravity.ForbiddenError
|
|||
|
|
if errors.As(err, &forbiddenErr) {
|
|||
|
|
now := time.Now()
|
|||
|
|
fbType := classifyForbiddenType(forbiddenErr.Body)
|
|||
|
|
return &QuotaResult{
|
|||
|
|
UsageInfo: &UsageInfo{
|
|||
|
|
UpdatedAt: &now,
|
|||
|
|
IsForbidden: true,
|
|||
|
|
ForbiddenReason: forbiddenErr.Body,
|
|||
|
|
ForbiddenType: fbType,
|
|||
|
|
ValidationURL: extractValidationURL(forbiddenErr.Body),
|
|||
|
|
NeedsVerify: fbType == forbiddenTypeValidation,
|
|||
|
|
IsBanned: fbType == forbiddenTypeViolation,
|
|||
|
|
ErrorCode: errorCodeForbidden,
|
|||
|
|
},
|
|||
|
|
}, nil
|
|||
|
|
}
|
|||
|
|
return nil, err
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
// 调用 LoadCodeAssist 获取订阅等级和 AI Credits 余额（非关键路径，失败不影响主流程）
|
|||
|
|
tierRaw, tierNormalized, loadResp := f.fetchSubscriptionTier(ctx, client, accessToken)
|
|||
|
|
|
|||
|
|
// 转换为 UsageInfo
|
|||
|
|
usageInfo := f.buildUsageInfo(modelsResp, tierRaw, tierNormalized, loadResp)
|
|||
|
|
|
|||
|
|
return &QuotaResult{
|
|||
|
|
UsageInfo: usageInfo,
|
|||
|
|
Raw: modelsRaw,
|
|||
|
|
}, nil
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
// fetchSubscriptionTier 获取账号订阅等级，失败返回空字符串。
|
|||
|
|
// 同时返回 LoadCodeAssistResponse，以便提取 AI Credits 余额。
|
|||
|
|
func (f *AntigravityQuotaFetcher) fetchSubscriptionTier(ctx context.Context, client *antigravity.Client, accessToken string) (raw, normalized string, loadResp *antigravity.LoadCodeAssistResponse) {
|
|||
|
|
loadResp, _, err := client.LoadCodeAssist(ctx, accessToken)
|
|||
|
|
if err != nil {
|
|||
|
|
slog.Warn("failed to fetch subscription tier", "error", err)
|
|||
|
|
return "", "", nil
|
|||
|
|
}
|
|||
|
|
if loadResp == nil {
|
|||
|
|
return "", "", nil
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
raw = loadResp.GetTier() // 已有方法：paidTier > currentTier
|
|||
|
|
normalized = normalizeTier(raw)
|
|||
|
|
return raw, normalized, loadResp
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
// normalizeTier 将原始 tier 字符串归一化为 FREE/PRO/ULTRA/UNKNOWN
|
|||
|
|
func normalizeTier(raw string) string {
|
|||
|
|
if raw == "" {
|
|||
|
|
return ""
|
|||
|
|
}
|
|||
|
|
lower := strings.ToLower(raw)
|
|||
|
|
switch {
|
|||
|
|
case strings.Contains(lower, "ultra"):
|
|||
|
|
return "ULTRA"
|
|||
|
|
case strings.Contains(lower, "pro"):
|
|||
|
|
return "PRO"
|
|||
|
|
case strings.Contains(lower, "free"):
|
|||
|
|
return "FREE"
|
|||
|
|
default:
|
|||
|
|
return "UNKNOWN"
|
|||
|
|
}
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
// buildUsageInfo 将 API 响应转换为 UsageInfo。
|
|||
|
|
func (f *AntigravityQuotaFetcher) buildUsageInfo(modelsResp *antigravity.FetchAvailableModelsResponse, tierRaw, tierNormalized string, loadResp *antigravity.LoadCodeAssistResponse) *UsageInfo {
|
|||
|
|
now := time.Now()
|
|||
|
|
info := &UsageInfo{
|
|||
|
|
UpdatedAt: &now,
|
|||
|
|
AntigravityQuota: make(map[string]*AntigravityModelQuota),
|
|||
|
|
AntigravityQuotaDetails: make(map[string]*AntigravityModelDetail),
|
|||
|
|
SubscriptionTier: tierNormalized,
|
|||
|
|
SubscriptionTierRaw: tierRaw,
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
// 遍历所有模型，填充 AntigravityQuota 和 AntigravityQuotaDetails
|
|||
|
|
for modelName, modelInfo := range modelsResp.Models {
|
|||
|
|
if modelInfo.QuotaInfo == nil {
|
|||
|
|
continue
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
// remainingFraction 是剩余比例 (0.0-1.0)，转换为使用率百分比
|
|||
|
|
utilization := int((1.0 - modelInfo.QuotaInfo.RemainingFraction) * 100)
|
|||
|
|
|
|||
|
|
info.AntigravityQuota[modelName] = &AntigravityModelQuota{
|
|||
|
|
Utilization: utilization,
|
|||
|
|
ResetTime: modelInfo.QuotaInfo.ResetTime,
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
// 填充模型详细能力信息
|
|||
|
|
detail := &AntigravityModelDetail{
|
|||
|
|
DisplayName: modelInfo.DisplayName,
|
|||
|
|
SupportsImages: modelInfo.SupportsImages,
|
|||
|
|
SupportsThinking: modelInfo.SupportsThinking,
|
|||
|
|
ThinkingBudget: modelInfo.ThinkingBudget,
|
|||
|
|
Recommended: modelInfo.Recommended,
|
|||
|
|
MaxTokens: modelInfo.MaxTokens,
|
|||
|
|
MaxOutputTokens: modelInfo.MaxOutputTokens,
|
|||
|
|
SupportedMimeTypes: modelInfo.SupportedMimeTypes,
|
|||
|
|
}
|
|||
|
|
info.AntigravityQuotaDetails[modelName] = detail
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
// 废弃模型转发规则
|
|||
|
|
if len(modelsResp.DeprecatedModelIDs) > 0 {
|
|||
|
|
info.ModelForwardingRules = make(map[string]string, len(modelsResp.DeprecatedModelIDs))
|
|||
|
|
for oldID, deprecated := range modelsResp.DeprecatedModelIDs {
|
|||
|
|
info.ModelForwardingRules[oldID] = deprecated.NewModelID
|
|||
|
|
}
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
// 同时设置 FiveHour 用于兼容展示（取主要模型）
|
|||
|
|
priorityModels := []string{"claude-sonnet-4-20250514", "claude-sonnet-4", "gemini-2.5-pro"}
|
|||
|
|
for _, modelName := range priorityModels {
|
|||
|
|
if modelInfo, ok := modelsResp.Models[modelName]; ok && modelInfo.QuotaInfo != nil {
|
|||
|
|
utilization := (1.0 - modelInfo.QuotaInfo.RemainingFraction) * 100
|
|||
|
|
progress := &UsageProgress{
|
|||
|
|
Utilization: utilization,
|
|||
|
|
}
|
|||
|
|
if modelInfo.QuotaInfo.ResetTime != "" {
|
|||
|
|
if resetTime, err := time.Parse(time.RFC3339, modelInfo.QuotaInfo.ResetTime); err == nil {
|
|||
|
|
progress.ResetsAt = &resetTime
|
|||
|
|
progress.RemainingSeconds = int(time.Until(resetTime).Seconds())
|
|||
|
|
}
|
|||
|
|
}
|
|||
|
|
info.FiveHour = progress
|
|||
|
|
break
|
|||
|
|
}
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
if loadResp != nil {
|
|||
|
|
for _, credit := range loadResp.GetAvailableCredits() {
|
|||
|
|
info.AICredits = append(info.AICredits, AICredit{
|
|||
|
|
CreditType: credit.CreditType,
|
|||
|
|
Amount: credit.GetAmount(),
|
|||
|
|
MinimumBalance: credit.GetMinimumAmount(),
|
|||
|
|
})
|
|||
|
|
}
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
return info
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
// GetProxyURL 获取账户的代理 URL
|
|||
|
|
func (f *AntigravityQuotaFetcher) GetProxyURL(ctx context.Context, account *Account) string {
|
|||
|
|
if account.ProxyID == nil || f.proxyRepo == nil {
|
|||
|
|
return ""
|
|||
|
|
}
|
|||
|
|
proxy, err := f.proxyRepo.GetByID(ctx, *account.ProxyID)
|
|||
|
|
if err != nil || proxy == nil {
|
|||
|
|
return ""
|
|||
|
|
}
|
|||
|
|
return proxy.URL()
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
// classifyForbiddenType 根据 403 响应体判断禁止类型
|
|||
|
|
func classifyForbiddenType(body string) string {
|
|||
|
|
lower := strings.ToLower(body)
|
|||
|
|
switch {
|
|||
|
|
case strings.Contains(lower, "validation_required") ||
|
|||
|
|
strings.Contains(lower, "verify your account") ||
|
|||
|
|
strings.Contains(lower, "validation_url"):
|
|||
|
|
return forbiddenTypeValidation
|
|||
|
|
case strings.Contains(lower, "terms of service") ||
|
|||
|
|
strings.Contains(lower, "violation"):
|
|||
|
|
return forbiddenTypeViolation
|
|||
|
|
default:
|
|||
|
|
return forbiddenTypeForbidden
|
|||
|
|
}
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
// urlPattern 用于从 403 响应体中提取 URL（降级方案）
|
|||
|
|
var urlPattern = regexp.MustCompile(`https://[^\s"'\\]+`)
|
|||
|
|
|
|||
|
|
// extractValidationURL 从 403 响应 JSON 中提取验证/申诉链接
|
|||
|
|
func extractValidationURL(body string) string {
|
|||
|
|
// 1. 尝试结构化 JSON 提取: /error/details[*]/metadata/validation_url 或 appeal_url
|
|||
|
|
var parsed struct {
|
|||
|
|
Error struct {
|
|||
|
|
Details []struct {
|
|||
|
|
Metadata map[string]string `json:"metadata"`
|
|||
|
|
} `json:"details"`
|
|||
|
|
} `json:"error"`
|
|||
|
|
}
|
|||
|
|
if json.Unmarshal([]byte(body), &parsed) == nil {
|
|||
|
|
for _, detail := range parsed.Error.Details {
|
|||
|
|
if u := detail.Metadata["validation_url"]; u != "" {
|
|||
|
|
return u
|
|||
|
|
}
|
|||
|
|
if u := detail.Metadata["appeal_url"]; u != "" {
|
|||
|
|
return u
|
|||
|
|
}
|
|||
|
|
}
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
// 2. 降级：正则匹配 URL
|
|||
|
|
lower := strings.ToLower(body)
|
|||
|
|
if !strings.Contains(lower, "validation") &&
|
|||
|
|
!strings.Contains(lower, "verify") &&
|
|||
|
|
!strings.Contains(lower, "appeal") {
|
|||
|
|
return ""
|
|||
|
|
}
|
|||
|
|
// 先解码常见转义再匹配
|
|||
|
|
normalized := strings.ReplaceAll(body, `\u0026`, "&")
|
|||
|
|
if m := urlPattern.FindString(normalized); m != "" {
|
|||
|
|
return m
|
|||
|
|
}
|
|||
|
|
return ""
|
|||
|
|
}
|