Unverified Commit 474165d7 authored by Wesley Liddick's avatar Wesley Liddick Committed by GitHub
Browse files

Merge pull request #1043 from touwaeriol/pr/antigravity-credits-overages

feat: Antigravity AI Credits overages handling & balance display
parents 94e067a2 552a4b99
...@@ -124,10 +124,68 @@ type IneligibleTier struct { ...@@ -124,10 +124,68 @@ type IneligibleTier struct {
type LoadCodeAssistResponse struct { type LoadCodeAssistResponse struct {
CloudAICompanionProject string `json:"cloudaicompanionProject"` CloudAICompanionProject string `json:"cloudaicompanionProject"`
CurrentTier *TierInfo `json:"currentTier,omitempty"` CurrentTier *TierInfo `json:"currentTier,omitempty"`
PaidTier *TierInfo `json:"paidTier,omitempty"` PaidTier *PaidTierInfo `json:"paidTier,omitempty"`
IneligibleTiers []*IneligibleTier `json:"ineligibleTiers,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 请求 // OnboardUserRequest onboardUser 请求
type OnboardUserRequest struct { type OnboardUserRequest struct {
TierID string `json:"tierId"` TierID string `json:"tierId"`
...@@ -157,6 +215,14 @@ func (r *LoadCodeAssistResponse) GetTier() string { ...@@ -157,6 +215,14 @@ func (r *LoadCodeAssistResponse) GetTier() string {
return "" return ""
} }
// GetAvailableCredits 返回 paid tier 中的 AI Credits 余额列表。
func (r *LoadCodeAssistResponse) GetAvailableCredits() []AvailableCredit {
if r.PaidTier == nil {
return nil
}
return r.PaidTier.AvailableCredits
}
// Client Antigravity API 客户端 // Client Antigravity API 客户端
type Client struct { type Client struct {
httpClient *http.Client httpClient *http.Client
......
...@@ -190,7 +190,7 @@ func TestTierInfo_UnmarshalJSON_通过JSON嵌套结构(t *testing.T) { ...@@ -190,7 +190,7 @@ func TestTierInfo_UnmarshalJSON_通过JSON嵌套结构(t *testing.T) {
func TestGetTier_PaidTier优先(t *testing.T) { func TestGetTier_PaidTier优先(t *testing.T) {
resp := &LoadCodeAssistResponse{ resp := &LoadCodeAssistResponse{
CurrentTier: &TierInfo{ID: "free-tier"}, 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" { if got := resp.GetTier(); got != "g1-pro-tier" {
t.Errorf("应返回 paidTier: got %s", got) t.Errorf("应返回 paidTier: got %s", got)
...@@ -209,7 +209,7 @@ func TestGetTier_回退到CurrentTier(t *testing.T) { ...@@ -209,7 +209,7 @@ func TestGetTier_回退到CurrentTier(t *testing.T) {
func TestGetTier_PaidTier为空ID(t *testing.T) { func TestGetTier_PaidTier为空ID(t *testing.T) {
resp := &LoadCodeAssistResponse{ resp := &LoadCodeAssistResponse{
CurrentTier: &TierInfo{ID: "free-tier"}, CurrentTier: &TierInfo{ID: "free-tier"},
PaidTier: &TierInfo{ID: ""}, PaidTier: &PaidTierInfo{ID: ""},
} }
// paidTier.ID 为空时应回退到 currentTier // paidTier.ID 为空时应回退到 currentTier
if got := resp.GetTier(); got != "free-tier" { if got := resp.GetTier(); got != "free-tier" {
...@@ -217,6 +217,32 @@ func TestGetTier_PaidTier为空ID(t *testing.T) { ...@@ -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) { func TestGetTier_两者都为nil(t *testing.T) {
resp := &LoadCodeAssistResponse{} resp := &LoadCodeAssistResponse{}
if got := resp.GetTier(); got != "" { if got := resp.GetTier(); got != "" {
......
...@@ -901,6 +901,22 @@ func (a *Account) IsMixedSchedulingEnabled() bool { ...@@ -901,6 +901,22 @@ func (a *Account) IsMixedSchedulingEnabled() bool {
return false 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 账号是否启用“自动透传(仅替换认证)”。 // IsOpenAIPassthroughEnabled 返回 OpenAI 账号是否启用“自动透传(仅替换认证)”。
// //
// 新字段:accounts.extra.openai_passthrough。 // 新字段:accounts.extra.openai_passthrough。
......
...@@ -166,6 +166,13 @@ type AntigravityModelDetail struct { ...@@ -166,6 +166,13 @@ type AntigravityModelDetail struct {
SupportedMimeTypes map[string]bool `json:"supported_mime_types,omitempty"` 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 账号使用量信息 // UsageInfo 账号使用量信息
type UsageInfo struct { type UsageInfo struct {
UpdatedAt *time.Time `json:"updated_at,omitempty"` // 更新时间 UpdatedAt *time.Time `json:"updated_at,omitempty"` // 更新时间
...@@ -189,6 +196,9 @@ type UsageInfo struct { ...@@ -189,6 +196,9 @@ type UsageInfo struct {
// Antigravity 模型详细能力信息(与 antigravity_quota 同 key) // Antigravity 模型详细能力信息(与 antigravity_quota 同 key)
AntigravityQuotaDetails map[string]*AntigravityModelDetail `json:"antigravity_quota_details,omitempty"` 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) // Antigravity 废弃模型转发规则 (old_model_id -> new_model_id)
ModelForwardingRules map[string]string `json:"model_forwarding_rules,omitempty"` ModelForwardingRules map[string]string `json:"model_forwarding_rules,omitempty"`
......
...@@ -368,6 +368,10 @@ type ProxyExitInfoProber interface { ...@@ -368,6 +368,10 @@ type ProxyExitInfoProber interface {
ProbeProxy(ctx context.Context, proxyURL string) (*ProxyExitInfo, int64, error) 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 { type proxyQualityTarget struct {
Target string Target string
URL string URL string
...@@ -445,10 +449,6 @@ type userGroupRateBatchReader interface { ...@@ -445,10 +449,6 @@ type userGroupRateBatchReader interface {
GetByUserIDs(ctx context.Context, userIDs []int64) (map[int64]map[int64]float64, error) 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 // NewAdminService creates a new AdminService
func NewAdminService( func NewAdminService(
userRepo UserRepository, userRepo UserRepository,
...@@ -1516,6 +1516,7 @@ func (s *adminServiceImpl) UpdateAccount(ctx context.Context, id int64, input *U ...@@ -1516,6 +1516,7 @@ func (s *adminServiceImpl) UpdateAccount(ctx context.Context, id int64, input *U
if err != nil { if err != nil {
return nil, err return nil, err
} }
wasOveragesEnabled := account.IsOveragesEnabled()
if input.Name != "" { if input.Name != "" {
account.Name = input.Name account.Name = input.Name
...@@ -1537,6 +1538,17 @@ func (s *adminServiceImpl) UpdateAccount(ctx context.Context, id int64, input *U ...@@ -1537,6 +1538,17 @@ func (s *adminServiceImpl) UpdateAccount(ctx context.Context, id int64, input *U
} }
} }
account.Extra = input.Extra 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 { if err := ValidateQuotaResetConfig(account.Extra); err != nil {
return nil, err return nil, err
......
//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 时应在持久化前清掉旧模型限流")
}
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))
}
}
//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, "普通模型限流应保留")
})
}
...@@ -188,9 +188,29 @@ func (s *AntigravityGatewayService) handleSmartRetry(p antigravityRetryLoopParam ...@@ -188,9 +188,29 @@ func (s *AntigravityGatewayService) handleSmartRetry(p antigravityRetryLoopParam
return &smartRetryResult{action: smartRetryActionContinueURL} return &smartRetryResult{action: smartRetryActionContinueURL}
} }
category := antigravity429Unknown
if resp.StatusCode == http.StatusTooManyRequests {
category = classifyAntigravity429(respBody)
}
// 判断是否触发智能重试 // 判断是否触发智能重试
shouldSmartRetry, shouldRateLimitModel, waitDuration, modelName, isModelCapacityExhausted := shouldTriggerAntigravitySmartRetry(p.account, 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 >= 阈值,限流模型并切换账号 // 情况1: retryDelay >= 阈值,限流模型并切换账号
if shouldRateLimitModel { if shouldRateLimitModel {
// 单账号 503 退避重试模式:不设限流、不切换账号,改为原地等待+重试 // 单账号 503 退避重试模式:不设限流、不切换账号,改为原地等待+重试
...@@ -532,14 +552,31 @@ func (s *AntigravityGatewayService) handleSingleAccountRetryInPlace( ...@@ -532,14 +552,31 @@ func (s *AntigravityGatewayService) handleSingleAccountRetryInPlace(
// antigravityRetryLoop 执行带 URL fallback 的重试循环 // antigravityRetryLoop 执行带 URL fallback 的重试循环
func (s *AntigravityGatewayService) antigravityRetryLoop(p antigravityRetryLoopParams) (*antigravityRetryLoopResult, error) { 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 p.requestedModel != "" {
if remaining := p.account.GetRateLimitRemainingTimeWithContext(p.ctx, p.requestedModel); remaining > 0 { if remaining := p.account.GetRateLimitRemainingTimeWithContext(p.ctx, p.requestedModel); remaining > 0 {
// 单账号 503 退避重试模式:跳过限流预检查,直接发请求。 // 已注入积分的请求不再受普通模型限流预检查阻断。
// 首次请求设的限流是为了多账号调度器跳过该账号,在单账号模式下无意义。 if overagesInjected {
// 如果上游确实还不可用,handleSmartRetry → handleSingleAccountRetryInPlace logger.LegacyPrintf("service.antigravity_gateway", "%s pre_check: credits_injected_ignore_rate_limit remaining=%v model=%s account=%d",
// 会在 Service 层原地等待+重试,不需要在预检查这里等。 p.prefix, remaining.Truncate(time.Millisecond), p.requestedModel, p.account.ID)
if isSingleAccountRetry(p.ctx) { } 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)", 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) p.prefix, remaining.Truncate(time.Millisecond), p.requestedModel, p.account.ID)
} else { } else {
...@@ -631,6 +668,15 @@ urlFallbackLoop: ...@@ -631,6 +668,15 @@ urlFallbackLoop:
respBody, _ := io.ReadAll(io.LimitReader(resp.Body, 2<<20)) respBody, _ := io.ReadAll(io.LimitReader(resp.Body, 2<<20))
_ = resp.Body.Close() _ = 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 handled, outStatus, policyErr := s.applyErrorPolicy(p, resp.StatusCode, resp.Header, respBody); handled {
if policyErr != nil { if policyErr != nil {
......
...@@ -78,11 +78,11 @@ func (f *AntigravityQuotaFetcher) FetchQuota(ctx context.Context, account *Accou ...@@ -78,11 +78,11 @@ func (f *AntigravityQuotaFetcher) FetchQuota(ctx context.Context, account *Accou
return nil, err return nil, err
} }
// 调用 LoadCodeAssist 获取订阅等级(非关键路径,失败不影响主流程) // 调用 LoadCodeAssist 获取订阅等级和 AI Credits 余额(非关键路径,失败不影响主流程)
tierRaw, tierNormalized := f.fetchSubscriptionTier(ctx, client, accessToken) tierRaw, tierNormalized, loadResp := f.fetchSubscriptionTier(ctx, client, accessToken)
// 转换为 UsageInfo // 转换为 UsageInfo
usageInfo := f.buildUsageInfo(modelsResp, tierRaw, tierNormalized) usageInfo := f.buildUsageInfo(modelsResp, tierRaw, tierNormalized, loadResp)
return &QuotaResult{ return &QuotaResult{
UsageInfo: usageInfo, UsageInfo: usageInfo,
...@@ -90,20 +90,21 @@ func (f *AntigravityQuotaFetcher) FetchQuota(ctx context.Context, account *Accou ...@@ -90,20 +90,21 @@ func (f *AntigravityQuotaFetcher) FetchQuota(ctx context.Context, account *Accou
}, nil }, nil
} }
// fetchSubscriptionTier 获取账号订阅等级,失败返回空字符串 // fetchSubscriptionTier 获取账号订阅等级,失败返回空字符串。
func (f *AntigravityQuotaFetcher) fetchSubscriptionTier(ctx context.Context, client *antigravity.Client, accessToken string) (raw, normalized string) { // 同时返回 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) loadResp, _, err := client.LoadCodeAssist(ctx, accessToken)
if err != nil { if err != nil {
slog.Warn("failed to fetch subscription tier", "error", err) slog.Warn("failed to fetch subscription tier", "error", err)
return "", "" return "", "", nil
} }
if loadResp == nil { if loadResp == nil {
return "", "" return "", "", nil
} }
raw = loadResp.GetTier() // 已有方法:paidTier > currentTier raw = loadResp.GetTier() // 已有方法:paidTier > currentTier
normalized = normalizeTier(raw) normalized = normalizeTier(raw)
return raw, normalized return raw, normalized, loadResp
} }
// normalizeTier 将原始 tier 字符串归一化为 FREE/PRO/ULTRA/UNKNOWN // normalizeTier 将原始 tier 字符串归一化为 FREE/PRO/ULTRA/UNKNOWN
...@@ -124,8 +125,8 @@ func normalizeTier(raw string) string { ...@@ -124,8 +125,8 @@ func normalizeTier(raw string) string {
} }
} }
// buildUsageInfo 将 API 响应转换为 UsageInfo // buildUsageInfo 将 API 响应转换为 UsageInfo
func (f *AntigravityQuotaFetcher) buildUsageInfo(modelsResp *antigravity.FetchAvailableModelsResponse, tierRaw, tierNormalized string) *UsageInfo { func (f *AntigravityQuotaFetcher) buildUsageInfo(modelsResp *antigravity.FetchAvailableModelsResponse, tierRaw, tierNormalized string, loadResp *antigravity.LoadCodeAssistResponse) *UsageInfo {
now := time.Now() now := time.Now()
info := &UsageInfo{ info := &UsageInfo{
UpdatedAt: &now, UpdatedAt: &now,
...@@ -190,6 +191,16 @@ func (f *AntigravityQuotaFetcher) buildUsageInfo(modelsResp *antigravity.FetchAv ...@@ -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 return info
} }
......
...@@ -81,7 +81,7 @@ func TestBuildUsageInfo_BasicModels(t *testing.T) { ...@@ -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") require.NotNil(t, info.UpdatedAt, "UpdatedAt should be set")
...@@ -141,7 +141,7 @@ func TestBuildUsageInfo_DeprecatedModels(t *testing.T) { ...@@ -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.Len(t, info.ModelForwardingRules, 2)
require.Equal(t, "claude-sonnet-4-20250514", info.ModelForwardingRules["claude-3-sonnet-20240229"]) require.Equal(t, "claude-sonnet-4-20250514", info.ModelForwardingRules["claude-3-sonnet-20240229"])
...@@ -159,7 +159,7 @@ func TestBuildUsageInfo_NoDeprecatedModels(t *testing.T) { ...@@ -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") require.Nil(t, info.ModelForwardingRules, "ModelForwardingRules should be nil when no deprecated models")
} }
...@@ -171,7 +171,7 @@ func TestBuildUsageInfo_EmptyModels(t *testing.T) { ...@@ -171,7 +171,7 @@ func TestBuildUsageInfo_EmptyModels(t *testing.T) {
Models: map[string]antigravity.ModelInfo{}, Models: map[string]antigravity.ModelInfo{},
} }
info := fetcher.buildUsageInfo(modelsResp, "", "") info := fetcher.buildUsageInfo(modelsResp, "", "", nil)
require.NotNil(t, info) require.NotNil(t, info)
require.NotNil(t, info.AntigravityQuota) require.NotNil(t, info.AntigravityQuota)
...@@ -193,7 +193,7 @@ func TestBuildUsageInfo_ModelWithNilQuotaInfo(t *testing.T) { ...@@ -193,7 +193,7 @@ func TestBuildUsageInfo_ModelWithNilQuotaInfo(t *testing.T) {
}, },
} }
info := fetcher.buildUsageInfo(modelsResp, "", "") info := fetcher.buildUsageInfo(modelsResp, "", "", nil)
require.NotNil(t, info) require.NotNil(t, info)
require.Empty(t, info.AntigravityQuota, "models with nil QuotaInfo should be skipped") require.Empty(t, info.AntigravityQuota, "models with nil QuotaInfo should be skipped")
...@@ -222,7 +222,7 @@ func TestBuildUsageInfo_FiveHourPriorityOrder(t *testing.T) { ...@@ -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") 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 // claude-sonnet-4-20250514 is first in priority list, so it should be used
...@@ -251,7 +251,7 @@ func TestBuildUsageInfo_FiveHourFallbackToClaude4(t *testing.T) { ...@@ -251,7 +251,7 @@ func TestBuildUsageInfo_FiveHourFallbackToClaude4(t *testing.T) {
}, },
} }
info := fetcher.buildUsageInfo(modelsResp, "", "") info := fetcher.buildUsageInfo(modelsResp, "", "", nil)
require.NotNil(t, info.FiveHour) require.NotNil(t, info.FiveHour)
expectedUtilization := (1.0 - 0.60) * 100 // 40 expectedUtilization := (1.0 - 0.60) * 100 // 40
...@@ -277,7 +277,7 @@ func TestBuildUsageInfo_FiveHourFallbackToGemini(t *testing.T) { ...@@ -277,7 +277,7 @@ func TestBuildUsageInfo_FiveHourFallbackToGemini(t *testing.T) {
}, },
} }
info := fetcher.buildUsageInfo(modelsResp, "", "") info := fetcher.buildUsageInfo(modelsResp, "", "", nil)
require.NotNil(t, info.FiveHour) require.NotNil(t, info.FiveHour)
expectedUtilization := (1.0 - 0.30) * 100 // 70 expectedUtilization := (1.0 - 0.30) * 100 // 70
...@@ -298,7 +298,7 @@ func TestBuildUsageInfo_FiveHourNoPriorityModel(t *testing.T) { ...@@ -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") require.Nil(t, info.FiveHour, "FiveHour should be nil when no priority model exists")
} }
...@@ -317,7 +317,7 @@ func TestBuildUsageInfo_FiveHourWithEmptyResetTime(t *testing.T) { ...@@ -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.NotNil(t, info.FiveHour)
require.Nil(t, info.FiveHour.ResetsAt, "ResetsAt should be nil when ResetTime is empty") require.Nil(t, info.FiveHour.ResetsAt, "ResetsAt should be nil when ResetTime is empty")
...@@ -338,7 +338,7 @@ func TestBuildUsageInfo_FullUtilization(t *testing.T) { ...@@ -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"] quota := info.AntigravityQuota["claude-sonnet-4-20250514"]
require.NotNil(t, quota) require.NotNil(t, quota)
...@@ -358,13 +358,38 @@ func TestBuildUsageInfo_ZeroUtilization(t *testing.T) { ...@@ -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"] quota := info.AntigravityQuota["claude-sonnet-4-20250514"]
require.NotNil(t, quota) require.NotNil(t, quota)
require.Equal(t, 0, quota.Utilization) 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) { func TestFetchQuota_ForbiddenReturnsIsForbidden(t *testing.T) {
// 模拟 FetchQuota 遇到 403 时的行为: // 模拟 FetchQuota 遇到 403 时的行为:
// FetchAvailableModels 返回 ForbiddenError → FetchQuota 应返回 is_forbidden=true // FetchAvailableModels 返回 ForbiddenError → FetchQuota 应返回 is_forbidden=true
......
...@@ -32,6 +32,10 @@ func (a *Account) IsSchedulableForModelWithContext(ctx context.Context, requeste ...@@ -32,6 +32,10 @@ func (a *Account) IsSchedulableForModelWithContext(ctx context.Context, requeste
return false return false
} }
if a.isModelRateLimitedWithContext(ctx, requestedModel) { if a.isModelRateLimitedWithContext(ctx, requestedModel) {
// Antigravity + overages 启用 + 积分未耗尽 → 放行(有积分可用)
if a.Platform == PlatformAntigravity && a.IsOveragesEnabled() && !a.isCreditsExhausted() {
return true
}
return false return false
} }
return true return true
......
...@@ -76,10 +76,16 @@ type modelRateLimitCall struct { ...@@ -76,10 +76,16 @@ type modelRateLimitCall struct {
resetAt time.Time resetAt time.Time
} }
type extraUpdateCall struct {
accountID int64
updates map[string]any
}
type stubAntigravityAccountRepo struct { type stubAntigravityAccountRepo struct {
AccountRepository AccountRepository
rateCalls []rateLimitCall rateCalls []rateLimitCall
modelRateLimitCalls []modelRateLimitCall modelRateLimitCalls []modelRateLimitCall
extraUpdateCalls []extraUpdateCall
} }
func (s *stubAntigravityAccountRepo) SetRateLimited(ctx context.Context, id int64, resetAt time.Time) error { 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 ...@@ -92,6 +98,11 @@ func (s *stubAntigravityAccountRepo) SetModelRateLimit(ctx context.Context, id i
return nil 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) { func TestAntigravityRetryLoop_NoURLFallback_UsesConfiguredBaseURL(t *testing.T) {
t.Setenv(antigravityForwardBaseURLEnv, "") t.Setenv(antigravityForwardBaseURLEnv, "")
......
...@@ -32,15 +32,23 @@ func (c *stubSmartRetryCache) DeleteSessionAccountID(_ context.Context, groupID ...@@ -32,15 +32,23 @@ func (c *stubSmartRetryCache) DeleteSessionAccountID(_ context.Context, groupID
// mockSmartRetryUpstream 用于 handleSmartRetry 测试的 mock upstream // mockSmartRetryUpstream 用于 handleSmartRetry 测试的 mock upstream
type mockSmartRetryUpstream struct { type mockSmartRetryUpstream struct {
responses []*http.Response responses []*http.Response
errors []error errors []error
callIdx int callIdx int
calls []string calls []string
requestBodies [][]byte
} }
func (m *mockSmartRetryUpstream) Do(req *http.Request, proxyURL string, accountID int64, accountConcurrency int) (*http.Response, error) { func (m *mockSmartRetryUpstream) Do(req *http.Request, proxyURL string, accountID int64, accountConcurrency int) (*http.Response, error) {
idx := m.callIdx idx := m.callIdx
m.calls = append(m.calls, req.URL.String()) 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++ m.callIdx++
if idx < len(m.responses) { if idx < len(m.responses) {
return m.responses[idx], m.errors[idx] return m.responses[idx], m.errors[idx]
......
...@@ -1174,7 +1174,8 @@ func hasRecoverableRuntimeState(account *Account) bool { ...@@ -1174,7 +1174,8 @@ func hasRecoverableRuntimeState(account *Account) bool {
if len(account.Extra) == 0 { if len(account.Extra) == 0 {
return false 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 { func hasNonEmptyMapValue(extra map[string]any, key string) bool {
......
...@@ -76,19 +76,39 @@ ...@@ -76,19 +76,39 @@
</div> </div>
</div> </div>
<!-- Model Rate Limit Indicators (Antigravity OAuth Smart Retry) --> <!-- Model Status Indicators (普通限流 / 超量请求中) -->
<div <div
v-if="activeModelRateLimits.length > 0" v-if="activeModelStatuses.length > 0"
:class="[ :class="[
activeModelRateLimits.length <= 4 activeModelStatuses.length <= 4
? 'flex flex-col gap-1' ? 'flex flex-col gap-1'
: activeModelRateLimits.length <= 8 : activeModelStatuses.length <= 8
? 'columns-2 gap-x-2' ? 'columns-2 gap-x-2'
: 'columns-3 gap-x-2' : 'columns-3 gap-x-2'
]" ]"
> >
<div v-for="item in activeModelRateLimits" :key="item.model" class="group relative mb-1 break-inside-avoid"> <div v-for="item in activeModelStatuses" :key="`${item.kind}-${item.model}`" class="group relative mb-1 break-inside-avoid">
<!-- 积分已用尽 -->
<span <span
v-if="item.kind === 'credits_exhausted'"
class="inline-flex items-center gap-1 rounded bg-red-100 px-1.5 py-0.5 text-xs font-medium text-red-700 dark:bg-red-900/30 dark:text-red-400"
>
<Icon name="exclamationTriangle" size="xs" :stroke-width="2" />
{{ t('admin.accounts.status.creditsExhausted') }}
<span class="text-[10px] opacity-70">{{ formatModelResetTime(item.reset_at) }}</span>
</span>
<!-- 正在走积分(模型限流但积分可用)-->
<span
v-else-if="item.kind === 'credits_active'"
class="inline-flex items-center gap-1 rounded bg-amber-100 px-1.5 py-0.5 text-xs font-medium text-amber-700 dark:bg-amber-900/30 dark:text-amber-400"
>
<span></span>
{{ formatScopeName(item.model) }}
<span class="text-[10px] opacity-70">{{ formatModelResetTime(item.reset_at) }}</span>
</span>
<!-- 普通模型限流 -->
<span
v-else
class="inline-flex items-center gap-1 rounded bg-purple-100 px-1.5 py-0.5 text-xs font-medium text-purple-700 dark:bg-purple-900/30 dark:text-purple-400" class="inline-flex items-center gap-1 rounded bg-purple-100 px-1.5 py-0.5 text-xs font-medium text-purple-700 dark:bg-purple-900/30 dark:text-purple-400"
> >
<Icon name="exclamationTriangle" size="xs" :stroke-width="2" /> <Icon name="exclamationTriangle" size="xs" :stroke-width="2" />
...@@ -99,7 +119,13 @@ ...@@ -99,7 +119,13 @@
<div <div
class="pointer-events-none absolute bottom-full left-1/2 z-50 mb-2 w-56 -translate-x-1/2 whitespace-normal rounded bg-gray-900 px-3 py-2 text-center text-xs leading-relaxed text-white opacity-0 transition-opacity group-hover:opacity-100 dark:bg-gray-700" class="pointer-events-none absolute bottom-full left-1/2 z-50 mb-2 w-56 -translate-x-1/2 whitespace-normal rounded bg-gray-900 px-3 py-2 text-center text-xs leading-relaxed text-white opacity-0 transition-opacity group-hover:opacity-100 dark:bg-gray-700"
> >
{{ 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) })
}}
<div <div
class="absolute left-1/2 top-full -translate-x-1/2 border-4 border-transparent border-t-gray-900 dark:border-t-gray-700" class="absolute left-1/2 top-full -translate-x-1/2 border-4 border-transparent border-t-gray-900 dark:border-t-gray-700"
></div> ></div>
...@@ -131,6 +157,7 @@ ...@@ -131,6 +157,7 @@
<script setup lang="ts"> <script setup lang="ts">
import { computed } from 'vue' import { computed } from 'vue'
import { useI18n } from 'vue-i18n' import { useI18n } from 'vue-i18n'
import Icon from '@/components/icons/Icon.vue'
import type { Account } from '@/types' import type { Account } from '@/types'
import { formatCountdown, formatDateTime, formatCountdownWithSuffix, formatTime } from '@/utils/format' import { formatCountdown, formatDateTime, formatCountdownWithSuffix, formatTime } from '@/utils/format'
...@@ -150,17 +177,44 @@ const isRateLimited = computed(() => { ...@@ -150,17 +177,44 @@ const isRateLimited = computed(() => {
return new Date(props.account.rate_limit_reset_at) > new Date() return new Date(props.account.rate_limit_reset_at) > new Date()
}) })
type AccountModelStatusItem = {
kind: 'rate_limit' | 'credits_exhausted' | 'credits_active'
model: string
reset_at: string
}
// Computed: active model rate limits (Antigravity OAuth Smart Retry) // Computed: active model statuses (普通模型限流 + 积分耗尽 + 走积分中)
const activeModelRateLimits = computed(() => { const activeModelStatuses = computed<AccountModelStatusItem[]>(() => {
const modelLimits = (props.account.extra as Record<string, unknown> | undefined)?.model_rate_limits as const extra = props.account.extra as Record<string, unknown> | undefined
const modelLimits = extra?.model_rate_limits as
| Record<string, { rate_limited_at: string; rate_limit_reset_at: string }> | Record<string, { rate_limited_at: string; rate_limit_reset_at: string }>
| undefined | undefined
if (!modelLimits) return []
const now = new Date() const now = new Date()
return Object.entries(modelLimits) const items: AccountModelStatusItem[] = []
.filter(([, info]) => new Date(info.rate_limit_reset_at) > now)
.map(([model, info]) => ({ model, reset_at: info.rate_limit_reset_at })) if (!modelLimits) return items
// 检查 AICredits key 是否生效(积分是否耗尽)
const aiCreditsEntry = modelLimits['AICredits']
const hasActiveAICredits = aiCreditsEntry && new Date(aiCreditsEntry.rate_limit_reset_at) > now
const allowOverages = !!(extra?.allow_overages)
for (const [model, info] of Object.entries(modelLimits)) {
if (new Date(info.rate_limit_reset_at) <= now) continue
if (model === 'AICredits') {
// AICredits key → 积分已用尽
items.push({ kind: 'credits_exhausted', model, reset_at: info.rate_limit_reset_at })
} else if (allowOverages && !hasActiveAICredits) {
// 普通模型限流 + overages 启用 + 积分可用 → 正在走积分
items.push({ kind: 'credits_active', model, reset_at: info.rate_limit_reset_at })
} else {
// 普通模型限流
items.push({ kind: 'rate_limit', model, reset_at: info.rate_limit_reset_at })
}
}
return items
}) })
const formatScopeName = (scope: string): string => { const formatScopeName = (scope: string): string => {
...@@ -182,7 +236,7 @@ const formatScopeName = (scope: string): string => { ...@@ -182,7 +236,7 @@ const formatScopeName = (scope: string): string => {
'gemini-3.1-pro-high': 'G3PH', 'gemini-3.1-pro-high': 'G3PH',
'gemini-3.1-pro-low': 'G3PL', 'gemini-3.1-pro-low': 'G3PL',
'gemini-3-pro-image': 'G3PI', 'gemini-3-pro-image': 'G3PI',
'gemini-3.1-flash-image': 'GImage', 'gemini-3.1-flash-image': 'G31FI',
// 其他 // 其他
'gpt-oss-120b-medium': 'GPT120', 'gpt-oss-120b-medium': 'GPT120',
'tab_flash_lite_preview': 'TabFL', 'tab_flash_lite_preview': 'TabFL',
......
...@@ -289,6 +289,13 @@ ...@@ -289,6 +289,13 @@
:resets-at="antigravityClaudeUsageFromAPI.resetTime" :resets-at="antigravityClaudeUsageFromAPI.resetTime"
color="amber" color="amber"
/> />
<div v-if="aiCreditsDisplay" class="mt-1 text-[10px] text-gray-500 dark:text-gray-400">
💳 {{ t('admin.accounts.aiCreditsBalance') }}: {{ aiCreditsDisplay }}
</div>
</div>
<div v-else-if="aiCreditsDisplay" class="text-[10px] text-gray-500 dark:text-gray-400">
💳 {{ t('admin.accounts.aiCreditsBalance') }}: {{ aiCreditsDisplay }}
</div> </div>
<div v-else class="text-xs text-gray-400">-</div> <div v-else class="text-xs text-gray-400">-</div>
</template> </template>
...@@ -581,6 +588,14 @@ const antigravityClaudeUsageFromAPI = computed(() => ...@@ -581,6 +588,14 @@ const antigravityClaudeUsageFromAPI = computed(() =>
]) ])
) )
const aiCreditsDisplay = computed(() => {
const credits = usageInfo.value?.ai_credits
if (!credits || credits.length === 0) return null
const total = credits.reduce((sum, credit) => sum + (credit.amount ?? 0), 0)
if (total <= 0) return null
return total.toFixed(0)
})
// Antigravity 账户类型(从 load_code_assist 响应中提取) // Antigravity 账户类型(从 load_code_assist 响应中提取)
const antigravityTier = computed(() => { const antigravityTier = computed(() => {
const extra = props.account.extra as Record<string, unknown> | undefined const extra = props.account.extra as Record<string, unknown> | undefined
......
...@@ -2449,6 +2449,33 @@ ...@@ -2449,6 +2449,33 @@
</div> </div>
</div> </div>
</div> </div>
<div v-if="form.platform === 'antigravity'" class="mt-3 flex items-center gap-2">
<label class="flex cursor-pointer items-center gap-2">
<input
type="checkbox"
v-model="allowOverages"
class="h-4 w-4 rounded border-gray-300 text-primary-500 focus:ring-primary-500 dark:border-dark-500"
/>
<span class="text-sm font-medium text-gray-700 dark:text-gray-300">
{{ t('admin.accounts.allowOverages') }}
</span>
</label>
<div class="group relative">
<span
class="inline-flex h-4 w-4 cursor-help items-center justify-center rounded-full bg-gray-200 text-xs text-gray-500 hover:bg-gray-300 dark:bg-dark-600 dark:text-gray-400 dark:hover:bg-dark-500"
>
?
</span>
<div
class="pointer-events-none absolute left-0 top-full z-[100] mt-1.5 w-72 rounded bg-gray-900 px-3 py-2 text-xs text-white opacity-0 transition-opacity group-hover:opacity-100 dark:bg-gray-700"
>
{{ t('admin.accounts.allowOveragesTooltip') }}
<div
class="absolute bottom-full left-3 border-4 border-transparent border-b-gray-900 dark:border-b-gray-700"
></div>
</div>
</div>
</div>
<!-- Group Selection - 仅标准模式显示 --> <!-- Group Selection - 仅标准模式显示 -->
<GroupSelector <GroupSelector
...@@ -2991,6 +3018,7 @@ const openaiAPIKeyResponsesWebSocketV2Mode = ref<OpenAIWSMode>(OPENAI_WS_MODE_OF ...@@ -2991,6 +3018,7 @@ const openaiAPIKeyResponsesWebSocketV2Mode = ref<OpenAIWSMode>(OPENAI_WS_MODE_OF
const codexCLIOnlyEnabled = ref(false) const codexCLIOnlyEnabled = ref(false)
const anthropicPassthroughEnabled = ref(false) const anthropicPassthroughEnabled = ref(false)
const mixedScheduling = ref(false) // For antigravity accounts: enable mixed scheduling const mixedScheduling = ref(false) // For antigravity accounts: enable mixed scheduling
const allowOverages = ref(false) // For antigravity accounts: enable AI Credits overages
const antigravityAccountType = ref<'oauth' | 'upstream'>('oauth') // For antigravity: oauth or upstream const antigravityAccountType = ref<'oauth' | 'upstream'>('oauth') // For antigravity: oauth or upstream
const soraAccountType = ref<'oauth' | 'apikey'>('oauth') // For sora: oauth or apikey (upstream) const soraAccountType = ref<'oauth' | 'apikey'>('oauth') // For sora: oauth or apikey (upstream)
const upstreamBaseUrl = ref('') // For upstream type: base URL const upstreamBaseUrl = ref('') // For upstream type: base URL
...@@ -3017,6 +3045,13 @@ const getTempUnschedRuleKey = createStableObjectKeyResolver<TempUnschedRuleForm> ...@@ -3017,6 +3045,13 @@ const getTempUnschedRuleKey = createStableObjectKeyResolver<TempUnschedRuleForm>
const geminiOAuthType = ref<'code_assist' | 'google_one' | 'ai_studio'>('google_one') const geminiOAuthType = ref<'code_assist' | 'google_one' | 'ai_studio'>('google_one')
const geminiAIStudioOAuthEnabled = ref(false) const geminiAIStudioOAuthEnabled = ref(false)
function buildAntigravityExtra(): Record<string, unknown> | undefined {
const extra: Record<string, unknown> = {}
if (mixedScheduling.value) extra.mixed_scheduling = true
if (allowOverages.value) extra.allow_overages = true
return Object.keys(extra).length > 0 ? extra : undefined
}
const showMixedChannelWarning = ref(false) const showMixedChannelWarning = ref(false)
const mixedChannelWarningDetails = ref<{ groupName: string; currentPlatform: string; otherPlatform: string } | null>( const mixedChannelWarningDetails = ref<{ groupName: string; currentPlatform: string; otherPlatform: string } | null>(
null null
...@@ -3282,6 +3317,7 @@ watch( ...@@ -3282,6 +3317,7 @@ watch(
accountCategory.value = 'oauth-based' accountCategory.value = 'oauth-based'
antigravityAccountType.value = 'oauth' antigravityAccountType.value = 'oauth'
} else { } else {
allowOverages.value = false
antigravityWhitelistModels.value = [] antigravityWhitelistModels.value = []
antigravityModelMappings.value = [] antigravityModelMappings.value = []
antigravityModelRestrictionMode.value = 'mapping' antigravityModelRestrictionMode.value = 'mapping'
...@@ -3712,6 +3748,7 @@ const resetForm = () => { ...@@ -3712,6 +3748,7 @@ const resetForm = () => {
sessionIdMaskingEnabled.value = false sessionIdMaskingEnabled.value = false
cacheTTLOverrideEnabled.value = false cacheTTLOverrideEnabled.value = false
cacheTTLOverrideTarget.value = '5m' cacheTTLOverrideTarget.value = '5m'
allowOverages.value = false
antigravityAccountType.value = 'oauth' antigravityAccountType.value = 'oauth'
upstreamBaseUrl.value = '' upstreamBaseUrl.value = ''
upstreamApiKey.value = '' upstreamApiKey.value = ''
...@@ -3960,7 +3997,7 @@ const handleSubmit = async () => { ...@@ -3960,7 +3997,7 @@ const handleSubmit = async () => {
applyInterceptWarmup(credentials, interceptWarmupRequests.value, 'create') applyInterceptWarmup(credentials, interceptWarmupRequests.value, 'create')
const extra = mixedScheduling.value ? { mixed_scheduling: true } : undefined const extra = buildAntigravityExtra()
await createAccountAndFinish(form.platform, 'apikey', credentials, extra) await createAccountAndFinish(form.platform, 'apikey', credentials, extra)
return return
} }
...@@ -4706,7 +4743,7 @@ const handleAntigravityExchange = async (authCode: string) => { ...@@ -4706,7 +4743,7 @@ const handleAntigravityExchange = async (authCode: string) => {
if (antigravityModelMapping) { if (antigravityModelMapping) {
credentials.model_mapping = antigravityModelMapping credentials.model_mapping = antigravityModelMapping
} }
const extra = mixedScheduling.value ? { mixed_scheduling: true } : undefined const extra = buildAntigravityExtra()
await createAccountAndFinish('antigravity', 'oauth', credentials, extra) await createAccountAndFinish('antigravity', 'oauth', credentials, extra)
} catch (error: any) { } catch (error: any) {
antigravityOAuth.error.value = error.response?.data?.detail || t('admin.accounts.oauth.authFailed') antigravityOAuth.error.value = error.response?.data?.detail || t('admin.accounts.oauth.authFailed')
......
...@@ -1610,6 +1610,33 @@ ...@@ -1610,6 +1610,33 @@
</div> </div>
</div> </div>
</div> </div>
<div v-if="account?.platform === 'antigravity'" class="mt-3 flex items-center gap-2">
<label class="flex cursor-pointer items-center gap-2">
<input
type="checkbox"
v-model="allowOverages"
class="h-4 w-4 rounded border-gray-300 text-primary-500 focus:ring-primary-500 dark:border-dark-500"
/>
<span class="text-sm font-medium text-gray-700 dark:text-gray-300">
{{ t('admin.accounts.allowOverages') }}
</span>
</label>
<div class="group relative">
<span
class="inline-flex h-4 w-4 cursor-help items-center justify-center rounded-full bg-gray-200 text-xs text-gray-500 hover:bg-gray-300 dark:bg-dark-600 dark:text-gray-400 dark:hover:bg-dark-500"
>
?
</span>
<div
class="pointer-events-none absolute left-0 top-full z-[100] mt-1.5 w-72 rounded bg-gray-900 px-3 py-2 text-xs text-white opacity-0 transition-opacity group-hover:opacity-100 dark:bg-gray-700"
>
{{ t('admin.accounts.allowOveragesTooltip') }}
<div
class="absolute bottom-full left-3 border-4 border-transparent border-b-gray-900 dark:border-b-gray-700"
></div>
</div>
</div>
</div>
</div> </div>
<!-- Group Selection - 仅标准模式显示 --> <!-- Group Selection - 仅标准模式显示 -->
...@@ -1778,6 +1805,7 @@ const customErrorCodeInput = ref<number | null>(null) ...@@ -1778,6 +1805,7 @@ const customErrorCodeInput = ref<number | null>(null)
const interceptWarmupRequests = ref(false) const interceptWarmupRequests = ref(false)
const autoPauseOnExpired = ref(false) const autoPauseOnExpired = ref(false)
const mixedScheduling = ref(false) // For antigravity accounts: enable mixed scheduling const mixedScheduling = ref(false) // For antigravity accounts: enable mixed scheduling
const allowOverages = ref(false) // For antigravity accounts: enable AI Credits overages
const antigravityModelRestrictionMode = ref<'whitelist' | 'mapping'>('whitelist') const antigravityModelRestrictionMode = ref<'whitelist' | 'mapping'>('whitelist')
const antigravityWhitelistModels = ref<string[]>([]) const antigravityWhitelistModels = ref<string[]>([])
const antigravityModelMappings = ref<ModelMapping[]>([]) const antigravityModelMappings = ref<ModelMapping[]>([])
...@@ -1980,8 +2008,11 @@ watch( ...@@ -1980,8 +2008,11 @@ watch(
autoPauseOnExpired.value = newAccount.auto_pause_on_expired === true autoPauseOnExpired.value = newAccount.auto_pause_on_expired === true
// Load mixed scheduling setting (only for antigravity accounts) // Load mixed scheduling setting (only for antigravity accounts)
mixedScheduling.value = false
allowOverages.value = false
const extra = newAccount.extra as Record<string, unknown> | undefined const extra = newAccount.extra as Record<string, unknown> | undefined
mixedScheduling.value = extra?.mixed_scheduling === true mixedScheduling.value = extra?.mixed_scheduling === true
allowOverages.value = extra?.allow_overages === true
// Load OpenAI passthrough toggle (OpenAI OAuth/API Key) // Load OpenAI passthrough toggle (OpenAI OAuth/API Key)
openaiPassthroughEnabled.value = false openaiPassthroughEnabled.value = false
...@@ -2822,7 +2853,7 @@ const handleSubmit = async () => { ...@@ -2822,7 +2853,7 @@ const handleSubmit = async () => {
updatePayload.credentials = newCredentials updatePayload.credentials = newCredentials
} }
// For antigravity accounts, handle mixed_scheduling in extra // For antigravity accounts, handle mixed_scheduling and allow_overages in extra
if (props.account.platform === 'antigravity') { if (props.account.platform === 'antigravity') {
const currentExtra = (props.account.extra as Record<string, unknown>) || {} const currentExtra = (props.account.extra as Record<string, unknown>) || {}
const newExtra: Record<string, unknown> = { ...currentExtra } const newExtra: Record<string, unknown> = { ...currentExtra }
...@@ -2831,6 +2862,11 @@ const handleSubmit = async () => { ...@@ -2831,6 +2862,11 @@ const handleSubmit = async () => {
} else { } else {
delete newExtra.mixed_scheduling delete newExtra.mixed_scheduling
} }
if (allowOverages.value) {
newExtra.allow_overages = true
} else {
delete newExtra.allow_overages
}
updatePayload.extra = newExtra updatePayload.extra = newExtra
} }
......
import { describe, expect, it, vi } from 'vitest'
import { mount } from '@vue/test-utils'
import AccountStatusIndicator from '../AccountStatusIndicator.vue'
import type { Account } from '@/types'
vi.mock('vue-i18n', async () => {
const actual = await vi.importActual<typeof import('vue-i18n')>('vue-i18n')
return {
...actual,
useI18n: () => ({
t: (key: string) => key
})
}
})
function makeAccount(overrides: Partial<Account>): Account {
return {
id: 1,
name: 'account',
platform: 'antigravity',
type: 'oauth',
proxy_id: null,
concurrency: 1,
priority: 1,
status: 'active',
error_message: null,
last_used_at: null,
expires_at: null,
auto_pause_on_expired: true,
created_at: '2026-03-15T00:00:00Z',
updated_at: '2026-03-15T00:00:00Z',
schedulable: true,
rate_limited_at: null,
rate_limit_reset_at: null,
overload_until: null,
temp_unschedulable_until: null,
temp_unschedulable_reason: null,
session_window_start: null,
session_window_end: null,
session_window_status: null,
...overrides,
}
}
describe('AccountStatusIndicator', () => {
it('模型限流 + overages 启用 + 无 AICredits key → 显示 ⚡ (credits_active)', () => {
const wrapper = mount(AccountStatusIndicator, {
props: {
account: makeAccount({
id: 1,
name: 'ag-1',
extra: {
allow_overages: true,
model_rate_limits: {
'claude-sonnet-4-5': {
rate_limited_at: '2026-03-15T00:00:00Z',
rate_limit_reset_at: '2099-03-15T00:00:00Z'
}
}
}
})
},
global: {
stubs: {
Icon: true
}
}
})
expect(wrapper.text()).toContain('')
expect(wrapper.text()).toContain('CSon45')
})
it('模型限流 + overages 未启用 → 普通限流样式(无 ⚡)', () => {
const wrapper = mount(AccountStatusIndicator, {
props: {
account: makeAccount({
id: 2,
name: 'ag-2',
extra: {
model_rate_limits: {
'claude-sonnet-4-5': {
rate_limited_at: '2026-03-15T00:00:00Z',
rate_limit_reset_at: '2099-03-15T00:00:00Z'
}
}
}
})
},
global: {
stubs: {
Icon: true
}
}
})
expect(wrapper.text()).toContain('CSon45')
expect(wrapper.text()).not.toContain('')
})
it('AICredits key 生效 → 显示积分已用尽 (credits_exhausted)', () => {
const wrapper = mount(AccountStatusIndicator, {
props: {
account: makeAccount({
id: 3,
name: 'ag-3',
extra: {
allow_overages: true,
model_rate_limits: {
'AICredits': {
rate_limited_at: '2026-03-15T00:00:00Z',
rate_limit_reset_at: '2099-03-15T00:00:00Z'
}
}
}
})
},
global: {
stubs: {
Icon: true
}
}
})
expect(wrapper.text()).toContain('account.creditsExhausted')
})
it('模型限流 + overages 启用 + AICredits key 生效 → 普通限流样式(积分耗尽,无 ⚡)', () => {
const wrapper = mount(AccountStatusIndicator, {
props: {
account: makeAccount({
id: 4,
name: 'ag-4',
extra: {
allow_overages: true,
model_rate_limits: {
'claude-sonnet-4-5': {
rate_limited_at: '2026-03-15T00:00:00Z',
rate_limit_reset_at: '2099-03-15T00:00:00Z'
},
'AICredits': {
rate_limited_at: '2026-03-15T00:00:00Z',
rate_limit_reset_at: '2099-03-15T00:00:00Z'
}
}
}
})
},
global: {
stubs: {
Icon: true
}
}
})
// 模型限流 + 积分耗尽 → 不应显示 ⚡
expect(wrapper.text()).toContain('CSon45')
expect(wrapper.text()).not.toContain('')
// AICredits 积分耗尽状态应显示
expect(wrapper.text()).toContain('account.creditsExhausted')
})
})
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment