diff --git a/backend/internal/repository/account_repo_integration_test.go b/backend/internal/repository/account_repo_integration_test.go index ae18f0afa6..6046f3ad86 100644 --- a/backend/internal/repository/account_repo_integration_test.go +++ b/backend/internal/repository/account_repo_integration_test.go @@ -1540,6 +1540,41 @@ func (s *AccountRepoSuite) TestUpdateExtra_SchedulerNeutralSkipsOutboxAndSyncsFr s.Require().Equal("2026-03-11T10:00:00Z", cacheRecorder.accounts[account.ID].Extra["codex_usage_updated_at"]) } +// Exercise the complete UpdateExtra -> PostgreSQL -> Redis metadata -> admission +// path. A recorder-only cache would miss fields discarded by the slim projection. +func (s *AccountRepoSuite) TestUpdateExtra_AnthropicThresholdRefreshesCandidateSnapshot() { + now := time.Now().UTC().Truncate(time.Second) + end := now.Add(time.Hour) + account := mustCreateAccount(s.T(), s.client, &service.Account{ + Name: "threshold-refresh", Platform: service.PlatformAnthropic, Type: service.AccountTypeOAuth, + Credentials: map[string]any{"account_scheduling_threshold": 60}, + Extra: map[string]any{"passive_usage_7d_utilization": .59, "passive_usage_7d_reset": end.Unix()}, + }) + cache := NewSchedulerCache(testRedis(s.T())) + s.repo.schedulerCache = cache + bucket := service.SchedulerBucket{GroupID: account.ID, Platform: service.PlatformAnthropic, Mode: service.SchedulerModeSingle} + token, err := cache.CaptureBucketWriteToken(s.ctx, bucket) + s.Require().NoError(err) + s.Require().NoError(cache.SetSnapshot(s.ctx, bucket, token, []service.Account{*account})) + for _, step := range []struct { + used float64 + reset time.Time + paused bool + }{ + {.59, end, false}, {.66, end, true}, {.66, now.Add(-time.Hour), false}, {.10, end, false}, + } { + s.Require().NoError(s.repo.UpdateExtra(s.ctx, account.ID, map[string]any{ + "passive_usage_7d_utilization": step.used, "passive_usage_7d_reset": step.reset.Unix(), + })) + candidates, hit, err := cache.GetSnapshot(s.ctx, bucket) + s.Require().NoError(err) + s.Require().True(hit) + s.Require().Len(candidates, 1) + decision := service.EvaluateAccountSchedulingThreshold(candidates[0], map[string]int{service.PlatformAnthropic: 100}, now) + s.Require().Equal(step.paused, decision.ShouldPause) + } +} + func (s *AccountRepoSuite) TestUpdateExtra_ExhaustedCodexSnapshotSyncsSchedulerCache() { account := mustCreateAccount(s.T(), s.client, &service.Account{ Name: "acc-extra-codex-exhausted", diff --git a/backend/internal/repository/scheduler_cache.go b/backend/internal/repository/scheduler_cache.go index e5ff4a9d9f..02a1fcf3be 100644 --- a/backend/internal/repository/scheduler_cache.go +++ b/backend/internal/repository/scheduler_cache.go @@ -954,7 +954,9 @@ func filterSchedulerCredentials(credentials map[string]any) map[string]any { if len(credentials) == 0 { return nil } - keys := []string{"model_mapping", "compact_model_mapping", "api_key", "project_id", "oauth_type", "plan_type"} + // Candidate-list admission evaluates the account override before hydrating + // the full account. Dropping it silently falls back to the platform threshold. + keys := []string{"model_mapping", "compact_model_mapping", "api_key", "project_id", "oauth_type", "plan_type", "account_scheduling_threshold"} filtered := make(map[string]any) for _, key := range keys { if value, ok := credentials[key]; ok && value != nil { @@ -972,6 +974,13 @@ func filterSchedulerExtra(extra map[string]any) map[string]any { return nil } keys := []string{ + // Anthropic shared-window and Fable-only threshold checks run on this + // projection. UpdateExtra refreshes both payloads without a bucket rebuild. + "session_window_utilization", + "passive_usage_7d_utilization", + "passive_usage_7d_reset", + "passive_usage_7d_oi_utilization", + "passive_usage_7d_oi_reset", "quota_limit", "quota_used", "quota_daily_limit", diff --git a/backend/internal/repository/scheduler_threshold_cache_unit_test.go b/backend/internal/repository/scheduler_threshold_cache_unit_test.go new file mode 100644 index 0000000000..1c081b43b7 --- /dev/null +++ b/backend/internal/repository/scheduler_threshold_cache_unit_test.go @@ -0,0 +1,132 @@ +//go:build unit + +package repository + +import ( + "context" + "testing" + "time" + + "github.com/Wei-Shaw/sub2api/internal/config" + "github.com/Wei-Shaw/sub2api/internal/service" + "github.com/stretchr/testify/require" +) + +// Only persistence is stubbed: the candidate is serialized through the real +// Redis snapshot projection before the production threshold service sees it. +type snapshotThresholdRepo struct { + service.AccountRepository + accountPauses int + modelPauses int + until time.Time +} + +func (r *snapshotThresholdRepo) SetTempUnschedulable(_ context.Context, _ int64, until time.Time, _ string) error { + r.accountPauses++ + r.until = until + return nil +} + +func (r *snapshotThresholdRepo) SetModelRateLimit(_ context.Context, _ int64, _ string, until time.Time, _ ...string) error { + r.modelPauses++ + r.until = until + return nil +} + +func TestSchedulerCacheAnthropicThresholdAdmission(t *testing.T) { + now := time.Now().UTC().Truncate(time.Second) + reset := now.Add(time.Hour) + cases := []struct { + name string + fiveHour, sevenDay, fable float64 + threshold int + expired bool + accountPaused, fablePaused bool + }{ + {name: "shared 5h", fiveHour: .60, threshold: 60, accountPaused: true}, + {name: "shared 7d", sevenDay: .66, threshold: 60, accountPaused: true}, + {name: "Fable only", sevenDay: .40, fable: .75, threshold: 60, fablePaused: true}, + {name: "below threshold", fiveHour: .59, sevenDay: .59, fable: .59, threshold: 60}, + {name: "expired windows", fiveHour: .75, sevenDay: .75, fable: .75, threshold: 60, expired: true}, + {name: "disabled override", fiveHour: .75, sevenDay: .75, fable: .75, threshold: 100}, + } + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + ctx := context.Background() + end := reset + if tc.expired { + end = now.Add(-time.Hour) + } + account := service.Account{ + ID: 3, Platform: service.PlatformAnthropic, Type: service.AccountTypeOAuth, + Status: service.StatusActive, Schedulable: true, SessionWindowEnd: &end, + Credentials: map[string]any{"account_scheduling_threshold": tc.threshold, "access_token": "test-only-token"}, + Extra: map[string]any{ + "session_window_utilization": tc.fiveHour, + "passive_usage_7d_utilization": tc.sevenDay, "passive_usage_7d_reset": end.Unix(), + "passive_usage_7d_oi_utilization": tc.fable, "passive_usage_7d_oi_reset": end.Unix(), + "unrelated_large_payload": "drop me", + }, + } + cache := newSchedulerCacheUnit(t) + bucket := service.SchedulerBucket{GroupID: 8, Platform: service.PlatformAnthropic, Mode: service.SchedulerModeSingle} + token, err := cache.CaptureBucketWriteToken(ctx, bucket) + require.NoError(t, err) + require.NoError(t, cache.SetSnapshot(ctx, bucket, token, []service.Account{account})) + candidates, hit, err := cache.GetSnapshot(ctx, bucket) + require.NoError(t, err) + require.True(t, hit) + require.Len(t, candidates, 1) + require.NotContains(t, candidates[0].Credentials, "access_token") + require.NotContains(t, candidates[0].Extra, "unrelated_large_payload") + full, err := cache.GetAccount(ctx, account.ID) + require.NoError(t, err) + require.NotNil(t, full) + for _, candidate := range []*service.Account{candidates[0], full} { + repo := &snapshotThresholdRepo{} + limiter := service.NewRateLimitService(repo, nil, &config.Config{}, nil, nil) + // Default platform thresholds are disabled, so the per-account override + // must survive both cache read paths to trigger admission control. + limiter.SetSettingService(service.NewSettingService(nil, &config.Config{})) + require.Equal(t, tc.accountPaused, limiter.ApplyAccountSchedulingThreshold(ctx, candidate)) + require.Equal(t, tc.accountPaused, repo.accountPauses == 1) + require.Equal(t, tc.fablePaused, repo.modelPauses == 1) + if tc.accountPaused || tc.fablePaused { + require.True(t, end.Equal(repo.until)) + } + if tc.fablePaused { + require.False(t, candidate.IsSchedulableForModel("claude-fable-5")) + require.True(t, candidate.IsSchedulableForModel("claude-opus-5")) + } + } + }) + } +} + +func TestSchedulerCacheAnthropicUsageRefresh(t *testing.T) { + ctx := context.Background() + now := time.Now().UTC().Truncate(time.Second) + end := now.Add(time.Hour) + cache := newSchedulerCacheUnit(t) + bucket := service.SchedulerBucket{GroupID: 8, Platform: service.PlatformAnthropic, Mode: service.SchedulerModeSingle} + account := service.Account{ + ID: 3, Platform: service.PlatformAnthropic, Type: service.AccountTypeOAuth, + Status: service.StatusActive, Schedulable: true, + Credentials: map[string]any{"account_scheduling_threshold": 60}, + Extra: map[string]any{"passive_usage_7d_utilization": .59, "passive_usage_7d_reset": end.Unix()}, + } + token, err := cache.CaptureBucketWriteToken(ctx, bucket) + require.NoError(t, err) + require.NoError(t, cache.SetSnapshot(ctx, bucket, token, []service.Account{account})) + for _, used := range []float64{.59, .66, .10} { + account.Extra["passive_usage_7d_utilization"] = used + // UpdateExtra refreshes account payloads without rebuilding bucket membership. + require.NoError(t, cache.SetAccount(ctx, &account)) + candidates, hit, err := cache.GetSnapshot(ctx, bucket) + require.NoError(t, err) + require.True(t, hit) + require.Len(t, candidates, 1) + decision := service.EvaluateAccountSchedulingThreshold(candidates[0], map[string]int{service.PlatformAnthropic: 100}, now) + require.Equal(t, used >= .60, decision.ShouldPause) + } +}