mirror of
https://github.com/Wei-Shaw/sub2api.git
synced 2026-09-24 16:05:44 +08:00
Merge remote-tracking branch 'origin/pr/antigravity-credits-overages' into release/custom-0.1.100
# Conflicts: # frontend/src/components/account/AccountStatusIndicator.vue
This commit is contained in:
@@ -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
|
||||
|
||||
@@ -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 != "" {
|
||||
|
||||
@@ -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。
|
||||
|
||||
@@ -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"`
|
||||
|
||||
|
||||
@@ -1535,6 +1535,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
|
||||
@@ -1556,6 +1557,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
|
||||
|
||||
@@ -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 时应在持久化前清掉旧模型限流")
|
||||
}
|
||||
@@ -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))
|
||||
}
|
||||
}
|
||||
@@ -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, "普通模型限流应保留")
|
||||
})
|
||||
}
|
||||
@@ -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 {
|
||||
|
||||
@@ -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
|
||||
}
|
||||
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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, "")
|
||||
|
||||
|
||||
@@ -51,15 +51,23 @@ func (c *stubSmartRetryCache) ClearAccountAffinity(_ context.Context, _ int64, _
|
||||
|
||||
// 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]
|
||||
|
||||
@@ -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 {
|
||||
|
||||
@@ -76,20 +76,40 @@
|
||||
</div>
|
||||
</div>
|
||||
|
||||
<!-- Model Rate Limit Indicators (Antigravity OAuth Smart Retry) -->
|
||||
<!-- Model Status Indicators (普通限流 / 超量请求中) -->
|
||||
<div
|
||||
v-if="activeModelRateLimits.length > 0"
|
||||
v-if="activeModelStatuses.length > 0"
|
||||
:class="[
|
||||
activeModelRateLimits.length <= 4
|
||||
? 'flex flex-col gap-0.5'
|
||||
: activeModelRateLimits.length <= 8
|
||||
activeModelStatuses.length <= 4
|
||||
? 'flex flex-col gap-1'
|
||||
: activeModelStatuses.length <= 8
|
||||
? 'columns-2 gap-x-2'
|
||||
: 'columns-3 gap-x-2'
|
||||
]"
|
||||
>
|
||||
<div v-for="item in activeModelRateLimits" :key="item.model" class="group relative mb-0.5 break-inside-avoid">
|
||||
<div v-for="item in activeModelStatuses" :key="`${item.kind}-${item.model}`" class="group relative mb-1 break-inside-avoid">
|
||||
<!-- 积分已用尽 -->
|
||||
<span
|
||||
class="inline-flex items-center gap-0.5 rounded bg-purple-100 px-1.5 py-px text-xs font-medium text-purple-700 dark:bg-purple-900/30 dark:text-purple-400"
|
||||
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"
|
||||
>
|
||||
<Icon name="exclamationTriangle" size="xs" :stroke-width="2" />
|
||||
{{ formatScopeName(item.model) }}
|
||||
@@ -99,7 +119,13 @@
|
||||
<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"
|
||||
>
|
||||
{{ 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
|
||||
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>
|
||||
@@ -131,6 +157,7 @@
|
||||
<script setup lang="ts">
|
||||
import { computed } from 'vue'
|
||||
import { useI18n } from 'vue-i18n'
|
||||
import Icon from '@/components/icons/Icon.vue'
|
||||
import type { Account } from '@/types'
|
||||
import { formatCountdown, formatDateTime, formatCountdownWithSuffix, formatTime } from '@/utils/format'
|
||||
|
||||
@@ -150,17 +177,44 @@ const isRateLimited = computed(() => {
|
||||
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)
|
||||
const activeModelRateLimits = computed(() => {
|
||||
const modelLimits = (props.account.extra as Record<string, unknown> | undefined)?.model_rate_limits as
|
||||
// Computed: active model statuses (普通模型限流 + 积分耗尽 + 走积分中)
|
||||
const activeModelStatuses = computed<AccountModelStatusItem[]>(() => {
|
||||
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 }>
|
||||
| undefined
|
||||
if (!modelLimits) return []
|
||||
const now = new Date()
|
||||
return Object.entries(modelLimits)
|
||||
.filter(([, info]) => new Date(info.rate_limit_reset_at) > now)
|
||||
.map(([model, info]) => ({ model, reset_at: info.rate_limit_reset_at }))
|
||||
const items: AccountModelStatusItem[] = []
|
||||
|
||||
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 => {
|
||||
|
||||
@@ -289,6 +289,13 @@
|
||||
:resets-at="antigravityClaudeUsageFromAPI.resetTime"
|
||||
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 v-else class="text-xs text-gray-400">-</div>
|
||||
</template>
|
||||
@@ -588,6 +595,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 响应中提取)
|
||||
const antigravityTier = computed(() => {
|
||||
const extra = props.account.extra as Record<string, unknown> | undefined
|
||||
|
||||
@@ -2502,6 +2502,33 @@
|
||||
</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 - 仅标准模式显示 -->
|
||||
<GroupSelector
|
||||
@@ -3045,6 +3072,7 @@ const openaiAPIKeyResponsesWebSocketV2Mode = ref<OpenAIWSMode>(OPENAI_WS_MODE_OF
|
||||
const codexCLIOnlyEnabled = ref(false)
|
||||
const anthropicPassthroughEnabled = ref(false)
|
||||
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 soraAccountType = ref<'oauth' | 'apikey'>('oauth') // For sora: oauth or apikey (upstream)
|
||||
const upstreamBaseUrl = ref('') // For upstream type: base URL
|
||||
@@ -3071,6 +3099,13 @@ const getTempUnschedRuleKey = createStableObjectKeyResolver<TempUnschedRuleForm>
|
||||
const geminiOAuthType = ref<'code_assist' | 'google_one' | 'ai_studio'>('google_one')
|
||||
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 mixedChannelWarningDetails = ref<{ groupName: string; currentPlatform: string; otherPlatform: string } | null>(
|
||||
null
|
||||
@@ -3341,6 +3376,7 @@ watch(
|
||||
accountCategory.value = 'oauth-based'
|
||||
antigravityAccountType.value = 'oauth'
|
||||
} else {
|
||||
allowOverages.value = false
|
||||
antigravityWhitelistModels.value = []
|
||||
antigravityModelMappings.value = []
|
||||
antigravityModelRestrictionMode.value = 'mapping'
|
||||
@@ -3774,6 +3810,7 @@ const resetForm = () => {
|
||||
sessionIdMaskingEnabled.value = false
|
||||
cacheTTLOverrideEnabled.value = false
|
||||
cacheTTLOverrideTarget.value = '5m'
|
||||
allowOverages.value = false
|
||||
antigravityAccountType.value = 'oauth'
|
||||
upstreamBaseUrl.value = ''
|
||||
upstreamApiKey.value = ''
|
||||
@@ -4044,7 +4081,7 @@ const handleSubmit = async () => {
|
||||
|
||||
applyInterceptWarmup(credentials, interceptWarmupRequests.value, 'create')
|
||||
|
||||
const extra = mixedScheduling.value ? { mixed_scheduling: true } : undefined
|
||||
const extra = buildAntigravityExtra()
|
||||
await createAccountAndFinish(form.platform, 'apikey', credentials, extra)
|
||||
return
|
||||
}
|
||||
@@ -4790,7 +4827,7 @@ const handleAntigravityExchange = async (authCode: string) => {
|
||||
if (antigravityModelMapping) {
|
||||
credentials.model_mapping = antigravityModelMapping
|
||||
}
|
||||
const extra = mixedScheduling.value ? { mixed_scheduling: true } : undefined
|
||||
const extra = buildAntigravityExtra()
|
||||
await createAccountAndFinish('antigravity', 'oauth', credentials, extra)
|
||||
} catch (error: any) {
|
||||
antigravityOAuth.error.value = error.response?.data?.detail || t('admin.accounts.oauth.authFailed')
|
||||
|
||||
@@ -1664,6 +1664,33 @@
|
||||
</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>
|
||||
|
||||
<!-- Group Selection - 仅标准模式显示 -->
|
||||
@@ -1833,6 +1860,7 @@ const customErrorCodeInput = ref<number | null>(null)
|
||||
const interceptWarmupRequests = ref(false)
|
||||
const autoPauseOnExpired = ref(false)
|
||||
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 antigravityWhitelistModels = ref<string[]>([])
|
||||
const antigravityModelMappings = ref<ModelMapping[]>([])
|
||||
@@ -2038,8 +2066,11 @@ watch(
|
||||
autoPauseOnExpired.value = newAccount.auto_pause_on_expired === true
|
||||
|
||||
// Load mixed scheduling setting (only for antigravity accounts)
|
||||
mixedScheduling.value = false
|
||||
allowOverages.value = false
|
||||
const extra = newAccount.extra as Record<string, unknown> | undefined
|
||||
mixedScheduling.value = extra?.mixed_scheduling === true
|
||||
allowOverages.value = extra?.allow_overages === true
|
||||
|
||||
// Load OpenAI passthrough toggle (OpenAI OAuth/API Key)
|
||||
openaiPassthroughEnabled.value = false
|
||||
@@ -2902,7 +2933,7 @@ const handleSubmit = async () => {
|
||||
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') {
|
||||
const currentExtra = (props.account.extra as Record<string, unknown>) || {}
|
||||
const newExtra: Record<string, unknown> = { ...currentExtra }
|
||||
@@ -2911,6 +2942,11 @@ const handleSubmit = async () => {
|
||||
} else {
|
||||
delete newExtra.mixed_scheduling
|
||||
}
|
||||
if (allowOverages.value) {
|
||||
newExtra.allow_overages = true
|
||||
} else {
|
||||
delete newExtra.allow_overages
|
||||
}
|
||||
updatePayload.extra = newExtra
|
||||
}
|
||||
|
||||
|
||||
@@ -0,0 +1,162 @@
|
||||
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')
|
||||
})
|
||||
})
|
||||
@@ -1,6 +1,7 @@
|
||||
import { describe, expect, it, vi, beforeEach } from 'vitest'
|
||||
import { flushPromises, mount } from '@vue/test-utils'
|
||||
import AccountUsageCell from '../AccountUsageCell.vue'
|
||||
import type { Account } from '@/types'
|
||||
|
||||
const { getUsage } = vi.hoisted(() => ({
|
||||
getUsage: vi.fn()
|
||||
@@ -28,6 +29,35 @@ vi.mock('vue-i18n', async () => {
|
||||
}
|
||||
})
|
||||
|
||||
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('AccountUsageCell', () => {
|
||||
beforeEach(() => {
|
||||
getUsage.mockReset()
|
||||
@@ -53,12 +83,12 @@ describe('AccountUsageCell', () => {
|
||||
|
||||
const wrapper = mount(AccountUsageCell, {
|
||||
props: {
|
||||
account: {
|
||||
account: makeAccount({
|
||||
id: 1001,
|
||||
platform: 'antigravity',
|
||||
type: 'oauth',
|
||||
extra: {}
|
||||
} as any
|
||||
})
|
||||
},
|
||||
global: {
|
||||
stubs: {
|
||||
@@ -76,6 +106,40 @@ describe('AccountUsageCell', () => {
|
||||
expect(wrapper.text()).toContain('admin.accounts.usageWindow.gemini3Image|70|2026-03-01T09:00:00Z')
|
||||
})
|
||||
|
||||
it('Antigravity 会显示 AI Credits 余额信息', async () => {
|
||||
getUsage.mockResolvedValue({
|
||||
ai_credits: [
|
||||
{
|
||||
credit_type: 'GOOGLE_ONE_AI',
|
||||
amount: 25,
|
||||
minimum_balance: 5
|
||||
}
|
||||
]
|
||||
})
|
||||
|
||||
const wrapper = mount(AccountUsageCell, {
|
||||
props: {
|
||||
account: makeAccount({
|
||||
id: 1002,
|
||||
platform: 'antigravity',
|
||||
type: 'oauth',
|
||||
extra: {}
|
||||
})
|
||||
},
|
||||
global: {
|
||||
stubs: {
|
||||
UsageProgressBar: true,
|
||||
AccountQuotaInfo: true
|
||||
}
|
||||
}
|
||||
})
|
||||
|
||||
await flushPromises()
|
||||
|
||||
expect(wrapper.text()).toContain('admin.accounts.aiCreditsBalance')
|
||||
expect(wrapper.text()).toContain('25')
|
||||
})
|
||||
|
||||
|
||||
it('OpenAI OAuth 快照已过期时首屏会重新请求 usage', async () => {
|
||||
getUsage.mockResolvedValue({
|
||||
@@ -107,7 +171,7 @@ describe('AccountUsageCell', () => {
|
||||
|
||||
const wrapper = mount(AccountUsageCell, {
|
||||
props: {
|
||||
account: {
|
||||
account: makeAccount({
|
||||
id: 2000,
|
||||
platform: 'openai',
|
||||
type: 'oauth',
|
||||
@@ -118,7 +182,7 @@ describe('AccountUsageCell', () => {
|
||||
codex_7d_used_percent: 34,
|
||||
codex_7d_reset_at: '2026-03-13T12:00:00Z'
|
||||
}
|
||||
} as any
|
||||
})
|
||||
},
|
||||
global: {
|
||||
stubs: {
|
||||
@@ -141,7 +205,7 @@ describe('AccountUsageCell', () => {
|
||||
it('OpenAI OAuth 有现成快照且未限额时不会首屏请求 usage', async () => {
|
||||
const wrapper = mount(AccountUsageCell, {
|
||||
props: {
|
||||
account: {
|
||||
account: makeAccount({
|
||||
id: 2001,
|
||||
platform: 'openai',
|
||||
type: 'oauth',
|
||||
@@ -152,7 +216,7 @@ describe('AccountUsageCell', () => {
|
||||
codex_7d_used_percent: 34,
|
||||
codex_7d_reset_at: '2099-03-13T12:00:00Z'
|
||||
}
|
||||
} as any
|
||||
})
|
||||
},
|
||||
global: {
|
||||
stubs: {
|
||||
@@ -200,15 +264,15 @@ describe('AccountUsageCell', () => {
|
||||
}
|
||||
})
|
||||
|
||||
const wrapper = mount(AccountUsageCell, {
|
||||
props: {
|
||||
account: {
|
||||
id: 2002,
|
||||
platform: 'openai',
|
||||
type: 'oauth',
|
||||
extra: {}
|
||||
} as any
|
||||
},
|
||||
const wrapper = mount(AccountUsageCell, {
|
||||
props: {
|
||||
account: makeAccount({
|
||||
id: 2002,
|
||||
platform: 'openai',
|
||||
type: 'oauth',
|
||||
extra: {}
|
||||
})
|
||||
},
|
||||
global: {
|
||||
stubs: {
|
||||
UsageProgressBar: {
|
||||
@@ -260,16 +324,16 @@ describe('AccountUsageCell', () => {
|
||||
seven_day: null
|
||||
})
|
||||
|
||||
const wrapper = mount(AccountUsageCell, {
|
||||
props: {
|
||||
account: {
|
||||
id: 2003,
|
||||
platform: 'openai',
|
||||
type: 'oauth',
|
||||
updated_at: '2026-03-07T10:00:00Z',
|
||||
extra: {}
|
||||
} as any
|
||||
},
|
||||
const wrapper = mount(AccountUsageCell, {
|
||||
props: {
|
||||
account: makeAccount({
|
||||
id: 2003,
|
||||
platform: 'openai',
|
||||
type: 'oauth',
|
||||
updated_at: '2026-03-07T10:00:00Z',
|
||||
extra: {}
|
||||
})
|
||||
},
|
||||
global: {
|
||||
stubs: {
|
||||
UsageProgressBar: {
|
||||
@@ -328,19 +392,19 @@ describe('AccountUsageCell', () => {
|
||||
}
|
||||
})
|
||||
|
||||
const wrapper = mount(AccountUsageCell, {
|
||||
props: {
|
||||
account: {
|
||||
id: 2004,
|
||||
platform: 'openai',
|
||||
type: 'oauth',
|
||||
rate_limit_reset_at: '2099-03-07T12:00:00Z',
|
||||
extra: {
|
||||
codex_5h_used_percent: 0,
|
||||
codex_7d_used_percent: 0
|
||||
}
|
||||
} as any
|
||||
},
|
||||
const wrapper = mount(AccountUsageCell, {
|
||||
props: {
|
||||
account: makeAccount({
|
||||
id: 2004,
|
||||
platform: 'openai',
|
||||
type: 'oauth',
|
||||
rate_limit_reset_at: '2099-03-07T12:00:00Z',
|
||||
extra: {
|
||||
codex_5h_used_percent: 0,
|
||||
codex_7d_used_percent: 0
|
||||
}
|
||||
})
|
||||
},
|
||||
global: {
|
||||
stubs: {
|
||||
UsageProgressBar: {
|
||||
|
||||
@@ -1876,6 +1876,9 @@ export default {
|
||||
rateLimitedUntil: 'Rate limited and removed from scheduling. Auto resumes at {time}',
|
||||
rateLimitedAutoResume: 'Auto resumes in {time}',
|
||||
modelRateLimitedUntil: '{model} rate limited until {time}',
|
||||
modelCreditOveragesUntil: '{model} using AI Credits until {time}',
|
||||
creditsExhausted: 'Credits Exhausted',
|
||||
creditsExhaustedUntil: 'AI Credits exhausted, expected recovery at {time}',
|
||||
overloadedUntil: 'Overloaded until {time}',
|
||||
viewTempUnschedDetails: 'View temp unschedulable details'
|
||||
},
|
||||
@@ -2266,6 +2269,10 @@ export default {
|
||||
mixedSchedulingHint: 'Enable to participate in Anthropic/Gemini group scheduling',
|
||||
mixedSchedulingTooltip:
|
||||
'!! WARNING !! Antigravity Claude and Anthropic Claude cannot be used in the same context. If you have both Anthropic and Antigravity accounts, enabling this option will cause frequent 400 errors. When enabled, please use the group feature to isolate Antigravity accounts from Anthropic accounts. Make sure you understand this before enabling!!',
|
||||
aiCreditsBalance: 'AI Credits',
|
||||
allowOverages: 'Allow Overages (AI Credits)',
|
||||
allowOveragesTooltip:
|
||||
'Only use AI Credits after free quota is explicitly exhausted. Ordinary concurrent 429 rate limits will not switch to overages.',
|
||||
creating: 'Creating...',
|
||||
updating: 'Updating...',
|
||||
accountCreated: 'Account created successfully',
|
||||
|
||||
@@ -2053,6 +2053,9 @@ export default {
|
||||
rateLimitedUntil: '限流中,当前不参与调度,预计 {time} 自动恢复',
|
||||
rateLimitedAutoResume: '{time} 自动恢复',
|
||||
modelRateLimitedUntil: '{model} 限流至 {time}',
|
||||
modelCreditOveragesUntil: '{model} 正在使用 AI Credits,至 {time}',
|
||||
creditsExhausted: '积分已用尽',
|
||||
creditsExhaustedUntil: 'AI Credits 已用尽,预计 {time} 恢复',
|
||||
overloadedUntil: '负载过重,重置时间:{time}',
|
||||
viewTempUnschedDetails: '查看临时不可调度详情'
|
||||
},
|
||||
@@ -2408,6 +2411,10 @@ export default {
|
||||
mixedSchedulingHint: '启用后可参与 Anthropic/Gemini 分组的调度',
|
||||
mixedSchedulingTooltip:
|
||||
'!!注意!! Antigravity Claude 和 Anthropic Claude 无法在同个上下文中使用,如果你同时有 Anthropic 账号和 Antigravity 账号,开启此选项会导致经常 400 报错。开启后,请用分组功能做好 Antigravity 账号和 Anthropic 账号的隔离。一定要弄明白再开启!!',
|
||||
aiCreditsBalance: 'AI Credits',
|
||||
allowOverages: '允许超量请求 (AI Credits)',
|
||||
allowOveragesTooltip:
|
||||
'仅在免费配额被明确判定为耗尽后才会使用 AI Credits。普通并发 429 限流不会切换到超量请求。',
|
||||
creating: '创建中...',
|
||||
updating: '更新中...',
|
||||
accountCreated: '账号创建成功',
|
||||
|
||||
@@ -790,6 +790,11 @@ export interface AccountUsageInfo {
|
||||
gemini_pro_minute?: UsageProgress | null
|
||||
gemini_flash_minute?: UsageProgress | null
|
||||
antigravity_quota?: Record<string, AntigravityModelQuota> | null
|
||||
ai_credits?: Array<{
|
||||
credit_type?: string
|
||||
amount?: number
|
||||
minimum_balance?: number
|
||||
}> | null
|
||||
// Antigravity 403 forbidden 状态
|
||||
is_forbidden?: boolean
|
||||
forbidden_reason?: string
|
||||
|
||||
Reference in New Issue
Block a user