diff --git a/backend/internal/pkg/antigravity/client.go b/backend/internal/pkg/antigravity/client.go index 1e63315b..af3a0bfc 100644 --- a/backend/internal/pkg/antigravity/client.go +++ b/backend/internal/pkg/antigravity/client.go @@ -124,10 +124,68 @@ type IneligibleTier struct { type LoadCodeAssistResponse struct { CloudAICompanionProject string `json:"cloudaicompanionProject"` CurrentTier *TierInfo `json:"currentTier,omitempty"` - PaidTier *TierInfo `json:"paidTier,omitempty"` + PaidTier *PaidTierInfo `json:"paidTier,omitempty"` IneligibleTiers []*IneligibleTier `json:"ineligibleTiers,omitempty"` } +// PaidTierInfo 付费等级信息，包含 AI Credits 余额。 +type PaidTierInfo struct { + ID string `json:"id"` + Name string `json:"name"` + Description string `json:"description"` + AvailableCredits []AvailableCredit `json:"availableCredits,omitempty"` +} + +// UnmarshalJSON 兼容 paidTier 既可能是字符串也可能是对象的情况。 +func (p *PaidTierInfo) UnmarshalJSON(data []byte) error { + data = bytes.TrimSpace(data) + if len(data) == 0 || string(data) == "null" { + return nil + } + if data[0] == '"' { + var id string + if err := json.Unmarshal(data, &id); err != nil { + return err + } + p.ID = id + return nil + } + type alias PaidTierInfo + var raw alias + if err := json.Unmarshal(data, &raw); err != nil { + return err + } + *p = PaidTierInfo(raw) + return nil +} + +// AvailableCredit 表示一条 AI Credits 余额记录。 +type AvailableCredit struct { + CreditType string `json:"creditType,omitempty"` + CreditAmount string `json:"creditAmount,omitempty"` + MinimumCreditAmountForUsage string `json:"minimumCreditAmountForUsage,omitempty"` +} + +// GetAmount 将 creditAmount 解析为浮点数。 +func (c *AvailableCredit) GetAmount() float64 { + if c.CreditAmount == "" { + return 0 + } + var value float64 + _, _ = fmt.Sscanf(c.CreditAmount, "%f", &value) + return value +} + +// GetMinimumAmount 将 minimumCreditAmountForUsage 解析为浮点数。 +func (c *AvailableCredit) GetMinimumAmount() float64 { + if c.MinimumCreditAmountForUsage == "" { + return 0 + } + var value float64 + _, _ = fmt.Sscanf(c.MinimumCreditAmountForUsage, "%f", &value) + return value +} + // OnboardUserRequest onboardUser 请求 type OnboardUserRequest struct { TierID string `json:"tierId"` @@ -157,6 +215,14 @@ func (r *LoadCodeAssistResponse) GetTier() string { return "" } +// GetAvailableCredits 返回 paid tier 中的 AI Credits 余额列表。 +func (r *LoadCodeAssistResponse) GetAvailableCredits() []AvailableCredit { + if r.PaidTier == nil { + return nil + } + return r.PaidTier.AvailableCredits +} + // Client Antigravity API 客户端 type Client struct { httpClient *http.Client diff --git a/backend/internal/pkg/antigravity/client_test.go b/backend/internal/pkg/antigravity/client_test.go index 20b57833..7d5bba93 100644 --- a/backend/internal/pkg/antigravity/client_test.go +++ b/backend/internal/pkg/antigravity/client_test.go @@ -190,7 +190,7 @@ func TestTierInfo_UnmarshalJSON_通过JSON嵌套结构(t *testing.T) { func TestGetTier_PaidTier优先(t *testing.T) { resp := &LoadCodeAssistResponse{ CurrentTier: &TierInfo{ID: "free-tier"}, - PaidTier: &TierInfo{ID: "g1-pro-tier"}, + PaidTier: &PaidTierInfo{ID: "g1-pro-tier"}, } if got := resp.GetTier(); got != "g1-pro-tier" { t.Errorf("应返回 paidTier: got %s", got) @@ -209,7 +209,7 @@ func TestGetTier_回退到CurrentTier(t *testing.T) { func TestGetTier_PaidTier为空ID(t *testing.T) { resp := &LoadCodeAssistResponse{ CurrentTier: &TierInfo{ID: "free-tier"}, - PaidTier: &TierInfo{ID: ""}, + PaidTier: &PaidTierInfo{ID: ""}, } // paidTier.ID 为空时应回退到 currentTier if got := resp.GetTier(); got != "free-tier" { @@ -217,6 +217,32 @@ func TestGetTier_PaidTier为空ID(t *testing.T) { } } +func TestGetAvailableCredits(t *testing.T) { + resp := &LoadCodeAssistResponse{ + PaidTier: &PaidTierInfo{ + ID: "g1-pro-tier", + AvailableCredits: []AvailableCredit{ + { + CreditType: "GOOGLE_ONE_AI", + CreditAmount: "25", + MinimumCreditAmountForUsage: "5", + }, + }, + }, + } + + credits := resp.GetAvailableCredits() + if len(credits) != 1 { + t.Fatalf("AI Credits 数量不匹配: got %d", len(credits)) + } + if credits[0].GetAmount() != 25 { + t.Errorf("CreditAmount 解析不正确: got %v", credits[0].GetAmount()) + } + if credits[0].GetMinimumAmount() != 5 { + t.Errorf("MinimumCreditAmountForUsage 解析不正确: got %v", credits[0].GetMinimumAmount()) + } +} + func TestGetTier_两者都为nil(t *testing.T) { resp := &LoadCodeAssistResponse{} if got := resp.GetTier(); got != "" { diff --git a/backend/internal/service/account.go b/backend/internal/service/account.go index 578d1da3..b6408f5f 100644 --- a/backend/internal/service/account.go +++ b/backend/internal/service/account.go @@ -901,6 +901,22 @@ func (a *Account) IsMixedSchedulingEnabled() bool { return false } +// IsOveragesEnabled 检查 Antigravity 账号是否启用 AI Credits 超量请求。 +func (a *Account) IsOveragesEnabled() bool { + if a.Platform != PlatformAntigravity { + return false + } + if a.Extra == nil { + return false + } + if v, ok := a.Extra["allow_overages"]; ok { + if enabled, ok := v.(bool); ok { + return enabled + } + } + return false +} + // IsOpenAIPassthroughEnabled 返回 OpenAI 账号是否启用“自动透传（仅替换认证）”。 // // 新字段：accounts.extra.openai_passthrough。 diff --git a/backend/internal/service/account_usage_service.go b/backend/internal/service/account_usage_service.go index b0a04d47..f117abfd 100644 --- a/backend/internal/service/account_usage_service.go +++ b/backend/internal/service/account_usage_service.go @@ -166,6 +166,13 @@ type AntigravityModelDetail struct { SupportedMimeTypes map[string]bool `json:"supported_mime_types,omitempty"` } +// AICredit 表示 Antigravity 账号的 AI Credits 余额信息。 +type AICredit struct { + CreditType string `json:"credit_type,omitempty"` + Amount float64 `json:"amount,omitempty"` + MinimumBalance float64 `json:"minimum_balance,omitempty"` +} + // UsageInfo 账号使用量信息 type UsageInfo struct { UpdatedAt *time.Time `json:"updated_at,omitempty"` // 更新时间 @@ -189,6 +196,9 @@ type UsageInfo struct { // Antigravity 模型详细能力信息（与 antigravity_quota 同 key） AntigravityQuotaDetails map[string]*AntigravityModelDetail `json:"antigravity_quota_details,omitempty"` + // Antigravity AI Credits 余额 + AICredits []AICredit `json:"ai_credits,omitempty"` + // Antigravity 废弃模型转发规则 (old_model_id -> new_model_id) ModelForwardingRules map[string]string `json:"model_forwarding_rules,omitempty"` diff --git a/backend/internal/service/admin_service.go b/backend/internal/service/admin_service.go index 7dc8bfbd..ea76e171 100644 --- a/backend/internal/service/admin_service.go +++ b/backend/internal/service/admin_service.go @@ -368,6 +368,10 @@ type ProxyExitInfoProber interface { ProbeProxy(ctx context.Context, proxyURL string) (*ProxyExitInfo, int64, error) } +type groupExistenceBatchReader interface { + ExistsByIDs(ctx context.Context, ids []int64) (map[int64]bool, error) +} + type proxyQualityTarget struct { Target string URL string @@ -445,10 +449,6 @@ type userGroupRateBatchReader interface { GetByUserIDs(ctx context.Context, userIDs []int64) (map[int64]map[int64]float64, error) } -type groupExistenceBatchReader interface { - ExistsByIDs(ctx context.Context, ids []int64) (map[int64]bool, error) -} - // NewAdminService creates a new AdminService func NewAdminService( userRepo UserRepository, @@ -1516,6 +1516,7 @@ func (s *adminServiceImpl) UpdateAccount(ctx context.Context, id int64, input *U if err != nil { return nil, err } + wasOveragesEnabled := account.IsOveragesEnabled() if input.Name != "" { account.Name = input.Name @@ -1537,6 +1538,17 @@ func (s *adminServiceImpl) UpdateAccount(ctx context.Context, id int64, input *U } } account.Extra = input.Extra + if account.Platform == PlatformAntigravity && wasOveragesEnabled && !account.IsOveragesEnabled() { + delete(account.Extra, "antigravity_credits_overages") // 清理旧版 overages 运行态 + // 清除 AICredits 限流 key + if rawLimits, ok := account.Extra[modelRateLimitsKey].(map[string]any); ok { + delete(rawLimits, creditsExhaustedKey) + } + } + if account.Platform == PlatformAntigravity && !wasOveragesEnabled && account.IsOveragesEnabled() { + delete(account.Extra, modelRateLimitsKey) + delete(account.Extra, "antigravity_credits_overages") // 清理旧版 overages 运行态 + } // 校验并预计算固定时间重置的下次重置时间 if err := ValidateQuotaResetConfig(account.Extra); err != nil { return nil, err diff --git a/backend/internal/service/admin_service_overages_test.go b/backend/internal/service/admin_service_overages_test.go new file mode 100644 index 00000000..779b08b9 --- /dev/null +++ b/backend/internal/service/admin_service_overages_test.go @@ -0,0 +1,123 @@ +//go:build unit + +package service + +import ( + "context" + "testing" + "time" + + "github.com/stretchr/testify/require" +) + +type updateAccountOveragesRepoStub struct { + mockAccountRepoForGemini + account *Account + updateCalls int +} + +func (r *updateAccountOveragesRepoStub) GetByID(ctx context.Context, id int64) (*Account, error) { + return r.account, nil +} + +func (r *updateAccountOveragesRepoStub) Update(ctx context.Context, account *Account) error { + r.updateCalls++ + r.account = account + return nil +} + +func TestUpdateAccount_DisableOveragesClearsAICreditsKey(t *testing.T) { + accountID := int64(101) + repo := &updateAccountOveragesRepoStub{ + account: &Account{ + ID: accountID, + Platform: PlatformAntigravity, + Type: AccountTypeOAuth, + Status: StatusActive, + Extra: map[string]any{ + "allow_overages": true, + "mixed_scheduling": true, + modelRateLimitsKey: map[string]any{ + "claude-sonnet-4-5": map[string]any{ + "rate_limited_at": "2026-03-15T00:00:00Z", + "rate_limit_reset_at": "2099-03-15T00:00:00Z", + }, + creditsExhaustedKey: map[string]any{ + "rate_limited_at": "2026-03-15T00:00:00Z", + "rate_limit_reset_at": time.Now().Add(5 * time.Hour).UTC().Format(time.RFC3339), + }, + }, + }, + }, + } + + svc := &adminServiceImpl{accountRepo: repo} + updated, err := svc.UpdateAccount(context.Background(), accountID, &UpdateAccountInput{ + Extra: map[string]any{ + "mixed_scheduling": true, + modelRateLimitsKey: map[string]any{ + "claude-sonnet-4-5": map[string]any{ + "rate_limited_at": "2026-03-15T00:00:00Z", + "rate_limit_reset_at": "2099-03-15T00:00:00Z", + }, + creditsExhaustedKey: map[string]any{ + "rate_limited_at": "2026-03-15T00:00:00Z", + "rate_limit_reset_at": time.Now().Add(5 * time.Hour).UTC().Format(time.RFC3339), + }, + }, + }, + }) + + require.NoError(t, err) + require.NotNil(t, updated) + require.Equal(t, 1, repo.updateCalls) + require.False(t, updated.IsOveragesEnabled()) + + // 关闭 overages 后，AICredits key 应被清除 + rawLimits, ok := repo.account.Extra[modelRateLimitsKey].(map[string]any) + if ok { + _, exists := rawLimits[creditsExhaustedKey] + require.False(t, exists, "关闭 overages 时应清除 AICredits 限流 key") + } + // 普通模型限流应保留 + require.True(t, ok) + _, exists := rawLimits["claude-sonnet-4-5"] + require.True(t, exists, "普通模型限流应保留") +} + +func TestUpdateAccount_EnableOveragesClearsModelRateLimitsBeforePersist(t *testing.T) { + accountID := int64(102) + repo := &updateAccountOveragesRepoStub{ + account: &Account{ + ID: accountID, + Platform: PlatformAntigravity, + Type: AccountTypeOAuth, + Status: StatusActive, + Extra: map[string]any{ + "mixed_scheduling": true, + modelRateLimitsKey: map[string]any{ + "claude-sonnet-4-5": map[string]any{ + "rate_limited_at": "2026-03-15T00:00:00Z", + "rate_limit_reset_at": "2099-03-15T00:00:00Z", + }, + }, + }, + }, + } + + svc := &adminServiceImpl{accountRepo: repo} + updated, err := svc.UpdateAccount(context.Background(), accountID, &UpdateAccountInput{ + Extra: map[string]any{ + "mixed_scheduling": true, + "allow_overages": true, + }, + }) + + require.NoError(t, err) + require.NotNil(t, updated) + require.Equal(t, 1, repo.updateCalls) + require.True(t, updated.IsOveragesEnabled()) + + _, exists := repo.account.Extra[modelRateLimitsKey] + require.False(t, exists, "开启 overages 时应在持久化前清掉旧模型限流") +} diff --git a/backend/internal/service/antigravity_credits_overages.go b/backend/internal/service/antigravity_credits_overages.go new file mode 100644 index 00000000..1521dfcd --- /dev/null +++ b/backend/internal/service/antigravity_credits_overages.go @@ -0,0 +1,234 @@ +package service + +import ( + "context" + "encoding/json" + "io" + "net/http" + "strings" + "time" + + "github.com/Wei-Shaw/sub2api/internal/pkg/antigravity" + "github.com/Wei-Shaw/sub2api/internal/pkg/logger" +) + +const ( + // creditsExhaustedKey 是 model_rate_limits 中标记积分耗尽的特殊 key。 + // 与普通模型限流完全同构：通过 SetModelRateLimit / isRateLimitActiveForKey 读写。 + creditsExhaustedKey = "AICredits" + creditsExhaustedDuration = 5 * time.Hour +) + +type antigravity429Category string + +const ( + antigravity429Unknown antigravity429Category = "unknown" + antigravity429RateLimited antigravity429Category = "rate_limited" + antigravity429QuotaExhausted antigravity429Category = "quota_exhausted" +) + +var ( + antigravityQuotaExhaustedKeywords = []string{ + "quota_exhausted", + "quota exhausted", + } + + creditsExhaustedKeywords = []string{ + "google_one_ai", + "insufficient credit", + "insufficient credits", + "not enough credit", + "not enough credits", + "credit exhausted", + "credits exhausted", + "credit balance", + "minimumcreditamountforusage", + "minimum credit amount for usage", + "minimum credit", + } +) + +// isCreditsExhausted 检查账号的 AICredits 限流 key 是否生效（积分是否耗尽）。 +func (a *Account) isCreditsExhausted() bool { + if a == nil { + return false + } + return a.isRateLimitActiveForKey(creditsExhaustedKey) +} + +// setCreditsExhausted 标记账号积分耗尽：写入 model_rate_limits["AICredits"] + 更新缓存。 +func (s *AntigravityGatewayService) setCreditsExhausted(ctx context.Context, account *Account) { + if account == nil || account.ID == 0 { + return + } + resetAt := time.Now().Add(creditsExhaustedDuration) + if err := s.accountRepo.SetModelRateLimit(ctx, account.ID, creditsExhaustedKey, resetAt); err != nil { + logger.LegacyPrintf("service.antigravity_gateway", "set credits exhausted failed: account=%d err=%v", account.ID, err) + return + } + s.updateAccountModelRateLimitInCache(ctx, account, creditsExhaustedKey, resetAt) + logger.LegacyPrintf("service.antigravity_gateway", "credits_exhausted_marked account=%d reset_at=%s", + account.ID, resetAt.UTC().Format(time.RFC3339)) +} + +// clearCreditsExhausted 清除账号的 AICredits 限流 key。 +func (s *AntigravityGatewayService) clearCreditsExhausted(ctx context.Context, account *Account) { + if account == nil || account.ID == 0 || account.Extra == nil { + return + } + rawLimits, ok := account.Extra[modelRateLimitsKey].(map[string]any) + if !ok { + return + } + if _, exists := rawLimits[creditsExhaustedKey]; !exists { + return + } + delete(rawLimits, creditsExhaustedKey) + account.Extra[modelRateLimitsKey] = rawLimits + if err := s.accountRepo.UpdateExtra(ctx, account.ID, map[string]any{ + modelRateLimitsKey: rawLimits, + }); err != nil { + logger.LegacyPrintf("service.antigravity_gateway", "clear credits exhausted failed: account=%d err=%v", account.ID, err) + } +} + +// classifyAntigravity429 将 Antigravity 的 429 响应归类为配额耗尽、限流或未知。 +func classifyAntigravity429(body []byte) antigravity429Category { + if len(body) == 0 { + return antigravity429Unknown + } + lowerBody := strings.ToLower(string(body)) + for _, keyword := range antigravityQuotaExhaustedKeywords { + if strings.Contains(lowerBody, keyword) { + return antigravity429QuotaExhausted + } + } + if info := parseAntigravitySmartRetryInfo(body); info != nil && !info.IsModelCapacityExhausted { + return antigravity429RateLimited + } + return antigravity429Unknown +} + +// injectEnabledCreditTypes 在已序列化的 v1internal JSON body 中注入 AI Credits 类型。 +func injectEnabledCreditTypes(body []byte) []byte { + var payload map[string]any + if err := json.Unmarshal(body, &payload); err != nil { + return nil + } + payload["enabledCreditTypes"] = []string{"GOOGLE_ONE_AI"} + result, err := json.Marshal(payload) + if err != nil { + return nil + } + return result +} + +// resolveCreditsOveragesModelKey 解析当前请求对应的 overages 状态模型 key。 +func resolveCreditsOveragesModelKey(ctx context.Context, account *Account, upstreamModelName, requestedModel string) string { + modelKey := strings.TrimSpace(upstreamModelName) + if modelKey != "" { + return modelKey + } + if account == nil { + return "" + } + modelKey = resolveFinalAntigravityModelKey(ctx, account, requestedModel) + if strings.TrimSpace(modelKey) != "" { + return modelKey + } + return resolveAntigravityModelKey(requestedModel) +} + +// shouldMarkCreditsExhausted 判断一次 credits 请求失败是否应标记为 credits 耗尽。 +func shouldMarkCreditsExhausted(resp *http.Response, respBody []byte, reqErr error) bool { + if reqErr != nil || resp == nil { + return false + } + if resp.StatusCode >= 500 || resp.StatusCode == http.StatusRequestTimeout { + return false + } + if isURLLevelRateLimit(respBody) { + return false + } + if info := parseAntigravitySmartRetryInfo(respBody); info != nil { + return false + } + bodyLower := strings.ToLower(string(respBody)) + for _, keyword := range creditsExhaustedKeywords { + if strings.Contains(bodyLower, keyword) { + return true + } + } + return false +} + +type creditsOveragesRetryResult struct { + handled bool + resp *http.Response +} + +// attemptCreditsOveragesRetry 在确认免费配额耗尽后，尝试注入 AI Credits 继续请求。 +func (s *AntigravityGatewayService) attemptCreditsOveragesRetry( + p antigravityRetryLoopParams, + baseURL string, + modelName string, + waitDuration time.Duration, + originalStatusCode int, + respBody []byte, +) *creditsOveragesRetryResult { + creditsBody := injectEnabledCreditTypes(p.body) + if creditsBody == nil { + return &creditsOveragesRetryResult{handled: false} + } + modelKey := resolveCreditsOveragesModelKey(p.ctx, p.account, modelName, p.requestedModel) + logger.LegacyPrintf("service.antigravity_gateway", "%s status=429 credit_overages_retry model=%s account=%d (injecting enabledCreditTypes)", + p.prefix, modelKey, p.account.ID) + + creditsReq, err := antigravity.NewAPIRequestWithURL(p.ctx, baseURL, p.action, p.accessToken, creditsBody) + if err != nil { + logger.LegacyPrintf("service.antigravity_gateway", "%s credit_overages_failed model=%s account=%d build_request_err=%v", + p.prefix, modelKey, p.account.ID, err) + return &creditsOveragesRetryResult{handled: true} + } + + creditsResp, err := p.httpUpstream.Do(creditsReq, p.proxyURL, p.account.ID, p.account.Concurrency) + if err == nil && creditsResp != nil && creditsResp.StatusCode < 400 { + s.clearCreditsExhausted(p.ctx, p.account) + logger.LegacyPrintf("service.antigravity_gateway", "%s status=%d credit_overages_success model=%s account=%d", + p.prefix, creditsResp.StatusCode, modelKey, p.account.ID) + return &creditsOveragesRetryResult{handled: true, resp: creditsResp} + } + + s.handleCreditsRetryFailure(p.ctx, p.prefix, modelKey, p.account, creditsResp, err) + return &creditsOveragesRetryResult{handled: true} +} + +func (s *AntigravityGatewayService) handleCreditsRetryFailure( + ctx context.Context, + prefix string, + modelKey string, + account *Account, + creditsResp *http.Response, + reqErr error, +) { + var creditsRespBody []byte + creditsStatusCode := 0 + if creditsResp != nil { + creditsStatusCode = creditsResp.StatusCode + if creditsResp.Body != nil { + creditsRespBody, _ = io.ReadAll(io.LimitReader(creditsResp.Body, 64<<10)) + _ = creditsResp.Body.Close() + } + } + + if shouldMarkCreditsExhausted(creditsResp, creditsRespBody, reqErr) && account != nil { + s.setCreditsExhausted(ctx, account) + logger.LegacyPrintf("service.antigravity_gateway", "%s credit_overages_failed model=%s account=%d marked_exhausted=true status=%d body=%s", + prefix, modelKey, account.ID, creditsStatusCode, truncateForLog(creditsRespBody, 200)) + return + } + if account != nil { + logger.LegacyPrintf("service.antigravity_gateway", "%s credit_overages_failed model=%s account=%d marked_exhausted=false status=%d err=%v body=%s", + prefix, modelKey, account.ID, creditsStatusCode, reqErr, truncateForLog(creditsRespBody, 200)) + } +} diff --git a/backend/internal/service/antigravity_credits_overages_test.go b/backend/internal/service/antigravity_credits_overages_test.go new file mode 100644 index 00000000..bc679494 --- /dev/null +++ b/backend/internal/service/antigravity_credits_overages_test.go @@ -0,0 +1,538 @@ +//go:build unit + +package service + +import ( + "bytes" + "context" + "io" + "net/http" + "strings" + "testing" + "time" + + "github.com/Wei-Shaw/sub2api/internal/pkg/antigravity" + "github.com/stretchr/testify/require" +) + +func TestClassifyAntigravity429(t *testing.T) { + t.Run("明确配额耗尽", func(t *testing.T) { + body := []byte(`{"error":{"status":"RESOURCE_EXHAUSTED","message":"QUOTA_EXHAUSTED"}}`) + require.Equal(t, antigravity429QuotaExhausted, classifyAntigravity429(body)) + }) + + t.Run("结构化限流", func(t *testing.T) { + body := []byte(`{ + "error": { + "status": "RESOURCE_EXHAUSTED", + "details": [ + {"@type": "type.googleapis.com/google.rpc.ErrorInfo", "metadata": {"model": "claude-sonnet-4-5"}, "reason": "RATE_LIMIT_EXCEEDED"}, + {"@type": "type.googleapis.com/google.rpc.RetryInfo", "retryDelay": "0.5s"} + ] + } + }`) + require.Equal(t, antigravity429RateLimited, classifyAntigravity429(body)) + }) + + t.Run("未知429", func(t *testing.T) { + body := []byte(`{"error":{"message":"too many requests"}}`) + require.Equal(t, antigravity429Unknown, classifyAntigravity429(body)) + }) +} + +func TestIsCreditsExhausted_UsesAICreditsKey(t *testing.T) { + t.Run("无 AICredits key 则积分可用", func(t *testing.T) { + account := &Account{ + ID: 1, + Platform: PlatformAntigravity, + Extra: map[string]any{ + "allow_overages": true, + }, + } + require.False(t, account.isCreditsExhausted()) + }) + + t.Run("AICredits key 生效则积分耗尽", func(t *testing.T) { + account := &Account{ + ID: 2, + Platform: PlatformAntigravity, + Extra: map[string]any{ + "allow_overages": true, + modelRateLimitsKey: map[string]any{ + creditsExhaustedKey: map[string]any{ + "rate_limited_at": time.Now().UTC().Format(time.RFC3339), + "rate_limit_reset_at": time.Now().Add(5 * time.Hour).UTC().Format(time.RFC3339), + }, + }, + }, + } + require.True(t, account.isCreditsExhausted()) + }) + + t.Run("AICredits key 过期则积分可用", func(t *testing.T) { + account := &Account{ + ID: 3, + Platform: PlatformAntigravity, + Extra: map[string]any{ + "allow_overages": true, + modelRateLimitsKey: map[string]any{ + creditsExhaustedKey: map[string]any{ + "rate_limited_at": time.Now().Add(-6 * time.Hour).UTC().Format(time.RFC3339), + "rate_limit_reset_at": time.Now().Add(-1 * time.Hour).UTC().Format(time.RFC3339), + }, + }, + }, + } + require.False(t, account.isCreditsExhausted()) + }) +} + +func TestHandleSmartRetry_QuotaExhausted_UsesCreditsAndStoresIndependentState(t *testing.T) { + successResp := &http.Response{ + StatusCode: http.StatusOK, + Header: http.Header{}, + Body: io.NopCloser(strings.NewReader(`{"ok":true}`)), + } + upstream := &mockSmartRetryUpstream{ + responses: []*http.Response{successResp}, + errors: []error{nil}, + } + repo := &stubAntigravityAccountRepo{} + account := &Account{ + ID: 101, + Name: "acc-101", + Type: AccountTypeOAuth, + Platform: PlatformAntigravity, + Extra: map[string]any{ + "allow_overages": true, + }, + Credentials: map[string]any{ + "model_mapping": map[string]any{ + "claude-opus-4-6": "claude-sonnet-4-5", + }, + }, + } + + respBody := []byte(`{"error":{"status":"RESOURCE_EXHAUSTED","message":"QUOTA_EXHAUSTED"}}`) + resp := &http.Response{ + StatusCode: http.StatusTooManyRequests, + Header: http.Header{}, + Body: io.NopCloser(bytes.NewReader(respBody)), + } + params := antigravityRetryLoopParams{ + ctx: context.Background(), + prefix: "[test]", + account: account, + accessToken: "token", + action: "generateContent", + body: []byte(`{"model":"claude-opus-4-6","request":{}}`), + httpUpstream: upstream, + accountRepo: repo, + requestedModel: "claude-opus-4-6", + handleError: func(ctx context.Context, prefix string, account *Account, statusCode int, headers http.Header, body []byte, requestedModel string, groupID int64, sessionHash string, isStickySession bool) *handleModelRateLimitResult { + return nil + }, + } + + svc := &AntigravityGatewayService{} + result := svc.handleSmartRetry(params, resp, respBody, "https://ag-1.test", 0, []string{"https://ag-1.test"}) + + require.NotNil(t, result) + require.Equal(t, smartRetryActionBreakWithResp, result.action) + require.NotNil(t, result.resp) + require.Nil(t, result.switchError) + require.Len(t, upstream.requestBodies, 1) + require.Contains(t, string(upstream.requestBodies[0]), "enabledCreditTypes") + require.Empty(t, repo.modelRateLimitCalls, "overages 成功后不应写入普通 model_rate_limits") +} + +func TestHandleSmartRetry_RateLimited_DoesNotUseCredits(t *testing.T) { + successResp := &http.Response{ + StatusCode: http.StatusOK, + Header: http.Header{}, + Body: io.NopCloser(strings.NewReader(`{"ok":true}`)), + } + upstream := &mockSmartRetryUpstream{ + responses: []*http.Response{successResp}, + errors: []error{nil}, + } + repo := &stubAntigravityAccountRepo{} + account := &Account{ + ID: 102, + Name: "acc-102", + Type: AccountTypeOAuth, + Platform: PlatformAntigravity, + Extra: map[string]any{ + "allow_overages": true, + }, + } + + respBody := []byte(`{ + "error": { + "status": "RESOURCE_EXHAUSTED", + "details": [ + {"@type": "type.googleapis.com/google.rpc.ErrorInfo", "metadata": {"model": "claude-sonnet-4-5"}, "reason": "RATE_LIMIT_EXCEEDED"}, + {"@type": "type.googleapis.com/google.rpc.RetryInfo", "retryDelay": "0.1s"} + ] + } + }`) + resp := &http.Response{ + StatusCode: http.StatusTooManyRequests, + Header: http.Header{}, + Body: io.NopCloser(bytes.NewReader(respBody)), + } + params := antigravityRetryLoopParams{ + ctx: context.Background(), + prefix: "[test]", + account: account, + accessToken: "token", + action: "generateContent", + body: []byte(`{"model":"claude-sonnet-4-5","request":{}}`), + httpUpstream: upstream, + accountRepo: repo, + handleError: func(ctx context.Context, prefix string, account *Account, statusCode int, headers http.Header, body []byte, requestedModel string, groupID int64, sessionHash string, isStickySession bool) *handleModelRateLimitResult { + return nil + }, + } + + svc := &AntigravityGatewayService{} + result := svc.handleSmartRetry(params, resp, respBody, "https://ag-1.test", 0, []string{"https://ag-1.test"}) + + require.NotNil(t, result) + require.Equal(t, smartRetryActionBreakWithResp, result.action) + require.NotNil(t, result.resp) + require.Len(t, upstream.requestBodies, 1) + require.NotContains(t, string(upstream.requestBodies[0]), "enabledCreditTypes") + require.Empty(t, repo.extraUpdateCalls) + require.Empty(t, repo.modelRateLimitCalls) +} + +func TestAntigravityRetryLoop_ModelRateLimited_InjectsCredits(t *testing.T) { + oldBaseURLs := append([]string(nil), antigravity.BaseURLs...) + oldAvailability := antigravity.DefaultURLAvailability + defer func() { + antigravity.BaseURLs = oldBaseURLs + antigravity.DefaultURLAvailability = oldAvailability + }() + + antigravity.BaseURLs = []string{"https://ag-1.test"} + antigravity.DefaultURLAvailability = antigravity.NewURLAvailability(time.Minute) + + upstream := &queuedHTTPUpstreamStub{ + responses: []*http.Response{ + { + StatusCode: http.StatusOK, + Header: http.Header{}, + Body: io.NopCloser(strings.NewReader(`{"ok":true}`)), + }, + }, + errors: []error{nil}, + } + // 模型已限流 + overages 启用 + 无 AICredits key → 应直接注入积分 + account := &Account{ + ID: 103, + Name: "acc-103", + Type: AccountTypeOAuth, + Platform: PlatformAntigravity, + Status: StatusActive, + Schedulable: true, + Extra: map[string]any{ + "allow_overages": true, + modelRateLimitsKey: map[string]any{ + "claude-sonnet-4-5": map[string]any{ + "rate_limited_at": time.Now().UTC().Format(time.RFC3339), + "rate_limit_reset_at": time.Now().Add(30 * time.Minute).UTC().Format(time.RFC3339), + }, + }, + }, + } + + svc := &AntigravityGatewayService{} + result, err := svc.antigravityRetryLoop(antigravityRetryLoopParams{ + ctx: context.Background(), + prefix: "[test]", + account: account, + accessToken: "token", + action: "generateContent", + body: []byte(`{"model":"claude-sonnet-4-5","request":{}}`), + httpUpstream: upstream, + requestedModel: "claude-sonnet-4-5", + handleError: func(ctx context.Context, prefix string, account *Account, statusCode int, headers http.Header, body []byte, requestedModel string, groupID int64, sessionHash string, isStickySession bool) *handleModelRateLimitResult { + return nil + }, + }) + + require.NoError(t, err) + require.NotNil(t, result) + require.Len(t, upstream.requestBodies, 1) + require.Contains(t, string(upstream.requestBodies[0]), "enabledCreditTypes") +} + +func TestAntigravityRetryLoop_CreditsExhausted_DoesNotInject(t *testing.T) { + oldBaseURLs := append([]string(nil), antigravity.BaseURLs...) + oldAvailability := antigravity.DefaultURLAvailability + defer func() { + antigravity.BaseURLs = oldBaseURLs + antigravity.DefaultURLAvailability = oldAvailability + }() + + antigravity.BaseURLs = []string{"https://ag-1.test"} + antigravity.DefaultURLAvailability = antigravity.NewURLAvailability(time.Minute) + + // 模型限流 + overages 启用 + AICredits key 生效 → 不应注入积分，应切号 + account := &Account{ + ID: 104, + Name: "acc-104", + Type: AccountTypeOAuth, + Platform: PlatformAntigravity, + Status: StatusActive, + Schedulable: true, + Extra: map[string]any{ + "allow_overages": true, + modelRateLimitsKey: map[string]any{ + "claude-sonnet-4-5": map[string]any{ + "rate_limited_at": time.Now().UTC().Format(time.RFC3339), + "rate_limit_reset_at": time.Now().Add(30 * time.Minute).UTC().Format(time.RFC3339), + }, + creditsExhaustedKey: map[string]any{ + "rate_limited_at": time.Now().UTC().Format(time.RFC3339), + "rate_limit_reset_at": time.Now().Add(5 * time.Hour).UTC().Format(time.RFC3339), + }, + }, + }, + } + + svc := &AntigravityGatewayService{} + _, err := svc.antigravityRetryLoop(antigravityRetryLoopParams{ + ctx: context.Background(), + prefix: "[test]", + account: account, + accessToken: "token", + action: "generateContent", + body: []byte(`{"model":"claude-sonnet-4-5","request":{}}`), + requestedModel: "claude-sonnet-4-5", + handleError: func(ctx context.Context, prefix string, account *Account, statusCode int, headers http.Header, body []byte, requestedModel string, groupID int64, sessionHash string, isStickySession bool) *handleModelRateLimitResult { + return nil + }, + }) + + // 模型限流 + 积分耗尽 → 应触发切号错误 + require.Error(t, err) + var switchErr *AntigravityAccountSwitchError + require.ErrorAs(t, err, &switchErr) +} + +func TestAntigravityRetryLoop_CreditErrorMarksExhausted(t *testing.T) { + oldBaseURLs := append([]string(nil), antigravity.BaseURLs...) + oldAvailability := antigravity.DefaultURLAvailability + defer func() { + antigravity.BaseURLs = oldBaseURLs + antigravity.DefaultURLAvailability = oldAvailability + }() + + antigravity.BaseURLs = []string{"https://ag-1.test"} + antigravity.DefaultURLAvailability = antigravity.NewURLAvailability(time.Minute) + + repo := &stubAntigravityAccountRepo{} + upstream := &queuedHTTPUpstreamStub{ + responses: []*http.Response{ + { + StatusCode: http.StatusForbidden, + Header: http.Header{}, + Body: io.NopCloser(strings.NewReader(`{"error":{"message":"Insufficient GOOGLE_ONE_AI credits"}}`)), + }, + }, + errors: []error{nil}, + } + // 模型限流 + overages 启用 + 积分可用 → 注入积分但上游返回积分不足 + account := &Account{ + ID: 105, + Name: "acc-105", + Type: AccountTypeOAuth, + Platform: PlatformAntigravity, + Status: StatusActive, + Schedulable: true, + Extra: map[string]any{ + "allow_overages": true, + modelRateLimitsKey: map[string]any{ + "claude-sonnet-4-5": map[string]any{ + "rate_limited_at": time.Now().UTC().Format(time.RFC3339), + "rate_limit_reset_at": time.Now().Add(30 * time.Minute).UTC().Format(time.RFC3339), + }, + }, + }, + } + + svc := &AntigravityGatewayService{accountRepo: repo} + result, err := svc.antigravityRetryLoop(antigravityRetryLoopParams{ + ctx: context.Background(), + prefix: "[test]", + account: account, + accessToken: "token", + action: "generateContent", + body: []byte(`{"model":"claude-sonnet-4-5","request":{}}`), + httpUpstream: upstream, + accountRepo: repo, + requestedModel: "claude-sonnet-4-5", + handleError: func(ctx context.Context, prefix string, account *Account, statusCode int, headers http.Header, body []byte, requestedModel string, groupID int64, sessionHash string, isStickySession bool) *handleModelRateLimitResult { + return nil + }, + }) + + require.NoError(t, err) + require.NotNil(t, result) + // 验证 AICredits key 已通过 SetModelRateLimit 写入数据库 + require.Len(t, repo.modelRateLimitCalls, 1, "应通过 SetModelRateLimit 写入 AICredits key") + require.Equal(t, creditsExhaustedKey, repo.modelRateLimitCalls[0].modelKey) +} + +func TestShouldMarkCreditsExhausted(t *testing.T) { + t.Run("reqErr 不为 nil 时不标记", func(t *testing.T) { + resp := &http.Response{StatusCode: http.StatusForbidden} + require.False(t, shouldMarkCreditsExhausted(resp, []byte(`{"error":"Insufficient credits"}`), io.ErrUnexpectedEOF)) + }) + + t.Run("resp 为 nil 时不标记", func(t *testing.T) { + require.False(t, shouldMarkCreditsExhausted(nil, []byte(`{"error":"Insufficient credits"}`), nil)) + }) + + t.Run("5xx 响应不标记", func(t *testing.T) { + resp := &http.Response{StatusCode: http.StatusInternalServerError} + require.False(t, shouldMarkCreditsExhausted(resp, []byte(`{"error":"Insufficient credits"}`), nil)) + }) + + t.Run("408 RequestTimeout 不标记", func(t *testing.T) { + resp := &http.Response{StatusCode: http.StatusRequestTimeout} + require.False(t, shouldMarkCreditsExhausted(resp, []byte(`{"error":"Insufficient credits"}`), nil)) + }) + + t.Run("URL 级限流不标记", func(t *testing.T) { + resp := &http.Response{StatusCode: http.StatusTooManyRequests} + body := []byte(`{"error":{"message":"Resource has been exhausted"}}`) + require.False(t, shouldMarkCreditsExhausted(resp, body, nil)) + }) + + t.Run("结构化限流不标记", func(t *testing.T) { + resp := &http.Response{StatusCode: http.StatusTooManyRequests} + body := []byte(`{"error":{"status":"RESOURCE_EXHAUSTED","details":[{"@type":"type.googleapis.com/google.rpc.ErrorInfo","reason":"RATE_LIMIT_EXCEEDED"},{"@type":"type.googleapis.com/google.rpc.RetryInfo","retryDelay":"0.5s"}]}}`) + require.False(t, shouldMarkCreditsExhausted(resp, body, nil)) + }) + + t.Run("含 credits 关键词时标记", func(t *testing.T) { + resp := &http.Response{StatusCode: http.StatusForbidden} + for _, keyword := range []string{ + "Insufficient GOOGLE_ONE_AI credits", + "insufficient credit balance", + "not enough credits for this request", + "Credits exhausted", + "minimumCreditAmountForUsage requirement not met", + } { + body := []byte(`{"error":{"message":"` + keyword + `"}}`) + require.True(t, shouldMarkCreditsExhausted(resp, body, nil), "should mark for keyword: %s", keyword) + } + }) + + t.Run("无 credits 关键词时不标记", func(t *testing.T) { + resp := &http.Response{StatusCode: http.StatusForbidden} + body := []byte(`{"error":{"message":"permission denied"}}`) + require.False(t, shouldMarkCreditsExhausted(resp, body, nil)) + }) +} + +func TestInjectEnabledCreditTypes(t *testing.T) { + t.Run("正常 JSON 注入成功", func(t *testing.T) { + body := []byte(`{"model":"claude-sonnet-4-5","request":{}}`) + result := injectEnabledCreditTypes(body) + require.NotNil(t, result) + require.Contains(t, string(result), `"enabledCreditTypes"`) + require.Contains(t, string(result), `GOOGLE_ONE_AI`) + }) + + t.Run("非法 JSON 返回 nil", func(t *testing.T) { + require.Nil(t, injectEnabledCreditTypes([]byte(`not json`))) + }) + + t.Run("空 body 返回 nil", func(t *testing.T) { + require.Nil(t, injectEnabledCreditTypes([]byte{})) + }) + + t.Run("已有 enabledCreditTypes 会被覆盖", func(t *testing.T) { + body := []byte(`{"enabledCreditTypes":["OLD"],"model":"test"}`) + result := injectEnabledCreditTypes(body) + require.NotNil(t, result) + require.Contains(t, string(result), `GOOGLE_ONE_AI`) + require.NotContains(t, string(result), `OLD`) + }) +} + +func TestClearCreditsExhausted(t *testing.T) { + t.Run("account 为 nil 不操作", func(t *testing.T) { + repo := &stubAntigravityAccountRepo{} + svc := &AntigravityGatewayService{accountRepo: repo} + svc.clearCreditsExhausted(context.Background(), nil) + require.Empty(t, repo.extraUpdateCalls) + }) + + t.Run("Extra 为 nil 不操作", func(t *testing.T) { + repo := &stubAntigravityAccountRepo{} + svc := &AntigravityGatewayService{accountRepo: repo} + svc.clearCreditsExhausted(context.Background(), &Account{ID: 1}) + require.Empty(t, repo.extraUpdateCalls) + }) + + t.Run("无 modelRateLimitsKey 不操作", func(t *testing.T) { + repo := &stubAntigravityAccountRepo{} + svc := &AntigravityGatewayService{accountRepo: repo} + svc.clearCreditsExhausted(context.Background(), &Account{ + ID: 1, + Extra: map[string]any{"some_key": "value"}, + }) + require.Empty(t, repo.extraUpdateCalls) + }) + + t.Run("无 AICredits key 不操作", func(t *testing.T) { + repo := &stubAntigravityAccountRepo{} + svc := &AntigravityGatewayService{accountRepo: repo} + svc.clearCreditsExhausted(context.Background(), &Account{ + ID: 1, + Extra: map[string]any{ + modelRateLimitsKey: map[string]any{ + "claude-sonnet-4-5": map[string]any{ + "rate_limited_at": "2026-03-15T00:00:00Z", + "rate_limit_reset_at": "2099-03-15T00:00:00Z", + }, + }, + }, + }) + require.Empty(t, repo.extraUpdateCalls) + }) + + t.Run("有 AICredits key 时删除并调用 UpdateExtra", func(t *testing.T) { + repo := &stubAntigravityAccountRepo{} + svc := &AntigravityGatewayService{accountRepo: repo} + account := &Account{ + ID: 1, + Extra: map[string]any{ + modelRateLimitsKey: map[string]any{ + "claude-sonnet-4-5": map[string]any{ + "rate_limited_at": "2026-03-15T00:00:00Z", + "rate_limit_reset_at": "2099-03-15T00:00:00Z", + }, + creditsExhaustedKey: map[string]any{ + "rate_limited_at": "2026-03-15T00:00:00Z", + "rate_limit_reset_at": time.Now().Add(5 * time.Hour).UTC().Format(time.RFC3339), + }, + }, + }, + } + svc.clearCreditsExhausted(context.Background(), account) + require.Len(t, repo.extraUpdateCalls, 1) + // AICredits key 应被删除 + rawLimits := account.Extra[modelRateLimitsKey].(map[string]any) + _, exists := rawLimits[creditsExhaustedKey] + require.False(t, exists, "AICredits key 应被删除") + // 普通模型限流应保留 + _, exists = rawLimits["claude-sonnet-4-5"] + require.True(t, exists, "普通模型限流应保留") + }) +} diff --git a/backend/internal/service/antigravity_gateway_service.go b/backend/internal/service/antigravity_gateway_service.go index f63802b8..cafc2a79 100644 --- a/backend/internal/service/antigravity_gateway_service.go +++ b/backend/internal/service/antigravity_gateway_service.go @@ -188,9 +188,29 @@ func (s *AntigravityGatewayService) handleSmartRetry(p antigravityRetryLoopParam return &smartRetryResult{action: smartRetryActionContinueURL} } + category := antigravity429Unknown + if resp.StatusCode == http.StatusTooManyRequests { + category = classifyAntigravity429(respBody) + } + // 判断是否触发智能重试 shouldSmartRetry, shouldRateLimitModel, waitDuration, modelName, isModelCapacityExhausted := shouldTriggerAntigravitySmartRetry(p.account, respBody) + // AI Credits 超量请求： + // 仅在上游明确返回免费配额耗尽时才允许切换到 credits。 + if resp.StatusCode == http.StatusTooManyRequests && + category == antigravity429QuotaExhausted && + p.account.IsOveragesEnabled() && + !p.account.isCreditsExhausted() { + result := s.attemptCreditsOveragesRetry(p, baseURL, modelName, waitDuration, resp.StatusCode, respBody) + if result.handled && result.resp != nil { + return &smartRetryResult{ + action: smartRetryActionBreakWithResp, + resp: result.resp, + } + } + } + // 情况1: retryDelay >= 阈值，限流模型并切换账号 if shouldRateLimitModel { // 单账号 503 退避重试模式：不设限流、不切换账号，改为原地等待+重试 @@ -532,14 +552,31 @@ func (s *AntigravityGatewayService) handleSingleAccountRetryInPlace( // antigravityRetryLoop 执行带 URL fallback 的重试循环 func (s *AntigravityGatewayService) antigravityRetryLoop(p antigravityRetryLoopParams) (*antigravityRetryLoopResult, error) { + // 预检查：模型限流 + overages 启用 + 积分未耗尽 → 直接注入 AI Credits + overagesInjected := false + if p.requestedModel != "" && p.account.Platform == PlatformAntigravity && + p.account.IsOveragesEnabled() && !p.account.isCreditsExhausted() && + p.account.isModelRateLimitedWithContext(p.ctx, p.requestedModel) { + if creditsBody := injectEnabledCreditTypes(p.body); creditsBody != nil { + p.body = creditsBody + overagesInjected = true + logger.LegacyPrintf("service.antigravity_gateway", "%s pre_check: model_rate_limited_credits_inject model=%s account=%d (injecting enabledCreditTypes)", + p.prefix, p.requestedModel, p.account.ID) + } + } + // 预检查：如果账号已限流，直接返回切换信号 if p.requestedModel != "" { if remaining := p.account.GetRateLimitRemainingTimeWithContext(p.ctx, p.requestedModel); remaining > 0 { - // 单账号 503 退避重试模式：跳过限流预检查，直接发请求。 - // 首次请求设的限流是为了多账号调度器跳过该账号，在单账号模式下无意义。 - // 如果上游确实还不可用，handleSmartRetry → handleSingleAccountRetryInPlace - // 会在 Service 层原地等待+重试，不需要在预检查这里等。 - if isSingleAccountRetry(p.ctx) { + // 已注入积分的请求不再受普通模型限流预检查阻断。 + if overagesInjected { + logger.LegacyPrintf("service.antigravity_gateway", "%s pre_check: credits_injected_ignore_rate_limit remaining=%v model=%s account=%d", + p.prefix, remaining.Truncate(time.Millisecond), p.requestedModel, p.account.ID) + } else if isSingleAccountRetry(p.ctx) { + // 单账号 503 退避重试模式：跳过限流预检查，直接发请求。 + // 首次请求设的限流是为了多账号调度器跳过该账号，在单账号模式下无意义。 + // 如果上游确实还不可用，handleSmartRetry → handleSingleAccountRetryInPlace + // 会在 Service 层原地等待+重试，不需要在预检查这里等。 logger.LegacyPrintf("service.antigravity_gateway", "%s pre_check: single_account_retry skipping rate_limit remaining=%v model=%s account=%d (will retry in-place if 503)", p.prefix, remaining.Truncate(time.Millisecond), p.requestedModel, p.account.ID) } else { @@ -631,6 +668,15 @@ urlFallbackLoop: respBody, _ := io.ReadAll(io.LimitReader(resp.Body, 2<<20)) _ = resp.Body.Close() + if overagesInjected && shouldMarkCreditsExhausted(resp, respBody, nil) { + modelKey := resolveCreditsOveragesModelKey(p.ctx, p.account, "", p.requestedModel) + s.handleCreditsRetryFailure(p.ctx, p.prefix, modelKey, p.account, &http.Response{ + StatusCode: resp.StatusCode, + Header: resp.Header.Clone(), + Body: io.NopCloser(bytes.NewReader(respBody)), + }, nil) + } + // ★ 统一入口：自定义错误码 + 临时不可调度 if handled, outStatus, policyErr := s.applyErrorPolicy(p, resp.StatusCode, resp.Header, respBody); handled { if policyErr != nil { diff --git a/backend/internal/service/antigravity_quota_fetcher.go b/backend/internal/service/antigravity_quota_fetcher.go index f8990b1a..9e09c904 100644 --- a/backend/internal/service/antigravity_quota_fetcher.go +++ b/backend/internal/service/antigravity_quota_fetcher.go @@ -78,11 +78,11 @@ func (f *AntigravityQuotaFetcher) FetchQuota(ctx context.Context, account *Accou return nil, err } - // 调用 LoadCodeAssist 获取订阅等级（非关键路径，失败不影响主流程） - tierRaw, tierNormalized := f.fetchSubscriptionTier(ctx, client, accessToken) + // 调用 LoadCodeAssist 获取订阅等级和 AI Credits 余额（非关键路径，失败不影响主流程） + tierRaw, tierNormalized, loadResp := f.fetchSubscriptionTier(ctx, client, accessToken) // 转换为 UsageInfo - usageInfo := f.buildUsageInfo(modelsResp, tierRaw, tierNormalized) + usageInfo := f.buildUsageInfo(modelsResp, tierRaw, tierNormalized, loadResp) return &QuotaResult{ UsageInfo: usageInfo, @@ -90,20 +90,21 @@ func (f *AntigravityQuotaFetcher) FetchQuota(ctx context.Context, account *Accou }, nil } -// fetchSubscriptionTier 获取账号订阅等级，失败返回空字符串 -func (f *AntigravityQuotaFetcher) fetchSubscriptionTier(ctx context.Context, client *antigravity.Client, accessToken string) (raw, normalized string) { +// 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 "", "" + return "", "", nil } if loadResp == nil { - return "", "" + return "", "", nil } raw = loadResp.GetTier() // 已有方法：paidTier > currentTier normalized = normalizeTier(raw) - return raw, normalized + return raw, normalized, loadResp } // normalizeTier 将原始 tier 字符串归一化为 FREE/PRO/ULTRA/UNKNOWN @@ -124,8 +125,8 @@ func normalizeTier(raw string) string { } } -// buildUsageInfo 将 API 响应转换为 UsageInfo -func (f *AntigravityQuotaFetcher) buildUsageInfo(modelsResp *antigravity.FetchAvailableModelsResponse, tierRaw, tierNormalized string) *UsageInfo { +// buildUsageInfo 将 API 响应转换为 UsageInfo。 +func (f *AntigravityQuotaFetcher) buildUsageInfo(modelsResp *antigravity.FetchAvailableModelsResponse, tierRaw, tierNormalized string, loadResp *antigravity.LoadCodeAssistResponse) *UsageInfo { now := time.Now() info := &UsageInfo{ UpdatedAt: &now, @@ -190,6 +191,16 @@ func (f *AntigravityQuotaFetcher) buildUsageInfo(modelsResp *antigravity.FetchAv } } + 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 } diff --git a/backend/internal/service/antigravity_quota_fetcher_test.go b/backend/internal/service/antigravity_quota_fetcher_test.go index 5ead8e60..e0f57051 100644 --- a/backend/internal/service/antigravity_quota_fetcher_test.go +++ b/backend/internal/service/antigravity_quota_fetcher_test.go @@ -81,7 +81,7 @@ func TestBuildUsageInfo_BasicModels(t *testing.T) { }, } - info := fetcher.buildUsageInfo(modelsResp, "g1-pro-tier", "PRO") + info := fetcher.buildUsageInfo(modelsResp, "g1-pro-tier", "PRO", nil) // 基本字段 require.NotNil(t, info.UpdatedAt, "UpdatedAt should be set") @@ -141,7 +141,7 @@ func TestBuildUsageInfo_DeprecatedModels(t *testing.T) { }, } - info := fetcher.buildUsageInfo(modelsResp, "", "") + info := fetcher.buildUsageInfo(modelsResp, "", "", nil) require.Len(t, info.ModelForwardingRules, 2) require.Equal(t, "claude-sonnet-4-20250514", info.ModelForwardingRules["claude-3-sonnet-20240229"]) @@ -159,7 +159,7 @@ func TestBuildUsageInfo_NoDeprecatedModels(t *testing.T) { }, } - info := fetcher.buildUsageInfo(modelsResp, "", "") + info := fetcher.buildUsageInfo(modelsResp, "", "", nil) require.Nil(t, info.ModelForwardingRules, "ModelForwardingRules should be nil when no deprecated models") } @@ -171,7 +171,7 @@ func TestBuildUsageInfo_EmptyModels(t *testing.T) { Models: map[string]antigravity.ModelInfo{}, } - info := fetcher.buildUsageInfo(modelsResp, "", "") + info := fetcher.buildUsageInfo(modelsResp, "", "", nil) require.NotNil(t, info) require.NotNil(t, info.AntigravityQuota) @@ -193,7 +193,7 @@ func TestBuildUsageInfo_ModelWithNilQuotaInfo(t *testing.T) { }, } - info := fetcher.buildUsageInfo(modelsResp, "", "") + info := fetcher.buildUsageInfo(modelsResp, "", "", nil) require.NotNil(t, info) require.Empty(t, info.AntigravityQuota, "models with nil QuotaInfo should be skipped") @@ -222,7 +222,7 @@ func TestBuildUsageInfo_FiveHourPriorityOrder(t *testing.T) { }, } - info := fetcher.buildUsageInfo(modelsResp, "", "") + info := fetcher.buildUsageInfo(modelsResp, "", "", nil) require.NotNil(t, info.FiveHour, "FiveHour should be set when a priority model exists") // claude-sonnet-4-20250514 is first in priority list, so it should be used @@ -251,7 +251,7 @@ func TestBuildUsageInfo_FiveHourFallbackToClaude4(t *testing.T) { }, } - info := fetcher.buildUsageInfo(modelsResp, "", "") + info := fetcher.buildUsageInfo(modelsResp, "", "", nil) require.NotNil(t, info.FiveHour) expectedUtilization := (1.0 - 0.60) * 100 // 40 @@ -277,7 +277,7 @@ func TestBuildUsageInfo_FiveHourFallbackToGemini(t *testing.T) { }, } - info := fetcher.buildUsageInfo(modelsResp, "", "") + info := fetcher.buildUsageInfo(modelsResp, "", "", nil) require.NotNil(t, info.FiveHour) expectedUtilization := (1.0 - 0.30) * 100 // 70 @@ -298,7 +298,7 @@ func TestBuildUsageInfo_FiveHourNoPriorityModel(t *testing.T) { }, } - info := fetcher.buildUsageInfo(modelsResp, "", "") + info := fetcher.buildUsageInfo(modelsResp, "", "", nil) require.Nil(t, info.FiveHour, "FiveHour should be nil when no priority model exists") } @@ -317,7 +317,7 @@ func TestBuildUsageInfo_FiveHourWithEmptyResetTime(t *testing.T) { }, } - info := fetcher.buildUsageInfo(modelsResp, "", "") + info := fetcher.buildUsageInfo(modelsResp, "", "", nil) require.NotNil(t, info.FiveHour) require.Nil(t, info.FiveHour.ResetsAt, "ResetsAt should be nil when ResetTime is empty") @@ -338,7 +338,7 @@ func TestBuildUsageInfo_FullUtilization(t *testing.T) { }, } - info := fetcher.buildUsageInfo(modelsResp, "", "") + info := fetcher.buildUsageInfo(modelsResp, "", "", nil) quota := info.AntigravityQuota["claude-sonnet-4-20250514"] require.NotNil(t, quota) @@ -358,13 +358,38 @@ func TestBuildUsageInfo_ZeroUtilization(t *testing.T) { }, } - info := fetcher.buildUsageInfo(modelsResp, "", "") - + info := fetcher.buildUsageInfo(modelsResp, "", "", nil) quota := info.AntigravityQuota["claude-sonnet-4-20250514"] require.NotNil(t, quota) require.Equal(t, 0, quota.Utilization) } +func TestBuildUsageInfo_AICredits(t *testing.T) { + fetcher := &AntigravityQuotaFetcher{} + modelsResp := &antigravity.FetchAvailableModelsResponse{ + Models: map[string]antigravity.ModelInfo{}, + } + loadResp := &antigravity.LoadCodeAssistResponse{ + PaidTier: &antigravity.PaidTierInfo{ + ID: "g1-pro-tier", + AvailableCredits: []antigravity.AvailableCredit{ + { + CreditType: "GOOGLE_ONE_AI", + CreditAmount: "25", + MinimumCreditAmountForUsage: "5", + }, + }, + }, + } + + info := fetcher.buildUsageInfo(modelsResp, "g1-pro-tier", "PRO", loadResp) + + require.Len(t, info.AICredits, 1) + require.Equal(t, "GOOGLE_ONE_AI", info.AICredits[0].CreditType) + require.Equal(t, 25.0, info.AICredits[0].Amount) + require.Equal(t, 5.0, info.AICredits[0].MinimumBalance) +} + func TestFetchQuota_ForbiddenReturnsIsForbidden(t *testing.T) { // 模拟 FetchQuota 遇到 403 时的行为： // FetchAvailableModels 返回 ForbiddenError → FetchQuota 应返回 is_forbidden=true diff --git a/backend/internal/service/antigravity_quota_scope.go b/backend/internal/service/antigravity_quota_scope.go index e181e7f8..b536d16c 100644 --- a/backend/internal/service/antigravity_quota_scope.go +++ b/backend/internal/service/antigravity_quota_scope.go @@ -32,6 +32,10 @@ func (a *Account) IsSchedulableForModelWithContext(ctx context.Context, requeste return false } if a.isModelRateLimitedWithContext(ctx, requestedModel) { + // Antigravity + overages 启用 + 积分未耗尽 → 放行（有积分可用） + if a.Platform == PlatformAntigravity && a.IsOveragesEnabled() && !a.isCreditsExhausted() { + return true + } return false } return true diff --git a/backend/internal/service/antigravity_rate_limit_test.go b/backend/internal/service/antigravity_rate_limit_test.go index dd8dd83f..df1ce9b9 100644 --- a/backend/internal/service/antigravity_rate_limit_test.go +++ b/backend/internal/service/antigravity_rate_limit_test.go @@ -76,10 +76,16 @@ type modelRateLimitCall struct { resetAt time.Time } +type extraUpdateCall struct { + accountID int64 + updates map[string]any +} + type stubAntigravityAccountRepo struct { AccountRepository rateCalls []rateLimitCall modelRateLimitCalls []modelRateLimitCall + extraUpdateCalls []extraUpdateCall } func (s *stubAntigravityAccountRepo) SetRateLimited(ctx context.Context, id int64, resetAt time.Time) error { @@ -92,6 +98,11 @@ func (s *stubAntigravityAccountRepo) SetModelRateLimit(ctx context.Context, id i return nil } +func (s *stubAntigravityAccountRepo) UpdateExtra(ctx context.Context, id int64, updates map[string]any) error { + s.extraUpdateCalls = append(s.extraUpdateCalls, extraUpdateCall{accountID: id, updates: updates}) + return nil +} + func TestAntigravityRetryLoop_NoURLFallback_UsesConfiguredBaseURL(t *testing.T) { t.Setenv(antigravityForwardBaseURLEnv, "") diff --git a/backend/internal/service/antigravity_smart_retry_test.go b/backend/internal/service/antigravity_smart_retry_test.go index 432c80e5..f569219f 100644 --- a/backend/internal/service/antigravity_smart_retry_test.go +++ b/backend/internal/service/antigravity_smart_retry_test.go @@ -32,15 +32,23 @@ func (c *stubSmartRetryCache) DeleteSessionAccountID(_ context.Context, groupID // mockSmartRetryUpstream 用于 handleSmartRetry 测试的 mock upstream type mockSmartRetryUpstream struct { - responses []*http.Response - errors []error - callIdx int - calls []string + responses []*http.Response + errors []error + callIdx int + calls []string + requestBodies [][]byte } func (m *mockSmartRetryUpstream) Do(req *http.Request, proxyURL string, accountID int64, accountConcurrency int) (*http.Response, error) { idx := m.callIdx m.calls = append(m.calls, req.URL.String()) + if req != nil && req.Body != nil { + body, _ := io.ReadAll(req.Body) + m.requestBodies = append(m.requestBodies, body) + req.Body = io.NopCloser(bytes.NewReader(body)) + } else { + m.requestBodies = append(m.requestBodies, nil) + } m.callIdx++ if idx < len(m.responses) { return m.responses[idx], m.errors[idx] diff --git a/backend/internal/service/ratelimit_service.go b/backend/internal/service/ratelimit_service.go index d410555d..5861a811 100644 --- a/backend/internal/service/ratelimit_service.go +++ b/backend/internal/service/ratelimit_service.go @@ -1174,7 +1174,8 @@ func hasRecoverableRuntimeState(account *Account) bool { if len(account.Extra) == 0 { return false } - return hasNonEmptyMapValue(account.Extra, "model_rate_limits") || hasNonEmptyMapValue(account.Extra, "antigravity_quota_scopes") + return hasNonEmptyMapValue(account.Extra, "model_rate_limits") || + hasNonEmptyMapValue(account.Extra, "antigravity_quota_scopes") } func hasNonEmptyMapValue(extra map[string]any, key string) bool { diff --git a/frontend/src/components/account/AccountStatusIndicator.vue b/frontend/src/components/account/AccountStatusIndicator.vue index 220b5c8b..fc2f7d0c 100644 --- a/frontend/src/components/account/AccountStatusIndicator.vue +++ b/frontend/src/components/account/AccountStatusIndicator.vue @@ -76,19 +76,39 @@ - +
-
+
+ + + {{ t('admin.accounts.status.creditsExhausted') }} + {{ formatModelResetTime(item.reset_at) }} + + + + ⚡ + {{ formatScopeName(item.model) }} + {{ formatModelResetTime(item.reset_at) }} + + + @@ -99,7 +119,13 @@
- {{ t('admin.accounts.status.modelRateLimitedUntil', { model: formatScopeName(item.model), time: formatTime(item.reset_at) }) }} + {{ + item.kind === 'credits_exhausted' + ? t('admin.accounts.status.creditsExhaustedUntil', { time: formatTime(item.reset_at) }) + : item.kind === 'credits_active' + ? t('admin.accounts.status.modelCreditOveragesUntil', { model: formatScopeName(item.model), time: formatTime(item.reset_at) }) + : t('admin.accounts.status.modelRateLimitedUntil', { model: formatScopeName(item.model), time: formatTime(item.reset_at) }) + }}
@@ -131,6 +157,7 @@