From 5f09442fced8a82c3eb5d58b608aa0a075c751ee Mon Sep 17 00:00:00 2001 From: zenor0 Date: Thu, 27 Aug 2026 00:17:58 +0800 Subject: [PATCH] fix(openai): refresh usage after quota reset --- .../handler/admin/openai_oauth_handler.go | 1 + .../openai_oauth_handler_reset_quota_test.go | 6 ++++ .../service/openai_quota_auto_reset.go | 1 + .../service/openai_quota_auto_reset_test.go | 4 +++ .../service/openai_quota_reset_workflow.go | 4 +-- .../internal/service/openai_quota_service.go | 25 +++++++++++++-- .../service/openai_quota_spark_window_test.go | 31 +++++++++++++++++-- .../components/account/AccountUsageCell.vue | 12 ------- .../__tests__/AccountUsageCell.spec.ts | 4 +-- 9 files changed, 66 insertions(+), 22 deletions(-) diff --git a/backend/internal/handler/admin/openai_oauth_handler.go b/backend/internal/handler/admin/openai_oauth_handler.go index 779e4455cd..031fe437fd 100644 --- a/backend/internal/handler/admin/openai_oauth_handler.go +++ b/backend/internal/handler/admin/openai_oauth_handler.go @@ -27,6 +27,7 @@ type OpenAIOAuthHandler struct { type openAIQuotaService interface { QueryUsage(ctx context.Context, accountID int64) (*service.OpenAIQuotaUsage, error) CacheResetCreditsSnapshot(ctx context.Context, accountID int64, credits *service.OpenAIRateLimitResetCredits) error + CachePostResetSnapshot(ctx context.Context, accountID int64, usage *service.OpenAIQuotaUsage) error ResetCredit(ctx context.Context, accountID int64) (*service.OpenAIQuotaResetResult, error) } diff --git a/backend/internal/handler/admin/openai_oauth_handler_reset_quota_test.go b/backend/internal/handler/admin/openai_oauth_handler_reset_quota_test.go index b9e8dac679..90677b3bf4 100644 --- a/backend/internal/handler/admin/openai_oauth_handler_reset_quota_test.go +++ b/backend/internal/handler/admin/openai_oauth_handler_reset_quota_test.go @@ -48,6 +48,12 @@ func (s *openAIQuotaWorkflowStub) CacheResetCreditsSnapshot(ctx context.Context, return s.cacheErr } +func (s *openAIQuotaWorkflowStub) CachePostResetSnapshot(ctx context.Context, _ int64, _ *service.OpenAIQuotaUsage) error { + s.cacheCalls++ + s.cacheCtxErr = ctx.Err() + return s.cacheErr +} + type openAIAccountStateRecovererStub struct { err error calls int diff --git a/backend/internal/service/openai_quota_auto_reset.go b/backend/internal/service/openai_quota_auto_reset.go index aafc5991d7..b7994e3cf5 100644 --- a/backend/internal/service/openai_quota_auto_reset.go +++ b/backend/internal/service/openai_quota_auto_reset.go @@ -54,6 +54,7 @@ type OpenAIAutoResetCreditState struct { type openAIAutoResetQuota interface { QueryUsage(ctx context.Context, accountID int64) (*OpenAIQuotaUsage, error) CacheResetCreditsSnapshot(ctx context.Context, accountID int64, credits *OpenAIRateLimitResetCredits) error + CachePostResetSnapshot(ctx context.Context, accountID int64, usage *OpenAIQuotaUsage) error ResetCreditTargeted(ctx context.Context, accountID int64, creditID, redeemRequestID string) (*OpenAIQuotaResetResult, error) } diff --git a/backend/internal/service/openai_quota_auto_reset_test.go b/backend/internal/service/openai_quota_auto_reset_test.go index ad2dc965aa..356224cad0 100644 --- a/backend/internal/service/openai_quota_auto_reset_test.go +++ b/backend/internal/service/openai_quota_auto_reset_test.go @@ -192,6 +192,10 @@ func (q *autoResetTestQuota) CacheResetCreditsSnapshot(context.Context, int64, * return nil } +func (q *autoResetTestQuota) CachePostResetSnapshot(context.Context, int64, *OpenAIQuotaUsage) error { + return nil +} + func (q *autoResetTestQuota) ResetCreditTargeted(_ context.Context, _ int64, creditID, redeemRequestID string) (*OpenAIQuotaResetResult, error) { if creditID == "" || redeemRequestID == "" { panic("targeted reset identifiers must be present") diff --git a/backend/internal/service/openai_quota_reset_workflow.go b/backend/internal/service/openai_quota_reset_workflow.go index 430b29618b..4dc0210cc7 100644 --- a/backend/internal/service/openai_quota_reset_workflow.go +++ b/backend/internal/service/openai_quota_reset_workflow.go @@ -15,7 +15,7 @@ const ( type openAIQuotaResetWorkflowQuota interface { QueryUsage(ctx context.Context, accountID int64) (*OpenAIQuotaUsage, error) - CacheResetCreditsSnapshot(ctx context.Context, accountID int64, credits *OpenAIRateLimitResetCredits) error + CachePostResetSnapshot(ctx context.Context, accountID int64, usage *OpenAIQuotaUsage) error } type openAIQuotaResetWorkflowRecoverer interface { @@ -60,7 +60,7 @@ func RunOpenAIQuotaResetPostProcess( slog.Warn("openai_quota_reset_cache_refresh_failed", "account_id", accountID, "error_code", infraerrors.Reason(usageErr)) result.WarningCode = OpenAIQuotaResetWarningCacheRefreshFailed default: - if err := quota.CacheResetCreditsSnapshot(ctx, accountID, usage.RateLimitResetCredits); err != nil { + if err := quota.CachePostResetSnapshot(ctx, accountID, usage); err != nil { slog.Warn("openai_quota_reset_cache_refresh_failed", "account_id", accountID, "error_code", infraerrors.Reason(err)) result.WarningCode = OpenAIQuotaResetWarningCacheRefreshFailed } else { diff --git a/backend/internal/service/openai_quota_service.go b/backend/internal/service/openai_quota_service.go index 6e40c497fc..b05241b222 100644 --- a/backend/internal/service/openai_quota_service.go +++ b/backend/internal/service/openai_quota_service.go @@ -225,6 +225,23 @@ func (s *OpenAIQuotaService) QueryUsage(ctx context.Context, accountID int64) (* // consume) credits that already expired. Callers must treat this rejection as a // partial success — the upstream read itself is still valid. func (s *OpenAIQuotaService) CacheResetCreditsSnapshot(ctx context.Context, accountID int64, credits *OpenAIRateLimitResetCredits) error { + return s.cacheResetCreditsSnapshot(ctx, accountID, credits, nil) +} + +// CachePostResetSnapshot persists the credits and usage windows observed after a reset. +func (s *OpenAIQuotaService) CachePostResetSnapshot(ctx context.Context, accountID int64, usage *OpenAIQuotaUsage) error { + if usage == nil { + return s.cacheResetCreditsSnapshot(ctx, accountID, nil, nil) + } + return s.cacheResetCreditsSnapshot( + ctx, + accountID, + usage.RateLimitResetCredits, + buildOpenAIAutoResetUsageUpdates(usage, time.Now()), + ) +} + +func (s *OpenAIQuotaService) cacheResetCreditsSnapshot(ctx context.Context, accountID int64, credits *OpenAIRateLimitResetCredits, updates map[string]any) error { if credits == nil || (credits.AvailableCount > 0 && len(credits.Credits) == 0) { return infraerrors.New( http.StatusBadGateway, @@ -232,9 +249,11 @@ func (s *OpenAIQuotaService) CacheResetCreditsSnapshot(ctx context.Context, acco "failed to refresh reset-credit expiration details; cached data was preserved", ) } - if err := s.accountRepo.UpdateExtra(ctx, accountID, map[string]any{ - openaiQuotaResetCreditsKey: credits, - }); err != nil { + if updates == nil { + updates = make(map[string]any, 1) + } + updates[openaiQuotaResetCreditsKey] = credits + if err := s.accountRepo.UpdateExtra(ctx, accountID, updates); err != nil { return infraerrors.New( http.StatusInternalServerError, "OPENAI_QUOTA_CACHE_WRITE_FAILED", diff --git a/backend/internal/service/openai_quota_spark_window_test.go b/backend/internal/service/openai_quota_spark_window_test.go index cb47a0c0c5..b1f636d416 100644 --- a/backend/internal/service/openai_quota_spark_window_test.go +++ b/backend/internal/service/openai_quota_spark_window_test.go @@ -27,9 +27,10 @@ import ( // stubQuotaAccountRepo 是多账号 AccountRepository stub，仅实现 GetByID。 type stubQuotaAccountRepo struct { AccountRepository - accounts map[int64]*Account - extraUpdates map[int64]map[string]any - extraUpdateErr error + accounts map[int64]*Account + extraUpdates map[int64]map[string]any + extraUpdateCalls int + extraUpdateErr error } func (r *stubQuotaAccountRepo) GetByID(_ context.Context, id int64) (*Account, error) { @@ -53,6 +54,7 @@ func (r *stubQuotaAccountRepo) UpdateExtra(_ context.Context, id int64, updates if r.extraUpdateErr != nil { return r.extraUpdateErr } + r.extraUpdateCalls++ if r.extraUpdates == nil { r.extraUpdates = make(map[int64]map[string]any) } @@ -694,6 +696,29 @@ func TestCacheResetCreditsSnapshot(t *testing.T) { }) } +func TestCachePostResetSnapshot(t *testing.T) { + repo := &stubQuotaAccountRepo{} + svc := &OpenAIQuotaService{accountRepo: repo} + credits := &OpenAIRateLimitResetCredits{AvailableCount: 0} + usage := &OpenAIQuotaUsage{ + RateLimitResetCredits: credits, + RateLimit: &OpenAIRateLimit{ + PrimaryWindow: &OpenAIRateLimitWindow{ + UsedPercent: 0, LimitWindowSeconds: 5 * 60 * 60, ResetAfterSeconds: 5 * 60 * 60, + }, + SecondaryWindow: &OpenAIRateLimitWindow{ + UsedPercent: 0, LimitWindowSeconds: 7 * 24 * 60 * 60, ResetAfterSeconds: 7 * 24 * 60 * 60, + }, + }, + } + + require.NoError(t, svc.CachePostResetSnapshot(context.Background(), 100, usage)) + require.Equal(t, 1, repo.extraUpdateCalls) + require.Equal(t, credits, repo.extraUpdates[100][openaiQuotaResetCreditsKey]) + require.Equal(t, 0.0, repo.extraUpdates[100]["codex_5h_used_percent"]) + require.Equal(t, 0.0, repo.extraUpdates[100]["codex_7d_used_percent"]) +} + // TestResetCreditGetByIDError_FailsClosed 验证守卫「失败关闭」语义： // 当守卫的 GetByID 发生瞬时错误时，ResetCredit 必须立即返回该错误， // 不得旁路进入 prepareUpstreamCall（否则影子账号会借 resolve 路径操作母账号）。 diff --git a/frontend/src/components/account/AccountUsageCell.vue b/frontend/src/components/account/AccountUsageCell.vue index 966b0a5e08..1ff84f3275 100644 --- a/frontend/src/components/account/AccountUsageCell.vue +++ b/frontend/src/components/account/AccountUsageCell.vue @@ -656,8 +656,6 @@ import { cnQuotaCellVisible as cnQuotaCellVisibleFn, cnBalanceCellVisible as cnB // Module-level cache shared across all AccountUsageCell instances const _usageCache = new Map() const USAGE_CACHE_TTL = 5 * 60 * 1000 // 5 minutes -// How long a quota-reset response may suppress the row-patch usage refetch. -const SUPPRESS_USAGE_REFRESH_WINDOW_MS = 5 * 1000 const props = withDefaults( defineProps<{ @@ -699,7 +697,6 @@ const usageInfo = ref(null) watch(usageInfo, (usage) => { if (usage) emit('usage-loaded', usage) }) -const suppressOpenAIUsageRefreshUntil = ref(0) const rootRef = ref(null) const isDesktopViewport = ref( typeof window === 'undefined' ? true : window.matchMedia(desktopViewportQuery).matches @@ -1525,11 +1522,6 @@ const quotaTotalBar = computed((): QuotaBarInfo | null => { }) const handleQuotaResetAccountUpdated = (account: Account) => { - // The reset response already carries authoritative quota and account data. - // Avoid turning the parent patch into a second automatic /usage request. - // The suppression is time-boxed so an unhandled emit (parent that ignores - // account-updated) cannot latch it and swallow a later, unrelated refresh. - suppressOpenAIUsageRefreshUntil.value = Date.now() + SUPPRESS_USAGE_REFRESH_WINDOW_MS emit('account-updated', account) } @@ -1619,10 +1611,6 @@ watch( watch(openAIUsageRefreshKey, (nextKey, prevKey) => { if (!prevKey || nextKey === prevKey) return if (props.account.platform !== 'openai' || props.account.type !== 'oauth') return - if (Date.now() < suppressOpenAIUsageRefreshUntil.value) { - suppressOpenAIUsageRefreshUntil.value = 0 - return - } if (isBatchManaged.value) { requestParentBatchUsage({ force: true }) diff --git a/frontend/src/components/account/__tests__/AccountUsageCell.spec.ts b/frontend/src/components/account/__tests__/AccountUsageCell.spec.ts index 1ad6866dd7..b87269f7a6 100644 --- a/frontend/src/components/account/__tests__/AccountUsageCell.spec.ts +++ b/frontend/src/components/account/__tests__/AccountUsageCell.spec.ts @@ -519,7 +519,7 @@ describe('AccountUsageCell', () => { expect(wrapper.text()).toContain('5h|0|200') }) - it('OpenAI 重置响应更新账号行时不会额外拉取 usage', async () => { + it('OpenAI 重置响应更新账号行后重新拉取 usage', async () => { getUsage.mockResolvedValue({ five_hour: { utilization: 0, @@ -560,7 +560,7 @@ describe('AccountUsageCell', () => { await wrapper.setProps({ account: updatedAccount as Account }) await flushPromises() - expect(getUsage).toHaveBeenCalledTimes(1) + expect(getUsage).toHaveBeenCalledTimes(2) }) it('OpenAI OAuth 已限额时显示 /usage API 返回的限额数据', async () => {