feat: optimize Antigravity AI Credits scheduling - use AICredits rate limit key

Replace sync.Map + per-model runtime state with a single AICredits key in
model_rate_limits. Rate-limited accounts with available credits are now
scheduled directly and credits injected without a 429 round trip.

- Scheduler: rate-limited + overages + credits available = allow scheduling
- Forwarding: inject credits proactively when model rate limited
- Storage: credits exhaustion stored as model_rate_limits["AICredits"]
- Frontend: show credits_active (yellow ⚡) vs credits_exhausted states
- Bump version to 0.1.100.4
This commit is contained in:
erio
2026-03-16 03:51:34 +08:00
parent 0158cb8aa8
commit 897c93528d
11 changed files with 239 additions and 246 deletions
+1 -1
View File
@@ -1 +1 @@
0.1.100.3
0.1.100.4
+6 -10
View File
@@ -1558,11 +1558,15 @@ 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, antigravityCreditsOveragesKey)
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, antigravityCreditsOveragesKey)
delete(account.Extra, "antigravity_credits_overages") // 清理旧版 overages 运行态
}
// 校验并预计算固定时间重置的下次重置时间
if err := ValidateQuotaResetConfig(account.Extra); err != nil {
@@ -1646,14 +1650,6 @@ func (s *adminServiceImpl) UpdateAccount(ctx context.Context, id int64, input *U
if err := s.accountRepo.Update(ctx, account); err != nil {
return nil, err
}
if account.Platform == PlatformAntigravity {
if !account.IsOveragesEnabled() && wasOveragesEnabled {
clearCreditsExhausted(account.ID)
}
if account.IsOveragesEnabled() && !wasOveragesEnabled {
clearCreditsExhausted(account.ID)
}
}
// 绑定分组
if input.GroupIDs != nil {
@@ -26,7 +26,7 @@ func (r *updateAccountOveragesRepoStub) Update(ctx context.Context, account *Acc
return nil
}
func TestUpdateAccount_DisableOveragesClearsRuntimeStateBeforePersist(t *testing.T) {
func TestUpdateAccount_DisableOveragesClearsAICreditsKey(t *testing.T) {
accountID := int64(101)
repo := &updateAccountOveragesRepoStub{
account: &Account{
@@ -37,24 +37,34 @@ func TestUpdateAccount_DisableOveragesClearsRuntimeStateBeforePersist(t *testing
Extra: map[string]any{
"allow_overages": true,
"mixed_scheduling": true,
antigravityCreditsOveragesKey: map[string]any{
modelRateLimitsKey: map[string]any{
"claude-sonnet-4-5": map[string]any{
"activated_at": "2026-03-15T00:00:00Z",
"active_until": "2099-03-15T00:00:00Z",
"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),
},
},
},
},
}
setCreditsExhausted(accountID, time.Now().Add(time.Minute))
t.Cleanup(func() {
clearCreditsExhausted(accountID)
})
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),
},
},
},
})
@@ -62,10 +72,17 @@ func TestUpdateAccount_DisableOveragesClearsRuntimeStateBeforePersist(t *testing
require.NotNil(t, updated)
require.Equal(t, 1, repo.updateCalls)
require.False(t, updated.IsOveragesEnabled())
require.False(t, isCreditsExhausted(accountID))
_, exists := repo.account.Extra[antigravityCreditsOveragesKey]
require.False(t, exists, "关闭 overages 时应在持久化前移除运行态")
// 关闭 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) {
@@ -87,10 +104,6 @@ func TestUpdateAccount_EnableOveragesClearsModelRateLimitsBeforePersist(t *testi
},
},
}
setCreditsExhausted(accountID, time.Now().Add(time.Minute))
t.Cleanup(func() {
clearCreditsExhausted(accountID)
})
svc := &adminServiceImpl{accountRepo: repo}
updated, err := svc.UpdateAccount(context.Background(), accountID, &UpdateAccountInput{
@@ -104,7 +117,6 @@ func TestUpdateAccount_EnableOveragesClearsModelRateLimitsBeforePersist(t *testi
require.NotNil(t, updated)
require.Equal(t, 1, repo.updateCalls)
require.True(t, updated.IsOveragesEnabled())
require.False(t, isCreditsExhausted(accountID))
_, exists := repo.account.Extra[modelRateLimitsKey]
require.False(t, exists, "开启 overages 时应在持久化前清掉旧模型限流")
@@ -6,14 +6,18 @@ import (
"io"
"net/http"
"strings"
"sync"
"time"
"github.com/Wei-Shaw/sub2api/internal/pkg/antigravity"
"github.com/Wei-Shaw/sub2api/internal/pkg/logger"
)
const antigravityCreditsOveragesKey = "antigravity_credits_overages"
const (
// creditsExhaustedKey 是 model_rate_limits 中标记积分耗尽的特殊 key。
// 与普通模型限流完全同构:通过 SetModelRateLimit / isRateLimitActiveForKey 读写。
creditsExhaustedKey = "AICredits"
creditsExhaustedDuration = 5 * time.Hour
)
type antigravity429Category string
@@ -24,8 +28,6 @@ const (
)
var (
creditsExhaustedCache sync.Map
antigravityQuotaExhaustedKeywords = []string{
"quota_exhausted",
"quota exhausted",
@@ -46,28 +48,48 @@ var (
}
)
// isCreditsExhausted 检查账号的 AI Credits 是否已被标记为耗尽。
func isCreditsExhausted(accountID int64) bool {
v, ok := creditsExhaustedCache.Load(accountID)
// 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 false
return
}
until, ok := v.(time.Time)
if !ok || time.Now().After(until) {
creditsExhaustedCache.Delete(accountID)
return false
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)
}
return true
}
// setCreditsExhausted 将账号标记为 AI Credits 已耗尽,直到指定时间。
func setCreditsExhausted(accountID int64, until time.Time) {
creditsExhaustedCache.Store(accountID, until)
}
// clearCreditsExhausted 清除账号的 AI Credits 耗尽标记。
func clearCreditsExhausted(accountID int64) {
creditsExhaustedCache.Delete(accountID)
}
// classifyAntigravity429 将 Antigravity 的 429 响应归类为配额耗尽、限流或未知。
@@ -117,105 +139,6 @@ func resolveCreditsOveragesModelKey(ctx context.Context, account *Account, upstr
return resolveAntigravityModelKey(requestedModel)
}
// canUseAntigravityCreditsOverages 判断当前请求是否应直接走已激活的 overages 链路。
func canUseAntigravityCreditsOverages(ctx context.Context, account *Account, requestedModel string) bool {
if account == nil || account.Platform != PlatformAntigravity {
return false
}
if !account.IsSchedulable() {
return false
}
if !account.IsOveragesEnabled() || isCreditsExhausted(account.ID) {
return false
}
return account.getAntigravityCreditsOveragesRemainingTimeWithContext(ctx, requestedModel) > 0
}
func (a *Account) getAntigravityCreditsOveragesRemainingTimeWithContext(ctx context.Context, requestedModel string) time.Duration {
if a == nil || a.Extra == nil {
return 0
}
modelKey := resolveFinalAntigravityModelKey(ctx, a, requestedModel)
modelKey = strings.TrimSpace(modelKey)
if modelKey == "" {
return 0
}
rawStates, ok := a.Extra[antigravityCreditsOveragesKey].(map[string]any)
if !ok {
return 0
}
rawState, ok := rawStates[modelKey].(map[string]any)
if !ok {
return 0
}
activeUntilRaw, ok := rawState["active_until"].(string)
if !ok || strings.TrimSpace(activeUntilRaw) == "" {
return 0
}
activeUntil, err := time.Parse(time.RFC3339, activeUntilRaw)
if err != nil {
return 0
}
remaining := time.Until(activeUntil)
if remaining > 0 {
return remaining
}
return 0
}
func setAntigravityCreditsOveragesActive(ctx context.Context, repo AccountRepository, account *Account, modelKey string, activeUntil time.Time) {
if repo == nil || account == nil || account.ID == 0 || strings.TrimSpace(modelKey) == "" {
return
}
stateMap := copyAntigravityCreditsOveragesState(account)
stateMap[modelKey] = map[string]any{
"activated_at": time.Now().UTC().Format(time.RFC3339),
"active_until": activeUntil.UTC().Format(time.RFC3339),
}
ensureAccountExtra(account)
account.Extra[antigravityCreditsOveragesKey] = stateMap
if err := repo.UpdateExtra(ctx, account.ID, map[string]any{antigravityCreditsOveragesKey: stateMap}); err != nil {
logger.LegacyPrintf("service.antigravity_gateway", "set overages state failed: account=%d model=%s err=%v", account.ID, modelKey, err)
}
}
func clearAntigravityCreditsOveragesState(ctx context.Context, repo AccountRepository, accountID int64) error {
if repo == nil || accountID == 0 {
return nil
}
return repo.UpdateExtra(ctx, accountID, map[string]any{
antigravityCreditsOveragesKey: map[string]any{},
})
}
func clearAntigravityCreditsOveragesStateForModel(ctx context.Context, repo AccountRepository, account *Account, modelKey string) {
if repo == nil || account == nil || account.ID == 0 {
return
}
stateMap := copyAntigravityCreditsOveragesState(account)
delete(stateMap, modelKey)
ensureAccountExtra(account)
account.Extra[antigravityCreditsOveragesKey] = stateMap
if err := repo.UpdateExtra(ctx, account.ID, map[string]any{antigravityCreditsOveragesKey: stateMap}); err != nil {
logger.LegacyPrintf("service.antigravity_gateway", "clear overages state failed: account=%d model=%s err=%v", account.ID, modelKey, err)
}
}
func copyAntigravityCreditsOveragesState(account *Account) map[string]any {
result := make(map[string]any)
if account == nil || account.Extra == nil {
return result
}
rawState, ok := account.Extra[antigravityCreditsOveragesKey].(map[string]any)
if !ok {
return result
}
for key, value := range rawState {
result[key] = value
}
return result
}
func ensureAccountExtra(account *Account) {
if account != nil && account.Extra == nil {
account.Extra = make(map[string]any)
@@ -276,23 +199,21 @@ func (s *AntigravityGatewayService) attemptCreditsOveragesRetry(
creditsResp, err := p.httpUpstream.Do(creditsReq, p.proxyURL, p.account.ID, p.account.Concurrency)
if err == nil && creditsResp != nil && creditsResp.StatusCode < 400 {
clearCreditsExhausted(p.account.ID)
activeUntil := s.resolveCreditsOveragesActiveUntil(respBody, waitDuration)
setAntigravityCreditsOveragesActive(p.ctx, p.accountRepo, p.account, modelKey, activeUntil)
logger.LegacyPrintf("service.antigravity_gateway", "%s status=%d credit_overages_success model=%s account=%d active_until=%s",
p.prefix, creditsResp.StatusCode, modelKey, p.account.ID, activeUntil.UTC().Format(time.RFC3339))
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.prefix, modelKey, p.account, waitDuration, creditsResp, err)
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,
waitDuration time.Duration,
creditsResp *http.Response,
reqErr error,
) {
@@ -307,11 +228,9 @@ func (s *AntigravityGatewayService) handleCreditsRetryFailure(
}
if shouldMarkCreditsExhausted(creditsResp, creditsRespBody, reqErr) && account != nil {
exhaustedUntil := s.resolveCreditsOveragesActiveUntil(creditsRespBody, waitDuration)
setCreditsExhausted(account.ID, exhaustedUntil)
clearAntigravityCreditsOveragesStateForModel(context.Background(), s.accountRepo, account, modelKey)
logger.LegacyPrintf("service.antigravity_gateway", "%s credit_overages_failed model=%s account=%d marked_exhausted=true status=%d exhausted_until=%v body=%s",
prefix, modelKey, account.ID, creditsStatusCode, exhaustedUntil, truncateForLog(creditsRespBody, 200))
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 {
@@ -40,58 +40,50 @@ func TestClassifyAntigravity429(t *testing.T) {
})
}
func TestCanUseAntigravityCreditsOverages(t *testing.T) {
activeUntil := time.Now().Add(10 * time.Minute).UTC().Format(time.RFC3339)
t.Run("必须有运行态才可直接走 overages", func(t *testing.T) {
func TestIsCreditsExhausted_UsesAICreditsKey(t *testing.T) {
t.Run("无 AICredits key 则积分可用", func(t *testing.T) {
account := &Account{
ID: 1,
Platform: PlatformAntigravity,
Status: StatusActive,
Schedulable: true,
ID: 1,
Platform: PlatformAntigravity,
Extra: map[string]any{
"allow_overages": true,
},
}
require.False(t, canUseAntigravityCreditsOverages(context.Background(), account, "claude-sonnet-4-5"))
require.False(t, account.isCreditsExhausted())
})
t.Run("运行态有效时允许使用 overages", func(t *testing.T) {
t.Run("AICredits key 生效则积分耗尽", func(t *testing.T) {
account := &Account{
ID: 2,
Platform: PlatformAntigravity,
Status: StatusActive,
Schedulable: true,
ID: 2,
Platform: PlatformAntigravity,
Extra: map[string]any{
"allow_overages": true,
antigravityCreditsOveragesKey: map[string]any{
"claude-sonnet-4-5": map[string]any{
"active_until": activeUntil,
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, canUseAntigravityCreditsOverages(context.Background(), account, "claude-sonnet-4-5"))
require.True(t, account.isCreditsExhausted())
})
t.Run("credits 耗尽后不可继续使用 overages", func(t *testing.T) {
t.Run("AICredits key 过期则积分可用", func(t *testing.T) {
account := &Account{
ID: 3,
Platform: PlatformAntigravity,
Status: StatusActive,
Schedulable: true,
ID: 3,
Platform: PlatformAntigravity,
Extra: map[string]any{
"allow_overages": true,
antigravityCreditsOveragesKey: map[string]any{
"claude-sonnet-4-5": map[string]any{
"active_until": activeUntil,
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),
},
},
},
}
setCreditsExhausted(account.ID, time.Now().Add(time.Minute))
t.Cleanup(func() { clearCreditsExhausted(account.ID) })
require.False(t, canUseAntigravityCreditsOverages(context.Background(), account, "claude-sonnet-4-5"))
require.False(t, account.isCreditsExhausted())
})
}
@@ -152,12 +144,6 @@ func TestHandleSmartRetry_QuotaExhausted_UsesCreditsAndStoresIndependentState(t
require.Len(t, upstream.requestBodies, 1)
require.Contains(t, string(upstream.requestBodies[0]), "enabledCreditTypes")
require.Empty(t, repo.modelRateLimitCalls, "overages 成功后不应写入普通 model_rate_limits")
require.Len(t, repo.extraUpdateCalls, 1)
state, ok := account.Extra[antigravityCreditsOveragesKey].(map[string]any)
require.True(t, ok)
_, exists := state["claude-sonnet-4-5"]
require.True(t, exists, "应使用最终映射模型写入独立 overages 运行态")
}
func TestHandleSmartRetry_RateLimited_DoesNotUseCredits(t *testing.T) {
@@ -221,7 +207,7 @@ func TestHandleSmartRetry_RateLimited_DoesNotUseCredits(t *testing.T) {
require.Empty(t, repo.modelRateLimitCalls)
}
func TestAntigravityRetryLoop_ActiveOverages_InjectsCreditsBody(t *testing.T) {
func TestAntigravityRetryLoop_ModelRateLimited_InjectsCredits(t *testing.T) {
oldBaseURLs := append([]string(nil), antigravity.BaseURLs...)
oldAvailability := antigravity.DefaultURLAvailability
defer func() {
@@ -232,7 +218,6 @@ func TestAntigravityRetryLoop_ActiveOverages_InjectsCreditsBody(t *testing.T) {
antigravity.BaseURLs = []string{"https://ag-1.test"}
antigravity.DefaultURLAvailability = antigravity.NewURLAvailability(time.Minute)
activeUntil := time.Now().Add(10 * time.Minute).UTC().Format(time.RFC3339)
upstream := &queuedHTTPUpstreamStub{
responses: []*http.Response{
{
@@ -243,6 +228,7 @@ func TestAntigravityRetryLoop_ActiveOverages_InjectsCreditsBody(t *testing.T) {
},
errors: []error{nil},
}
// 模型已限流 + overages 启用 + 无 AICredits key → 应直接注入积分
account := &Account{
ID: 103,
Name: "acc-103",
@@ -252,9 +238,10 @@ func TestAntigravityRetryLoop_ActiveOverages_InjectsCreditsBody(t *testing.T) {
Schedulable: true,
Extra: map[string]any{
"allow_overages": true,
antigravityCreditsOveragesKey: map[string]any{
modelRateLimitsKey: map[string]any{
"claude-sonnet-4-5": map[string]any{
"active_until": activeUntil,
"rate_limited_at": time.Now().UTC().Format(time.RFC3339),
"rate_limit_reset_at": time.Now().Add(30 * time.Minute).UTC().Format(time.RFC3339),
},
},
},
@@ -281,7 +268,61 @@ func TestAntigravityRetryLoop_ActiveOverages_InjectsCreditsBody(t *testing.T) {
require.Contains(t, string(upstream.requestBodies[0]), "enabledCreditTypes")
}
func TestAntigravityRetryLoop_ActiveOverages_ExplicitCreditErrorMarksExhausted(t *testing.T) {
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() {
@@ -292,8 +333,6 @@ func TestAntigravityRetryLoop_ActiveOverages_ExplicitCreditErrorMarksExhausted(t
antigravity.BaseURLs = []string{"https://ag-1.test"}
antigravity.DefaultURLAvailability = antigravity.NewURLAvailability(time.Minute)
accountID := int64(104)
activeUntil := time.Now().Add(10 * time.Minute).UTC().Format(time.RFC3339)
repo := &stubAntigravityAccountRepo{}
upstream := &queuedHTTPUpstreamStub{
responses: []*http.Response{
@@ -305,24 +344,24 @@ func TestAntigravityRetryLoop_ActiveOverages_ExplicitCreditErrorMarksExhausted(t
},
errors: []error{nil},
}
// 模型限流 + overages 启用 + 积分可用 → 注入积分但上游返回积分不足
account := &Account{
ID: accountID,
Name: "acc-104",
ID: 105,
Name: "acc-105",
Type: AccountTypeOAuth,
Platform: PlatformAntigravity,
Status: StatusActive,
Schedulable: true,
Extra: map[string]any{
"allow_overages": true,
antigravityCreditsOveragesKey: map[string]any{
modelRateLimitsKey: map[string]any{
"claude-sonnet-4-5": map[string]any{
"active_until": activeUntil,
"rate_limited_at": time.Now().UTC().Format(time.RFC3339),
"rate_limit_reset_at": time.Now().Add(30 * time.Minute).UTC().Format(time.RFC3339),
},
},
},
}
clearCreditsExhausted(accountID)
t.Cleanup(func() { clearCreditsExhausted(accountID) })
svc := &AntigravityGatewayService{accountRepo: repo}
result, err := svc.antigravityRetryLoop(antigravityRetryLoopParams{
@@ -333,6 +372,7 @@ func TestAntigravityRetryLoop_ActiveOverages_ExplicitCreditErrorMarksExhausted(t
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
@@ -341,6 +381,7 @@ func TestAntigravityRetryLoop_ActiveOverages_ExplicitCreditErrorMarksExhausted(t
require.NoError(t, err)
require.NotNil(t, result)
require.True(t, isCreditsExhausted(accountID))
require.Len(t, repo.extraUpdateCalls, 1, "应清理对应模型的 overages 运行态")
// 验证 AICredits key 已通过 SetModelRateLimit 写入数据库
require.Len(t, repo.modelRateLimitCalls, 1, "应通过 SetModelRateLimit 写入 AICredits key")
require.Equal(t, creditsExhaustedKey, repo.modelRateLimitCalls[0].modelKey)
}
@@ -201,7 +201,7 @@ func (s *AntigravityGatewayService) handleSmartRetry(p antigravityRetryLoopParam
if resp.StatusCode == http.StatusTooManyRequests &&
category == antigravity429QuotaExhausted &&
p.account.IsOveragesEnabled() &&
!isCreditsExhausted(p.account.ID) {
!p.account.isCreditsExhausted() {
result := s.attemptCreditsOveragesRetry(p, baseURL, modelName, waitDuration, resp.StatusCode, respBody)
if result.handled && result.resp != nil {
return &smartRetryResult{
@@ -552,13 +552,15 @@ func (s *AntigravityGatewayService) handleSingleAccountRetryInPlace(
// antigravityRetryLoop 执行带 URL fallback 的重试循环
func (s *AntigravityGatewayService) antigravityRetryLoop(p antigravityRetryLoopParams) (*antigravityRetryLoopResult, error) {
// 预检查:如果模型已进入 overages 运行态,则直接注入 AI Credits。
overagesActive := false
if p.requestedModel != "" && canUseAntigravityCreditsOverages(p.ctx, p.account, p.requestedModel) {
// 预检查:模型限流 + 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
overagesActive = true
logger.LegacyPrintf("service.antigravity_gateway", "%s pre_check: credit_overages_active model=%s account=%d (injecting enabledCreditTypes)",
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)
}
}
@@ -566,9 +568,9 @@ func (s *AntigravityGatewayService) antigravityRetryLoop(p antigravityRetryLoopP
// 预检查:如果账号已限流,直接返回切换信号
if p.requestedModel != "" {
if remaining := p.account.GetRateLimitRemainingTimeWithContext(p.ctx, p.requestedModel); remaining > 0 {
// 进入 overages 运行态的模型不再受普通模型限流预检查阻断。
if overagesActive {
logger.LegacyPrintf("service.antigravity_gateway", "%s pre_check: credit_overages_ignore_rate_limit remaining=%v model=%s account=%d",
// 已注入积分的请求不再受普通模型限流预检查阻断。
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 退避重试模式:跳过限流预检查,直接发请求。
@@ -666,9 +668,9 @@ urlFallbackLoop:
respBody, _ := io.ReadAll(io.LimitReader(resp.Body, 2<<20))
_ = resp.Body.Close()
if overagesActive && shouldMarkCreditsExhausted(resp, respBody, nil) {
if overagesInjected && shouldMarkCreditsExhausted(resp, respBody, nil) {
modelKey := resolveCreditsOveragesModelKey(p.ctx, p.account, "", p.requestedModel)
s.handleCreditsRetryFailure(p.prefix, modelKey, p.account, 0, &http.Response{
s.handleCreditsRetryFailure(p.ctx, p.prefix, modelKey, p.account, &http.Response{
StatusCode: resp.StatusCode,
Header: resp.Header.Clone(),
Body: io.NopCloser(bytes.NewReader(respBody)),
@@ -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
@@ -1100,9 +1100,6 @@ func (s *RateLimitService) ClearRateLimit(ctx context.Context, accountID int64)
if err := s.accountRepo.ClearModelRateLimits(ctx, accountID); err != nil {
return err
}
if err := clearAntigravityCreditsOveragesState(ctx, s.accountRepo, accountID); err != nil {
return err
}
// 清除限流时一并清理临时不可调度状态,避免周限/窗口重置后仍被本地临时状态阻断。
if err := s.accountRepo.ClearTempUnschedulable(ctx, accountID); err != nil {
return err
@@ -1112,7 +1109,6 @@ func (s *RateLimitService) ClearRateLimit(ctx context.Context, accountID int64)
slog.Warn("temp_unsched_cache_delete_failed", "account_id", accountID, "error", err)
}
}
clearCreditsExhausted(accountID)
return nil
}
@@ -1179,8 +1175,7 @@ func hasRecoverableRuntimeState(account *Account) bool {
return false
}
return hasNonEmptyMapValue(account.Extra, "model_rate_limits") ||
hasNonEmptyMapValue(account.Extra, "antigravity_quota_scopes") ||
hasNonEmptyMapValue(account.Extra, antigravityCreditsOveragesKey)
hasNonEmptyMapValue(account.Extra, "antigravity_quota_scopes")
}
func hasNonEmptyMapValue(extra map[string]any, key string) bool {
@@ -88,14 +88,25 @@
]"
>
<div v-for="item in activeModelStatuses" :key="`${item.kind}-${item.model}`" class="group relative mb-1 break-inside-avoid">
<!-- 积分已用尽 -->
<span
v-if="item.kind === 'overages'"
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"
@@ -109,9 +120,11 @@
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"
>
{{
item.kind === 'overages'
? 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) })
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"
@@ -165,12 +178,12 @@ const isRateLimited = computed(() => {
})
type AccountModelStatusItem = {
kind: 'rate_limit' | 'overages'
kind: 'rate_limit' | 'credits_exhausted' | 'credits_active'
model: string
reset_at: string
}
// Computed: active model statuses (普通模型限流 + 超量请求运行态)
// 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
@@ -179,19 +192,26 @@ const activeModelStatuses = computed<AccountModelStatusItem[]>(() => {
const now = new Date()
const items: AccountModelStatusItem[] = []
if (modelLimits) {
items.push(...Object.entries(modelLimits)
.filter(([, info]) => new Date(info.rate_limit_reset_at) > now)
.map(([model, info]) => ({ kind: 'rate_limit' as const, model, reset_at: info.rate_limit_reset_at })))
}
if (!modelLimits) return items
const overagesStates = extra?.antigravity_credits_overages as
| Record<string, { activated_at?: string; active_until: string }>
| undefined
if (overagesStates) {
items.push(...Object.entries(overagesStates)
.filter(([, info]) => new Date(info.active_until) > now)
.map(([model, info]) => ({ kind: 'overages' as const, model, reset_at: info.active_until })))
// 检查 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
+2
View File
@@ -1877,6 +1877,8 @@ export default {
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'
},
+2
View File
@@ -2054,6 +2054,8 @@ export default {
rateLimitedAutoResume: '{time} 自动恢复',
modelRateLimitedUntil: '{model} 限流至 {time}',
modelCreditOveragesUntil: '{model} 正在使用 AI Credits,至 {time}',
creditsExhausted: '积分已用尽',
creditsExhaustedUntil: 'AI Credits 已用尽,预计 {time} 恢复',
overloadedUntil: '负载过重,重置时间:{time}',
viewTempUnschedDetails: '查看临时不可调度详情'
},