Merge pull request #7555 from Yukinoshita-Yukino/fix/openai-auto-reset-credit-sched-projection

fix(scheduler): keep auto reset-credit fields in the scheduler projection
This commit is contained in:
Wesley Liddick
2026-09-28 10:27:00 +08:00
committed by GitHub
5 changed files with 256 additions and 4 deletions
@@ -1034,6 +1034,13 @@ func filterSchedulerExtra(extra map[string]any) map[string]any {
"auto_pause_7d_threshold",
"auto_pause_5h_disabled",
"auto_pause_7d_disabled",
// 自动用卡:卡可用的 OpenAI 号在暂停阈值与用卡阈值之间继续调度。
// 候选过滤读的是本投影,缺这几个键时放行分支永远不会生效,
// 账号会在暂停阈值处被一刀切停调,直到窗口自然重置。
service.OpenAIAutoResetCreditEnabledExtraKey,
service.OpenAIAutoResetCredit5hThresholdExtraKey,
service.OpenAIAutoResetCredit7dThresholdExtraKey,
service.OpenAIAutoResetCreditStateExtraKey,
"model_rate_limits",
service.UpstreamBillingProbeExtraKey,
service.GrokMediaEligibleExtraKey,
@@ -446,6 +446,47 @@ func TestBuildSchedulerMetadataAccount_KeepsQuotaAutoPauseFields(t *testing.T) {
require.Equal(t, false, got.Extra["auto_pause_7d_disabled"])
}
// 候选过滤读的是 Redis 元数据投影;自动用卡的开关、阈值和卡状态缺失时,
// 卡可用的 OpenAI 号会在暂停阈值处被一刀切停调,放行到用卡阈值的逻辑永远不生效。
func TestSchedulerMetadataPayload_KeepsOpenAIAutoResetCreditFields(t *testing.T) {
checkedAt := time.Now().UTC().Format(time.RFC3339)
account := service.Account{
ID: 40,
Platform: service.PlatformOpenAI,
Type: service.AccountTypeOAuth,
Status: service.StatusActive,
Schedulable: true,
Extra: map[string]any{
"codex_7d_used_percent": 90.0,
service.OpenAIAutoResetCreditEnabledExtraKey: true,
service.OpenAIAutoResetCredit5hThresholdExtraKey: 0.95,
service.OpenAIAutoResetCredit7dThresholdExtraKey: 0.98,
service.OpenAIAutoResetCreditStateExtraKey: map[string]any{
"status": service.OpenAIAutoResetStatusAvailable,
"available_count": 3,
"checked_at": checkedAt,
"trigger_window": "7d",
},
},
}
_, metaPayload, err := marshalSchedulerCacheAccount(account)
require.NoError(t, err)
var cached service.Account
require.NoError(t, json.Unmarshal(metaPayload, &cached))
config := service.ResolveOpenAIAutoResetCreditConfig(&cached)
require.True(t, config.Enabled)
require.Equal(t, 0.95, config.Threshold5h)
require.Equal(t, 0.98, config.Threshold7d)
state, ok := cached.Extra[service.OpenAIAutoResetCreditStateExtraKey].(map[string]any)
require.True(t, ok, "卡状态必须进入调度投影")
require.Equal(t, service.OpenAIAutoResetStatusAvailable, state["status"])
require.EqualValues(t, 3, state["available_count"])
require.Equal(t, checkedAt, state["checked_at"])
}
func TestBuildSchedulerMetadataAccount_KeepsQuotaStateForCachedAccounts(t *testing.T) {
now := time.Now().UTC()
activeStart := now.Add(-time.Hour).Format(time.RFC3339)
@@ -529,11 +529,11 @@ func shouldAutoPauseOpenAIAccountByQuota(ctx context.Context, account *Account)
utilization5h, has5h := resolveOpenAIQuotaUtilization(account.Extra, "5h", now)
utilization7d, has7d := resolveOpenAIQuotaUtilization(account.Extra, "7d", now)
if has5h && utilization5h >= config.Threshold5h {
notifyOpenAIAutoReset(account.ID)
notifyOpenAIAutoResetFromScheduler(account.ID)
return true, openAIQuotaAutoPauseDecision{window: "5h", threshold: config.Threshold5h, utilization: utilization5h, reason: "quota_auto_reset_pending_5h"}
}
if has7d && utilization7d >= config.Threshold7d {
notifyOpenAIAutoReset(account.ID)
notifyOpenAIAutoResetFromScheduler(account.ID)
return true, openAIQuotaAutoPauseDecision{window: "7d", threshold: config.Threshold7d, utilization: utilization7d, reason: "quota_auto_reset_pending_7d"}
}
@@ -547,7 +547,7 @@ func shouldAutoPauseOpenAIAccountByQuota(ctx context.Context, account *Account)
if state != nil && state.Status == OpenAIAutoResetStatusAvailable && state.AvailableCount > 0 && !openAIAutoResetStateStale(state, now) {
return false, openAIQuotaAutoPauseDecision{}
}
notifyOpenAIAutoReset(account.ID)
notifyOpenAIAutoResetFromScheduler(account.ID)
if pauseReached5h {
return true, openAIQuotaAutoPauseDecision{window: "5h", threshold: pause5h, utilization: utilization5h, reason: "quota_auto_reset_credit_check_5h"}
}
@@ -27,6 +27,13 @@ const (
openAIAutoResetQueueCapacity = 1024
openAIAutoResetAttemptTTL = 8 * 24 * time.Hour
openAIAutoResetLeaderLockKey = "jobs:openai-auto-reset-credit"
// 查询阶段失败后的最短重试间隔。调度热路径和扫描都会触发评估,不设间隔时上游
// 故障期间同一账号会被十几秒重查一次;用卡阶段的失败不受此限,超时后仍立即用
// 同一幂等键重试。
openAIAutoResetQueryFailureRetryAfter = time.Minute
// 调度热路径每次过滤候选都会评估暂停,对同一账号的通知按此冷却合并;
// 后台每分钟全量扫描兜底,冷却不会让账号漏检。
openAIAutoResetSchedulerNotifyCooldown = 30 * time.Second
)
const (
@@ -288,10 +295,16 @@ func (s *OpenAIQuotaAutoResetService) evaluateAccount(ctx context.Context, accou
now := time.Now()
assessment := s.assessExtra(account, config, now)
state := openAIAutoResetStateFromExtra(account.Extra)
needsQuery := openAIAutoResetSnapshotStale(account.Extra, now) || assessment.resetReached
// 达到用卡阈值本应立即查询以便用卡;但 10 分钟内已确认无卡时,重查不会改变结论,
// 只会让调度热路径的通知把同一账号的上游额度接口打到十几秒一次。
needsQuery := openAIAutoResetSnapshotStale(account.Extra, now) ||
(assessment.resetReached && !openAIAutoResetNoCreditConfirmed(state, now))
if assessment.pauseReached && !assessment.resetReached {
needsQuery = needsQuery || state == nil || state.Status == OpenAIAutoResetStatusChecking || state.Status == OpenAIAutoResetStatusFailed || openAIAutoResetStateStale(state, now)
}
if needsQuery && openAIAutoResetQueryFailureBackoffActive(state, now) {
needsQuery = false
}
if !needsQuery {
if !assessment.pauseReached && state != nil && state.TriggerWindow != "" {
state.TriggerWindow = ""
@@ -690,6 +703,30 @@ func openAIAutoResetStateFromExtra(extra map[string]any) *OpenAIAutoResetCreditS
return &state
}
// openAIAutoResetNoCreditConfirmed 报告账号在状态有效期内已被确认无卡。
func openAIAutoResetNoCreditConfirmed(state *OpenAIAutoResetCreditState, now time.Time) bool {
return state != nil && state.Status == OpenAIAutoResetStatusNoCredit && !openAIAutoResetStateStale(state, now)
}
// openAIAutoResetQueryFailureBackoffActive 报告查询阶段刚失败、仍在重试间隔内。
// 只覆盖查询、快照写入、卡明细缺失三类失败;用卡阶段的失败需要立即按原幂等键重试。
func openAIAutoResetQueryFailureBackoffActive(state *OpenAIAutoResetCreditState, now time.Time) bool {
if state == nil || state.Status != OpenAIAutoResetStatusFailed {
return false
}
switch state.ErrorCode {
case "RESET_CREDIT_QUERY_FAILED", "USAGE_SNAPSHOT_WRITE_FAILED", "RESET_CREDIT_DETAILS_UNAVAILABLE":
default:
return false
}
failedAt, err := time.Parse(time.RFC3339, state.LastResultAt)
if err != nil {
return false
}
elapsed := now.Sub(failedAt)
return elapsed >= 0 && elapsed < openAIAutoResetQueryFailureRetryAfter
}
func openAIAutoResetStateStale(state *OpenAIAutoResetCreditState, now time.Time) bool {
if state == nil || state.CheckedAt == "" {
return true
@@ -808,3 +845,29 @@ func notifyOpenAIAutoReset(accountID int64) {
func NotifyOpenAIAutoResetCredit(accountID int64) {
notifyOpenAIAutoReset(accountID)
}
// openAIAutoResetSchedulerNotifiedAt 记录调度热路径最近一次为某账号发出通知的时间。
var openAIAutoResetSchedulerNotifiedAt sync.Map // accountID(int64) -> time.Time
// notifyOpenAIAutoResetFromScheduler 供调度候选过滤使用。候选过滤按请求逐账号执行,
// 每条通知都会让后台读一次账号,不做冷却时无卡或待用卡的账号会持续占满后台协程。
func notifyOpenAIAutoResetFromScheduler(accountID int64) {
notifyOpenAIAutoResetFromSchedulerAt(accountID, time.Now())
}
func notifyOpenAIAutoResetFromSchedulerAt(accountID int64, now time.Time) bool {
if accountID <= 0 {
return false
}
if last, ok := openAIAutoResetSchedulerNotifiedAt.Load(accountID); ok {
if lastAt, ok := last.(time.Time); ok {
elapsed := now.Sub(lastAt)
if elapsed >= 0 && elapsed < openAIAutoResetSchedulerNotifyCooldown {
return false
}
}
}
openAIAutoResetSchedulerNotifiedAt.Store(accountID, now)
notifyOpenAIAutoReset(accountID)
return true
}
@@ -3,6 +3,7 @@ package service
import (
"context"
"encoding/json"
"errors"
"sync"
"sync/atomic"
"testing"
@@ -100,6 +101,43 @@ func TestShouldAutoPauseOpenAIAccountByQuota_AutoResetCreditStates(t *testing.T)
})
}
func TestNotifyOpenAIAutoResetFromScheduler_CoolsDownPerAccount(t *testing.T) {
svc := &OpenAIQuotaAutoResetService{ctx: context.Background(), queue: make(chan int64, 8)}
setOpenAIAutoResetNotifier(svc)
t.Cleanup(func() { clearOpenAIAutoResetNotifier(svc) })
const accountA, accountB int64 = 9_900_001, 9_900_002
t.Cleanup(func() {
openAIAutoResetSchedulerNotifiedAt.Delete(accountA)
openAIAutoResetSchedulerNotifiedAt.Delete(accountB)
})
drain := func() []int64 {
var got []int64
for {
select {
case id := <-svc.queue:
svc.pending.Delete(id)
got = append(got, id)
default:
return got
}
}
}
base := time.Now()
require.True(t, notifyOpenAIAutoResetFromSchedulerAt(accountA, base))
require.Equal(t, []int64{accountA}, drain())
require.False(t, notifyOpenAIAutoResetFromSchedulerAt(accountA, base.Add(10*time.Second)), "冷却期内同一账号不重复通知")
require.Empty(t, drain())
require.True(t, notifyOpenAIAutoResetFromSchedulerAt(accountB, base.Add(10*time.Second)), "冷却按账号独立")
require.Equal(t, []int64{accountB}, drain())
require.True(t, notifyOpenAIAutoResetFromSchedulerAt(accountA, base.Add(openAIAutoResetSchedulerNotifyCooldown)))
require.Equal(t, []int64{accountA}, drain())
}
func TestSelectOpenAIAutoResetCandidate_FailsClosed(t *testing.T) {
candidates := []openAIAutoResetCreditCandidate{
{ID: "later", ExpiresAt: "2026-09-02T00:00:00Z"},
@@ -181,9 +219,15 @@ type autoResetTestQuota struct {
mu sync.Mutex
resetArgs [][2]string
failFirst bool
queryCalls atomic.Int32
queryErr error
}
func (q *autoResetTestQuota) QueryUsage(context.Context, int64) (*OpenAIQuotaUsage, error) {
q.queryCalls.Add(1)
if q.queryErr != nil {
return nil, q.queryErr
}
copy := *q.usage
return &copy, nil
}
@@ -334,3 +378,100 @@ func TestOpenAIQuotaAutoResetService_TimeoutRetryReusesRequestBody(t *testing.T)
require.Len(t, args, 2)
require.Equal(t, args[0], args[1], "超时重试必须复用相同 credit_id 与 redeem_request_id")
}
// newResetThresholdTestFixture 构造一个 7d 用量已过用卡阈值、处于停调中的账号。
func newResetThresholdTestFixture(t *testing.T, state *OpenAIAutoResetCreditState, availableCredits int) (*OpenAIQuotaAutoResetService, *autoResetTestAccountRepo, *autoResetTestQuota) {
t.Helper()
now := time.Now().UTC()
extra := map[string]any{
OpenAIAutoResetCreditEnabledExtraKey: true,
OpenAIAutoResetCredit5hThresholdExtraKey: 0.95,
OpenAIAutoResetCredit7dThresholdExtraKey: 0.95,
"codex_7d_used_percent": 99.0,
"codex_usage_updated_at": now.Format(time.RFC3339),
"codex_7d_reset_at": now.Add(48 * time.Hour).Format(time.RFC3339),
}
if state != nil {
extra[OpenAIAutoResetCreditStateExtraKey] = state
}
account := &Account{ID: 501, Platform: PlatformOpenAI, Type: AccountTypeOAuth, Status: StatusActive, Schedulable: true, Extra: extra}
repo := &autoResetTestAccountRepo{account: account}
quota := &autoResetTestQuota{usage: &OpenAIQuotaUsage{
FetchedAt: now.Unix(),
RateLimit: &OpenAIRateLimit{
SecondaryWindow: &OpenAIRateLimitWindow{UsedPercent: 99, LimitWindowSeconds: 7 * 24 * 60 * 60, ResetAfterSeconds: 48 * 3600, ResetAt: now.Add(48 * time.Hour).Unix()},
},
RateLimitResetCredits: &OpenAIRateLimitResetCredits{AvailableCount: availableCredits},
}}
idempotencyConfig := DefaultIdempotencyConfig()
idempotencyConfig.ObserveOnly = false
service := NewOpenAIQuotaAutoResetService(repo, quota, autoResetTestRecoverer{},
NewIdempotencyCoordinator(newInMemoryIdempotencyRepo(), idempotencyConfig), nil, nil, nil)
return service, repo, quota
}
func (r *autoResetTestAccountRepo) setExtraForTest(key string, value any) {
r.mu.Lock()
defer r.mu.Unlock()
r.account.Extra[key] = value
}
func (r *autoResetTestAccountRepo) stateForTest() *OpenAIAutoResetCreditState {
r.mu.Lock()
defer r.mu.Unlock()
return openAIAutoResetStateFromExtra(r.account.Extra)
}
func TestOpenAIQuotaAutoResetService_ResetThresholdSkipsRequeryWhileNoCreditIsFresh(t *testing.T) {
now := time.Now().UTC()
service, repo, quota := newResetThresholdTestFixture(t, &OpenAIAutoResetCreditState{
Status: OpenAIAutoResetStatusNoCredit, TriggerWindow: "7d", CheckedAt: now.Format(time.RFC3339), ErrorCode: "NO_RESET_CREDIT",
}, 0)
for i := 0; i < 3; i++ {
require.NoError(t, service.evaluateAccount(context.Background(), 501))
}
require.Zero(t, quota.queryCalls.Load(), "10 分钟内已确认无卡,调度通知不应再触发上游查询")
stale := now.Add(-openAIAutoResetSnapshotTTL - time.Minute).Format(time.RFC3339)
repo.setExtraForTest(OpenAIAutoResetCreditStateExtraKey, &OpenAIAutoResetCreditState{
Status: OpenAIAutoResetStatusNoCredit, TriggerWindow: "7d", CheckedAt: stale, ErrorCode: "NO_RESET_CREDIT",
})
repo.setExtraForTest("codex_usage_updated_at", stale)
require.NoError(t, service.evaluateAccount(context.Background(), 501))
require.Equal(t, int32(1), quota.queryCalls.Load(), "无卡结论过期后应重新查询,以便发现新发放的卡")
require.Equal(t, OpenAIAutoResetStatusNoCredit, repo.stateForTest().Status)
}
func TestOpenAIQuotaAutoResetService_ResetThresholdStillQueriesWhenCreditAvailable(t *testing.T) {
now := time.Now().UTC()
service, _, quota := newResetThresholdTestFixture(t, &OpenAIAutoResetCreditState{
Status: OpenAIAutoResetStatusAvailable, AvailableCount: 1, TriggerWindow: "7d", CheckedAt: now.Format(time.RFC3339),
}, 0)
require.NoError(t, service.evaluateAccount(context.Background(), 501))
require.Equal(t, int32(1), quota.queryCalls.Load(), "有卡时到达用卡阈值必须立即查询并进入用卡流程")
}
func TestOpenAIQuotaAutoResetService_QueryFailureBacksOffBeforeRetry(t *testing.T) {
service, repo, quota := newResetThresholdTestFixture(t, nil, 0)
quota.queryErr = errors.New("upstream EOF")
require.Error(t, service.evaluateAccount(context.Background(), 501))
require.Equal(t, int32(1), quota.queryCalls.Load())
state := repo.stateForTest()
require.Equal(t, OpenAIAutoResetStatusFailed, state.Status)
require.Equal(t, "RESET_CREDIT_QUERY_FAILED", state.ErrorCode)
for i := 0; i < 3; i++ {
require.NoError(t, service.evaluateAccount(context.Background(), 501))
}
require.Equal(t, int32(1), quota.queryCalls.Load(), "查询失败后一分钟内不应重复打上游")
state.LastResultAt = time.Now().UTC().Add(-openAIAutoResetQueryFailureRetryAfter - time.Second).Format(time.RFC3339)
repo.setExtraForTest(OpenAIAutoResetCreditStateExtraKey, state)
quota.queryErr = nil
require.NoError(t, service.evaluateAccount(context.Background(), 501))
require.Equal(t, int32(2), quota.queryCalls.Load(), "重试间隔过后应恢复查询")
require.Equal(t, OpenAIAutoResetStatusNoCredit, repo.stateForTest().Status)
}