diff --git a/backend/internal/application/account/quota_refresh_test.go b/backend/internal/application/account/quota_refresh_test.go index a459ed568..986c33675 100644 --- a/backend/internal/application/account/quota_refresh_test.go +++ b/backend/internal/application/account/quota_refresh_test.go @@ -3,6 +3,7 @@ package account import ( "context" "errors" + "fmt" "path/filepath" "sync" "sync/atomic" @@ -411,6 +412,57 @@ func TestRefreshWebImagineQuotaModeAtomicallyReplacesGroup(t *testing.T) { } } +func TestRefreshPaidWebImagineFallsBackToSharedWeeklyQuota(t *testing.T) { + ctx := context.Background() + database, err := relational.OpenSQLite(ctx, filepath.Join(t.TempDir(), "imagine-shared-weekly.db")) + if err != nil { + t.Fatal(err) + } + t.Cleanup(func() { _ = database.Close() }) + if err := database.InitializeSchema(ctx); err != nil { + t.Fatal(err) + } + accounts := relational.NewAccountRepository(database) + credential, _, err := accounts.UpsertByIdentity(ctx, accountdomain.Credential{ + Provider: accountdomain.ProviderWeb, AuthType: accountdomain.AuthTypeSSO, WebTier: accountdomain.WebTierSuper, + Name: "web-imagine-weekly", SourceKey: "web-imagine-weekly", EncryptedAccessToken: "encrypted", + Enabled: true, AuthStatus: accountdomain.AuthStatusActive, + }) + if err != nil { + t.Fatal(err) + } + now := time.Now().UTC() + if err := accounts.SaveQuotaWindows(ctx, credential.ID, accountdomain.WebTierSuper, now, []accountdomain.QuotaWindow{ + {AccountID: credential.ID, Mode: "weekly", Remaining: 50, Total: 100, UpdatedAt: now}, + {AccountID: credential.ID, Mode: accountdomain.QuotaModeWebImagePro, Remaining: 4, UpdatedAt: now}, + }); err != nil { + t.Fatal(err) + } + adapter := &sharedWeeklyImagineAdapter{} + service := NewService(accounts, nil, nil, nil, provider.NewRegistry(adapter), nil, nil) + window, err := service.RefreshQuotaMode(ctx, credential.ID, accountdomain.QuotaModeWebImagePro) + if err != nil { + t.Fatal(err) + } + if window.Mode != "weekly" || window.Remaining != 80 || adapter.groupCalls.Load() != 1 || adapter.modeCalls.Load() != 1 { + t.Fatalf("window = %#v, group calls = %d, mode calls = %d", window, adapter.groupCalls.Load(), adapter.modeCalls.Load()) + } + stored, err := accounts.GetQuotaWindows(ctx, []uint64{credential.ID}) + if err != nil { + t.Fatal(err) + } + byMode := make(map[string]accountdomain.QuotaWindow) + for _, current := range stored[credential.ID] { + byMode[current.Mode] = current + } + if byMode["weekly"].Remaining != 80 { + t.Fatalf("weekly quota was not refreshed: %#v", stored[credential.ID]) + } + if _, exists := byMode[accountdomain.QuotaModeWebImagePro]; exists { + t.Fatalf("stale product quota survived availability-only refresh: %#v", stored[credential.ID]) + } +} + func TestRefreshConsoleQuotaModePersistsCompleteUsageSnapshot(t *testing.T) { ctx := context.Background() database, err := relational.OpenSQLite(ctx, filepath.Join(t.TempDir(), "console-quota-mode.db")) @@ -900,6 +952,48 @@ type imagineQuotaGroupAdapter struct { calls atomic.Int64 } +type sharedWeeklyImagineAdapter struct { + groupCalls atomic.Int64 + modeCalls atomic.Int64 +} + +func (a *sharedWeeklyImagineAdapter) Provider() accountdomain.Provider { + return accountdomain.ProviderWeb +} + +func (a *sharedWeeklyImagineAdapter) Definition() provider.Definition { + return provider.Definition{ + Provider: accountdomain.ProviderWeb, ModelNamespace: accountdomain.ProviderWeb.ModelNamespace(), + Quota: provider.QuotaRemoteWindow, Credential: provider.CredentialSurface{AuthType: accountdomain.AuthTypeSSO}, + } +} + +func (a *sharedWeeklyImagineAdapter) SyncQuota(context.Context, accountdomain.Credential) (provider.QuotaSnapshot, error) { + return provider.QuotaSnapshot{}, errors.New("unexpected full quota sync") +} + +func (a *sharedWeeklyImagineAdapter) SyncQuotaMode(_ context.Context, credential accountdomain.Credential, mode string) (accountdomain.QuotaWindow, error) { + a.modeCalls.Add(1) + if mode != "weekly" { + return accountdomain.QuotaWindow{}, fmt.Errorf("unexpected quota mode %q", mode) + } + now := time.Now().UTC() + return accountdomain.QuotaWindow{ + AccountID: credential.ID, Mode: mode, Remaining: 80, Total: 100, + SyncedAt: &now, Source: accountdomain.QuotaSourceUpstream, UpdatedAt: now, + }, nil +} + +func (a *sharedWeeklyImagineAdapter) SyncQuotaGroup(_ context.Context, _ accountdomain.Credential, group string) (provider.QuotaGroupSnapshot, error) { + a.groupCalls.Add(1) + if group != accountdomain.QuotaGroupWebImagine { + return provider.QuotaGroupSnapshot{}, fmt.Errorf("unexpected quota group %q", group) + } + return provider.QuotaGroupSnapshot{ + Group: group, Modes: accountdomain.WebImagineQuotaModes(), SyncedAt: time.Now().UTC(), + }, nil +} + func (a *imagineQuotaGroupAdapter) Provider() accountdomain.Provider { return accountdomain.ProviderWeb } diff --git a/backend/internal/application/account/service.go b/backend/internal/application/account/service.go index 126461370..6fff12db4 100644 --- a/backend/internal/application/account/service.go +++ b/backend/internal/application/account/service.go @@ -3064,9 +3064,9 @@ func (s *Service) RefreshQuotaMode(ctx context.Context, id uint64, mode string) return accountdomain.QuotaWindow{}, err } } - window, ok := quotaWindowByMode(refreshed.Windows, mode) - if !ok { - return accountdomain.QuotaWindow{}, fmt.Errorf("Provider usage 响应缺少 %s 额度", mode) + window, err := s.resolveRefreshedQuotaWindow(ctx, id, mode, refreshed) + if err != nil { + return accountdomain.QuotaWindow{}, err } if len(refreshed.Modes) == 0 && refreshed.Credential.Provider == accountdomain.ProviderConsole { // One Console request refreshes all three authoritative windows. Reconcile @@ -3079,6 +3079,13 @@ func (s *Service) RefreshQuotaMode(ctx context.Context, id uint64, mode string) if err := s.reconcileQuotaRecoveryWindow(ctx, refreshed.Credential.Provider, id, window); err != nil { return window, err } + } else if window.Mode == "weekly" { + // The requested Imagine product was availability-only and resolved to + // the paid shared pool. Reconcile the authoritative weekly window too; + // the group reconciliation above only covers product-specific modes. + if err := s.reconcileQuotaRecoveryWindow(ctx, refreshed.Credential.Provider, id, window); err != nil { + return window, err + } } return window, nil } @@ -3102,11 +3109,26 @@ func (s *Service) ProbeQuotaMode(ctx context.Context, id uint64, mode string) (a if !ok { return accountdomain.QuotaWindow{}, fmt.Errorf("Provider 模式额度探测返回类型无效") } - window, ok := quotaWindowByMode(refreshed.Windows, mode) - if !ok { - return accountdomain.QuotaWindow{}, fmt.Errorf("Provider usage 响应缺少 %s 额度", mode) + return s.resolveRefreshedQuotaWindow(ctx, id, mode, refreshed) +} + +func (s *Service) resolveRefreshedQuotaWindow(ctx context.Context, id uint64, mode string, refreshed quotaRefreshResult) (accountdomain.QuotaWindow, error) { + if window, ok := quotaWindowByMode(refreshed.Windows, mode); ok { + return window, nil } - return window, nil + credential := refreshed.Credential + paidWebImagine := credential.Provider == accountdomain.ProviderWeb && isWebImagineQuotaMode(mode) && + (credential.WebTier == accountdomain.WebTierSuper || credential.WebTier == accountdomain.WebTierHeavy) + if paidWebImagine { + weekly, err := s.refreshQuotaMode(ctx, id, "weekly") + if err != nil { + return accountdomain.QuotaWindow{}, err + } + if window, ok := quotaWindowByMode(weekly.Windows, "weekly"); ok { + return window, nil + } + } + return accountdomain.QuotaWindow{}, fmt.Errorf("Provider usage 响应缺少 %s 额度", mode) } func (s *Service) refreshQuotaGroup(ctx context.Context, id uint64, group string) (quotaRefreshResult, error) { diff --git a/backend/internal/application/gateway/image.go b/backend/internal/application/gateway/image.go index 258a05da4..4ac4bac5f 100644 --- a/backend/internal/application/gateway/image.go +++ b/backend/internal/application/gateway/image.go @@ -325,7 +325,7 @@ func (s *Service) executeImage( } } quotaKind, _ := s.providers.QuotaKind(route.Provider) - refreshMode, decrementMode := quotaFinalizationModes(effectiveQuotaMode, quotaRefreshGroup) + refreshMode, decrementMode, availabilityMode := quotaFinalizationModes(effectiveQuotaMode, quotaRefreshGroup) if successful && quotaKind == provider.QuotaRemoteWindow && refreshMode != "" { if decrementMode != "" && decrementMode != "weekly" { units := max(1, response.QuotaUnits) @@ -342,6 +342,9 @@ func (s *Service) executeImage( } } s.accounts.QueueQuotaRefresh(accountID, refreshMode) + if availabilityMode != "" && availabilityMode != refreshMode { + s.accounts.QueueQuotaRefresh(accountID, availabilityMode) + } } if err := budget.run("audit", finalizationAuditBudget, func(stageCtx context.Context) error { return s.audits.Create(stageCtx, record) @@ -359,10 +362,17 @@ func (s *Service) executeImage( // upstream windows atomically, while the local fence must charge the exact // window selected for this account so concurrent media requests cannot // over-allocate during the short refresh delay. -func quotaFinalizationModes(effectiveMode, refreshGroup string) (refreshMode, decrementMode string) { +func quotaFinalizationModes(effectiveMode, refreshGroup string) (refreshMode, decrementMode, availabilityMode string) { + // Availability-only Imagine products on paid Web tiers are governed by the + // shared weekly pool. Refresh its numeric counter and also re-read the + // product group so available=false/nextAvailableAt can install an exact + // product fence that overrides weekly routing. + if effectiveMode == "weekly" { + return effectiveMode, effectiveMode, refreshGroup + } refreshMode = effectiveMode if refreshGroup != "" { refreshMode = refreshGroup } - return refreshMode, effectiveMode + return refreshMode, effectiveMode, "" } diff --git a/backend/internal/application/gateway/selector_test.go b/backend/internal/application/gateway/selector_test.go index ea2e267de..7ea3ffba0 100644 --- a/backend/internal/application/gateway/selector_test.go +++ b/backend/internal/application/gateway/selector_test.go @@ -745,9 +745,13 @@ func TestSelectorWebCatalogCapabilityStillEnforcesTier(t *testing.T) { } func TestImageQuotaFinalizationKeepsEffectiveConsumptionFence(t *testing.T) { - refreshMode, decrementMode := quotaFinalizationModes(account.QuotaModeWebImagePro, account.QuotaGroupWebImagine) - if refreshMode != account.QuotaGroupWebImagine || decrementMode != account.QuotaModeWebImagePro { - t.Fatalf("refresh=%q decrement=%q", refreshMode, decrementMode) + refreshMode, decrementMode, availabilityMode := quotaFinalizationModes(account.QuotaModeWebImagePro, account.QuotaGroupWebImagine) + if refreshMode != account.QuotaGroupWebImagine || decrementMode != account.QuotaModeWebImagePro || availabilityMode != "" { + t.Fatalf("refresh=%q decrement=%q availability=%q", refreshMode, decrementMode, availabilityMode) + } + refreshMode, decrementMode, availabilityMode = quotaFinalizationModes("weekly", account.QuotaGroupWebImagine) + if refreshMode != "weekly" || decrementMode != "weekly" || availabilityMode != account.QuotaGroupWebImagine { + t.Fatalf("shared weekly refresh=%q decrement=%q availability=%q", refreshMode, decrementMode, availabilityMode) } } diff --git a/backend/internal/application/gateway/video.go b/backend/internal/application/gateway/video.go index 15e555590..ccdbdf0d1 100644 --- a/backend/internal/application/gateway/video.go +++ b/backend/internal/application/gateway/video.go @@ -741,7 +741,7 @@ func (s *Service) runVideoJob(parent context.Context, job media.Job, route model return } s.selector.MarkSuccess(context.Background(), lease.Credential) - refreshMode, decrementMode := quotaFinalizationModes(lease.QuotaMode, quotaRefreshGroup) + refreshMode, decrementMode, availabilityMode := quotaFinalizationModes(lease.QuotaMode, quotaRefreshGroup) if decrementMode != "" && decrementMode != "weekly" { quotaCtx, quotaCancel := context.WithTimeout(context.Background(), accountStateWriteTimeout) updated, quotaErr := s.accounts.DecrementQuota(quotaCtx, job.AccountID, decrementMode, 1) @@ -757,6 +757,9 @@ func (s *Service) runVideoJob(parent context.Context, job media.Job, route model } if quotaKind, _ := s.providers.QuotaKind(route.Provider); quotaKind == provider.QuotaRemoteWindow && refreshMode != "" { s.accounts.QueueQuotaRefresh(job.AccountID, refreshMode) + if availabilityMode != "" && availabilityMode != refreshMode { + s.accounts.QueueQuotaRefresh(job.AccountID, availabilityMode) + } } // 输入回收放在账号状态、计费和审计收尾之后，存储抖动不得延迟关键终态逻辑。 s.releaseVideoInputs(job) diff --git a/backend/internal/application/gateway/video_test.go b/backend/internal/application/gateway/video_test.go index 5c607bf57..5658131c3 100644 --- a/backend/internal/application/gateway/video_test.go +++ b/backend/internal/application/gateway/video_test.go @@ -47,9 +47,13 @@ func TestVideoQuotaModeUsesWeb720pProduct(t *testing.T) { } func TestVideoQuotaFinalizationKeepsEffectiveConsumptionFence(t *testing.T) { - refreshMode, decrementMode := quotaFinalizationModes(account.QuotaModeWebVideo720p, account.QuotaGroupWebImagine) - if refreshMode != account.QuotaGroupWebImagine || decrementMode != account.QuotaModeWebVideo720p { - t.Fatalf("refresh=%q decrement=%q", refreshMode, decrementMode) + refreshMode, decrementMode, availabilityMode := quotaFinalizationModes(account.QuotaModeWebVideo720p, account.QuotaGroupWebImagine) + if refreshMode != account.QuotaGroupWebImagine || decrementMode != account.QuotaModeWebVideo720p || availabilityMode != "" { + t.Fatalf("refresh=%q decrement=%q availability=%q", refreshMode, decrementMode, availabilityMode) + } + refreshMode, decrementMode, availabilityMode = quotaFinalizationModes("weekly", account.QuotaGroupWebImagine) + if refreshMode != "weekly" || decrementMode != "weekly" || availabilityMode != account.QuotaGroupWebImagine { + t.Fatalf("shared weekly refresh=%q decrement=%q availability=%q", refreshMode, decrementMode, availabilityMode) } } diff --git a/backend/internal/infra/persistence/relational/account_repository.go b/backend/internal/infra/persistence/relational/account_repository.go index e8d8d0c51..7f49c9adf 100644 --- a/backend/internal/infra/persistence/relational/account_repository.go +++ b/backend/internal/infra/persistence/relational/account_repository.go @@ -680,11 +680,13 @@ func (r *AccountRepository) getRoutingQuotaWindows(ctx context.Context, provider if provider != account.ProviderWeb && quotaMode == "" { return result, nil } - modes := make([]string, 0, 2) - // Paid Web chat routes are governed by the shared weekly pool. Imagine - // products have independent authoritative windows and must not be hidden by - // a weekly row merely because the same account also has paid chat access. - if provider == account.ProviderWeb && !account.IsWebImagineQuotaMode(quotaMode) { + modes := make([]string, 0, 3) + webImagineMode := provider == account.ProviderWeb && account.IsWebImagineQuotaMode(quotaMode) + // Paid Web routes use the shared weekly pool. Imagine may additionally + // expose a product-specific remainingQueries window (notably for Basic and + // older response shapes); load both and prefer the exact product window + // below, falling back to weekly only for confirmed Super/Heavy accounts. + if provider == account.ProviderWeb { modes = append(modes, "weekly") } if provider == account.ProviderWeb && quotaMode == account.QuotaModeWebImageEdit { @@ -711,10 +713,23 @@ func (r *AccountRepository) getRoutingQuotaWindows(ctx context.Context, provider webTiers[credential.ID] = credential.WebTier } for _, row := range rows { - if provider == account.ProviderWeb && quotaMode == account.QuotaModeWebImageEdit { - if row.Mode != webImageEditRoutingQuotaMode(webTiers[row.AccountID]) { - continue + if webImagineMode { + tier := webTiers[row.AccountID] + productMode := quotaMode + if quotaMode == account.QuotaModeWebImageEdit { + productMode = webImageEditRoutingQuotaMode(tier) } + switch { + case row.Mode == productMode: + // An explicit product window is more precise than the shared pool, + // regardless of query order. + result[row.AccountID] = toRoutingQuotaWindowDomain(row) + case row.Mode == "weekly" && (tier == account.WebTierSuper || tier == account.WebTierHeavy): + if existing, exists := result[row.AccountID]; !exists || existing.Mode != productMode { + result[row.AccountID] = toRoutingQuotaWindowDomain(row) + } + } + continue } if _, exists := result[row.AccountID]; !exists { result[row.AccountID] = toRoutingQuotaWindowDomain(row) diff --git a/backend/internal/infra/persistence/relational/routing_projection_test.go b/backend/internal/infra/persistence/relational/routing_projection_test.go index 1b9ad68ff..7452b0553 100644 --- a/backend/internal/infra/persistence/relational/routing_projection_test.go +++ b/backend/internal/infra/persistence/relational/routing_projection_test.go @@ -130,6 +130,7 @@ func TestRoutingProjectionMapsWebImageEditQuotaByTier(t *testing.T) { t.Fatal(createErr) } if saveErr := accounts.SaveQuotaWindows(ctx, value.ID, tier, now, []account.QuotaWindow{ + {AccountID: value.ID, Mode: "weekly", Remaining: 11, SyncedAt: &now, Source: account.QuotaSourceUpstream}, {AccountID: value.ID, Mode: account.QuotaModeWebImagePro, Remaining: 3, SyncedAt: &now, Source: account.QuotaSourceUpstream}, {AccountID: value.ID, Mode: account.QuotaModeWebImageEdit, Remaining: 7, SyncedAt: &now, Source: account.QuotaSourceUpstream}, }); saveErr != nil { @@ -155,6 +156,56 @@ func TestRoutingProjectionMapsWebImageEditQuotaByTier(t *testing.T) { } } +func TestRoutingProjectionFallsBackToWeeklyForPaidWebImagine(t *testing.T) { + ctx := context.Background() + database, err := OpenSQLite(ctx, filepath.Join(t.TempDir(), "routing-web-imagine-weekly.db")) + if err != nil { + t.Fatal(err) + } + defer database.Close() + if err := database.InitializeSchema(ctx); err != nil { + t.Fatal(err) + } + accounts := NewAccountRepository(database) + now := time.Now().UTC() + create := func(name string, tier account.WebTier) account.Credential { + value, _, createErr := accounts.UpsertByIdentity(ctx, account.Credential{ + Provider: account.ProviderWeb, AuthType: account.AuthTypeSSO, Name: name, SourceKey: name, + EncryptedAccessToken: "encrypted", Enabled: true, AuthStatus: account.AuthStatusActive, WebTier: tier, + }) + if createErr != nil { + t.Fatal(createErr) + } + if saveErr := accounts.SaveQuotaWindows(ctx, value.ID, tier, now, []account.QuotaWindow{{ + AccountID: value.ID, Mode: "weekly", Remaining: 9, Total: 10, + SyncedAt: &now, Source: account.QuotaSourceUpstream, + }}); saveErr != nil { + t.Fatal(saveErr) + } + return value + } + basic := create("basic-weekly", account.WebTierBasic) + super := create("super-weekly", account.WebTierSuper) + heavy := create("heavy-weekly", account.WebTierHeavy) + + bases, err := accounts.ListRoutingAccountBases(ctx, account.ProviderWeb, account.QuotaModeWebImagePro) + if err != nil { + t.Fatal(err) + } + byID := make(map[uint64]account.RoutingAccountBase, len(bases)) + for _, base := range bases { + byID[base.Credential.ID] = base + } + if got := byID[basic.ID].QuotaWindow; got != nil { + t.Fatalf("Basic must not inherit paid weekly Imagine quota: %#v", got) + } + for _, value := range []account.Credential{super, heavy} { + if got := byID[value.ID].QuotaWindow; got == nil || got.Mode != "weekly" || got.Remaining != 9 { + t.Fatalf("paid Imagine weekly fallback for %d = %#v", value.ID, got) + } + } +} + func TestGetCredentialMaterialHydratesOneAccountAndMapsNotFound(t *testing.T) { ctx := context.Background() database, err := OpenSQLite(ctx, filepath.Join(t.TempDir(), "credential-material.db")) diff --git a/backend/internal/infra/provider/web/quota.go b/backend/internal/infra/provider/web/quota.go index 0762fd54b..c617601e3 100644 --- a/backend/internal/infra/provider/web/quota.go +++ b/backend/internal/infra/provider/web/quota.go @@ -49,13 +49,18 @@ func (a *Adapter) SyncQuota(ctx context.Context, credential account.Credential) // Basic/未知账号没有付费周池，避免为每次完整同步额外访问付费端点。 // 只有模式额度已经确认付费等级时才读取 weekly 作为权威额度。 if tier == account.WebTierSuper || tier == account.WebTierHeavy { - if weekly, weeklyErr := a.syncWeeklyCredits(ctx, credential); weeklyErr == nil { - // 周池覆盖 chat 模式窗口，但保留 imagine 窗口供前端展示与触顶判定。 - kept := make([]account.QuotaWindow, 0, 1+len(imagineSnapshot.Windows)) - kept = append(kept, weekly) - kept = append(kept, imagineSnapshot.Windows...) - windows = kept + weekly, weeklyErr := a.syncWeeklyCredits(ctx, credential) + if weeklyErr != nil { + // Paid Web routing is governed by the shared weekly pool. Returning a + // partial successful snapshot would make the application replace and + // erase the last authoritative weekly window. + return provider.QuotaSnapshot{}, weeklyErr } + // 周池覆盖 chat 模式窗口，但保留 imagine 窗口供前端展示与触顶判定。 + kept := make([]account.QuotaWindow, 0, 1+len(imagineSnapshot.Windows)) + kept = append(kept, weekly) + kept = append(kept, imagineSnapshot.Windows...) + windows = kept } if windows == nil { windows = append(chatWindows, imagineSnapshot.Windows...) @@ -182,11 +187,26 @@ func decodeImagineQuotaSnapshot(body []byte, accountID uint64, now time.Time) ([ if item.mode == "" { continue } - if *product.Available && (product.RemainingQueries == nil || product.WindowSizeSeconds == nil) { + // Paid Web tiers can use the shared weekly pool. For those accounts the + // Imagine endpoint reports only product availability and a window size, + // without an independent remainingQueries counter. Absence of that counter + // means "no product-specific window", not zero remaining quota. Omitting the + // row lets routing use the paid account's weekly window and atomically + // removes any stale per-product counter from an older response shape. + if *product.Available && product.RemainingQueries == nil { + if product.WindowSizeSeconds == nil { + return nil, fmt.Errorf("Grok Web Imagine 配额字段 %s 结构不完整", item.field) + } + if *product.WindowSizeSeconds <= 0 { + return nil, fmt.Errorf("Grok Web Imagine 配额字段 %s 的 windowSizeSeconds 无效", item.field) + } + continue + } + if *product.Available && product.WindowSizeSeconds == nil { return nil, fmt.Errorf("Grok Web Imagine 配额字段 %s 结构不完整", item.field) } remaining := 0 - if product.RemainingQueries != nil { + if *product.Available && product.RemainingQueries != nil { remaining = max(0, *product.RemainingQueries) } windowSeconds := 86400 @@ -291,6 +311,9 @@ func (a *Adapter) SyncQuotaMode(ctx context.Context, credential account.Credenti return w, nil } } + if credential.WebTier == account.WebTierSuper || credential.WebTier == account.WebTierHeavy { + return a.syncWeeklyCredits(ctx, credential) + } return account.QuotaWindow{}, fmt.Errorf("imagine 配额响应缺少 %s", mode) } cfg := a.config() diff --git a/backend/internal/infra/provider/web/quota_test.go b/backend/internal/infra/provider/web/quota_test.go index e8c36d90b..927cdafce 100644 --- a/backend/internal/infra/provider/web/quota_test.go +++ b/backend/internal/infra/provider/web/quota_test.go @@ -79,7 +79,14 @@ func TestSyncQuotaFetchesWeeklyOnlyAfterPaidTierIsConfirmed(t *testing.T) { server := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) { switch request.URL.Path { case "/rest/media/imagine/quota_info": - writeEmptyImagineQuota(writer) + writer.Header().Set("Content-Type", "application/json") + _, _ = writer.Write([]byte(`{ + "image":{"available":true,"windowSizeSeconds":64800}, + "imagePro":{"available":true,"windowSizeSeconds":64800}, + "imageEdit":{"available":true,"windowSizeSeconds":64800}, + "video":{"available":true,"windowSizeSeconds":64800}, + "video720p":{"available":true,"windowSizeSeconds":64800} + }`)) case "/grok_api_v2.GrokBuildBilling/GetGrokCreditsConfig": weeklyCalls.Add(1) writer.Header().Set("Content-Type", "application/grpc-web+proto") @@ -121,6 +128,58 @@ func TestSyncQuotaFetchesWeeklyOnlyAfterPaidTierIsConfirmed(t *testing.T) { } } +func TestSyncQuotaFailsWhenPaidWeeklySnapshotIsUnavailable(t *testing.T) { + var weeklyCalls atomic.Int64 + server := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) { + switch request.URL.Path { + case "/rest/media/imagine/quota_info": + writer.Header().Set("Content-Type", "application/json") + _, _ = writer.Write([]byte(`{ + "image":{"available":true,"windowSizeSeconds":64800}, + "imagePro":{"available":true,"windowSizeSeconds":64800}, + "imageEdit":{"available":true,"windowSizeSeconds":64800}, + "video":{"available":true,"windowSizeSeconds":64800}, + "video720p":{"available":true,"windowSizeSeconds":64800} + }`)) + case "/grok_api_v2.GrokBuildBilling/GetGrokCreditsConfig": + weeklyCalls.Add(1) + http.Error(writer, "temporary weekly failure", http.StatusServiceUnavailable) + case "/rest/rate-limits": + var payload struct { + ModelName string `json:"modelName"` + } + if err := json.NewDecoder(request.Body).Decode(&payload); err != nil { + t.Errorf("quota payload: %v", err) + } + total := map[string]int{"auto": 50, "fast": 140}[payload.ModelName] + _ = json.NewEncoder(writer).Encode(map[string]any{ + "windowSizeSeconds": 7200, "remainingQueries": total, "totalQueries": total, + }) + default: + http.NotFound(writer, request) + } + })) + defer server.Close() + + cipher, err := security.NewCipher(base64.StdEncoding.EncodeToString(make([]byte, 32))) + if err != nil { + t.Fatal(err) + } + encrypted, err := cipher.Encrypt("test-sso") + if err != nil { + t.Fatal(err) + } + adapter := NewAdapter(Config{ + BaseURL: server.URL, StatsigMode: "manual", StatsigManualValue: "test-signature", + }, infraegress.NewManager(egressRepositoryStub{}, cipher), cipher, nil, nil) + if _, err := adapter.SyncQuota(context.Background(), account.Credential{ID: 2, WebTier: account.WebTierAuto, EncryptedAccessToken: encrypted}); err == nil { + t.Fatal("expected the paid weekly failure to reject the partial snapshot") + } + if weeklyCalls.Load() != 1 { + t.Fatalf("weekly calls = %d", weeklyCalls.Load()) + } +} + func TestSyncQuotaStopsAfterFirstUnauthorizedMode(t *testing.T) { var calls atomic.Int64 server := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) { @@ -514,10 +573,41 @@ func TestDecodeImagineQuotaSnapshotAcceptsExplicitUnavailableProduct(t *testing. } } -func TestDecodeImagineQuotaSnapshotRejectsIncompleteAvailableProduct(t *testing.T) { +func TestDecodeImagineQuotaSnapshotAcceptsAvailabilityOnlySharedWeeklyProducts(t *testing.T) { + now := time.Now().UTC() + windows, err := decodeImagineQuotaSnapshot([]byte(`{ + "image":{"available":true,"windowSizeSeconds":64800}, + "imagePro":{"available":true,"windowSizeSeconds":64800}, + "imageEdit":{"available":true,"windowSizeSeconds":64800}, + "video":{"available":true,"windowSizeSeconds":64800}, + "video720p":{"available":true,"windowSizeSeconds":64800} + }`), 42, now) + if err != nil { + t.Fatal(err) + } + if len(windows) != 0 { + t.Fatalf("availability-only products must use the shared weekly pool, got %#v", windows) + } +} + +func TestDecodeImagineQuotaSnapshotRequiresWindowForIndependentCounter(t *testing.T) { + now := time.Now().UTC() + _, err := decodeImagineQuotaSnapshot([]byte(`{ + "image":null,"imageEdit":null, + "imagePro":{"available":true,"remainingQueries":2}, + "video":null,"video720p":null + }`), 42, now) + if err == nil || !strings.Contains(err.Error(), "imagePro") { + t.Fatalf("err = %v", err) + } +} + +func TestDecodeImagineQuotaSnapshotRequiresWindowForAvailabilityOnlyProduct(t *testing.T) { now := time.Now().UTC() _, err := decodeImagineQuotaSnapshot([]byte(`{ - "image":null,"imageEdit":null,"imagePro":{"available":true},"video":null,"video720p":null + "image":null,"imageEdit":null, + "imagePro":{"available":true}, + "video":null,"video720p":null }`), 42, now) if err == nil || !strings.Contains(err.Error(), "imagePro") { t.Fatalf("err = %v", err)