scheduler_cache_unit_test.go 45 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768697071727374757677787980818283848586878889909192939495969798991001011021031041051061071081091101111121131141151161171181191201211221231241251261271281291301311321331341351361371381391401411421431441451461471481491501511521531541551561571581591601611621631641651661671681691701711721731741751761771781791801811821831841851861871881891901911921931941951961971981992002012022032042052062072082092102112122132142152162172182192202212222232242252262272282292302312322332342352362372382392402412422432442452462472482492502512522532542552562572582592602612622632642652662672682692702712722732742752762772782792802812822832842852862872882892902912922932942952962972982993003013023033043053063073083093103113123133143153163173183193203213223233243253263273283293303313323333343353363373383393403413423433443453463473483493503513523533543553563573583593603613623633643653663673683693703713723733743753763773783793803813823833843853863873883893903913923933943953963973983994004014024034044054064074084094104114124134144154164174184194204214224234244254264274284294304314324334344354364374384394404414424434444454464474484494504514524534544554564574584594604614624634644654664674684694704714724734744754764774784794804814824834844854864874884894904914924934944954964974984995005015025035045055065075085095105115125135145155165175185195205215225235245255265275285295305315325335345355365375385395405415425435445455465475485495505515525535545555565575585595605615625635645655665675685695705715725735745755765775785795805815825835845855865875885895905915925935945955965975985996006016026036046056066076086096106116126136146156166176186196206216226236246256266276286296306316326336346356366376386396406416426436446456466476486496506516526536546556566576586596606616626636646656666676686696706716726736746756766776786796806816826836846856866876886896906916926936946956966976986997007017027037047057067077087097107117127137147157167177187197207217227237247257267277287297307317327337347357367377387397407417427437447457467477487497507517527537547557567577587597607617627637647657667677687697707717727737747757767777787797807817827837847857867877887897907917927937947957967977987998008018028038048058068078088098108118128138148158168178188198208218228238248258268278288298308318328338348358368378388398408418428438448458468478488498508518528538548558568578588598608618628638648658668678688698708718728738748758768778788798808818828838848858868878888898908918928938948958968978988999009019029039049059069079089099109119129139149159169179189199209219229239249259269279289299309319329339349359369379389399409419429439449459469479489499509519529539549559569579589599609619629639649659669679689699709719729739749759769779789799809819829839849859869879889899909919929939949959969979989991000100110021003100410051006100710081009101010111012101310141015101610171018101910201021102210231024102510261027102810291030103110321033103410351036103710381039104010411042104310441045104610471048104910501051105210531054105510561057105810591060106110621063106410651066106710681069107010711072107310741075107610771078107910801081108210831084108510861087108810891090109110921093109410951096109710981099110011011102110311041105110611071108110911101111111211131114111511161117111811191120112111221123
  1. //go:build unit
  2. package repository
  3. import (
  4. "context"
  5. "encoding/hex"
  6. "encoding/json"
  7. "fmt"
  8. "strconv"
  9. "strings"
  10. "testing"
  11. "time"
  12. "github.com/Wei-Shaw/sub2api/internal/service"
  13. "github.com/alicebob/miniredis/v2"
  14. "github.com/redis/go-redis/v9"
  15. "github.com/stretchr/testify/require"
  16. )
  17. func newSchedulerCacheUnit(t *testing.T) *schedulerCache {
  18. cache, _ := newSchedulerCacheUnitWithRedis(t)
  19. return cache
  20. }
  21. func newSchedulerCacheUnitWithRedis(t *testing.T) (*schedulerCache, *miniredis.Miniredis) {
  22. t.Helper()
  23. mr := miniredis.RunT(t)
  24. rdb := redis.NewClient(&redis.Options{Addr: mr.Addr()})
  25. t.Cleanup(func() { _ = rdb.Close() })
  26. cache, ok := newSchedulerCacheWithChunkSizes(rdb, defaultSchedulerSnapshotMGetChunkSize, defaultSchedulerSnapshotWriteChunkSize).(*schedulerCache)
  27. require.True(t, ok)
  28. return cache, mr
  29. }
  30. func TestSchedulerCacheWriteAccountIDsSkipsUnencodableTimes(t *testing.T) {
  31. ctx := context.Background()
  32. cache := newSchedulerCacheUnit(t)
  33. invalidTime := time.Date(10000, time.January, 1, 0, 0, 0, 0, time.UTC)
  34. accountIDs, err := cache.writeAccountIDs(ctx, []service.Account{
  35. {ID: 111, Platform: service.PlatformOpenAI, Type: service.AccountTypeAPIKey},
  36. {ID: 112, Platform: service.PlatformOpenAI, Type: service.AccountTypeAPIKey, ExpiresAt: &invalidTime},
  37. })
  38. require.NoError(t, err)
  39. require.Equal(t, []int64{111}, accountIDs)
  40. cached, err := cache.GetAccount(ctx, 111)
  41. require.NoError(t, err)
  42. require.NotNil(t, cached)
  43. invalid, err := cache.GetAccount(ctx, 112)
  44. require.NoError(t, err)
  45. require.Nil(t, invalid)
  46. }
  47. func TestSchedulerCacheSetAccountClearsUnencodablePayload(t *testing.T) {
  48. ctx := context.Background()
  49. cache := newSchedulerCacheUnit(t)
  50. account := service.Account{ID: 113, Platform: service.PlatformOpenAI, Type: service.AccountTypeAPIKey}
  51. require.NoError(t, cache.SetAccount(ctx, &account))
  52. invalidTime := time.Date(10000, time.January, 1, 0, 0, 0, 0, time.UTC)
  53. account.ExpiresAt = &invalidTime
  54. require.NoError(t, cache.SetAccount(ctx, &account))
  55. cached, err := cache.GetAccount(ctx, account.ID)
  56. require.NoError(t, err)
  57. require.Nil(t, cached)
  58. }
  59. func TestSchedulerCacheUpdateLastUsedClearsUnencodablePayload(t *testing.T) {
  60. ctx := context.Background()
  61. cache := newSchedulerCacheUnit(t)
  62. account := service.Account{ID: 114, Platform: service.PlatformOpenAI, Type: service.AccountTypeAPIKey}
  63. require.NoError(t, cache.SetAccount(ctx, &account))
  64. invalidTime := time.Date(10000, time.January, 1, 0, 0, 0, 0, time.UTC)
  65. require.NoError(t, cache.UpdateLastUsed(ctx, map[int64]time.Time{account.ID: invalidTime}))
  66. cached, err := cache.GetAccount(ctx, account.ID)
  67. require.NoError(t, err)
  68. require.Nil(t, cached)
  69. }
  70. func TestSchedulerCacheSnapshotAccountIDReusePreservesPayloadAndMembers(t *testing.T) {
  71. ctx := context.Background()
  72. cache, _ := newSchedulerCacheUnitWithRedis(t)
  73. invalidTime := time.Date(10000, time.January, 1, 0, 0, 0, 0, time.UTC)
  74. validOne := service.Account{
  75. ID: 701,
  76. Name: "first",
  77. Platform: service.PlatformOpenAI,
  78. Type: service.AccountTypeOAuth,
  79. Credentials: map[string]any{"model_mapping": map[string]any{"z": "last", "a": "first"}},
  80. Extra: map[string]any{"mixed_scheduling": true},
  81. GroupIDs: []int64{17},
  82. }
  83. validTwo := service.Account{ID: 702, Name: "second", Platform: service.PlatformOpenAI, Type: service.AccountTypeAPIKey}
  84. invalid := service.Account{ID: 799, Platform: service.PlatformOpenAI, Type: service.AccountTypeAPIKey, ExpiresAt: &invalidTime}
  85. accounts := []service.Account{validOne, invalid, validTwo, validOne}
  86. single := service.SchedulerBucket{GroupID: 17, Platform: service.PlatformOpenAI, Mode: service.SchedulerModeSingle}
  87. singleToken, err := cache.CaptureBucketWriteToken(ctx, single)
  88. require.NoError(t, err)
  89. accountIDs, err := cache.SetSnapshotAndReturnAccountIDs(ctx, single, singleToken, accounts)
  90. require.NoError(t, err)
  91. require.Equal(t, []int64{701, 702, 701}, accountIDs, "应保留可编码账号的原顺序和重复项")
  92. wantFull, err := json.Marshal(validOne)
  93. require.NoError(t, err)
  94. wantMeta, err := json.Marshal(buildSchedulerMetadataAccount(validOne))
  95. require.NoError(t, err)
  96. fullBefore, err := cache.rdb.Get(ctx, schedulerAccountKey("701")).Bytes()
  97. require.NoError(t, err)
  98. metaBefore, err := cache.rdb.Get(ctx, schedulerAccountMetaKey("701")).Bytes()
  99. require.NoError(t, err)
  100. require.Equal(t, wantFull, fullBefore)
  101. require.Equal(t, wantMeta, metaBefore)
  102. forced := service.SchedulerBucket{GroupID: 17, Platform: service.PlatformOpenAI, Mode: service.SchedulerModeForced}
  103. forcedToken, err := cache.CaptureBucketWriteToken(ctx, forced)
  104. require.NoError(t, err)
  105. require.NoError(t, cache.SetSnapshotByAccountIDs(ctx, forced, forcedToken, accountIDs))
  106. fullAfter, err := cache.rdb.Get(ctx, schedulerAccountKey("701")).Bytes()
  107. require.NoError(t, err)
  108. metaAfter, err := cache.rdb.Get(ctx, schedulerAccountMetaKey("701")).Bytes()
  109. require.NoError(t, err)
  110. require.Equal(t, fullBefore, fullAfter, "ID-only 路径不得重写完整账号键")
  111. require.Equal(t, metaBefore, metaAfter, "ID-only 路径不得重写调度元数据键")
  112. for _, bucket := range []service.SchedulerBucket{single, forced} {
  113. version, err := cache.rdb.Get(ctx, schedulerBucketKey(schedulerActivePrefix, bucket)).Result()
  114. require.NoError(t, err)
  115. members, err := cache.rdb.ZRange(ctx, schedulerSnapshotKey(bucket, version), 0, -1).Result()
  116. require.NoError(t, err)
  117. require.Equal(t, []string{"702", "701"}, members, bucket.String())
  118. }
  119. missing, err := cache.GetAccount(ctx, invalid.ID)
  120. require.NoError(t, err)
  121. require.Nil(t, missing)
  122. }
  123. func TestSchedulerCacheSetSnapshotMatchesIDPublishing(t *testing.T) {
  124. ctx := context.Background()
  125. cache, _ := newSchedulerCacheUnitWithRedis(t)
  126. invalidTime := time.Date(10000, time.January, 1, 0, 0, 0, 0, time.UTC)
  127. validOne := service.Account{
  128. ID: 721,
  129. Name: "first",
  130. Platform: service.PlatformOpenAI,
  131. Type: service.AccountTypeOAuth,
  132. Credentials: map[string]any{"model_mapping": map[string]any{"source": "target"}},
  133. Extra: map[string]any{"mixed_scheduling": true},
  134. GroupIDs: []int64{21},
  135. }
  136. validTwo := service.Account{ID: 722, Name: "second", Platform: service.PlatformOpenAI, Type: service.AccountTypeAPIKey}
  137. invalid := service.Account{ID: 799, Platform: service.PlatformOpenAI, Type: service.AccountTypeAPIKey, ExpiresAt: &invalidTime}
  138. accounts := []service.Account{validOne, invalid, validTwo, validOne}
  139. normal := service.SchedulerBucket{GroupID: 21, Platform: service.PlatformOpenAI, Mode: service.SchedulerModeSingle}
  140. normalToken, err := cache.CaptureBucketWriteToken(ctx, normal)
  141. require.NoError(t, err)
  142. require.NoError(t, cache.SetSnapshot(ctx, normal, normalToken, accounts))
  143. fullBefore, err := cache.rdb.Get(ctx, schedulerAccountKey("721")).Bytes()
  144. require.NoError(t, err)
  145. metaBefore, err := cache.rdb.Get(ctx, schedulerAccountMetaKey("721")).Bytes()
  146. require.NoError(t, err)
  147. idOnly := service.SchedulerBucket{GroupID: 21, Platform: service.PlatformOpenAI, Mode: service.SchedulerModeForced}
  148. idOnlyToken, err := cache.CaptureBucketWriteToken(ctx, idOnly)
  149. require.NoError(t, err)
  150. accountIDs, err := cache.SetSnapshotAndReturnAccountIDs(ctx, idOnly, idOnlyToken, accounts)
  151. require.NoError(t, err)
  152. require.Equal(t, []int64{721, 722, 721}, accountIDs)
  153. fullAfter, err := cache.rdb.Get(ctx, schedulerAccountKey("721")).Bytes()
  154. require.NoError(t, err)
  155. metaAfter, err := cache.rdb.Get(ctx, schedulerAccountMetaKey("721")).Bytes()
  156. require.NoError(t, err)
  157. require.Equal(t, fullBefore, fullAfter, "普通快照和 ID 发布必须写入相同完整账号 payload")
  158. require.Equal(t, metaBefore, metaAfter, "普通快照和 ID 发布必须写入相同元数据 payload")
  159. for _, bucket := range []service.SchedulerBucket{normal, idOnly} {
  160. version, err := cache.rdb.Get(ctx, schedulerBucketKey(schedulerActivePrefix, bucket)).Result()
  161. require.NoError(t, err)
  162. members, err := cache.rdb.ZRange(ctx, schedulerSnapshotKey(bucket, version), 0, -1).Result()
  163. require.NoError(t, err)
  164. require.Equal(t, []string{"722", "721"}, members, bucket.String())
  165. }
  166. }
  167. func TestSchedulerCacheSnapshotAccountIDReuseKeepsEmptySnapshotSemantics(t *testing.T) {
  168. ctx := context.Background()
  169. cache := newSchedulerCacheUnit(t)
  170. invalidTime := time.Date(10000, time.January, 1, 0, 0, 0, 0, time.UTC)
  171. accounts := []service.Account{{ID: 811, Platform: service.PlatformOpenAI, ExpiresAt: &invalidTime}}
  172. single := service.SchedulerBucket{GroupID: 18, Platform: service.PlatformOpenAI, Mode: service.SchedulerModeSingle}
  173. singleToken, err := cache.CaptureBucketWriteToken(ctx, single)
  174. require.NoError(t, err)
  175. accountIDs, err := cache.SetSnapshotAndReturnAccountIDs(ctx, single, singleToken, accounts)
  176. require.NoError(t, err)
  177. require.Empty(t, accountIDs)
  178. forced := service.SchedulerBucket{GroupID: 18, Platform: service.PlatformOpenAI, Mode: service.SchedulerModeForced}
  179. forcedToken, err := cache.CaptureBucketWriteToken(ctx, forced)
  180. require.NoError(t, err)
  181. require.NoError(t, cache.SetSnapshotByAccountIDs(ctx, forced, forcedToken, accountIDs))
  182. for _, bucket := range []service.SchedulerBucket{single, forced} {
  183. ready, err := cache.rdb.Get(ctx, schedulerBucketKey(schedulerReadyPrefix, bucket)).Result()
  184. require.NoError(t, err)
  185. require.Equal(t, "1", ready)
  186. snapshot, hit, err := cache.GetSnapshot(ctx, bucket)
  187. require.NoError(t, err)
  188. require.False(t, hit, bucket.String())
  189. require.Nil(t, snapshot)
  190. }
  191. }
  192. func TestSchedulerCacheSetSnapshotByAccountIDsKeepsFencing(t *testing.T) {
  193. ctx := context.Background()
  194. cache := newSchedulerCacheUnit(t)
  195. bucket := service.SchedulerBucket{GroupID: 19, Platform: service.PlatformOpenAI, Mode: service.SchedulerModeForced}
  196. err := cache.SetSnapshotByAccountIDs(ctx, bucket, service.SchedulerBucketWriteToken{}, []int64{901})
  197. require.ErrorIs(t, err, service.ErrSchedulerBucketWriteFenced)
  198. _, err = cache.rdb.Get(ctx, schedulerBucketKey(schedulerVersionPrefix, bucket)).Result()
  199. require.ErrorIs(t, err, redis.Nil)
  200. token, err := cache.CaptureBucketWriteToken(ctx, bucket)
  201. require.NoError(t, err)
  202. require.NoError(t, cache.RetireBucket(ctx, bucket))
  203. err = cache.SetSnapshotByAccountIDs(ctx, bucket, token, []int64{901})
  204. require.ErrorIs(t, err, service.ErrSchedulerBucketRetired)
  205. }
  206. func TestSchedulerCacheSetSnapshotByAccountIDsDoesNotResurrectDeletedAccount(t *testing.T) {
  207. ctx := context.Background()
  208. cache := newSchedulerCacheUnit(t)
  209. account := service.Account{ID: 902, Platform: service.PlatformOpenAI, Type: service.AccountTypeOAuth}
  210. single := service.SchedulerBucket{GroupID: 20, Platform: service.PlatformOpenAI, Mode: service.SchedulerModeSingle}
  211. singleToken, err := cache.CaptureBucketWriteToken(ctx, single)
  212. require.NoError(t, err)
  213. accountIDs, err := cache.SetSnapshotAndReturnAccountIDs(ctx, single, singleToken, []service.Account{account})
  214. require.NoError(t, err)
  215. require.Equal(t, []int64{account.ID}, accountIDs)
  216. require.NoError(t, cache.DeleteAccount(ctx, account.ID))
  217. forced := service.SchedulerBucket{GroupID: 20, Platform: service.PlatformOpenAI, Mode: service.SchedulerModeForced}
  218. forcedToken, err := cache.CaptureBucketWriteToken(ctx, forced)
  219. require.NoError(t, err)
  220. require.NoError(t, cache.SetSnapshotByAccountIDs(ctx, forced, forcedToken, accountIDs))
  221. full, err := cache.GetAccount(ctx, account.ID)
  222. require.NoError(t, err)
  223. require.Nil(t, full, "ID-only 发布不得复活已删除的完整账号键")
  224. snapshot, hit, err := cache.GetSnapshot(ctx, forced)
  225. require.NoError(t, err)
  226. require.False(t, hit, "元数据缺失时必须安全回源,而不是返回残缺快照")
  227. require.Nil(t, snapshot)
  228. }
  229. func TestMarshalSchedulerCacheAccountKeepsEncodingJSONWireFormat(t *testing.T) {
  230. cases := []struct {
  231. name string
  232. account service.Account
  233. }{
  234. {name: "nil collections", account: service.Account{ID: 801}},
  235. {name: "empty collections", account: service.Account{
  236. ID: 802,
  237. Credentials: map[string]any{},
  238. Extra: map[string]any{},
  239. GroupIDs: []int64{},
  240. Groups: []*service.Group{},
  241. }},
  242. {name: "nested maps and escaping", account: service.Account{
  243. ID: 803,
  244. Credentials: map[string]any{"model_mapping": map[string]any{"z": "<last>", "a": "&first"}},
  245. Extra: map[string]any{"mixed_scheduling": true},
  246. }},
  247. }
  248. for _, tc := range cases {
  249. t.Run(tc.name, func(t *testing.T) {
  250. full, meta, err := marshalSchedulerCacheAccount(tc.account)
  251. require.NoError(t, err)
  252. wantFull, err := json.Marshal(tc.account)
  253. require.NoError(t, err)
  254. wantMeta, err := json.Marshal(buildSchedulerMetadataAccount(tc.account))
  255. require.NoError(t, err)
  256. require.Equal(t, wantFull, full)
  257. require.Equal(t, wantMeta, meta)
  258. })
  259. }
  260. }
  261. func TestBuildSchedulerMetadataAccount_KeepsOpenAIWSFlags(t *testing.T) {
  262. account := service.Account{
  263. ID: 42,
  264. Platform: service.PlatformOpenAI,
  265. Type: service.AccountTypeOAuth,
  266. Extra: map[string]any{
  267. "openai_oauth_responses_websockets_v2_enabled": true,
  268. "openai_oauth_responses_websockets_v2_mode": service.OpenAIWSIngressModePassthrough,
  269. "openai_ws_force_http": true,
  270. "openai_responses_mode": "force_chat_completions",
  271. "openai_responses_supported": false,
  272. "codex_fingerprint_mode": "session",
  273. "codex_fingerprint_seed": "11111111-1111-4111-8111-111111111111",
  274. "mixed_scheduling": true,
  275. "unused_large_field": "drop-me",
  276. },
  277. }
  278. got := buildSchedulerMetadataAccount(account)
  279. require.Equal(t, true, got.Extra["openai_oauth_responses_websockets_v2_enabled"])
  280. require.Equal(t, service.OpenAIWSIngressModePassthrough, got.Extra["openai_oauth_responses_websockets_v2_mode"])
  281. require.Equal(t, true, got.Extra["openai_ws_force_http"])
  282. require.Equal(t, "force_chat_completions", got.Extra["openai_responses_mode"])
  283. require.Equal(t, false, got.Extra["openai_responses_supported"])
  284. require.Equal(t, "session", got.Extra["codex_fingerprint_mode"])
  285. require.Equal(t, "11111111-1111-4111-8111-111111111111", got.Extra["codex_fingerprint_seed"])
  286. require.Equal(t, true, got.Extra["mixed_scheduling"])
  287. require.Nil(t, got.Extra["unused_large_field"])
  288. }
  289. func TestBuildSchedulerMetadataAccount_KeepsGrokMediaEligibility(t *testing.T) {
  290. t.Run("explicit override", func(t *testing.T) {
  291. account := service.Account{
  292. ID: 43,
  293. Platform: service.PlatformGrok,
  294. Type: service.AccountTypeOAuth,
  295. Extra: map[string]any{
  296. service.GrokMediaEligibleExtraKey: false,
  297. "unused_large_field": "drop-me",
  298. },
  299. }
  300. got := buildSchedulerMetadataAccount(account)
  301. eligible, reason := got.GrokMediaGenerationEligibility()
  302. require.False(t, eligible)
  303. require.Equal(t, "override_disabled", reason)
  304. require.Equal(t, false, got.Extra[service.GrokMediaEligibleExtraKey])
  305. require.Nil(t, got.Extra["unused_large_field"])
  306. })
  307. t.Run("forbidden billing observation", func(t *testing.T) {
  308. account := service.Account{
  309. ID: 44,
  310. Platform: service.PlatformGrok,
  311. Type: service.AccountTypeOAuth,
  312. Extra: map[string]any{
  313. "grok_billing_snapshot": map[string]any{
  314. "status_code": 200,
  315. "weekly_status_code": 403,
  316. "monthly_status_code": 200,
  317. },
  318. },
  319. }
  320. got := buildSchedulerMetadataAccount(account)
  321. eligible, reason := got.GrokMediaGenerationEligibility()
  322. require.False(t, eligible)
  323. require.Equal(t, "billing_forbidden", reason)
  324. require.NotNil(t, got.Extra["grok_billing_snapshot"])
  325. })
  326. }
  327. func TestBuildSchedulerMetadataAccount_KeepsSlimGroupMembership(t *testing.T) {
  328. account := service.Account{
  329. ID: 42,
  330. Platform: service.PlatformAnthropic,
  331. GroupIDs: []int64{7, 9, 7, 0},
  332. AccountGroups: []service.AccountGroup{
  333. {
  334. AccountID: 42,
  335. GroupID: 7,
  336. Priority: 2,
  337. Account: &service.Account{ID: 42, Name: "drop-from-metadata"},
  338. Group: &service.Group{ID: 7, Name: "drop-from-metadata"},
  339. },
  340. {
  341. AccountID: 42,
  342. GroupID: 11,
  343. Priority: 3,
  344. Group: &service.Group{ID: 11, Name: "drop-from-metadata"},
  345. },
  346. {
  347. AccountID: 42,
  348. GroupID: 0,
  349. Priority: 4,
  350. },
  351. },
  352. }
  353. got := buildSchedulerMetadataAccount(account)
  354. require.Equal(t, []int64{7, 9, 11}, got.GroupIDs)
  355. require.Len(t, got.AccountGroups, 2)
  356. require.Equal(t, int64(42), got.AccountGroups[0].AccountID)
  357. require.Equal(t, int64(7), got.AccountGroups[0].GroupID)
  358. require.Equal(t, 2, got.AccountGroups[0].Priority)
  359. require.Nil(t, got.AccountGroups[0].Account)
  360. require.Nil(t, got.AccountGroups[0].Group)
  361. require.Equal(t, int64(11), got.AccountGroups[1].GroupID)
  362. require.Nil(t, got.Groups)
  363. }
  364. func TestBuildSchedulerMetadataAccount_KeepsQuotaAutoPauseFields(t *testing.T) {
  365. account := service.Account{
  366. ID: 88,
  367. Extra: map[string]any{
  368. "codex_5h_used_percent": 12.34,
  369. "codex_7d_used_percent": 56.78,
  370. "codex_5h_reset_at": "2026-05-29T10:00:00Z",
  371. "codex_7d_reset_at": "2026-06-01T10:00:00Z",
  372. "codex_5h_reset_after_seconds": 300,
  373. "codex_7d_reset_after_seconds": 600,
  374. "codex_usage_updated_at": "2026-05-29T09:00:00Z",
  375. "auto_pause_5h_threshold": 0.95,
  376. "auto_pause_7d_threshold": 0.96,
  377. "auto_pause_5h_disabled": true,
  378. "auto_pause_7d_disabled": false,
  379. },
  380. }
  381. got := buildSchedulerMetadataAccount(account)
  382. require.Equal(t, 12.34, got.Extra["codex_5h_used_percent"])
  383. require.Equal(t, 56.78, got.Extra["codex_7d_used_percent"])
  384. require.Equal(t, "2026-05-29T10:00:00Z", got.Extra["codex_5h_reset_at"])
  385. require.Equal(t, "2026-06-01T10:00:00Z", got.Extra["codex_7d_reset_at"])
  386. require.Equal(t, 300, got.Extra["codex_5h_reset_after_seconds"])
  387. require.Equal(t, 600, got.Extra["codex_7d_reset_after_seconds"])
  388. require.Equal(t, "2026-05-29T09:00:00Z", got.Extra["codex_usage_updated_at"])
  389. require.Equal(t, 0.95, got.Extra["auto_pause_5h_threshold"])
  390. require.Equal(t, 0.96, got.Extra["auto_pause_7d_threshold"])
  391. require.Equal(t, true, got.Extra["auto_pause_5h_disabled"])
  392. require.Equal(t, false, got.Extra["auto_pause_7d_disabled"])
  393. }
  394. func TestBuildSchedulerMetadataAccount_KeepsQuotaStateForCachedAccounts(t *testing.T) {
  395. now := time.Now().UTC()
  396. activeStart := now.Add(-time.Hour).Format(time.RFC3339)
  397. expiredDailyStart := now.Add(-25 * time.Hour).Format(time.RFC3339)
  398. expiredWeeklyStart := now.Add(-8 * 24 * time.Hour).Format(time.RFC3339)
  399. weeklyResetDay := float64(now.AddDate(0, 0, 1).Weekday())
  400. cases := []struct {
  401. name string
  402. platform string
  403. typ string
  404. extra map[string]any
  405. quotaExceeded bool
  406. }{
  407. {
  408. name: "anthropic api key total quota exhausted", platform: service.PlatformAnthropic, typ: service.AccountTypeAPIKey,
  409. extra: map[string]any{"quota_limit": 10.0, "quota_used": 10.0}, quotaExceeded: true,
  410. },
  411. {
  412. name: "gemini api key rolling daily quota exhausted", platform: service.PlatformGemini, typ: service.AccountTypeAPIKey,
  413. extra: map[string]any{
  414. "quota_daily_limit": 20.0, "quota_daily_used": 20.0,
  415. "quota_daily_start": activeStart, "quota_daily_reset_mode": "rolling",
  416. }, quotaExceeded: true,
  417. },
  418. {
  419. name: "gemini api key expired rolling daily window", platform: service.PlatformGemini, typ: service.AccountTypeAPIKey,
  420. extra: map[string]any{
  421. "quota_daily_limit": 20.0, "quota_daily_used": 20.0,
  422. "quota_daily_start": expiredDailyStart, "quota_daily_reset_mode": "rolling",
  423. },
  424. },
  425. {
  426. name: "bedrock fixed weekly quota exhausted", platform: service.PlatformAnthropic, typ: service.AccountTypeBedrock,
  427. extra: map[string]any{
  428. "quota_weekly_limit": 30.0, "quota_weekly_used": 30.0, "quota_weekly_start": activeStart,
  429. "quota_weekly_reset_mode": "fixed", "quota_weekly_reset_day": weeklyResetDay,
  430. "quota_weekly_reset_hour": 0.0, "quota_reset_timezone": "UTC",
  431. }, quotaExceeded: true,
  432. },
  433. {
  434. name: "bedrock expired fixed weekly window", platform: service.PlatformAnthropic, typ: service.AccountTypeBedrock,
  435. extra: map[string]any{
  436. "quota_weekly_limit": 30.0, "quota_weekly_used": 30.0, "quota_weekly_start": expiredWeeklyStart,
  437. "quota_weekly_reset_mode": "fixed", "quota_weekly_reset_day": weeklyResetDay,
  438. "quota_weekly_reset_hour": 0.0, "quota_reset_timezone": "UTC",
  439. },
  440. },
  441. }
  442. for i, tc := range cases {
  443. t.Run(tc.name, func(t *testing.T) {
  444. extra := make(map[string]any, len(tc.extra)+1)
  445. for key, value := range tc.extra {
  446. extra[key] = value
  447. }
  448. extra["unrelated"] = "drop me"
  449. account := service.Account{
  450. ID: int64(46690 + i), Platform: tc.platform, Type: tc.typ, Extra: extra,
  451. Status: service.StatusActive, Schedulable: true,
  452. }
  453. cache := newSchedulerCacheUnit(t)
  454. ctx := context.Background()
  455. bucket := service.SchedulerBucket{GroupID: int64(46690 + i), Platform: tc.platform, Mode: service.SchedulerModeSingle}
  456. token, err := cache.CaptureBucketWriteToken(ctx, bucket)
  457. require.NoError(t, err)
  458. require.NoError(t, cache.SetSnapshot(ctx, bucket, token, []service.Account{account}))
  459. snapshot, hit, err := cache.GetSnapshot(ctx, bucket)
  460. require.NoError(t, err)
  461. require.True(t, hit)
  462. require.Len(t, snapshot, 1)
  463. cached := snapshot[0]
  464. require.Equal(t, tc.extra, cached.Extra)
  465. require.NotContains(t, cached.Extra, "unrelated")
  466. require.Equal(t, tc.quotaExceeded, cached.IsQuotaExceeded())
  467. require.Equal(t, !tc.quotaExceeded, cached.IsSchedulable())
  468. })
  469. }
  470. }
  471. func TestBuildSchedulerMetadataAccount_KeepsModelRateLimits(t *testing.T) {
  472. account := service.Account{
  473. ID: 90,
  474. Platform: service.PlatformAntigravity,
  475. Extra: map[string]any{
  476. "model_rate_limits": map[string]any{
  477. "gemini-3-flash": map[string]any{
  478. "rate_limit_reset_at": "2026-05-30T10:10:00Z",
  479. },
  480. "antigravity:gemini": map[string]any{
  481. "rate_limit_reset_at": "2026-05-30T10:10:00Z",
  482. },
  483. },
  484. "unused_large_field": "drop-me",
  485. },
  486. }
  487. got := buildSchedulerMetadataAccount(account)
  488. limits, ok := got.Extra["model_rate_limits"].(map[string]any)
  489. require.True(t, ok)
  490. require.Contains(t, limits, "gemini-3-flash")
  491. require.Contains(t, limits, "antigravity:gemini")
  492. require.Nil(t, got.Extra["unused_large_field"])
  493. }
  494. func TestBuildSchedulerMetadataAccount_KeepsSparkShadowRoutingIdentity(t *testing.T) {
  495. parentID := int64(100)
  496. account := service.Account{
  497. ID: 200,
  498. Platform: service.PlatformOpenAI,
  499. Type: service.AccountTypeOAuth,
  500. ParentAccountID: &parentID,
  501. QuotaDimension: service.QuotaDimensionSpark,
  502. Credentials: map[string]any{
  503. "model_mapping": map[string]any{
  504. "gpt-5.3-codex-spark": "gpt-5.3-codex-spark",
  505. },
  506. "compact_model_mapping": map[string]any{
  507. "gpt-5.4": "gpt-5.4-openai-compact",
  508. },
  509. "access_token": "drop-me",
  510. },
  511. }
  512. got := buildSchedulerMetadataAccount(account)
  513. require.NotNil(t, got.ParentAccountID)
  514. require.Equal(t, parentID, *got.ParentAccountID)
  515. require.Equal(t, service.QuotaDimensionSpark, got.QuotaDimension)
  516. require.Equal(t, map[string]any{"gpt-5.3-codex-spark": "gpt-5.3-codex-spark"}, got.Credentials["model_mapping"])
  517. require.Equal(t, map[string]any{"gpt-5.4": "gpt-5.4-openai-compact"}, got.Credentials["compact_model_mapping"])
  518. require.Nil(t, got.Credentials["access_token"])
  519. }
  520. func TestSchedulerCacheBucketRetirementFencesWritersAndReopen(t *testing.T) {
  521. ctx := context.Background()
  522. cache, mr := newSchedulerCacheUnitWithRedis(t)
  523. bucket := service.SchedulerBucket{GroupID: 41, Platform: service.PlatformOpenAI, Mode: service.SchedulerModeSingle}
  524. otherBucket := service.SchedulerBucket{GroupID: 42, Platform: service.PlatformOpenAI, Mode: service.SchedulerModeSingle}
  525. account := service.Account{ID: 4101, Platform: service.PlatformOpenAI, Type: service.AccountTypeAPIKey}
  526. token, err := cache.CaptureBucketWriteToken(ctx, bucket)
  527. require.NoError(t, err)
  528. require.True(t, token.ValidFor(bucket))
  529. require.NoError(t, cache.SetSnapshot(ctx, bucket, token, []service.Account{account}))
  530. // A token is bound to the full bucket identity, not just an epoch number.
  531. err = cache.SetSnapshot(ctx, otherBucket, token, []service.Account{account})
  532. require.ErrorIs(t, err, service.ErrSchedulerBucketWriteFenced)
  533. _, err = cache.rdb.Get(ctx, schedulerBucketKey(schedulerVersionPrefix, otherBucket)).Result()
  534. require.ErrorIs(t, err, redis.Nil)
  535. otherAccount := service.Account{ID: 4201, Platform: service.PlatformOpenAI, Type: service.AccountTypeAPIKey}
  536. otherToken, err := cache.CaptureBucketWriteToken(ctx, otherBucket)
  537. require.NoError(t, err)
  538. require.NoError(t, cache.SetSnapshot(ctx, otherBucket, otherToken, []service.Account{otherAccount}))
  539. otherEpoch := otherToken.Epoch
  540. activeVersion, err := cache.rdb.Get(ctx, schedulerBucketKey(schedulerActivePrefix, bucket)).Result()
  541. require.NoError(t, err)
  542. require.NoError(t, cache.RetireBucket(ctx, bucket))
  543. retiredEpoch, err := cache.rdb.Get(ctx, schedulerBucketKey(schedulerEpochPrefix, bucket)).Int64()
  544. require.NoError(t, err)
  545. require.Greater(t, retiredEpoch, token.Epoch)
  546. // Retirement is idempotent and does not advance the epoch again.
  547. require.NoError(t, cache.RetireBucket(ctx, bucket))
  548. retiredEpochAgain, err := cache.rdb.Get(ctx, schedulerBucketKey(schedulerEpochPrefix, bucket)).Int64()
  549. require.NoError(t, err)
  550. require.Equal(t, retiredEpoch, retiredEpochAgain)
  551. // New readers miss because ready/active were removed atomically. A reader that
  552. // captured activeVersion before retirement may still finish against that version.
  553. _, hit, err := cache.GetSnapshot(ctx, bucket)
  554. require.NoError(t, err)
  555. require.False(t, hit)
  556. ids, err := cache.rdb.ZRange(ctx, schedulerSnapshotKey(bucket, activeVersion), 0, -1).Result()
  557. require.NoError(t, err)
  558. require.Equal(t, []string{"4101"}, ids)
  559. ttl, err := cache.rdb.TTL(ctx, schedulerSnapshotKey(bucket, activeVersion)).Result()
  560. require.NoError(t, err)
  561. require.Positive(t, ttl)
  562. require.LessOrEqual(t, ttl, time.Duration(snapshotGraceTTLSeconds)*time.Second)
  563. buckets, err := cache.ListBuckets(ctx)
  564. require.NoError(t, err)
  565. require.NotContains(t, buckets, bucket)
  566. require.Contains(t, buckets, otherBucket)
  567. otherSnapshot, otherHit, err := cache.GetSnapshot(ctx, otherBucket)
  568. require.NoError(t, err)
  569. require.True(t, otherHit)
  570. require.Len(t, otherSnapshot, 1)
  571. require.Equal(t, otherAccount.ID, otherSnapshot[0].ID)
  572. otherEpochAfter, err := cache.rdb.Get(ctx, schedulerBucketKey(schedulerEpochPrefix, otherBucket)).Int64()
  573. require.NoError(t, err)
  574. require.Equal(t, otherEpoch, otherEpochAfter)
  575. _, err = cache.CaptureBucketWriteToken(ctx, bucket)
  576. require.ErrorIs(t, err, service.ErrSchedulerBucketRetired)
  577. versionBeforeRejectedWrite, err := cache.rdb.Get(ctx, schedulerBucketKey(schedulerVersionPrefix, bucket)).Int64()
  578. require.NoError(t, err)
  579. err = cache.SetSnapshot(ctx, bucket, token, []service.Account{account})
  580. require.ErrorIs(t, err, service.ErrSchedulerBucketRetired)
  581. versionAfterRejectedWrite, err := cache.rdb.Get(ctx, schedulerBucketKey(schedulerVersionPrefix, bucket)).Int64()
  582. require.NoError(t, err)
  583. require.Equal(t, versionBeforeRejectedWrite, versionAfterRejectedWrite, "fenced writers must not allocate a new version")
  584. retired, err := cache.rdb.Exists(ctx, schedulerBucketKey(schedulerRetiredPrefix, bucket)).Result()
  585. require.NoError(t, err)
  586. require.EqualValues(t, 1, retired, "ordinary writers must never clear the tombstone")
  587. mr.FastForward(time.Duration(snapshotGraceTTLSeconds+1) * time.Second)
  588. exists, err := cache.rdb.Exists(ctx, schedulerSnapshotKey(bucket, activeVersion)).Result()
  589. require.NoError(t, err)
  590. require.Zero(t, exists, "retired active snapshot must expire after the in-flight grace period")
  591. newToken, err := cache.ReopenBucket(ctx, bucket)
  592. require.NoError(t, err)
  593. require.True(t, newToken.ValidFor(bucket))
  594. require.Equal(t, retiredEpoch, newToken.Epoch)
  595. reopenedAgain, err := cache.ReopenBucket(ctx, bucket)
  596. require.NoError(t, err)
  597. require.Equal(t, newToken, reopenedAgain, "reopen must be idempotent within one retirement generation")
  598. err = cache.SetSnapshot(ctx, bucket, token, []service.Account{account})
  599. require.ErrorIs(t, err, service.ErrSchedulerBucketWriteFenced)
  600. require.NoError(t, cache.SetSnapshot(ctx, bucket, newToken, []service.Account{account}))
  601. reopenedWhileOpen, err := cache.ReopenBucket(ctx, bucket)
  602. require.NoError(t, err)
  603. require.Equal(t, newToken, reopenedWhileOpen)
  604. snapshot, hit, err := cache.GetSnapshot(ctx, bucket)
  605. require.NoError(t, err)
  606. require.True(t, hit)
  607. require.Len(t, snapshot, 1)
  608. require.Equal(t, account.ID, snapshot[0].ID)
  609. }
  610. func TestSchedulerCacheActivationIsFencedAfterRetire(t *testing.T) {
  611. ctx := context.Background()
  612. cache := newSchedulerCacheUnit(t)
  613. bucket := service.SchedulerBucket{GroupID: 51, Platform: service.PlatformAnthropic, Mode: service.SchedulerModeMixed}
  614. account := service.Account{ID: 5101, Platform: service.PlatformAnthropic, Type: service.AccountTypeAPIKey}
  615. token, err := cache.CaptureBucketWriteToken(ctx, bucket)
  616. require.NoError(t, err)
  617. version, err := cache.allocateSnapshotVersion(ctx, bucket, token)
  618. require.NoError(t, err)
  619. _, err = cache.writeSnapshotVersionAndReturnAccountIDs(ctx, bucket, version, []service.Account{account})
  620. require.NoError(t, err)
  621. // Deterministic race C: retirement and authoritative reopen both happen after
  622. // INCR/write but before the old writer activates.
  623. require.NoError(t, cache.RetireBucket(ctx, bucket))
  624. _, err = cache.ReopenBucket(ctx, bucket)
  625. require.NoError(t, err)
  626. err = cache.activateSnapshotVersion(ctx, bucket, token, version)
  627. require.ErrorIs(t, err, service.ErrSchedulerBucketWriteFenced)
  628. exists, err := cache.rdb.Exists(ctx, schedulerSnapshotKey(bucket, version)).Result()
  629. require.NoError(t, err)
  630. require.Zero(t, exists, "fenced activation must delete its unpublished snapshot")
  631. exists, err = cache.rdb.Exists(
  632. ctx,
  633. schedulerBucketKey(schedulerReadyPrefix, bucket),
  634. schedulerBucketKey(schedulerActivePrefix, bucket),
  635. ).Result()
  636. require.NoError(t, err)
  637. require.Zero(t, exists)
  638. buckets, err := cache.ListBuckets(ctx)
  639. require.NoError(t, err)
  640. require.NotContains(t, buckets, bucket)
  641. }
  642. func TestSchedulerCacheConcurrentReopenReturnsSameToken(t *testing.T) {
  643. ctx := context.Background()
  644. cache := newSchedulerCacheUnit(t)
  645. bucket := service.SchedulerBucket{GroupID: 53, Platform: service.PlatformOpenAI, Mode: service.SchedulerModeForced}
  646. account := service.Account{ID: 5301, Platform: service.PlatformOpenAI, Type: service.AccountTypeAPIKey}
  647. oldToken, err := cache.CaptureBucketWriteToken(ctx, bucket)
  648. require.NoError(t, err)
  649. require.NoError(t, cache.RetireBucket(ctx, bucket))
  650. type reopenResult struct {
  651. token service.SchedulerBucketWriteToken
  652. err error
  653. }
  654. start := make(chan struct{})
  655. results := make(chan reopenResult, 2)
  656. for range 2 {
  657. go func() {
  658. <-start
  659. token, err := cache.ReopenBucket(ctx, bucket)
  660. results <- reopenResult{token: token, err: err}
  661. }()
  662. }
  663. close(start)
  664. first := <-results
  665. second := <-results
  666. require.NoError(t, first.err)
  667. require.NoError(t, second.err)
  668. require.Equal(t, first.token, second.token)
  669. require.Greater(t, first.token.Epoch, oldToken.Epoch)
  670. require.ErrorIs(t, cache.SetSnapshot(ctx, bucket, oldToken, []service.Account{account}), service.ErrSchedulerBucketWriteFenced)
  671. require.NoError(t, cache.SetSnapshot(ctx, bucket, first.token, []service.Account{account}))
  672. }
  673. func TestSchedulerCacheReopenExpiresPreviousActiveSnapshot(t *testing.T) {
  674. ctx := context.Background()
  675. cache, mr := newSchedulerCacheUnitWithRedis(t)
  676. bucket := service.SchedulerBucket{GroupID: 52, Platform: service.PlatformGemini, Mode: service.SchedulerModeForced}
  677. account := service.Account{ID: 5201, Platform: service.PlatformGemini, Type: service.AccountTypeAPIKey}
  678. oldToken, err := cache.CaptureBucketWriteToken(ctx, bucket)
  679. require.NoError(t, err)
  680. require.NoError(t, cache.SetSnapshot(ctx, bucket, oldToken, []service.Account{account}))
  681. oldVersion, err := cache.rdb.Get(ctx, schedulerBucketKey(schedulerActivePrefix, bucket)).Result()
  682. require.NoError(t, err)
  683. retiredEpoch := oldToken.Epoch + 1
  684. require.NoError(t, cache.rdb.Set(ctx, schedulerBucketKey(schedulerEpochPrefix, bucket), retiredEpoch, 0).Err())
  685. require.NoError(t, cache.rdb.Set(ctx, schedulerBucketKey(schedulerRetiredPrefix, bucket), retiredEpoch, 0).Err())
  686. newToken, err := cache.ReopenBucket(ctx, bucket)
  687. require.NoError(t, err)
  688. require.Equal(t, retiredEpoch, newToken.Epoch)
  689. _, hit, err := cache.GetSnapshot(ctx, bucket)
  690. require.NoError(t, err)
  691. require.False(t, hit)
  692. ttl, err := cache.rdb.TTL(ctx, schedulerSnapshotKey(bucket, oldVersion)).Result()
  693. require.NoError(t, err)
  694. require.Positive(t, ttl)
  695. require.LessOrEqual(t, ttl, time.Duration(snapshotGraceTTLSeconds)*time.Second)
  696. require.ErrorIs(t, cache.SetSnapshot(ctx, bucket, oldToken, []service.Account{account}), service.ErrSchedulerBucketWriteFenced)
  697. mr.FastForward(time.Duration(snapshotGraceTTLSeconds+1) * time.Second)
  698. exists, err := cache.rdb.Exists(ctx, schedulerSnapshotKey(bucket, oldVersion)).Result()
  699. require.NoError(t, err)
  700. require.Zero(t, exists)
  701. require.NoError(t, cache.SetSnapshot(ctx, bucket, newToken, []service.Account{account}))
  702. }
  703. func TestSchedulerCacheGroupLifecycleLeaseConcurrentAcquireSingleOwner(t *testing.T) {
  704. ctx := context.Background()
  705. cache := newSchedulerCacheUnit(t)
  706. const groupID int64 = 71
  707. type result struct {
  708. lease service.SchedulerGroupLifecycleLease
  709. acquired bool
  710. err error
  711. }
  712. start := make(chan struct{})
  713. results := make(chan result, 32)
  714. for range 32 {
  715. go func() {
  716. <-start
  717. lease, acquired, err := cache.TryAcquireGroupLifecycleLease(ctx, groupID, time.Minute)
  718. results <- result{lease: lease, acquired: acquired, err: err}
  719. }()
  720. }
  721. close(start)
  722. var owner service.SchedulerGroupLifecycleLease
  723. acquiredCount := 0
  724. for range 32 {
  725. got := <-results
  726. require.NoError(t, got.err)
  727. if got.acquired {
  728. acquiredCount++
  729. owner = got.lease
  730. require.True(t, got.lease.ValidFor(groupID))
  731. } else {
  732. require.Equal(t, service.SchedulerGroupLifecycleLease{}, got.lease)
  733. }
  734. }
  735. require.Equal(t, 1, acquiredCount)
  736. require.Len(t, owner.OwnerToken, schedulerGroupLifecycleOwnerTokenBytes*2)
  737. require.Equal(t, strings.ToLower(owner.OwnerToken), owner.OwnerToken)
  738. decodedOwner, err := hex.DecodeString(owner.OwnerToken)
  739. require.NoError(t, err)
  740. require.Len(t, decodedOwner, schedulerGroupLifecycleOwnerTokenBytes)
  741. require.NoError(t, cache.ReleaseGroupLifecycleLease(ctx, owner))
  742. next, acquired, err := cache.TryAcquireGroupLifecycleLease(ctx, groupID, time.Minute)
  743. require.NoError(t, err)
  744. require.True(t, acquired)
  745. require.True(t, next.ValidFor(groupID))
  746. require.NotEqual(t, owner.OwnerToken, next.OwnerToken)
  747. require.NoError(t, cache.ReleaseGroupLifecycleLease(ctx, next))
  748. }
  749. func TestSchedulerCacheGroupLifecycleLeaseStaleReleaseCannotDeleteSuccessor(t *testing.T) {
  750. ctx := context.Background()
  751. cache, mr := newSchedulerCacheUnitWithRedis(t)
  752. const groupID int64 = 72
  753. const ttl = time.Minute
  754. first, acquired, err := cache.TryAcquireGroupLifecycleLease(ctx, groupID, ttl)
  755. require.NoError(t, err)
  756. require.True(t, acquired)
  757. mr.FastForward(ttl + time.Second)
  758. second, acquired, err := cache.TryAcquireGroupLifecycleLease(ctx, groupID, ttl)
  759. require.NoError(t, err)
  760. require.True(t, acquired)
  761. require.NotEqual(t, first.OwnerToken, second.OwnerToken)
  762. require.ErrorIs(t, cache.ReleaseGroupLifecycleLease(ctx, first), service.ErrSchedulerGroupLifecycleLeaseLost)
  763. owner, err := cache.rdb.Get(ctx, schedulerGroupLifecycleLockKey(groupID)).Result()
  764. require.NoError(t, err)
  765. require.Equal(t, second.OwnerToken, owner)
  766. _, acquired, err = cache.TryAcquireGroupLifecycleLease(ctx, groupID, ttl)
  767. require.NoError(t, err)
  768. require.False(t, acquired)
  769. require.NoError(t, cache.ReleaseGroupLifecycleLease(ctx, second))
  770. require.ErrorIs(t, cache.ReleaseGroupLifecycleLease(ctx, second), service.ErrSchedulerGroupLifecycleLeaseLost)
  771. }
  772. func TestSchedulerCacheGroupLifecycleLeaseExpiredReleaseIsLost(t *testing.T) {
  773. ctx := context.Background()
  774. cache, mr := newSchedulerCacheUnitWithRedis(t)
  775. const groupID int64 = 73
  776. const ttl = time.Minute
  777. lease, acquired, err := cache.TryAcquireGroupLifecycleLease(ctx, groupID, ttl)
  778. require.NoError(t, err)
  779. require.True(t, acquired)
  780. mr.FastForward(ttl + time.Second)
  781. require.ErrorIs(t, cache.ReleaseGroupLifecycleLease(ctx, lease), service.ErrSchedulerGroupLifecycleLeaseLost)
  782. }
  783. func TestSchedulerCacheGroupLifecycleLeaseWrongOwnerAndCrossGroupAreLost(t *testing.T) {
  784. ctx := context.Background()
  785. cache := newSchedulerCacheUnit(t)
  786. const firstGroupID int64 = 74
  787. const secondGroupID int64 = 75
  788. first, acquired, err := cache.TryAcquireGroupLifecycleLease(ctx, firstGroupID, time.Minute)
  789. require.NoError(t, err)
  790. require.True(t, acquired)
  791. second, acquired, err := cache.TryAcquireGroupLifecycleLease(ctx, secondGroupID, time.Minute)
  792. require.NoError(t, err)
  793. require.True(t, acquired, "different groups must acquire independently")
  794. require.NotEqual(t, first.OwnerToken, second.OwnerToken)
  795. wrongOwner := first
  796. wrongOwner.OwnerToken = strings.Repeat("0", schedulerGroupLifecycleOwnerTokenBytes*2)
  797. if wrongOwner.OwnerToken == first.OwnerToken {
  798. wrongOwner.OwnerToken = strings.Repeat("1", schedulerGroupLifecycleOwnerTokenBytes*2)
  799. }
  800. require.ErrorIs(t, cache.ReleaseGroupLifecycleLease(ctx, wrongOwner), service.ErrSchedulerGroupLifecycleLeaseLost)
  801. crossGroup := first
  802. crossGroup.GroupID = secondGroupID
  803. require.ErrorIs(t, cache.ReleaseGroupLifecycleLease(ctx, crossGroup), service.ErrSchedulerGroupLifecycleLeaseLost)
  804. firstOwner, err := cache.rdb.Get(ctx, schedulerGroupLifecycleLockKey(firstGroupID)).Result()
  805. require.NoError(t, err)
  806. require.Equal(t, first.OwnerToken, firstOwner)
  807. secondOwner, err := cache.rdb.Get(ctx, schedulerGroupLifecycleLockKey(secondGroupID)).Result()
  808. require.NoError(t, err)
  809. require.Equal(t, second.OwnerToken, secondOwner)
  810. require.NoError(t, cache.ReleaseGroupLifecycleLease(ctx, first))
  811. require.NoError(t, cache.ReleaseGroupLifecycleLease(ctx, second))
  812. }
  813. func TestSchedulerCacheGroupLifecycleLeaseCanceledContextFailsClosed(t *testing.T) {
  814. cache := newSchedulerCacheUnit(t)
  815. canceledCtx, cancel := context.WithCancel(context.Background())
  816. cancel()
  817. lease, acquired, err := cache.TryAcquireGroupLifecycleLease(canceledCtx, 76, time.Minute)
  818. require.ErrorIs(t, err, context.Canceled)
  819. require.False(t, acquired)
  820. require.Equal(t, service.SchedulerGroupLifecycleLease{}, lease)
  821. ctx := context.Background()
  822. lease, acquired, err = cache.TryAcquireGroupLifecycleLease(ctx, 76, time.Minute)
  823. require.NoError(t, err)
  824. require.True(t, acquired)
  825. require.ErrorIs(t, cache.ReleaseGroupLifecycleLease(canceledCtx, lease), context.Canceled)
  826. owner, err := cache.rdb.Get(ctx, schedulerGroupLifecycleLockKey(lease.GroupID)).Result()
  827. require.NoError(t, err)
  828. require.Equal(t, lease.OwnerToken, owner)
  829. require.NoError(t, cache.ReleaseGroupLifecycleLease(ctx, lease))
  830. }
  831. func TestSchedulerCacheGroupLifecycleLeaseRejectsInvalidInput(t *testing.T) {
  832. ctx := context.Background()
  833. cache := newSchedulerCacheUnit(t)
  834. lease, acquired, err := cache.TryAcquireGroupLifecycleLease(ctx, 0, time.Minute)
  835. require.ErrorIs(t, err, service.ErrSchedulerGroupLifecycleLeaseInvalid)
  836. require.False(t, acquired)
  837. require.Equal(t, service.SchedulerGroupLifecycleLease{}, lease)
  838. lease, acquired, err = cache.TryAcquireGroupLifecycleLease(ctx, 73, 0)
  839. require.ErrorIs(t, err, service.ErrSchedulerGroupLifecycleLeaseInvalid)
  840. require.False(t, acquired)
  841. require.Equal(t, service.SchedulerGroupLifecycleLease{}, lease)
  842. canceledCtx, cancel := context.WithCancel(ctx)
  843. cancel()
  844. require.ErrorIs(t, cache.ReleaseGroupLifecycleLease(canceledCtx, service.SchedulerGroupLifecycleLease{}), service.ErrSchedulerGroupLifecycleLeaseInvalid)
  845. keys, err := cache.rdb.DBSize(ctx).Result()
  846. require.NoError(t, err)
  847. require.Zero(t, keys)
  848. }
  849. var schedulerCachePayloadBenchmarkSink int
  850. func BenchmarkSchedulerCacheAccountPayloadReuse(b *testing.B) {
  851. for _, size := range []int{1, 100, 10_000} {
  852. accounts := schedulerCacheBenchmarkAccounts(size)
  853. b.Run(fmt.Sprintf("pair_baseline_%d_accounts", size), func(b *testing.B) {
  854. b.ReportAllocs()
  855. for i := 0; i < b.N; i++ {
  856. first, err := benchmarkSchedulerLegacySnapshotPayload(accounts)
  857. if err != nil {
  858. b.Fatal(err)
  859. }
  860. second, err := benchmarkSchedulerLegacySnapshotPayload(accounts)
  861. if err != nil {
  862. b.Fatal(err)
  863. }
  864. schedulerCachePayloadBenchmarkSink = first + second
  865. }
  866. })
  867. b.Run(fmt.Sprintf("pair_reuse_%d_accounts", size), func(b *testing.B) {
  868. b.ReportAllocs()
  869. for i := 0; i < b.N; i++ {
  870. ids, total, err := benchmarkSchedulerReusableSnapshotPayload(accounts)
  871. if err != nil {
  872. b.Fatal(err)
  873. }
  874. // 第二个桶仍构造成员,只跳过账号 JSON 与全局账号键。
  875. total += len(schedulerSnapshotMembers(ids))
  876. schedulerCachePayloadBenchmarkSink = total
  877. }
  878. })
  879. b.Run(fmt.Sprintf("first_baseline_%d_accounts", size), func(b *testing.B) {
  880. b.ReportAllocs()
  881. for i := 0; i < b.N; i++ {
  882. total, err := benchmarkSchedulerLegacySnapshotPayload(accounts)
  883. if err != nil {
  884. b.Fatal(err)
  885. }
  886. schedulerCachePayloadBenchmarkSink = total
  887. }
  888. })
  889. b.Run(fmt.Sprintf("first_reuse_%d_accounts", size), func(b *testing.B) {
  890. b.ReportAllocs()
  891. for i := 0; i < b.N; i++ {
  892. ids, total, err := benchmarkSchedulerReusableSnapshotPayload(accounts)
  893. if err != nil {
  894. b.Fatal(err)
  895. }
  896. total += len(ids)
  897. schedulerCachePayloadBenchmarkSink = total
  898. }
  899. })
  900. }
  901. }
  902. func benchmarkSchedulerLegacySnapshotPayload(accounts []service.Account) (int, error) {
  903. cacheable := make([]service.Account, 0, len(accounts))
  904. total := 0
  905. for _, account := range accounts {
  906. full, meta, err := marshalSchedulerCacheAccount(account)
  907. if err != nil {
  908. continue
  909. }
  910. total += len(full) + len(meta)
  911. cacheable = append(cacheable, account)
  912. }
  913. members := make([]redis.Z, 0, len(cacheable))
  914. for idx, account := range cacheable {
  915. members = append(members, redis.Z{Score: float64(idx), Member: strconv.FormatInt(account.ID, 10)})
  916. }
  917. return total + len(members), nil
  918. }
  919. func benchmarkSchedulerReusableSnapshotPayload(accounts []service.Account) ([]int64, int, error) {
  920. accountIDs := make([]int64, 0, len(accounts))
  921. total := 0
  922. for _, account := range accounts {
  923. full, meta, err := marshalSchedulerCacheAccount(account)
  924. if err != nil {
  925. continue
  926. }
  927. total += len(full) + len(meta)
  928. accountIDs = append(accountIDs, account.ID)
  929. }
  930. total += len(schedulerSnapshotMembers(accountIDs))
  931. return accountIDs, total, nil
  932. }
  933. func schedulerCacheBenchmarkAccounts(size int) []service.Account {
  934. largeValue := strings.Repeat("x", 4096)
  935. credentials := map[string]any{
  936. "api_key": "benchmark-key",
  937. "model_mapping": map[string]any{"z-model": "z-target", "a-model": "a-target"},
  938. "large_value": largeValue,
  939. }
  940. extra := map[string]any{
  941. "mixed_scheduling": true,
  942. "model_rate_limits": map[string]any{
  943. "z-model": map[string]any{"rate_limit_reset_at": "2026-07-16T00:00:00Z"},
  944. "a-model": map[string]any{"rate_limit_reset_at": "2026-07-16T00:00:00Z"},
  945. },
  946. "large_value": largeValue,
  947. }
  948. accounts := make([]service.Account, size)
  949. for i := range accounts {
  950. id := int64(i + 1)
  951. accounts[i] = service.Account{
  952. ID: id,
  953. Name: "benchmark-account",
  954. Platform: service.PlatformOpenAI,
  955. Type: service.AccountTypeOAuth,
  956. Credentials: credentials,
  957. Extra: extra,
  958. GroupIDs: []int64{7, 9},
  959. AccountGroups: []service.AccountGroup{
  960. {AccountID: id, GroupID: 7, Priority: 1},
  961. {AccountID: id, GroupID: 9, Priority: 2},
  962. },
  963. }
  964. }
  965. return accounts
  966. }
  967. // 调度投影必须保留 OpenAI 透传开关。
  968. //
  969. // 候选过滤走 ListSchedulableAccounts,读的是 buildSchedulerMetadataAccount 产出的精简投影;
  970. // Account.IsModelSupported 又靠 extra 上的透传开关短路 model_mapping 白名单(#4936)。
  971. // 一旦投影把开关裁掉、却保留了白名单,透传账号在选号阶段就会退回白名单判定并被误判成
  972. // model_not_supported,而转发阶段(读完整账号)仍按透传工作 —— 表现为"单独测这个账号能通、
  973. // 走网关却报 no available accounts"。#4936 修的是判定逻辑,这里守的是喂给判定的输入。
  974. func TestBuildSchedulerMetadataAccount_KeepsOpenAIPassthroughForModelGate(t *testing.T) {
  975. for _, key := range []string{"openai_passthrough", "openai_oauth_passthrough"} {
  976. t.Run(key, func(t *testing.T) {
  977. account := service.Account{
  978. ID: 383,
  979. Platform: service.PlatformOpenAI,
  980. Type: service.AccountTypeOAuth,
  981. Credentials: map[string]any{
  982. // 账号从白名单模式切到透传后常见的残留映射,未列出请求的模型。
  983. "model_mapping": map[string]any{"gpt-5.5": "gpt-5.5"},
  984. "access_token": "drop-me",
  985. },
  986. Extra: map[string]any{key: true},
  987. }
  988. require.True(t, account.IsModelSupported("gpt-5.6-sol"),
  989. "前置条件:透传账号本应放行白名单外的模型")
  990. meta := buildSchedulerMetadataAccount(account)
  991. // 走一遍真实的序列化/反序列化路径(写入 sched:meta 再由 decodeCachedAccount 读回)。
  992. payload, err := json.Marshal(meta)
  993. require.NoError(t, err)
  994. var restored service.Account
  995. require.NoError(t, json.Unmarshal(payload, &restored))
  996. require.Equal(t, true, restored.Extra[key])
  997. require.True(t, restored.IsOpenAIPassthroughEnabled())
  998. require.True(t, restored.IsModelSupported("gpt-5.6-sol"),
  999. "投影裁掉透传开关会让透传账号在候选过滤阶段被误判为 model_not_supported")
  1000. // 白名单本身仍需保留:非透传账号依赖它做模型门。
  1001. require.Equal(t, map[string]any{"gpt-5.5": "gpt-5.5"}, restored.Credentials["model_mapping"])
  1002. })
  1003. }
  1004. }