| 1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768697071727374757677787980818283848586878889909192939495969798991001011021031041051061071081091101111121131141151161171181191201211221231241251261271281291301311321331341351361371381391401411421431441451461471481491501511521531541551561571581591601611621631641651661671681691701711721731741751761771781791801811821831841851861871881891901911921931941951961971981992002012022032042052062072082092102112122132142152162172182192202212222232242252262272282292302312322332342352362372382392402412422432442452462472482492502512522532542552562572582592602612622632642652662672682692702712722732742752762772782792802812822832842852862872882892902912922932942952962972982993003013023033043053063073083093103113123133143153163173183193203213223233243253263273283293303313323333343353363373383393403413423433443453463473483493503513523533543553563573583593603613623633643653663673683693703713723733743753763773783793803813823833843853863873883893903913923933943953963973983994004014024034044054064074084094104114124134144154164174184194204214224234244254264274284294304314324334344354364374384394404414424434444454464474484494504514524534544554564574584594604614624634644654664674684694704714724734744754764774784794804814824834844854864874884894904914924934944954964974984995005015025035045055065075085095105115125135145155165175185195205215225235245255265275285295305315325335345355365375385395405415425435445455465475485495505515525535545555565575585595605615625635645655665675685695705715725735745755765775785795805815825835845855865875885895905915925935945955965975985996006016026036046056066076086096106116126136146156166176186196206216226236246256266276286296306316326336346356366376386396406416426436446456466476486496506516526536546556566576586596606616626636646656666676686696706716726736746756766776786796806816826836846856866876886896906916926936946956966976986997007017027037047057067077087097107117127137147157167177187197207217227237247257267277287297307317327337347357367377387397407417427437447457467477487497507517527537547557567577587597607617627637647657667677687697707717727737747757767777787797807817827837847857867877887897907917927937947957967977987998008018028038048058068078088098108118128138148158168178188198208218228238248258268278288298308318328338348358368378388398408418428438448458468478488498508518528538548558568578588598608618628638648658668678688698708718728738748758768778788798808818828838848858868878888898908918928938948958968978988999009019029039049059069079089099109119129139149159169179189199209219229239249259269279289299309319329339349359369379389399409419429439449459469479489499509519529539549559569579589599609619629639649659669679689699709719729739749759769779789799809819829839849859869879889899909919929939949959969979989991000100110021003100410051006100710081009101010111012101310141015101610171018101910201021102210231024102510261027102810291030103110321033103410351036103710381039104010411042104310441045104610471048104910501051105210531054105510561057105810591060106110621063106410651066106710681069107010711072107310741075107610771078107910801081108210831084108510861087108810891090109110921093109410951096109710981099110011011102110311041105110611071108110911101111111211131114111511161117111811191120112111221123 |
- //go:build unit
- package repository
- import (
- "context"
- "encoding/hex"
- "encoding/json"
- "fmt"
- "strconv"
- "strings"
- "testing"
- "time"
- "github.com/Wei-Shaw/sub2api/internal/service"
- "github.com/alicebob/miniredis/v2"
- "github.com/redis/go-redis/v9"
- "github.com/stretchr/testify/require"
- )
- func newSchedulerCacheUnit(t *testing.T) *schedulerCache {
- cache, _ := newSchedulerCacheUnitWithRedis(t)
- return cache
- }
- func newSchedulerCacheUnitWithRedis(t *testing.T) (*schedulerCache, *miniredis.Miniredis) {
- t.Helper()
- mr := miniredis.RunT(t)
- rdb := redis.NewClient(&redis.Options{Addr: mr.Addr()})
- t.Cleanup(func() { _ = rdb.Close() })
- cache, ok := newSchedulerCacheWithChunkSizes(rdb, defaultSchedulerSnapshotMGetChunkSize, defaultSchedulerSnapshotWriteChunkSize).(*schedulerCache)
- require.True(t, ok)
- return cache, mr
- }
- func TestSchedulerCacheWriteAccountIDsSkipsUnencodableTimes(t *testing.T) {
- ctx := context.Background()
- cache := newSchedulerCacheUnit(t)
- invalidTime := time.Date(10000, time.January, 1, 0, 0, 0, 0, time.UTC)
- accountIDs, err := cache.writeAccountIDs(ctx, []service.Account{
- {ID: 111, Platform: service.PlatformOpenAI, Type: service.AccountTypeAPIKey},
- {ID: 112, Platform: service.PlatformOpenAI, Type: service.AccountTypeAPIKey, ExpiresAt: &invalidTime},
- })
- require.NoError(t, err)
- require.Equal(t, []int64{111}, accountIDs)
- cached, err := cache.GetAccount(ctx, 111)
- require.NoError(t, err)
- require.NotNil(t, cached)
- invalid, err := cache.GetAccount(ctx, 112)
- require.NoError(t, err)
- require.Nil(t, invalid)
- }
- func TestSchedulerCacheSetAccountClearsUnencodablePayload(t *testing.T) {
- ctx := context.Background()
- cache := newSchedulerCacheUnit(t)
- account := service.Account{ID: 113, Platform: service.PlatformOpenAI, Type: service.AccountTypeAPIKey}
- require.NoError(t, cache.SetAccount(ctx, &account))
- invalidTime := time.Date(10000, time.January, 1, 0, 0, 0, 0, time.UTC)
- account.ExpiresAt = &invalidTime
- require.NoError(t, cache.SetAccount(ctx, &account))
- cached, err := cache.GetAccount(ctx, account.ID)
- require.NoError(t, err)
- require.Nil(t, cached)
- }
- func TestSchedulerCacheUpdateLastUsedClearsUnencodablePayload(t *testing.T) {
- ctx := context.Background()
- cache := newSchedulerCacheUnit(t)
- account := service.Account{ID: 114, Platform: service.PlatformOpenAI, Type: service.AccountTypeAPIKey}
- require.NoError(t, cache.SetAccount(ctx, &account))
- invalidTime := time.Date(10000, time.January, 1, 0, 0, 0, 0, time.UTC)
- require.NoError(t, cache.UpdateLastUsed(ctx, map[int64]time.Time{account.ID: invalidTime}))
- cached, err := cache.GetAccount(ctx, account.ID)
- require.NoError(t, err)
- require.Nil(t, cached)
- }
- func TestSchedulerCacheSnapshotAccountIDReusePreservesPayloadAndMembers(t *testing.T) {
- ctx := context.Background()
- cache, _ := newSchedulerCacheUnitWithRedis(t)
- invalidTime := time.Date(10000, time.January, 1, 0, 0, 0, 0, time.UTC)
- validOne := service.Account{
- ID: 701,
- Name: "first",
- Platform: service.PlatformOpenAI,
- Type: service.AccountTypeOAuth,
- Credentials: map[string]any{"model_mapping": map[string]any{"z": "last", "a": "first"}},
- Extra: map[string]any{"mixed_scheduling": true},
- GroupIDs: []int64{17},
- }
- validTwo := service.Account{ID: 702, Name: "second", Platform: service.PlatformOpenAI, Type: service.AccountTypeAPIKey}
- invalid := service.Account{ID: 799, Platform: service.PlatformOpenAI, Type: service.AccountTypeAPIKey, ExpiresAt: &invalidTime}
- accounts := []service.Account{validOne, invalid, validTwo, validOne}
- single := service.SchedulerBucket{GroupID: 17, Platform: service.PlatformOpenAI, Mode: service.SchedulerModeSingle}
- singleToken, err := cache.CaptureBucketWriteToken(ctx, single)
- require.NoError(t, err)
- accountIDs, err := cache.SetSnapshotAndReturnAccountIDs(ctx, single, singleToken, accounts)
- require.NoError(t, err)
- require.Equal(t, []int64{701, 702, 701}, accountIDs, "应保留可编码账号的原顺序和重复项")
- wantFull, err := json.Marshal(validOne)
- require.NoError(t, err)
- wantMeta, err := json.Marshal(buildSchedulerMetadataAccount(validOne))
- require.NoError(t, err)
- fullBefore, err := cache.rdb.Get(ctx, schedulerAccountKey("701")).Bytes()
- require.NoError(t, err)
- metaBefore, err := cache.rdb.Get(ctx, schedulerAccountMetaKey("701")).Bytes()
- require.NoError(t, err)
- require.Equal(t, wantFull, fullBefore)
- require.Equal(t, wantMeta, metaBefore)
- forced := service.SchedulerBucket{GroupID: 17, Platform: service.PlatformOpenAI, Mode: service.SchedulerModeForced}
- forcedToken, err := cache.CaptureBucketWriteToken(ctx, forced)
- require.NoError(t, err)
- require.NoError(t, cache.SetSnapshotByAccountIDs(ctx, forced, forcedToken, accountIDs))
- fullAfter, err := cache.rdb.Get(ctx, schedulerAccountKey("701")).Bytes()
- require.NoError(t, err)
- metaAfter, err := cache.rdb.Get(ctx, schedulerAccountMetaKey("701")).Bytes()
- require.NoError(t, err)
- require.Equal(t, fullBefore, fullAfter, "ID-only 路径不得重写完整账号键")
- require.Equal(t, metaBefore, metaAfter, "ID-only 路径不得重写调度元数据键")
- for _, bucket := range []service.SchedulerBucket{single, forced} {
- version, err := cache.rdb.Get(ctx, schedulerBucketKey(schedulerActivePrefix, bucket)).Result()
- require.NoError(t, err)
- members, err := cache.rdb.ZRange(ctx, schedulerSnapshotKey(bucket, version), 0, -1).Result()
- require.NoError(t, err)
- require.Equal(t, []string{"702", "701"}, members, bucket.String())
- }
- missing, err := cache.GetAccount(ctx, invalid.ID)
- require.NoError(t, err)
- require.Nil(t, missing)
- }
- func TestSchedulerCacheSetSnapshotMatchesIDPublishing(t *testing.T) {
- ctx := context.Background()
- cache, _ := newSchedulerCacheUnitWithRedis(t)
- invalidTime := time.Date(10000, time.January, 1, 0, 0, 0, 0, time.UTC)
- validOne := service.Account{
- ID: 721,
- Name: "first",
- Platform: service.PlatformOpenAI,
- Type: service.AccountTypeOAuth,
- Credentials: map[string]any{"model_mapping": map[string]any{"source": "target"}},
- Extra: map[string]any{"mixed_scheduling": true},
- GroupIDs: []int64{21},
- }
- validTwo := service.Account{ID: 722, Name: "second", Platform: service.PlatformOpenAI, Type: service.AccountTypeAPIKey}
- invalid := service.Account{ID: 799, Platform: service.PlatformOpenAI, Type: service.AccountTypeAPIKey, ExpiresAt: &invalidTime}
- accounts := []service.Account{validOne, invalid, validTwo, validOne}
- normal := service.SchedulerBucket{GroupID: 21, Platform: service.PlatformOpenAI, Mode: service.SchedulerModeSingle}
- normalToken, err := cache.CaptureBucketWriteToken(ctx, normal)
- require.NoError(t, err)
- require.NoError(t, cache.SetSnapshot(ctx, normal, normalToken, accounts))
- fullBefore, err := cache.rdb.Get(ctx, schedulerAccountKey("721")).Bytes()
- require.NoError(t, err)
- metaBefore, err := cache.rdb.Get(ctx, schedulerAccountMetaKey("721")).Bytes()
- require.NoError(t, err)
- idOnly := service.SchedulerBucket{GroupID: 21, Platform: service.PlatformOpenAI, Mode: service.SchedulerModeForced}
- idOnlyToken, err := cache.CaptureBucketWriteToken(ctx, idOnly)
- require.NoError(t, err)
- accountIDs, err := cache.SetSnapshotAndReturnAccountIDs(ctx, idOnly, idOnlyToken, accounts)
- require.NoError(t, err)
- require.Equal(t, []int64{721, 722, 721}, accountIDs)
- fullAfter, err := cache.rdb.Get(ctx, schedulerAccountKey("721")).Bytes()
- require.NoError(t, err)
- metaAfter, err := cache.rdb.Get(ctx, schedulerAccountMetaKey("721")).Bytes()
- require.NoError(t, err)
- require.Equal(t, fullBefore, fullAfter, "普通快照和 ID 发布必须写入相同完整账号 payload")
- require.Equal(t, metaBefore, metaAfter, "普通快照和 ID 发布必须写入相同元数据 payload")
- for _, bucket := range []service.SchedulerBucket{normal, idOnly} {
- version, err := cache.rdb.Get(ctx, schedulerBucketKey(schedulerActivePrefix, bucket)).Result()
- require.NoError(t, err)
- members, err := cache.rdb.ZRange(ctx, schedulerSnapshotKey(bucket, version), 0, -1).Result()
- require.NoError(t, err)
- require.Equal(t, []string{"722", "721"}, members, bucket.String())
- }
- }
- func TestSchedulerCacheSnapshotAccountIDReuseKeepsEmptySnapshotSemantics(t *testing.T) {
- ctx := context.Background()
- cache := newSchedulerCacheUnit(t)
- invalidTime := time.Date(10000, time.January, 1, 0, 0, 0, 0, time.UTC)
- accounts := []service.Account{{ID: 811, Platform: service.PlatformOpenAI, ExpiresAt: &invalidTime}}
- single := service.SchedulerBucket{GroupID: 18, Platform: service.PlatformOpenAI, Mode: service.SchedulerModeSingle}
- singleToken, err := cache.CaptureBucketWriteToken(ctx, single)
- require.NoError(t, err)
- accountIDs, err := cache.SetSnapshotAndReturnAccountIDs(ctx, single, singleToken, accounts)
- require.NoError(t, err)
- require.Empty(t, accountIDs)
- forced := service.SchedulerBucket{GroupID: 18, Platform: service.PlatformOpenAI, Mode: service.SchedulerModeForced}
- forcedToken, err := cache.CaptureBucketWriteToken(ctx, forced)
- require.NoError(t, err)
- require.NoError(t, cache.SetSnapshotByAccountIDs(ctx, forced, forcedToken, accountIDs))
- for _, bucket := range []service.SchedulerBucket{single, forced} {
- ready, err := cache.rdb.Get(ctx, schedulerBucketKey(schedulerReadyPrefix, bucket)).Result()
- require.NoError(t, err)
- require.Equal(t, "1", ready)
- snapshot, hit, err := cache.GetSnapshot(ctx, bucket)
- require.NoError(t, err)
- require.False(t, hit, bucket.String())
- require.Nil(t, snapshot)
- }
- }
- func TestSchedulerCacheSetSnapshotByAccountIDsKeepsFencing(t *testing.T) {
- ctx := context.Background()
- cache := newSchedulerCacheUnit(t)
- bucket := service.SchedulerBucket{GroupID: 19, Platform: service.PlatformOpenAI, Mode: service.SchedulerModeForced}
- err := cache.SetSnapshotByAccountIDs(ctx, bucket, service.SchedulerBucketWriteToken{}, []int64{901})
- require.ErrorIs(t, err, service.ErrSchedulerBucketWriteFenced)
- _, err = cache.rdb.Get(ctx, schedulerBucketKey(schedulerVersionPrefix, bucket)).Result()
- require.ErrorIs(t, err, redis.Nil)
- token, err := cache.CaptureBucketWriteToken(ctx, bucket)
- require.NoError(t, err)
- require.NoError(t, cache.RetireBucket(ctx, bucket))
- err = cache.SetSnapshotByAccountIDs(ctx, bucket, token, []int64{901})
- require.ErrorIs(t, err, service.ErrSchedulerBucketRetired)
- }
- func TestSchedulerCacheSetSnapshotByAccountIDsDoesNotResurrectDeletedAccount(t *testing.T) {
- ctx := context.Background()
- cache := newSchedulerCacheUnit(t)
- account := service.Account{ID: 902, Platform: service.PlatformOpenAI, Type: service.AccountTypeOAuth}
- single := service.SchedulerBucket{GroupID: 20, Platform: service.PlatformOpenAI, Mode: service.SchedulerModeSingle}
- singleToken, err := cache.CaptureBucketWriteToken(ctx, single)
- require.NoError(t, err)
- accountIDs, err := cache.SetSnapshotAndReturnAccountIDs(ctx, single, singleToken, []service.Account{account})
- require.NoError(t, err)
- require.Equal(t, []int64{account.ID}, accountIDs)
- require.NoError(t, cache.DeleteAccount(ctx, account.ID))
- forced := service.SchedulerBucket{GroupID: 20, Platform: service.PlatformOpenAI, Mode: service.SchedulerModeForced}
- forcedToken, err := cache.CaptureBucketWriteToken(ctx, forced)
- require.NoError(t, err)
- require.NoError(t, cache.SetSnapshotByAccountIDs(ctx, forced, forcedToken, accountIDs))
- full, err := cache.GetAccount(ctx, account.ID)
- require.NoError(t, err)
- require.Nil(t, full, "ID-only 发布不得复活已删除的完整账号键")
- snapshot, hit, err := cache.GetSnapshot(ctx, forced)
- require.NoError(t, err)
- require.False(t, hit, "元数据缺失时必须安全回源,而不是返回残缺快照")
- require.Nil(t, snapshot)
- }
- func TestMarshalSchedulerCacheAccountKeepsEncodingJSONWireFormat(t *testing.T) {
- cases := []struct {
- name string
- account service.Account
- }{
- {name: "nil collections", account: service.Account{ID: 801}},
- {name: "empty collections", account: service.Account{
- ID: 802,
- Credentials: map[string]any{},
- Extra: map[string]any{},
- GroupIDs: []int64{},
- Groups: []*service.Group{},
- }},
- {name: "nested maps and escaping", account: service.Account{
- ID: 803,
- Credentials: map[string]any{"model_mapping": map[string]any{"z": "<last>", "a": "&first"}},
- Extra: map[string]any{"mixed_scheduling": true},
- }},
- }
- for _, tc := range cases {
- t.Run(tc.name, func(t *testing.T) {
- full, meta, err := marshalSchedulerCacheAccount(tc.account)
- require.NoError(t, err)
- wantFull, err := json.Marshal(tc.account)
- require.NoError(t, err)
- wantMeta, err := json.Marshal(buildSchedulerMetadataAccount(tc.account))
- require.NoError(t, err)
- require.Equal(t, wantFull, full)
- require.Equal(t, wantMeta, meta)
- })
- }
- }
- func TestBuildSchedulerMetadataAccount_KeepsOpenAIWSFlags(t *testing.T) {
- account := service.Account{
- ID: 42,
- Platform: service.PlatformOpenAI,
- Type: service.AccountTypeOAuth,
- Extra: map[string]any{
- "openai_oauth_responses_websockets_v2_enabled": true,
- "openai_oauth_responses_websockets_v2_mode": service.OpenAIWSIngressModePassthrough,
- "openai_ws_force_http": true,
- "openai_responses_mode": "force_chat_completions",
- "openai_responses_supported": false,
- "codex_fingerprint_mode": "session",
- "codex_fingerprint_seed": "11111111-1111-4111-8111-111111111111",
- "mixed_scheduling": true,
- "unused_large_field": "drop-me",
- },
- }
- got := buildSchedulerMetadataAccount(account)
- require.Equal(t, true, got.Extra["openai_oauth_responses_websockets_v2_enabled"])
- require.Equal(t, service.OpenAIWSIngressModePassthrough, got.Extra["openai_oauth_responses_websockets_v2_mode"])
- require.Equal(t, true, got.Extra["openai_ws_force_http"])
- require.Equal(t, "force_chat_completions", got.Extra["openai_responses_mode"])
- require.Equal(t, false, got.Extra["openai_responses_supported"])
- require.Equal(t, "session", got.Extra["codex_fingerprint_mode"])
- require.Equal(t, "11111111-1111-4111-8111-111111111111", got.Extra["codex_fingerprint_seed"])
- require.Equal(t, true, got.Extra["mixed_scheduling"])
- require.Nil(t, got.Extra["unused_large_field"])
- }
- func TestBuildSchedulerMetadataAccount_KeepsGrokMediaEligibility(t *testing.T) {
- t.Run("explicit override", func(t *testing.T) {
- account := service.Account{
- ID: 43,
- Platform: service.PlatformGrok,
- Type: service.AccountTypeOAuth,
- Extra: map[string]any{
- service.GrokMediaEligibleExtraKey: false,
- "unused_large_field": "drop-me",
- },
- }
- got := buildSchedulerMetadataAccount(account)
- eligible, reason := got.GrokMediaGenerationEligibility()
- require.False(t, eligible)
- require.Equal(t, "override_disabled", reason)
- require.Equal(t, false, got.Extra[service.GrokMediaEligibleExtraKey])
- require.Nil(t, got.Extra["unused_large_field"])
- })
- t.Run("forbidden billing observation", func(t *testing.T) {
- account := service.Account{
- ID: 44,
- Platform: service.PlatformGrok,
- Type: service.AccountTypeOAuth,
- Extra: map[string]any{
- "grok_billing_snapshot": map[string]any{
- "status_code": 200,
- "weekly_status_code": 403,
- "monthly_status_code": 200,
- },
- },
- }
- got := buildSchedulerMetadataAccount(account)
- eligible, reason := got.GrokMediaGenerationEligibility()
- require.False(t, eligible)
- require.Equal(t, "billing_forbidden", reason)
- require.NotNil(t, got.Extra["grok_billing_snapshot"])
- })
- }
- func TestBuildSchedulerMetadataAccount_KeepsSlimGroupMembership(t *testing.T) {
- account := service.Account{
- ID: 42,
- Platform: service.PlatformAnthropic,
- GroupIDs: []int64{7, 9, 7, 0},
- AccountGroups: []service.AccountGroup{
- {
- AccountID: 42,
- GroupID: 7,
- Priority: 2,
- Account: &service.Account{ID: 42, Name: "drop-from-metadata"},
- Group: &service.Group{ID: 7, Name: "drop-from-metadata"},
- },
- {
- AccountID: 42,
- GroupID: 11,
- Priority: 3,
- Group: &service.Group{ID: 11, Name: "drop-from-metadata"},
- },
- {
- AccountID: 42,
- GroupID: 0,
- Priority: 4,
- },
- },
- }
- got := buildSchedulerMetadataAccount(account)
- require.Equal(t, []int64{7, 9, 11}, got.GroupIDs)
- require.Len(t, got.AccountGroups, 2)
- require.Equal(t, int64(42), got.AccountGroups[0].AccountID)
- require.Equal(t, int64(7), got.AccountGroups[0].GroupID)
- require.Equal(t, 2, got.AccountGroups[0].Priority)
- require.Nil(t, got.AccountGroups[0].Account)
- require.Nil(t, got.AccountGroups[0].Group)
- require.Equal(t, int64(11), got.AccountGroups[1].GroupID)
- require.Nil(t, got.Groups)
- }
- func TestBuildSchedulerMetadataAccount_KeepsQuotaAutoPauseFields(t *testing.T) {
- account := service.Account{
- ID: 88,
- Extra: map[string]any{
- "codex_5h_used_percent": 12.34,
- "codex_7d_used_percent": 56.78,
- "codex_5h_reset_at": "2026-05-29T10:00:00Z",
- "codex_7d_reset_at": "2026-06-01T10:00:00Z",
- "codex_5h_reset_after_seconds": 300,
- "codex_7d_reset_after_seconds": 600,
- "codex_usage_updated_at": "2026-05-29T09:00:00Z",
- "auto_pause_5h_threshold": 0.95,
- "auto_pause_7d_threshold": 0.96,
- "auto_pause_5h_disabled": true,
- "auto_pause_7d_disabled": false,
- },
- }
- got := buildSchedulerMetadataAccount(account)
- require.Equal(t, 12.34, got.Extra["codex_5h_used_percent"])
- require.Equal(t, 56.78, got.Extra["codex_7d_used_percent"])
- require.Equal(t, "2026-05-29T10:00:00Z", got.Extra["codex_5h_reset_at"])
- require.Equal(t, "2026-06-01T10:00:00Z", got.Extra["codex_7d_reset_at"])
- require.Equal(t, 300, got.Extra["codex_5h_reset_after_seconds"])
- require.Equal(t, 600, got.Extra["codex_7d_reset_after_seconds"])
- require.Equal(t, "2026-05-29T09:00:00Z", got.Extra["codex_usage_updated_at"])
- require.Equal(t, 0.95, got.Extra["auto_pause_5h_threshold"])
- require.Equal(t, 0.96, got.Extra["auto_pause_7d_threshold"])
- require.Equal(t, true, got.Extra["auto_pause_5h_disabled"])
- require.Equal(t, false, got.Extra["auto_pause_7d_disabled"])
- }
- func TestBuildSchedulerMetadataAccount_KeepsQuotaStateForCachedAccounts(t *testing.T) {
- now := time.Now().UTC()
- activeStart := now.Add(-time.Hour).Format(time.RFC3339)
- expiredDailyStart := now.Add(-25 * time.Hour).Format(time.RFC3339)
- expiredWeeklyStart := now.Add(-8 * 24 * time.Hour).Format(time.RFC3339)
- weeklyResetDay := float64(now.AddDate(0, 0, 1).Weekday())
- cases := []struct {
- name string
- platform string
- typ string
- extra map[string]any
- quotaExceeded bool
- }{
- {
- name: "anthropic api key total quota exhausted", platform: service.PlatformAnthropic, typ: service.AccountTypeAPIKey,
- extra: map[string]any{"quota_limit": 10.0, "quota_used": 10.0}, quotaExceeded: true,
- },
- {
- name: "gemini api key rolling daily quota exhausted", platform: service.PlatformGemini, typ: service.AccountTypeAPIKey,
- extra: map[string]any{
- "quota_daily_limit": 20.0, "quota_daily_used": 20.0,
- "quota_daily_start": activeStart, "quota_daily_reset_mode": "rolling",
- }, quotaExceeded: true,
- },
- {
- name: "gemini api key expired rolling daily window", platform: service.PlatformGemini, typ: service.AccountTypeAPIKey,
- extra: map[string]any{
- "quota_daily_limit": 20.0, "quota_daily_used": 20.0,
- "quota_daily_start": expiredDailyStart, "quota_daily_reset_mode": "rolling",
- },
- },
- {
- name: "bedrock fixed weekly quota exhausted", platform: service.PlatformAnthropic, typ: service.AccountTypeBedrock,
- extra: map[string]any{
- "quota_weekly_limit": 30.0, "quota_weekly_used": 30.0, "quota_weekly_start": activeStart,
- "quota_weekly_reset_mode": "fixed", "quota_weekly_reset_day": weeklyResetDay,
- "quota_weekly_reset_hour": 0.0, "quota_reset_timezone": "UTC",
- }, quotaExceeded: true,
- },
- {
- name: "bedrock expired fixed weekly window", platform: service.PlatformAnthropic, typ: service.AccountTypeBedrock,
- extra: map[string]any{
- "quota_weekly_limit": 30.0, "quota_weekly_used": 30.0, "quota_weekly_start": expiredWeeklyStart,
- "quota_weekly_reset_mode": "fixed", "quota_weekly_reset_day": weeklyResetDay,
- "quota_weekly_reset_hour": 0.0, "quota_reset_timezone": "UTC",
- },
- },
- }
- for i, tc := range cases {
- t.Run(tc.name, func(t *testing.T) {
- extra := make(map[string]any, len(tc.extra)+1)
- for key, value := range tc.extra {
- extra[key] = value
- }
- extra["unrelated"] = "drop me"
- account := service.Account{
- ID: int64(46690 + i), Platform: tc.platform, Type: tc.typ, Extra: extra,
- Status: service.StatusActive, Schedulable: true,
- }
- cache := newSchedulerCacheUnit(t)
- ctx := context.Background()
- bucket := service.SchedulerBucket{GroupID: int64(46690 + i), Platform: tc.platform, Mode: service.SchedulerModeSingle}
- token, err := cache.CaptureBucketWriteToken(ctx, bucket)
- require.NoError(t, err)
- require.NoError(t, cache.SetSnapshot(ctx, bucket, token, []service.Account{account}))
- snapshot, hit, err := cache.GetSnapshot(ctx, bucket)
- require.NoError(t, err)
- require.True(t, hit)
- require.Len(t, snapshot, 1)
- cached := snapshot[0]
- require.Equal(t, tc.extra, cached.Extra)
- require.NotContains(t, cached.Extra, "unrelated")
- require.Equal(t, tc.quotaExceeded, cached.IsQuotaExceeded())
- require.Equal(t, !tc.quotaExceeded, cached.IsSchedulable())
- })
- }
- }
- func TestBuildSchedulerMetadataAccount_KeepsModelRateLimits(t *testing.T) {
- account := service.Account{
- ID: 90,
- Platform: service.PlatformAntigravity,
- Extra: map[string]any{
- "model_rate_limits": map[string]any{
- "gemini-3-flash": map[string]any{
- "rate_limit_reset_at": "2026-05-30T10:10:00Z",
- },
- "antigravity:gemini": map[string]any{
- "rate_limit_reset_at": "2026-05-30T10:10:00Z",
- },
- },
- "unused_large_field": "drop-me",
- },
- }
- got := buildSchedulerMetadataAccount(account)
- limits, ok := got.Extra["model_rate_limits"].(map[string]any)
- require.True(t, ok)
- require.Contains(t, limits, "gemini-3-flash")
- require.Contains(t, limits, "antigravity:gemini")
- require.Nil(t, got.Extra["unused_large_field"])
- }
- func TestBuildSchedulerMetadataAccount_KeepsSparkShadowRoutingIdentity(t *testing.T) {
- parentID := int64(100)
- account := service.Account{
- ID: 200,
- Platform: service.PlatformOpenAI,
- Type: service.AccountTypeOAuth,
- ParentAccountID: &parentID,
- QuotaDimension: service.QuotaDimensionSpark,
- Credentials: map[string]any{
- "model_mapping": map[string]any{
- "gpt-5.3-codex-spark": "gpt-5.3-codex-spark",
- },
- "compact_model_mapping": map[string]any{
- "gpt-5.4": "gpt-5.4-openai-compact",
- },
- "access_token": "drop-me",
- },
- }
- got := buildSchedulerMetadataAccount(account)
- require.NotNil(t, got.ParentAccountID)
- require.Equal(t, parentID, *got.ParentAccountID)
- require.Equal(t, service.QuotaDimensionSpark, got.QuotaDimension)
- require.Equal(t, map[string]any{"gpt-5.3-codex-spark": "gpt-5.3-codex-spark"}, got.Credentials["model_mapping"])
- require.Equal(t, map[string]any{"gpt-5.4": "gpt-5.4-openai-compact"}, got.Credentials["compact_model_mapping"])
- require.Nil(t, got.Credentials["access_token"])
- }
- func TestSchedulerCacheBucketRetirementFencesWritersAndReopen(t *testing.T) {
- ctx := context.Background()
- cache, mr := newSchedulerCacheUnitWithRedis(t)
- bucket := service.SchedulerBucket{GroupID: 41, Platform: service.PlatformOpenAI, Mode: service.SchedulerModeSingle}
- otherBucket := service.SchedulerBucket{GroupID: 42, Platform: service.PlatformOpenAI, Mode: service.SchedulerModeSingle}
- account := service.Account{ID: 4101, Platform: service.PlatformOpenAI, Type: service.AccountTypeAPIKey}
- token, err := cache.CaptureBucketWriteToken(ctx, bucket)
- require.NoError(t, err)
- require.True(t, token.ValidFor(bucket))
- require.NoError(t, cache.SetSnapshot(ctx, bucket, token, []service.Account{account}))
- // A token is bound to the full bucket identity, not just an epoch number.
- err = cache.SetSnapshot(ctx, otherBucket, token, []service.Account{account})
- require.ErrorIs(t, err, service.ErrSchedulerBucketWriteFenced)
- _, err = cache.rdb.Get(ctx, schedulerBucketKey(schedulerVersionPrefix, otherBucket)).Result()
- require.ErrorIs(t, err, redis.Nil)
- otherAccount := service.Account{ID: 4201, Platform: service.PlatformOpenAI, Type: service.AccountTypeAPIKey}
- otherToken, err := cache.CaptureBucketWriteToken(ctx, otherBucket)
- require.NoError(t, err)
- require.NoError(t, cache.SetSnapshot(ctx, otherBucket, otherToken, []service.Account{otherAccount}))
- otherEpoch := otherToken.Epoch
- activeVersion, err := cache.rdb.Get(ctx, schedulerBucketKey(schedulerActivePrefix, bucket)).Result()
- require.NoError(t, err)
- require.NoError(t, cache.RetireBucket(ctx, bucket))
- retiredEpoch, err := cache.rdb.Get(ctx, schedulerBucketKey(schedulerEpochPrefix, bucket)).Int64()
- require.NoError(t, err)
- require.Greater(t, retiredEpoch, token.Epoch)
- // Retirement is idempotent and does not advance the epoch again.
- require.NoError(t, cache.RetireBucket(ctx, bucket))
- retiredEpochAgain, err := cache.rdb.Get(ctx, schedulerBucketKey(schedulerEpochPrefix, bucket)).Int64()
- require.NoError(t, err)
- require.Equal(t, retiredEpoch, retiredEpochAgain)
- // New readers miss because ready/active were removed atomically. A reader that
- // captured activeVersion before retirement may still finish against that version.
- _, hit, err := cache.GetSnapshot(ctx, bucket)
- require.NoError(t, err)
- require.False(t, hit)
- ids, err := cache.rdb.ZRange(ctx, schedulerSnapshotKey(bucket, activeVersion), 0, -1).Result()
- require.NoError(t, err)
- require.Equal(t, []string{"4101"}, ids)
- ttl, err := cache.rdb.TTL(ctx, schedulerSnapshotKey(bucket, activeVersion)).Result()
- require.NoError(t, err)
- require.Positive(t, ttl)
- require.LessOrEqual(t, ttl, time.Duration(snapshotGraceTTLSeconds)*time.Second)
- buckets, err := cache.ListBuckets(ctx)
- require.NoError(t, err)
- require.NotContains(t, buckets, bucket)
- require.Contains(t, buckets, otherBucket)
- otherSnapshot, otherHit, err := cache.GetSnapshot(ctx, otherBucket)
- require.NoError(t, err)
- require.True(t, otherHit)
- require.Len(t, otherSnapshot, 1)
- require.Equal(t, otherAccount.ID, otherSnapshot[0].ID)
- otherEpochAfter, err := cache.rdb.Get(ctx, schedulerBucketKey(schedulerEpochPrefix, otherBucket)).Int64()
- require.NoError(t, err)
- require.Equal(t, otherEpoch, otherEpochAfter)
- _, err = cache.CaptureBucketWriteToken(ctx, bucket)
- require.ErrorIs(t, err, service.ErrSchedulerBucketRetired)
- versionBeforeRejectedWrite, err := cache.rdb.Get(ctx, schedulerBucketKey(schedulerVersionPrefix, bucket)).Int64()
- require.NoError(t, err)
- err = cache.SetSnapshot(ctx, bucket, token, []service.Account{account})
- require.ErrorIs(t, err, service.ErrSchedulerBucketRetired)
- versionAfterRejectedWrite, err := cache.rdb.Get(ctx, schedulerBucketKey(schedulerVersionPrefix, bucket)).Int64()
- require.NoError(t, err)
- require.Equal(t, versionBeforeRejectedWrite, versionAfterRejectedWrite, "fenced writers must not allocate a new version")
- retired, err := cache.rdb.Exists(ctx, schedulerBucketKey(schedulerRetiredPrefix, bucket)).Result()
- require.NoError(t, err)
- require.EqualValues(t, 1, retired, "ordinary writers must never clear the tombstone")
- mr.FastForward(time.Duration(snapshotGraceTTLSeconds+1) * time.Second)
- exists, err := cache.rdb.Exists(ctx, schedulerSnapshotKey(bucket, activeVersion)).Result()
- require.NoError(t, err)
- require.Zero(t, exists, "retired active snapshot must expire after the in-flight grace period")
- newToken, err := cache.ReopenBucket(ctx, bucket)
- require.NoError(t, err)
- require.True(t, newToken.ValidFor(bucket))
- require.Equal(t, retiredEpoch, newToken.Epoch)
- reopenedAgain, err := cache.ReopenBucket(ctx, bucket)
- require.NoError(t, err)
- require.Equal(t, newToken, reopenedAgain, "reopen must be idempotent within one retirement generation")
- err = cache.SetSnapshot(ctx, bucket, token, []service.Account{account})
- require.ErrorIs(t, err, service.ErrSchedulerBucketWriteFenced)
- require.NoError(t, cache.SetSnapshot(ctx, bucket, newToken, []service.Account{account}))
- reopenedWhileOpen, err := cache.ReopenBucket(ctx, bucket)
- require.NoError(t, err)
- require.Equal(t, newToken, reopenedWhileOpen)
- snapshot, hit, err := cache.GetSnapshot(ctx, bucket)
- require.NoError(t, err)
- require.True(t, hit)
- require.Len(t, snapshot, 1)
- require.Equal(t, account.ID, snapshot[0].ID)
- }
- func TestSchedulerCacheActivationIsFencedAfterRetire(t *testing.T) {
- ctx := context.Background()
- cache := newSchedulerCacheUnit(t)
- bucket := service.SchedulerBucket{GroupID: 51, Platform: service.PlatformAnthropic, Mode: service.SchedulerModeMixed}
- account := service.Account{ID: 5101, Platform: service.PlatformAnthropic, Type: service.AccountTypeAPIKey}
- token, err := cache.CaptureBucketWriteToken(ctx, bucket)
- require.NoError(t, err)
- version, err := cache.allocateSnapshotVersion(ctx, bucket, token)
- require.NoError(t, err)
- _, err = cache.writeSnapshotVersionAndReturnAccountIDs(ctx, bucket, version, []service.Account{account})
- require.NoError(t, err)
- // Deterministic race C: retirement and authoritative reopen both happen after
- // INCR/write but before the old writer activates.
- require.NoError(t, cache.RetireBucket(ctx, bucket))
- _, err = cache.ReopenBucket(ctx, bucket)
- require.NoError(t, err)
- err = cache.activateSnapshotVersion(ctx, bucket, token, version)
- require.ErrorIs(t, err, service.ErrSchedulerBucketWriteFenced)
- exists, err := cache.rdb.Exists(ctx, schedulerSnapshotKey(bucket, version)).Result()
- require.NoError(t, err)
- require.Zero(t, exists, "fenced activation must delete its unpublished snapshot")
- exists, err = cache.rdb.Exists(
- ctx,
- schedulerBucketKey(schedulerReadyPrefix, bucket),
- schedulerBucketKey(schedulerActivePrefix, bucket),
- ).Result()
- require.NoError(t, err)
- require.Zero(t, exists)
- buckets, err := cache.ListBuckets(ctx)
- require.NoError(t, err)
- require.NotContains(t, buckets, bucket)
- }
- func TestSchedulerCacheConcurrentReopenReturnsSameToken(t *testing.T) {
- ctx := context.Background()
- cache := newSchedulerCacheUnit(t)
- bucket := service.SchedulerBucket{GroupID: 53, Platform: service.PlatformOpenAI, Mode: service.SchedulerModeForced}
- account := service.Account{ID: 5301, Platform: service.PlatformOpenAI, Type: service.AccountTypeAPIKey}
- oldToken, err := cache.CaptureBucketWriteToken(ctx, bucket)
- require.NoError(t, err)
- require.NoError(t, cache.RetireBucket(ctx, bucket))
- type reopenResult struct {
- token service.SchedulerBucketWriteToken
- err error
- }
- start := make(chan struct{})
- results := make(chan reopenResult, 2)
- for range 2 {
- go func() {
- <-start
- token, err := cache.ReopenBucket(ctx, bucket)
- results <- reopenResult{token: token, err: err}
- }()
- }
- close(start)
- first := <-results
- second := <-results
- require.NoError(t, first.err)
- require.NoError(t, second.err)
- require.Equal(t, first.token, second.token)
- require.Greater(t, first.token.Epoch, oldToken.Epoch)
- require.ErrorIs(t, cache.SetSnapshot(ctx, bucket, oldToken, []service.Account{account}), service.ErrSchedulerBucketWriteFenced)
- require.NoError(t, cache.SetSnapshot(ctx, bucket, first.token, []service.Account{account}))
- }
- func TestSchedulerCacheReopenExpiresPreviousActiveSnapshot(t *testing.T) {
- ctx := context.Background()
- cache, mr := newSchedulerCacheUnitWithRedis(t)
- bucket := service.SchedulerBucket{GroupID: 52, Platform: service.PlatformGemini, Mode: service.SchedulerModeForced}
- account := service.Account{ID: 5201, Platform: service.PlatformGemini, Type: service.AccountTypeAPIKey}
- oldToken, err := cache.CaptureBucketWriteToken(ctx, bucket)
- require.NoError(t, err)
- require.NoError(t, cache.SetSnapshot(ctx, bucket, oldToken, []service.Account{account}))
- oldVersion, err := cache.rdb.Get(ctx, schedulerBucketKey(schedulerActivePrefix, bucket)).Result()
- require.NoError(t, err)
- retiredEpoch := oldToken.Epoch + 1
- require.NoError(t, cache.rdb.Set(ctx, schedulerBucketKey(schedulerEpochPrefix, bucket), retiredEpoch, 0).Err())
- require.NoError(t, cache.rdb.Set(ctx, schedulerBucketKey(schedulerRetiredPrefix, bucket), retiredEpoch, 0).Err())
- newToken, err := cache.ReopenBucket(ctx, bucket)
- require.NoError(t, err)
- require.Equal(t, retiredEpoch, newToken.Epoch)
- _, hit, err := cache.GetSnapshot(ctx, bucket)
- require.NoError(t, err)
- require.False(t, hit)
- ttl, err := cache.rdb.TTL(ctx, schedulerSnapshotKey(bucket, oldVersion)).Result()
- require.NoError(t, err)
- require.Positive(t, ttl)
- require.LessOrEqual(t, ttl, time.Duration(snapshotGraceTTLSeconds)*time.Second)
- require.ErrorIs(t, cache.SetSnapshot(ctx, bucket, oldToken, []service.Account{account}), service.ErrSchedulerBucketWriteFenced)
- mr.FastForward(time.Duration(snapshotGraceTTLSeconds+1) * time.Second)
- exists, err := cache.rdb.Exists(ctx, schedulerSnapshotKey(bucket, oldVersion)).Result()
- require.NoError(t, err)
- require.Zero(t, exists)
- require.NoError(t, cache.SetSnapshot(ctx, bucket, newToken, []service.Account{account}))
- }
- func TestSchedulerCacheGroupLifecycleLeaseConcurrentAcquireSingleOwner(t *testing.T) {
- ctx := context.Background()
- cache := newSchedulerCacheUnit(t)
- const groupID int64 = 71
- type result struct {
- lease service.SchedulerGroupLifecycleLease
- acquired bool
- err error
- }
- start := make(chan struct{})
- results := make(chan result, 32)
- for range 32 {
- go func() {
- <-start
- lease, acquired, err := cache.TryAcquireGroupLifecycleLease(ctx, groupID, time.Minute)
- results <- result{lease: lease, acquired: acquired, err: err}
- }()
- }
- close(start)
- var owner service.SchedulerGroupLifecycleLease
- acquiredCount := 0
- for range 32 {
- got := <-results
- require.NoError(t, got.err)
- if got.acquired {
- acquiredCount++
- owner = got.lease
- require.True(t, got.lease.ValidFor(groupID))
- } else {
- require.Equal(t, service.SchedulerGroupLifecycleLease{}, got.lease)
- }
- }
- require.Equal(t, 1, acquiredCount)
- require.Len(t, owner.OwnerToken, schedulerGroupLifecycleOwnerTokenBytes*2)
- require.Equal(t, strings.ToLower(owner.OwnerToken), owner.OwnerToken)
- decodedOwner, err := hex.DecodeString(owner.OwnerToken)
- require.NoError(t, err)
- require.Len(t, decodedOwner, schedulerGroupLifecycleOwnerTokenBytes)
- require.NoError(t, cache.ReleaseGroupLifecycleLease(ctx, owner))
- next, acquired, err := cache.TryAcquireGroupLifecycleLease(ctx, groupID, time.Minute)
- require.NoError(t, err)
- require.True(t, acquired)
- require.True(t, next.ValidFor(groupID))
- require.NotEqual(t, owner.OwnerToken, next.OwnerToken)
- require.NoError(t, cache.ReleaseGroupLifecycleLease(ctx, next))
- }
- func TestSchedulerCacheGroupLifecycleLeaseStaleReleaseCannotDeleteSuccessor(t *testing.T) {
- ctx := context.Background()
- cache, mr := newSchedulerCacheUnitWithRedis(t)
- const groupID int64 = 72
- const ttl = time.Minute
- first, acquired, err := cache.TryAcquireGroupLifecycleLease(ctx, groupID, ttl)
- require.NoError(t, err)
- require.True(t, acquired)
- mr.FastForward(ttl + time.Second)
- second, acquired, err := cache.TryAcquireGroupLifecycleLease(ctx, groupID, ttl)
- require.NoError(t, err)
- require.True(t, acquired)
- require.NotEqual(t, first.OwnerToken, second.OwnerToken)
- require.ErrorIs(t, cache.ReleaseGroupLifecycleLease(ctx, first), service.ErrSchedulerGroupLifecycleLeaseLost)
- owner, err := cache.rdb.Get(ctx, schedulerGroupLifecycleLockKey(groupID)).Result()
- require.NoError(t, err)
- require.Equal(t, second.OwnerToken, owner)
- _, acquired, err = cache.TryAcquireGroupLifecycleLease(ctx, groupID, ttl)
- require.NoError(t, err)
- require.False(t, acquired)
- require.NoError(t, cache.ReleaseGroupLifecycleLease(ctx, second))
- require.ErrorIs(t, cache.ReleaseGroupLifecycleLease(ctx, second), service.ErrSchedulerGroupLifecycleLeaseLost)
- }
- func TestSchedulerCacheGroupLifecycleLeaseExpiredReleaseIsLost(t *testing.T) {
- ctx := context.Background()
- cache, mr := newSchedulerCacheUnitWithRedis(t)
- const groupID int64 = 73
- const ttl = time.Minute
- lease, acquired, err := cache.TryAcquireGroupLifecycleLease(ctx, groupID, ttl)
- require.NoError(t, err)
- require.True(t, acquired)
- mr.FastForward(ttl + time.Second)
- require.ErrorIs(t, cache.ReleaseGroupLifecycleLease(ctx, lease), service.ErrSchedulerGroupLifecycleLeaseLost)
- }
- func TestSchedulerCacheGroupLifecycleLeaseWrongOwnerAndCrossGroupAreLost(t *testing.T) {
- ctx := context.Background()
- cache := newSchedulerCacheUnit(t)
- const firstGroupID int64 = 74
- const secondGroupID int64 = 75
- first, acquired, err := cache.TryAcquireGroupLifecycleLease(ctx, firstGroupID, time.Minute)
- require.NoError(t, err)
- require.True(t, acquired)
- second, acquired, err := cache.TryAcquireGroupLifecycleLease(ctx, secondGroupID, time.Minute)
- require.NoError(t, err)
- require.True(t, acquired, "different groups must acquire independently")
- require.NotEqual(t, first.OwnerToken, second.OwnerToken)
- wrongOwner := first
- wrongOwner.OwnerToken = strings.Repeat("0", schedulerGroupLifecycleOwnerTokenBytes*2)
- if wrongOwner.OwnerToken == first.OwnerToken {
- wrongOwner.OwnerToken = strings.Repeat("1", schedulerGroupLifecycleOwnerTokenBytes*2)
- }
- require.ErrorIs(t, cache.ReleaseGroupLifecycleLease(ctx, wrongOwner), service.ErrSchedulerGroupLifecycleLeaseLost)
- crossGroup := first
- crossGroup.GroupID = secondGroupID
- require.ErrorIs(t, cache.ReleaseGroupLifecycleLease(ctx, crossGroup), service.ErrSchedulerGroupLifecycleLeaseLost)
- firstOwner, err := cache.rdb.Get(ctx, schedulerGroupLifecycleLockKey(firstGroupID)).Result()
- require.NoError(t, err)
- require.Equal(t, first.OwnerToken, firstOwner)
- secondOwner, err := cache.rdb.Get(ctx, schedulerGroupLifecycleLockKey(secondGroupID)).Result()
- require.NoError(t, err)
- require.Equal(t, second.OwnerToken, secondOwner)
- require.NoError(t, cache.ReleaseGroupLifecycleLease(ctx, first))
- require.NoError(t, cache.ReleaseGroupLifecycleLease(ctx, second))
- }
- func TestSchedulerCacheGroupLifecycleLeaseCanceledContextFailsClosed(t *testing.T) {
- cache := newSchedulerCacheUnit(t)
- canceledCtx, cancel := context.WithCancel(context.Background())
- cancel()
- lease, acquired, err := cache.TryAcquireGroupLifecycleLease(canceledCtx, 76, time.Minute)
- require.ErrorIs(t, err, context.Canceled)
- require.False(t, acquired)
- require.Equal(t, service.SchedulerGroupLifecycleLease{}, lease)
- ctx := context.Background()
- lease, acquired, err = cache.TryAcquireGroupLifecycleLease(ctx, 76, time.Minute)
- require.NoError(t, err)
- require.True(t, acquired)
- require.ErrorIs(t, cache.ReleaseGroupLifecycleLease(canceledCtx, lease), context.Canceled)
- owner, err := cache.rdb.Get(ctx, schedulerGroupLifecycleLockKey(lease.GroupID)).Result()
- require.NoError(t, err)
- require.Equal(t, lease.OwnerToken, owner)
- require.NoError(t, cache.ReleaseGroupLifecycleLease(ctx, lease))
- }
- func TestSchedulerCacheGroupLifecycleLeaseRejectsInvalidInput(t *testing.T) {
- ctx := context.Background()
- cache := newSchedulerCacheUnit(t)
- lease, acquired, err := cache.TryAcquireGroupLifecycleLease(ctx, 0, time.Minute)
- require.ErrorIs(t, err, service.ErrSchedulerGroupLifecycleLeaseInvalid)
- require.False(t, acquired)
- require.Equal(t, service.SchedulerGroupLifecycleLease{}, lease)
- lease, acquired, err = cache.TryAcquireGroupLifecycleLease(ctx, 73, 0)
- require.ErrorIs(t, err, service.ErrSchedulerGroupLifecycleLeaseInvalid)
- require.False(t, acquired)
- require.Equal(t, service.SchedulerGroupLifecycleLease{}, lease)
- canceledCtx, cancel := context.WithCancel(ctx)
- cancel()
- require.ErrorIs(t, cache.ReleaseGroupLifecycleLease(canceledCtx, service.SchedulerGroupLifecycleLease{}), service.ErrSchedulerGroupLifecycleLeaseInvalid)
- keys, err := cache.rdb.DBSize(ctx).Result()
- require.NoError(t, err)
- require.Zero(t, keys)
- }
- var schedulerCachePayloadBenchmarkSink int
- func BenchmarkSchedulerCacheAccountPayloadReuse(b *testing.B) {
- for _, size := range []int{1, 100, 10_000} {
- accounts := schedulerCacheBenchmarkAccounts(size)
- b.Run(fmt.Sprintf("pair_baseline_%d_accounts", size), func(b *testing.B) {
- b.ReportAllocs()
- for i := 0; i < b.N; i++ {
- first, err := benchmarkSchedulerLegacySnapshotPayload(accounts)
- if err != nil {
- b.Fatal(err)
- }
- second, err := benchmarkSchedulerLegacySnapshotPayload(accounts)
- if err != nil {
- b.Fatal(err)
- }
- schedulerCachePayloadBenchmarkSink = first + second
- }
- })
- b.Run(fmt.Sprintf("pair_reuse_%d_accounts", size), func(b *testing.B) {
- b.ReportAllocs()
- for i := 0; i < b.N; i++ {
- ids, total, err := benchmarkSchedulerReusableSnapshotPayload(accounts)
- if err != nil {
- b.Fatal(err)
- }
- // 第二个桶仍构造成员,只跳过账号 JSON 与全局账号键。
- total += len(schedulerSnapshotMembers(ids))
- schedulerCachePayloadBenchmarkSink = total
- }
- })
- b.Run(fmt.Sprintf("first_baseline_%d_accounts", size), func(b *testing.B) {
- b.ReportAllocs()
- for i := 0; i < b.N; i++ {
- total, err := benchmarkSchedulerLegacySnapshotPayload(accounts)
- if err != nil {
- b.Fatal(err)
- }
- schedulerCachePayloadBenchmarkSink = total
- }
- })
- b.Run(fmt.Sprintf("first_reuse_%d_accounts", size), func(b *testing.B) {
- b.ReportAllocs()
- for i := 0; i < b.N; i++ {
- ids, total, err := benchmarkSchedulerReusableSnapshotPayload(accounts)
- if err != nil {
- b.Fatal(err)
- }
- total += len(ids)
- schedulerCachePayloadBenchmarkSink = total
- }
- })
- }
- }
- func benchmarkSchedulerLegacySnapshotPayload(accounts []service.Account) (int, error) {
- cacheable := make([]service.Account, 0, len(accounts))
- total := 0
- for _, account := range accounts {
- full, meta, err := marshalSchedulerCacheAccount(account)
- if err != nil {
- continue
- }
- total += len(full) + len(meta)
- cacheable = append(cacheable, account)
- }
- members := make([]redis.Z, 0, len(cacheable))
- for idx, account := range cacheable {
- members = append(members, redis.Z{Score: float64(idx), Member: strconv.FormatInt(account.ID, 10)})
- }
- return total + len(members), nil
- }
- func benchmarkSchedulerReusableSnapshotPayload(accounts []service.Account) ([]int64, int, error) {
- accountIDs := make([]int64, 0, len(accounts))
- total := 0
- for _, account := range accounts {
- full, meta, err := marshalSchedulerCacheAccount(account)
- if err != nil {
- continue
- }
- total += len(full) + len(meta)
- accountIDs = append(accountIDs, account.ID)
- }
- total += len(schedulerSnapshotMembers(accountIDs))
- return accountIDs, total, nil
- }
- func schedulerCacheBenchmarkAccounts(size int) []service.Account {
- largeValue := strings.Repeat("x", 4096)
- credentials := map[string]any{
- "api_key": "benchmark-key",
- "model_mapping": map[string]any{"z-model": "z-target", "a-model": "a-target"},
- "large_value": largeValue,
- }
- extra := map[string]any{
- "mixed_scheduling": true,
- "model_rate_limits": map[string]any{
- "z-model": map[string]any{"rate_limit_reset_at": "2026-07-16T00:00:00Z"},
- "a-model": map[string]any{"rate_limit_reset_at": "2026-07-16T00:00:00Z"},
- },
- "large_value": largeValue,
- }
- accounts := make([]service.Account, size)
- for i := range accounts {
- id := int64(i + 1)
- accounts[i] = service.Account{
- ID: id,
- Name: "benchmark-account",
- Platform: service.PlatformOpenAI,
- Type: service.AccountTypeOAuth,
- Credentials: credentials,
- Extra: extra,
- GroupIDs: []int64{7, 9},
- AccountGroups: []service.AccountGroup{
- {AccountID: id, GroupID: 7, Priority: 1},
- {AccountID: id, GroupID: 9, Priority: 2},
- },
- }
- }
- return accounts
- }
- // 调度投影必须保留 OpenAI 透传开关。
- //
- // 候选过滤走 ListSchedulableAccounts,读的是 buildSchedulerMetadataAccount 产出的精简投影;
- // Account.IsModelSupported 又靠 extra 上的透传开关短路 model_mapping 白名单(#4936)。
- // 一旦投影把开关裁掉、却保留了白名单,透传账号在选号阶段就会退回白名单判定并被误判成
- // model_not_supported,而转发阶段(读完整账号)仍按透传工作 —— 表现为"单独测这个账号能通、
- // 走网关却报 no available accounts"。#4936 修的是判定逻辑,这里守的是喂给判定的输入。
- func TestBuildSchedulerMetadataAccount_KeepsOpenAIPassthroughForModelGate(t *testing.T) {
- for _, key := range []string{"openai_passthrough", "openai_oauth_passthrough"} {
- t.Run(key, func(t *testing.T) {
- account := service.Account{
- ID: 383,
- Platform: service.PlatformOpenAI,
- Type: service.AccountTypeOAuth,
- Credentials: map[string]any{
- // 账号从白名单模式切到透传后常见的残留映射,未列出请求的模型。
- "model_mapping": map[string]any{"gpt-5.5": "gpt-5.5"},
- "access_token": "drop-me",
- },
- Extra: map[string]any{key: true},
- }
- require.True(t, account.IsModelSupported("gpt-5.6-sol"),
- "前置条件:透传账号本应放行白名单外的模型")
- meta := buildSchedulerMetadataAccount(account)
- // 走一遍真实的序列化/反序列化路径(写入 sched:meta 再由 decodeCachedAccount 读回)。
- payload, err := json.Marshal(meta)
- require.NoError(t, err)
- var restored service.Account
- require.NoError(t, json.Unmarshal(payload, &restored))
- require.Equal(t, true, restored.Extra[key])
- require.True(t, restored.IsOpenAIPassthroughEnabled())
- require.True(t, restored.IsModelSupported("gpt-5.6-sol"),
- "投影裁掉透传开关会让透传账号在候选过滤阶段被误判为 model_not_supported")
- // 白名单本身仍需保留:非透传账号依赖它做模型门。
- require.Equal(t, map[string]any{"gpt-5.5": "gpt-5.5"}, restored.Credentials["model_mapping"])
- })
- }
- }
|