fix: 补齐旧账号的 OpenAI 限流补偿

This commit is contained in:
神乐
2026-03-08 00:14:15 +08:00
parent 45d57018eb
commit 1307d604e7
4 changed files with 141 additions and 2 deletions

View File

@@ -1447,10 +1447,20 @@ func (s *OpenAIGatewayService) resolveFreshSchedulableOpenAIAccount(ctx context.
}
func (s *OpenAIGatewayService) getSchedulableAccount(ctx context.Context, accountID int64) (*Account, error) {
var (
account *Account
err error
)
if s.schedulerSnapshot != nil {
return s.schedulerSnapshot.GetAccount(ctx, accountID)
account, err = s.schedulerSnapshot.GetAccount(ctx, accountID)
} else {
account, err = s.accountRepo.GetByID(ctx, accountID)
}
return s.accountRepo.GetByID(ctx, accountID)
if err != nil || account == nil {
return account, err
}
syncOpenAICodexRateLimitFromExtra(ctx, s.accountRepo, account, time.Now())
return account, nil
}
func (s *OpenAIGatewayService) schedulingConfig() config.GatewaySchedulingConfig {
@@ -3923,6 +3933,45 @@ func codexRateLimitResetAtFromSnapshot(snapshot *OpenAICodexUsageSnapshot, fallb
return nil
}
func codexRateLimitResetAtFromExtra(extra map[string]any, now time.Time) *time.Time {
if len(extra) == 0 {
return nil
}
if progress := buildCodexUsageProgressFromExtra(extra, "7d", now); progress != nil && codexUsagePercentExhausted(&progress.Utilization) && progress.ResetsAt != nil && now.Before(*progress.ResetsAt) {
resetAt := progress.ResetsAt.UTC()
return &resetAt
}
if progress := buildCodexUsageProgressFromExtra(extra, "5h", now); progress != nil && codexUsagePercentExhausted(&progress.Utilization) && progress.ResetsAt != nil && now.Before(*progress.ResetsAt) {
resetAt := progress.ResetsAt.UTC()
return &resetAt
}
return nil
}
func applyOpenAICodexRateLimitFromExtra(account *Account, now time.Time) (*time.Time, bool) {
if account == nil || !account.IsOpenAI() {
return nil, false
}
resetAt := codexRateLimitResetAtFromExtra(account.Extra, now)
if resetAt == nil {
return nil, false
}
if account.RateLimitResetAt != nil && now.Before(*account.RateLimitResetAt) && !account.RateLimitResetAt.Before(*resetAt) {
return account.RateLimitResetAt, false
}
account.RateLimitResetAt = resetAt
return resetAt, true
}
func syncOpenAICodexRateLimitFromExtra(ctx context.Context, repo AccountRepository, account *Account, now time.Time) *time.Time {
resetAt, changed := applyOpenAICodexRateLimitFromExtra(account, now)
if !changed || resetAt == nil || repo == nil || account == nil || account.ID <= 0 {
return resetAt
}
_ = repo.SetRateLimited(ctx, account.ID, *resetAt)
return resetAt
}
// updateCodexUsageSnapshot saves the Codex usage snapshot to account's Extra field
func (s *OpenAIGatewayService) updateCodexUsageSnapshot(ctx context.Context, accountID int64, snapshot *OpenAICodexUsageSnapshot) {
if snapshot == nil {