refactor: replace sync.Map credits state with AICredits rate limit key

Replace process-memory sync.Map + per-model runtime state with a single
"AICredits" key in model_rate_limits, making credits exhaustion fully
isomorphic with model-level rate limiting.

Scheduler: rate-limited accounts with overages enabled + credits available
are now scheduled instead of excluded.

Forwarding: when model is rate-limited + credits available, inject credits
proactively without waiting for a 429 round trip.

Storage: credits exhaustion stored as model_rate_limits["AICredits"] with
5h duration, reusing SetModelRateLimit/isRateLimitActiveForKey.

Frontend: show credits_active (yellow ) when model rate-limited but
credits available, credits_exhausted (red) when AICredits key active.

Tests: add unit tests for shouldMarkCreditsExhausted, injectEnabledCreditTypes,
clearCreditsExhausted, and update existing overages tests.
This commit is contained in:
erio
2026-03-16 04:31:22 +08:00
parent e14c87597a
commit 8a260defc2
12 changed files with 692 additions and 327 deletions

View File

@@ -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 时应在持久化前清掉旧模型限流")