From 4543a6f0432aaa6da21d9474e631becf109a66b1 Mon Sep 17 00:00:00 2001 From: song Date: Fri, 2 Jan 2026 22:41:55 +0800 Subject: [PATCH] =?UTF-8?q?refactor(antigravity):=20=E7=BB=9F=E4=B8=80?= =?UTF-8?q?=E9=A2=9D=E5=BA=A6=E5=88=B7=E6=96=B0=E6=9C=BA=E5=88=B6=E4=B8=8E?= =?UTF-8?q?=20Claude=20=E4=B8=80=E8=87=B4?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 将 Antigravity 的额度刷新从后台定时刷新改为按需获取模式，与 Claude 统一： - 删除 AntigravityQuotaRefresher 后台服务 - 新增 QuotaFetcher 接口和 AntigravityQuotaFetcher 实现 - 前端改为调用 usage API 获取额度，支持 loading/error 状态 - 统一使用内存缓存（10 分钟 TTL） --- backend/cmd/server/wire.go | 5 - backend/cmd/server/wire_gen.go | 11 +- .../internal/service/account_usage_service.go | 85 ++++++- .../service/antigravity_quota_fetcher.go | 111 +++++++++ .../service/antigravity_quota_refresher.go | 222 ------------------ backend/internal/service/quota_fetcher.go | 19 ++ backend/internal/service/wire.go | 14 +- .../components/account/AccountUsageCell.vue | 107 +++++---- frontend/src/types/index.ts | 7 + 9 files changed, 274 insertions(+), 307 deletions(-) create mode 100644 backend/internal/service/antigravity_quota_fetcher.go delete mode 100644 backend/internal/service/antigravity_quota_refresher.go create mode 100644 backend/internal/service/quota_fetcher.go diff --git a/backend/cmd/server/wire.go b/backend/cmd/server/wire.go index 8596b8ba..ff6ab4e6 100644 --- a/backend/cmd/server/wire.go +++ b/backend/cmd/server/wire.go @@ -70,7 +70,6 @@ func provideCleanup( openaiOAuth *service.OpenAIOAuthService, geminiOAuth *service.GeminiOAuthService, antigravityOAuth *service.AntigravityOAuthService, - antigravityQuota *service.AntigravityQuotaRefresher, ) func() { return func() { ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second) @@ -113,10 +112,6 @@ func provideCleanup( antigravityOAuth.Stop() return nil }}, - {"AntigravityQuotaRefresher", func() error { - antigravityQuota.Stop() - return nil - }}, {"Redis", func() error { return rdb.Close() }}, diff --git a/backend/cmd/server/wire_gen.go b/backend/cmd/server/wire_gen.go index 5cbc774d..ac4e23ce 100644 --- a/backend/cmd/server/wire_gen.go +++ b/backend/cmd/server/wire_gen.go @@ -90,7 +90,8 @@ func initializeApplication(buildInfo handler.BuildInfo) (*Application, error) { geminiQuotaService := service.NewGeminiQuotaService(configConfig, settingRepository) rateLimitService := service.NewRateLimitService(accountRepository, usageLogRepository, configConfig, geminiQuotaService) claudeUsageFetcher := repository.NewClaudeUsageFetcher() - accountUsageService := service.NewAccountUsageService(accountRepository, usageLogRepository, claudeUsageFetcher, geminiQuotaService) + antigravityQuotaFetcher := service.NewAntigravityQuotaFetcher(proxyRepository) + accountUsageService := service.NewAccountUsageService(accountRepository, usageLogRepository, claudeUsageFetcher, geminiQuotaService, antigravityQuotaFetcher) geminiTokenCache := repository.NewGeminiTokenCache(redisClient) geminiTokenProvider := service.NewGeminiTokenProvider(accountRepository, geminiTokenCache, geminiOAuthService) gatewayCache := repository.NewGatewayCache(redisClient) @@ -145,8 +146,7 @@ func initializeApplication(buildInfo handler.BuildInfo) (*Application, error) { engine := server.ProvideRouter(configConfig, handlers, jwtAuthMiddleware, adminAuthMiddleware, apiKeyAuthMiddleware, apiKeyService, subscriptionService) httpServer := server.ProvideHTTPServer(configConfig, engine) tokenRefreshService := service.ProvideTokenRefreshService(accountRepository, oAuthService, openAIOAuthService, geminiOAuthService, antigravityOAuthService, configConfig) - antigravityQuotaRefresher := service.ProvideAntigravityQuotaRefresher(accountRepository, proxyRepository, antigravityOAuthService, configConfig) - v := provideCleanup(client, redisClient, tokenRefreshService, pricingService, emailQueueService, billingCacheService, oAuthService, openAIOAuthService, geminiOAuthService, antigravityOAuthService, antigravityQuotaRefresher) + v := provideCleanup(client, redisClient, tokenRefreshService, pricingService, emailQueueService, billingCacheService, oAuthService, openAIOAuthService, geminiOAuthService, antigravityOAuthService) application := &Application{ Server: httpServer, Cleanup: v, @@ -179,7 +179,6 @@ func provideCleanup( openaiOAuth *service.OpenAIOAuthService, geminiOAuth *service.GeminiOAuthService, antigravityOAuth *service.AntigravityOAuthService, - antigravityQuota *service.AntigravityQuotaRefresher, ) func() { return func() { ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second) @@ -221,10 +220,6 @@ func provideCleanup( antigravityOAuth.Stop() return nil }}, - {"AntigravityQuotaRefresher", func() error { - antigravityQuota.Stop() - return nil - }}, {"Redis", func() error { return rdb.Close() }}, diff --git a/backend/internal/service/account_usage_service.go b/backend/internal/service/account_usage_service.go index dfceac07..d851af6e 100644 --- a/backend/internal/service/account_usage_service.go +++ b/backend/internal/service/account_usage_service.go @@ -91,6 +91,12 @@ type UsageProgress struct { WindowStats *WindowStats `json:"window_stats,omitempty"` // 窗口期统计（从窗口开始到当前的使用量） } +// AntigravityModelQuota Antigravity 单个模型的配额信息 +type AntigravityModelQuota struct { + Utilization int `json:"utilization"` // 使用率 0-100 + ResetTime string `json:"reset_time"` // 重置时间 ISO8601 +} + // UsageInfo 账号使用量信息 type UsageInfo struct { UpdatedAt *time.Time `json:"updated_at,omitempty"` // 更新时间 @@ -99,6 +105,9 @@ type UsageInfo struct { SevenDaySonnet *UsageProgress `json:"seven_day_sonnet,omitempty"` // 7天Sonnet窗口 GeminiProDaily *UsageProgress `json:"gemini_pro_daily,omitempty"` // Gemini Pro 日配额 GeminiFlashDaily *UsageProgress `json:"gemini_flash_daily,omitempty"` // Gemini Flash 日配额 + + // Antigravity 多模型配额 + AntigravityQuota map[string]*AntigravityModelQuota `json:"antigravity_quota,omitempty"` } // ClaudeUsageResponse Anthropic API返回的usage结构 @@ -124,19 +133,27 @@ type ClaudeUsageFetcher interface { // AccountUsageService 账号使用量查询服务 type AccountUsageService struct { - accountRepo AccountRepository - usageLogRepo UsageLogRepository - usageFetcher ClaudeUsageFetcher - geminiQuotaService *GeminiQuotaService + accountRepo AccountRepository + usageLogRepo UsageLogRepository + usageFetcher ClaudeUsageFetcher + geminiQuotaService *GeminiQuotaService + antigravityQuotaFetcher *AntigravityQuotaFetcher } // NewAccountUsageService 创建AccountUsageService实例 -func NewAccountUsageService(accountRepo AccountRepository, usageLogRepo UsageLogRepository, usageFetcher ClaudeUsageFetcher, geminiQuotaService *GeminiQuotaService) *AccountUsageService { +func NewAccountUsageService( + accountRepo AccountRepository, + usageLogRepo UsageLogRepository, + usageFetcher ClaudeUsageFetcher, + geminiQuotaService *GeminiQuotaService, + antigravityQuotaFetcher *AntigravityQuotaFetcher, +) *AccountUsageService { return &AccountUsageService{ - accountRepo: accountRepo, - usageLogRepo: usageLogRepo, - usageFetcher: usageFetcher, - geminiQuotaService: geminiQuotaService, + accountRepo: accountRepo, + usageLogRepo: usageLogRepo, + usageFetcher: usageFetcher, + geminiQuotaService: geminiQuotaService, + antigravityQuotaFetcher: antigravityQuotaFetcher, } } @@ -154,6 +171,11 @@ func (s *AccountUsageService) GetUsage(ctx context.Context, accountID int64) (*U return s.getGeminiUsage(ctx, account) } + // Antigravity 平台：使用 AntigravityQuotaFetcher 获取额度 + if account.Platform == PlatformAntigravity { + return s.getAntigravityUsage(ctx, account) + } + // 只有oauth类型账号可以通过API获取usage（有profile scope） if account.CanGetUsage() { var apiResp *ClaudeUsageResponse @@ -230,6 +252,51 @@ func (s *AccountUsageService) getGeminiUsage(ctx context.Context, account *Accou return usage, nil } +// antigravityUsageCache 缓存 Antigravity 额度数据 +type antigravityUsageCache struct { + usageInfo *UsageInfo + timestamp time.Time +} + +var antigravityCacheMap = sync.Map{} + +// getAntigravityUsage 获取 Antigravity 账户额度 +func (s *AccountUsageService) getAntigravityUsage(ctx context.Context, account *Account) (*UsageInfo, error) { + if s.antigravityQuotaFetcher == nil || !s.antigravityQuotaFetcher.CanFetch(account) { + now := time.Now() + return &UsageInfo{UpdatedAt: &now}, nil + } + + // 1. 检查缓存（10 分钟） + if cached, ok := antigravityCacheMap.Load(account.ID); ok { + if cache, ok := cached.(*antigravityUsageCache); ok && time.Since(cache.timestamp) < apiCacheTTL { + // 重新计算 RemainingSeconds + usage := cache.usageInfo + if usage.FiveHour != nil && usage.FiveHour.ResetsAt != nil { + usage.FiveHour.RemainingSeconds = int(time.Until(*usage.FiveHour.ResetsAt).Seconds()) + } + return usage, nil + } + } + + // 2. 获取代理 URL + proxyURL := s.antigravityQuotaFetcher.GetProxyURL(ctx, account) + + // 3. 调用 API 获取额度 + result, err := s.antigravityQuotaFetcher.FetchQuota(ctx, account, proxyURL) + if err != nil { + return nil, fmt.Errorf("fetch antigravity quota failed: %w", err) + } + + // 4. 缓存结果 + antigravityCacheMap.Store(account.ID, &antigravityUsageCache{ + usageInfo: result.UsageInfo, + timestamp: time.Now(), + }) + + return result.UsageInfo, nil +} + // addWindowStats 为 usage 数据添加窗口期统计 // 使用独立缓存（1 分钟），与 API 缓存分离 func (s *AccountUsageService) addWindowStats(ctx context.Context, account *Account, usage *UsageInfo) { diff --git a/backend/internal/service/antigravity_quota_fetcher.go b/backend/internal/service/antigravity_quota_fetcher.go new file mode 100644 index 00000000..c9024e33 --- /dev/null +++ b/backend/internal/service/antigravity_quota_fetcher.go @@ -0,0 +1,111 @@ +package service + +import ( + "context" + "time" + + "github.com/Wei-Shaw/sub2api/internal/pkg/antigravity" +) + +// 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") + + // 如果没有 project_id，生成一个随机的 + if projectID == "" { + projectID = antigravity.GenerateMockProjectID() + } + + client := antigravity.NewClient(proxyURL) + + // 调用 API 获取配额 + modelsResp, modelsRaw, err := client.FetchAvailableModels(ctx, accessToken, projectID) + if err != nil { + return nil, err + } + + // 转换为 UsageInfo + usageInfo := f.buildUsageInfo(modelsResp) + + return &QuotaResult{ + UsageInfo: usageInfo, + Raw: modelsRaw, + }, nil +} + +// buildUsageInfo 将 API 响应转换为 UsageInfo +func (f *AntigravityQuotaFetcher) buildUsageInfo(modelsResp *antigravity.FetchAvailableModelsResponse) *UsageInfo { + now := time.Now() + info := &UsageInfo{ + UpdatedAt: &now, + AntigravityQuota: make(map[string]*AntigravityModelQuota), + } + + // 遍历所有模型，填充 AntigravityQuota + 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, + } + } + + // 同时设置 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 + } + } + + 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() +} diff --git a/backend/internal/service/antigravity_quota_refresher.go b/backend/internal/service/antigravity_quota_refresher.go deleted file mode 100644 index c4b11d73..00000000 --- a/backend/internal/service/antigravity_quota_refresher.go +++ /dev/null @@ -1,222 +0,0 @@ -package service - -import ( - "context" - "log" - "sync" - "time" - - "github.com/Wei-Shaw/sub2api/internal/config" - "github.com/Wei-Shaw/sub2api/internal/pkg/antigravity" -) - -// AntigravityQuotaRefresher 定时刷新 Antigravity 账户的配额信息 -type AntigravityQuotaRefresher struct { - accountRepo AccountRepository - proxyRepo ProxyRepository - cfg *config.TokenRefreshConfig - - stopCh chan struct{} - wg sync.WaitGroup -} - -// NewAntigravityQuotaRefresher 创建配额刷新器 -func NewAntigravityQuotaRefresher( - accountRepo AccountRepository, - proxyRepo ProxyRepository, - _ *AntigravityOAuthService, - cfg *config.Config, -) *AntigravityQuotaRefresher { - return &AntigravityQuotaRefresher{ - accountRepo: accountRepo, - proxyRepo: proxyRepo, - cfg: &cfg.TokenRefresh, - stopCh: make(chan struct{}), - } -} - -// Start 启动后台配额刷新服务 -func (r *AntigravityQuotaRefresher) Start() { - if !r.cfg.Enabled { - log.Println("[AntigravityQuota] Service disabled by configuration") - return - } - - r.wg.Add(1) - go r.refreshLoop() - - log.Printf("[AntigravityQuota] Service started (check every %d minutes)", r.cfg.CheckIntervalMinutes) -} - -// Stop 停止服务 -func (r *AntigravityQuotaRefresher) Stop() { - close(r.stopCh) - r.wg.Wait() - log.Println("[AntigravityQuota] Service stopped") -} - -// refreshLoop 刷新循环 -func (r *AntigravityQuotaRefresher) refreshLoop() { - defer r.wg.Done() - - checkInterval := time.Duration(r.cfg.CheckIntervalMinutes) * time.Minute - if checkInterval < time.Minute { - checkInterval = 5 * time.Minute - } - - ticker := time.NewTicker(checkInterval) - defer ticker.Stop() - - // 启动时立即执行一次 - r.processRefresh() - - for { - select { - case <-ticker.C: - r.processRefresh() - case <-r.stopCh: - return - } - } -} - -// processRefresh 执行一次刷新 -func (r *AntigravityQuotaRefresher) processRefresh() { - ctx := context.Background() - - // 查询所有 active 的账户，然后过滤 antigravity 平台 - allAccounts, err := r.accountRepo.ListActive(ctx) - if err != nil { - log.Printf("[AntigravityQuota] Failed to list accounts: %v", err) - return - } - - // 过滤 antigravity 平台账户 - var accounts []Account - for _, acc := range allAccounts { - if acc.Platform == PlatformAntigravity { - accounts = append(accounts, acc) - } - } - - if len(accounts) == 0 { - return - } - - refreshed, failed := 0, 0 - - for i := range accounts { - account := &accounts[i] - - if err := r.refreshAccountQuota(ctx, account); err != nil { - log.Printf("[AntigravityQuota] Account %d (%s) failed: %v", account.ID, account.Name, err) - failed++ - } else { - refreshed++ - } - } - - log.Printf("[AntigravityQuota] Cycle complete: total=%d, refreshed=%d, failed=%d", - len(accounts), refreshed, failed) -} - -// refreshAccountQuota 刷新单个账户的配额 -func (r *AntigravityQuotaRefresher) refreshAccountQuota(ctx context.Context, account *Account) error { - accessToken := account.GetCredential("access_token") - projectID := account.GetCredential("project_id") - - if accessToken == "" { - return nil // 没有 access_token，跳过 - } - - // token 过期则跳过，由 TokenRefreshService 负责刷新 - if r.isTokenExpired(account) { - return nil - } - - // 获取代理 URL - var proxyURL string - if account.ProxyID != nil { - proxy, err := r.proxyRepo.GetByID(ctx, *account.ProxyID) - if err == nil && proxy != nil { - proxyURL = proxy.URL() - } - } - - client := antigravity.NewClient(proxyURL) - - if account.Extra == nil { - account.Extra = make(map[string]any) - } - - // 获取账户信息（tier、project_id 等） - loadResp, loadRaw, _ := client.LoadCodeAssist(ctx, accessToken) - if loadRaw != nil { - account.Extra["load_code_assist"] = loadRaw - } - if loadResp != nil { - // 尝试从 API 获取 project_id - if projectID == "" && loadResp.CloudAICompanionProject != "" { - projectID = loadResp.CloudAICompanionProject - account.Credentials["project_id"] = projectID - } - } - - // 如果仍然没有 project_id，随机生成一个并保存 - if projectID == "" { - projectID = antigravity.GenerateMockProjectID() - account.Credentials["project_id"] = projectID - log.Printf("[AntigravityQuotaRefresher] 为账户 %d 生成随机 project_id: %s", account.ID, projectID) - } - - // 调用 API 获取配额 - modelsResp, modelsRaw, err := client.FetchAvailableModels(ctx, accessToken, projectID) - if err != nil { - return r.accountRepo.Update(ctx, account) // 保存已有的 load_code_assist 信息 - } - - // 保存完整的配额响应 - if modelsRaw != nil { - account.Extra["available_models"] = modelsRaw - } - - // 解析配额数据为前端使用的格式 - r.updateAccountQuota(account, modelsResp) - - account.Extra["last_refresh"] = time.Now().Format(time.RFC3339) - - // 保存到数据库 - return r.accountRepo.Update(ctx, account) -} - -// isTokenExpired 检查 token 是否过期 -func (r *AntigravityQuotaRefresher) isTokenExpired(account *Account) bool { - expiresAt := account.GetCredentialAsTime("expires_at") - if expiresAt == nil { - return false - } - - // 提前 5 分钟认为过期 - return time.Now().Add(5 * time.Minute).After(*expiresAt) -} - -// updateAccountQuota 更新账户的配额信息（前端使用的格式） -func (r *AntigravityQuotaRefresher) updateAccountQuota(account *Account, modelsResp *antigravity.FetchAvailableModelsResponse) { - quota := make(map[string]any) - - for modelName, modelInfo := range modelsResp.Models { - if modelInfo.QuotaInfo == nil { - continue - } - - // 转换 remainingFraction (0.0-1.0) 为百分比 (0-100) - remaining := int(modelInfo.QuotaInfo.RemainingFraction * 100) - - quota[modelName] = map[string]any{ - "remaining": remaining, - "reset_time": modelInfo.QuotaInfo.ResetTime, - } - } - - account.Extra["quota"] = quota -} diff --git a/backend/internal/service/quota_fetcher.go b/backend/internal/service/quota_fetcher.go new file mode 100644 index 00000000..40d8572c --- /dev/null +++ b/backend/internal/service/quota_fetcher.go @@ -0,0 +1,19 @@ +package service + +import ( + "context" +) + +// QuotaFetcher 额度获取接口，各平台实现此接口 +type QuotaFetcher interface { + // CanFetch 检查是否可以获取此账户的额度 + CanFetch(account *Account) bool + // FetchQuota 获取账户额度信息 + FetchQuota(ctx context.Context, account *Account, proxyURL string) (*QuotaResult, error) +} + +// QuotaResult 额度获取结果 +type QuotaResult struct { + UsageInfo *UsageInfo // 转换后的使用信息 + Raw map[string]any // 原始响应，可存入 account.Extra +} diff --git a/backend/internal/service/wire.go b/backend/internal/service/wire.go index 7971f041..74ed2f93 100644 --- a/backend/internal/service/wire.go +++ b/backend/internal/service/wire.go @@ -54,18 +54,6 @@ func ProvideTimingWheelService() *TimingWheelService { return svc } -// ProvideAntigravityQuotaRefresher creates and starts AntigravityQuotaRefresher -func ProvideAntigravityQuotaRefresher( - accountRepo AccountRepository, - proxyRepo ProxyRepository, - oauthSvc *AntigravityOAuthService, - cfg *config.Config, -) *AntigravityQuotaRefresher { - svc := NewAntigravityQuotaRefresher(accountRepo, proxyRepo, oauthSvc, cfg) - svc.Start() - return svc -} - // ProvideDeferredService creates and starts DeferredService func ProvideDeferredService(accountRepo AccountRepository, timingWheel *TimingWheelService) *DeferredService { svc := NewDeferredService(accountRepo, timingWheel, 10*time.Second) @@ -124,6 +112,6 @@ var ProviderSet = wire.NewSet( ProvideTokenRefreshService, ProvideTimingWheelService, ProvideDeferredService, - ProvideAntigravityQuotaRefresher, + NewAntigravityQuotaFetcher, NewUserAttributeService, ) diff --git a/frontend/src/components/account/AccountUsageCell.vue b/frontend/src/components/account/AccountUsageCell.vue index 8dfb9f38..b0bc6c32 100644 --- a/frontend/src/components/account/AccountUsageCell.vue +++ b/frontend/src/components/account/AccountUsageCell.vue @@ -93,7 +93,7 @@
-
- +