fix(account): prevent quota-exceeded API key/Bedrock accounts from being scheduled
Add quota exceeded check to IsSchedulable() and refactor shouldClearStickySession to delegate to IsSchedulable(), eliminating duplicated logic and fixing missed overload/rate-limit/expired checks. Frontend displays quota exceeded status independently via quota fields.
This commit is contained in:
@@ -121,6 +121,9 @@ func (a *Account) IsSchedulable() bool {
|
|||||||
if a.TempUnschedulableUntil != nil && now.Before(*a.TempUnschedulableUntil) {
|
if a.TempUnschedulableUntil != nil && now.Before(*a.TempUnschedulableUntil) {
|
||||||
return false
|
return false
|
||||||
}
|
}
|
||||||
|
if a.IsAPIKeyOrBedrock() && a.IsQuotaExceeded() {
|
||||||
|
return false
|
||||||
|
}
|
||||||
return true
|
return true
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|||||||
123
backend/internal/service/account_quota_schedulable_test.go
Normal file
123
backend/internal/service/account_quota_schedulable_test.go
Normal file
@@ -0,0 +1,123 @@
|
|||||||
|
//go:build unit
|
||||||
|
|
||||||
|
package service
|
||||||
|
|
||||||
|
import (
|
||||||
|
"testing"
|
||||||
|
"time"
|
||||||
|
|
||||||
|
"github.com/stretchr/testify/require"
|
||||||
|
)
|
||||||
|
|
||||||
|
func TestAccountIsSchedulable_QuotaExceeded(t *testing.T) {
|
||||||
|
now := time.Now()
|
||||||
|
|
||||||
|
tests := []struct {
|
||||||
|
name string
|
||||||
|
account *Account
|
||||||
|
want bool
|
||||||
|
}{
|
||||||
|
{
|
||||||
|
name: "apikey daily quota exceeded",
|
||||||
|
account: &Account{
|
||||||
|
Status: StatusActive,
|
||||||
|
Schedulable: true,
|
||||||
|
Type: AccountTypeAPIKey,
|
||||||
|
Extra: map[string]any{
|
||||||
|
"quota_daily_limit": 10.0,
|
||||||
|
"quota_daily_used": 10.0,
|
||||||
|
"quota_daily_start": now.Add(-1 * time.Hour).Format(time.RFC3339),
|
||||||
|
},
|
||||||
|
},
|
||||||
|
want: false,
|
||||||
|
},
|
||||||
|
{
|
||||||
|
name: "apikey weekly quota exceeded",
|
||||||
|
account: &Account{
|
||||||
|
Status: StatusActive,
|
||||||
|
Schedulable: true,
|
||||||
|
Type: AccountTypeAPIKey,
|
||||||
|
Extra: map[string]any{
|
||||||
|
"quota_weekly_limit": 50.0,
|
||||||
|
"quota_weekly_used": 50.0,
|
||||||
|
"quota_weekly_start": now.Add(-2 * 24 * time.Hour).Format(time.RFC3339),
|
||||||
|
},
|
||||||
|
},
|
||||||
|
want: false,
|
||||||
|
},
|
||||||
|
{
|
||||||
|
name: "apikey total quota exceeded",
|
||||||
|
account: &Account{
|
||||||
|
Status: StatusActive,
|
||||||
|
Schedulable: true,
|
||||||
|
Type: AccountTypeAPIKey,
|
||||||
|
Extra: map[string]any{
|
||||||
|
"quota_limit": 100.0,
|
||||||
|
"quota_used": 100.0,
|
||||||
|
},
|
||||||
|
},
|
||||||
|
want: false,
|
||||||
|
},
|
||||||
|
{
|
||||||
|
name: "apikey quota not exceeded",
|
||||||
|
account: &Account{
|
||||||
|
Status: StatusActive,
|
||||||
|
Schedulable: true,
|
||||||
|
Type: AccountTypeAPIKey,
|
||||||
|
Extra: map[string]any{
|
||||||
|
"quota_daily_limit": 10.0,
|
||||||
|
"quota_daily_used": 5.0,
|
||||||
|
"quota_daily_start": now.Add(-1 * time.Hour).Format(time.RFC3339),
|
||||||
|
},
|
||||||
|
},
|
||||||
|
want: true,
|
||||||
|
},
|
||||||
|
{
|
||||||
|
name: "apikey expired daily period restores schedulable",
|
||||||
|
account: &Account{
|
||||||
|
Status: StatusActive,
|
||||||
|
Schedulable: true,
|
||||||
|
Type: AccountTypeAPIKey,
|
||||||
|
Extra: map[string]any{
|
||||||
|
"quota_daily_limit": 10.0,
|
||||||
|
"quota_daily_used": 10.0,
|
||||||
|
"quota_daily_start": now.Add(-25 * time.Hour).Format(time.RFC3339),
|
||||||
|
},
|
||||||
|
},
|
||||||
|
want: true,
|
||||||
|
},
|
||||||
|
{
|
||||||
|
name: "oauth ignores quota exceeded",
|
||||||
|
account: &Account{
|
||||||
|
Status: StatusActive,
|
||||||
|
Schedulable: true,
|
||||||
|
Type: AccountTypeOAuth,
|
||||||
|
Extra: map[string]any{
|
||||||
|
"quota_daily_limit": 10.0,
|
||||||
|
"quota_daily_used": 10.0,
|
||||||
|
"quota_daily_start": now.Add(-1 * time.Hour).Format(time.RFC3339),
|
||||||
|
},
|
||||||
|
},
|
||||||
|
want: true,
|
||||||
|
},
|
||||||
|
{
|
||||||
|
name: "bedrock quota exceeded",
|
||||||
|
account: &Account{
|
||||||
|
Status: StatusActive,
|
||||||
|
Schedulable: true,
|
||||||
|
Type: AccountTypeBedrock,
|
||||||
|
Extra: map[string]any{
|
||||||
|
"quota_limit": 200.0,
|
||||||
|
"quota_used": 200.0,
|
||||||
|
},
|
||||||
|
},
|
||||||
|
want: false,
|
||||||
|
},
|
||||||
|
}
|
||||||
|
|
||||||
|
for _, tt := range tests {
|
||||||
|
t.Run(tt.name, func(t *testing.T) {
|
||||||
|
require.Equal(t, tt.want, tt.account.IsSchedulable())
|
||||||
|
})
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -435,26 +435,19 @@ func prefetchedStickyAccountIDFromContext(ctx context.Context, groupID *int64) i
|
|||||||
}
|
}
|
||||||
|
|
||||||
// shouldClearStickySession 检查账号是否处于不可调度状态,需要清理粘性会话绑定。
|
// shouldClearStickySession 检查账号是否处于不可调度状态,需要清理粘性会话绑定。
|
||||||
// 当账号状态为错误、禁用、不可调度、处于临时不可调度期间,
|
// 委托 IsSchedulable() 判断账号级可调度性(状态、配额、过载、限流等),
|
||||||
// 或请求的模型处于限流状态时,返回 true。
|
// 额外检查模型级限流。
|
||||||
// 这确保后续请求不会继续使用不可用的账号。
|
|
||||||
//
|
//
|
||||||
// shouldClearStickySession checks if an account is in an unschedulable state
|
// shouldClearStickySession checks if an account is in an unschedulable state
|
||||||
// and the sticky session binding should be cleared.
|
// and the sticky session binding should be cleared.
|
||||||
// Returns true when account status is error/disabled, schedulable is false,
|
// Delegates to IsSchedulable() for account-level checks, plus model-level rate limiting.
|
||||||
// within temporary unschedulable period, or the requested model is rate-limited.
|
|
||||||
// This ensures subsequent requests won't continue using unavailable accounts.
|
|
||||||
func shouldClearStickySession(account *Account, requestedModel string) bool {
|
func shouldClearStickySession(account *Account, requestedModel string) bool {
|
||||||
if account == nil {
|
if account == nil {
|
||||||
return false
|
return false
|
||||||
}
|
}
|
||||||
if account.Status == StatusError || account.Status == StatusDisabled || !account.Schedulable {
|
if !account.IsSchedulable() {
|
||||||
return true
|
return true
|
||||||
}
|
}
|
||||||
if account.TempUnschedulableUntil != nil && time.Now().Before(*account.TempUnschedulableUntil) {
|
|
||||||
return true
|
|
||||||
}
|
|
||||||
// 检查模型限流和 scope 限流,有限流即清除粘性会话
|
|
||||||
if remaining := account.GetRateLimitRemainingTimeWithContext(context.Background(), requestedModel); remaining > 0 {
|
if remaining := account.GetRateLimitRemainingTimeWithContext(context.Background(), requestedModel); remaining > 0 {
|
||||||
return true
|
return true
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -15,20 +15,8 @@ import (
|
|||||||
"github.com/stretchr/testify/require"
|
"github.com/stretchr/testify/require"
|
||||||
)
|
)
|
||||||
|
|
||||||
// TestShouldClearStickySession 测试粘性会话清理判断逻辑。
|
// TestShouldClearStickySession tests sticky session clearing via IsSchedulable() delegation
|
||||||
// 验证在以下情况下是否正确判断需要清理粘性会话:
|
// plus model-level rate limiting.
|
||||||
// - nil 账号:不清理(返回 false)
|
|
||||||
// - 状态为错误或禁用:清理
|
|
||||||
// - 不可调度:清理
|
|
||||||
// - 临时不可调度且未过期:清理
|
|
||||||
// - 临时不可调度已过期:不清理
|
|
||||||
// - 正常可调度状态:不清理
|
|
||||||
// - 模型限流(任意时长):清理
|
|
||||||
//
|
|
||||||
// TestShouldClearStickySession tests the sticky session clearing logic.
|
|
||||||
// Verifies correct behavior for various account states including:
|
|
||||||
// nil account, error/disabled status, unschedulable, temporary unschedulable,
|
|
||||||
// and model rate limiting scenarios.
|
|
||||||
func TestShouldClearStickySession(t *testing.T) {
|
func TestShouldClearStickySession(t *testing.T) {
|
||||||
now := time.Now()
|
now := time.Now()
|
||||||
future := now.Add(1 * time.Hour)
|
future := now.Add(1 * time.Hour)
|
||||||
@@ -101,6 +89,56 @@ func TestShouldClearStickySession(t *testing.T) {
|
|||||||
requestedModel: "claude-opus-4", // 请求不同模型
|
requestedModel: "claude-opus-4", // 请求不同模型
|
||||||
want: false, // 不同模型不受影响
|
want: false, // 不同模型不受影响
|
||||||
},
|
},
|
||||||
|
{
|
||||||
|
name: "apikey quota exceeded",
|
||||||
|
account: &Account{
|
||||||
|
Status: StatusActive,
|
||||||
|
Schedulable: true,
|
||||||
|
Type: AccountTypeAPIKey,
|
||||||
|
Extra: map[string]any{
|
||||||
|
"quota_daily_limit": 10.0,
|
||||||
|
"quota_daily_used": 10.0,
|
||||||
|
"quota_daily_start": now.Add(-1 * time.Hour).Format(time.RFC3339),
|
||||||
|
},
|
||||||
|
},
|
||||||
|
requestedModel: "",
|
||||||
|
want: true,
|
||||||
|
},
|
||||||
|
{
|
||||||
|
name: "oauth quota exceeded not cleared",
|
||||||
|
account: &Account{
|
||||||
|
Status: StatusActive,
|
||||||
|
Schedulable: true,
|
||||||
|
Type: AccountTypeOAuth,
|
||||||
|
Extra: map[string]any{
|
||||||
|
"quota_daily_limit": 10.0,
|
||||||
|
"quota_daily_used": 10.0,
|
||||||
|
"quota_daily_start": now.Add(-1 * time.Hour).Format(time.RFC3339),
|
||||||
|
},
|
||||||
|
},
|
||||||
|
requestedModel: "",
|
||||||
|
want: false,
|
||||||
|
},
|
||||||
|
{
|
||||||
|
name: "overloaded account",
|
||||||
|
account: &Account{
|
||||||
|
Status: StatusActive,
|
||||||
|
Schedulable: true,
|
||||||
|
OverloadUntil: &future,
|
||||||
|
},
|
||||||
|
requestedModel: "",
|
||||||
|
want: true,
|
||||||
|
},
|
||||||
|
{
|
||||||
|
name: "account-level rate limited",
|
||||||
|
account: &Account{
|
||||||
|
Status: StatusActive,
|
||||||
|
Schedulable: true,
|
||||||
|
RateLimitResetAt: &future,
|
||||||
|
},
|
||||||
|
requestedModel: "",
|
||||||
|
want: true,
|
||||||
|
},
|
||||||
}
|
}
|
||||||
|
|
||||||
for _, tt := range tests {
|
for _, tt := range tests {
|
||||||
|
|||||||
@@ -284,6 +284,16 @@ const hasError = computed(() => {
|
|||||||
return props.account.status === 'error'
|
return props.account.status === 'error'
|
||||||
})
|
})
|
||||||
|
|
||||||
|
const isQuotaExceeded = computed(() => {
|
||||||
|
const exceeded = (used?: number | null, limit?: number | null) =>
|
||||||
|
typeof limit === 'number' && limit > 0 && typeof used === 'number' && used >= limit
|
||||||
|
return (
|
||||||
|
exceeded(props.account.quota_used, props.account.quota_limit) ||
|
||||||
|
exceeded(props.account.quota_daily_used, props.account.quota_daily_limit) ||
|
||||||
|
exceeded(props.account.quota_weekly_used, props.account.quota_weekly_limit)
|
||||||
|
)
|
||||||
|
})
|
||||||
|
|
||||||
// Computed: countdown text for rate limit (429)
|
// Computed: countdown text for rate limit (429)
|
||||||
const rateLimitCountdown = computed(() => {
|
const rateLimitCountdown = computed(() => {
|
||||||
return formatCountdown(props.account.rate_limit_reset_at)
|
return formatCountdown(props.account.rate_limit_reset_at)
|
||||||
@@ -307,19 +317,16 @@ const statusClass = computed(() => {
|
|||||||
if (isTempUnschedulable.value) {
|
if (isTempUnschedulable.value) {
|
||||||
return 'badge-warning'
|
return 'badge-warning'
|
||||||
}
|
}
|
||||||
|
if (props.account.status !== 'active') {
|
||||||
|
return props.account.status === 'error' ? 'badge-danger' : 'badge-gray'
|
||||||
|
}
|
||||||
|
if (isQuotaExceeded.value) {
|
||||||
|
return 'badge-warning'
|
||||||
|
}
|
||||||
if (!props.account.schedulable) {
|
if (!props.account.schedulable) {
|
||||||
return 'badge-gray'
|
return 'badge-gray'
|
||||||
}
|
}
|
||||||
switch (props.account.status) {
|
return 'badge-success'
|
||||||
case 'active':
|
|
||||||
return 'badge-success'
|
|
||||||
case 'inactive':
|
|
||||||
return 'badge-gray'
|
|
||||||
case 'error':
|
|
||||||
return 'badge-danger'
|
|
||||||
default:
|
|
||||||
return 'badge-gray'
|
|
||||||
}
|
|
||||||
})
|
})
|
||||||
|
|
||||||
// Computed: status text
|
// Computed: status text
|
||||||
@@ -330,6 +337,12 @@ const statusText = computed(() => {
|
|||||||
if (isTempUnschedulable.value) {
|
if (isTempUnschedulable.value) {
|
||||||
return t('admin.accounts.status.tempUnschedulable')
|
return t('admin.accounts.status.tempUnschedulable')
|
||||||
}
|
}
|
||||||
|
if (props.account.status !== 'active') {
|
||||||
|
return t(`admin.accounts.status.${props.account.status}`)
|
||||||
|
}
|
||||||
|
if (isQuotaExceeded.value) {
|
||||||
|
return t('admin.accounts.status.quotaExceeded')
|
||||||
|
}
|
||||||
if (!props.account.schedulable) {
|
if (!props.account.schedulable) {
|
||||||
return t('admin.accounts.status.paused')
|
return t('admin.accounts.status.paused')
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -2126,6 +2126,7 @@ export default {
|
|||||||
rateLimited: 'Rate Limited',
|
rateLimited: 'Rate Limited',
|
||||||
overloaded: 'Overloaded',
|
overloaded: 'Overloaded',
|
||||||
tempUnschedulable: 'Temp Unschedulable',
|
tempUnschedulable: 'Temp Unschedulable',
|
||||||
|
quotaExceeded: 'Quota Exceeded',
|
||||||
unschedulable: 'Unschedulable',
|
unschedulable: 'Unschedulable',
|
||||||
rateLimitedUntil: 'Rate limited and removed from scheduling. Auto resumes at {time}',
|
rateLimitedUntil: 'Rate limited and removed from scheduling. Auto resumes at {time}',
|
||||||
rateLimitedAutoResume: 'Auto resumes in {time}',
|
rateLimitedAutoResume: 'Auto resumes in {time}',
|
||||||
|
|||||||
@@ -2315,6 +2315,7 @@ export default {
|
|||||||
rateLimited: '限流中',
|
rateLimited: '限流中',
|
||||||
overloaded: '过载中',
|
overloaded: '过载中',
|
||||||
tempUnschedulable: '临时不可调度',
|
tempUnschedulable: '临时不可调度',
|
||||||
|
quotaExceeded: '配额超限',
|
||||||
unschedulable: '不可调度',
|
unschedulable: '不可调度',
|
||||||
rateLimitedUntil: '限流中,当前不参与调度,预计 {time} 自动恢复',
|
rateLimitedUntil: '限流中,当前不参与调度,预计 {time} 自动恢复',
|
||||||
rateLimitedAutoResume: '{time} 自动恢复',
|
rateLimitedAutoResume: '{time} 自动恢复',
|
||||||
|
|||||||
Reference in New Issue
Block a user