openai_account_scheduler_test.go 134 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667686970717273747576777879808182838485868788899091929394959697989910010110210310410510610710810911011111211311411511611711811912012112212312412512612712812913013113213313413513613713813914014114214314414514614714814915015115215315415515615715815916016116216316416516616716816917017117217317417517617717817918018118218318418518618718818919019119219319419519619719819920020120220320420520620720820921021121221321421521621721821922022122222322422522622722822923023123223323423523623723823924024124224324424524624724824925025125225325425525625725825926026126226326426526626726826927027127227327427527627727827928028128228328428528628728828929029129229329429529629729829930030130230330430530630730830931031131231331431531631731831932032132232332432532632732832933033133233333433533633733833934034134234334434534634734834935035135235335435535635735835936036136236336436536636736836937037137237337437537637737837938038138238338438538638738838939039139239339439539639739839940040140240340440540640740840941041141241341441541641741841942042142242342442542642742842943043143243343443543643743843944044144244344444544644744844945045145245345445545645745845946046146246346446546646746846947047147247347447547647747847948048148248348448548648748848949049149249349449549649749849950050150250350450550650750850951051151251351451551651751851952052152252352452552652752852953053153253353453553653753853954054154254354454554654754854955055155255355455555655755855956056156256356456556656756856957057157257357457557657757857958058158258358458558658758858959059159259359459559659759859960060160260360460560660760860961061161261361461561661761861962062162262362462562662762862963063163263363463563663763863964064164264364464564664764864965065165265365465565665765865966066166266366466566666766866967067167267367467567667767867968068168268368468568668768868969069169269369469569669769869970070170270370470570670770870971071171271371471571671771871972072172272372472572672772872973073173273373473573673773873974074174274374474574674774874975075175275375475575675775875976076176276376476576676776876977077177277377477577677777877978078178278378478578678778878979079179279379479579679779879980080180280380480580680780880981081181281381481581681781881982082182282382482582682782882983083183283383483583683783883984084184284384484584684784884985085185285385485585685785885986086186286386486586686786886987087187287387487587687787887988088188288388488588688788888989089189289389489589689789889990090190290390490590690790890991091191291391491591691791891992092192292392492592692792892993093193293393493593693793893994094194294394494594694794894995095195295395495595695795895996096196296396496596696796896997097197297397497597697797897998098198298398498598698798898999099199299399499599699799899910001001100210031004100510061007100810091010101110121013101410151016101710181019102010211022102310241025102610271028102910301031103210331034103510361037103810391040104110421043104410451046104710481049105010511052105310541055105610571058105910601061106210631064106510661067106810691070107110721073107410751076107710781079108010811082108310841085108610871088108910901091109210931094109510961097109810991100110111021103110411051106110711081109111011111112111311141115111611171118111911201121112211231124112511261127112811291130113111321133113411351136113711381139114011411142114311441145114611471148114911501151115211531154115511561157115811591160116111621163116411651166116711681169117011711172117311741175117611771178117911801181118211831184118511861187118811891190119111921193119411951196119711981199120012011202120312041205120612071208120912101211121212131214121512161217121812191220122112221223122412251226122712281229123012311232123312341235123612371238123912401241124212431244124512461247124812491250125112521253125412551256125712581259126012611262126312641265126612671268126912701271127212731274127512761277127812791280128112821283128412851286128712881289129012911292129312941295129612971298129913001301130213031304130513061307130813091310131113121313131413151316131713181319132013211322132313241325132613271328132913301331133213331334133513361337133813391340134113421343134413451346134713481349135013511352135313541355135613571358135913601361136213631364136513661367136813691370137113721373137413751376137713781379138013811382138313841385138613871388138913901391139213931394139513961397139813991400140114021403140414051406140714081409141014111412141314141415141614171418141914201421142214231424142514261427142814291430143114321433143414351436143714381439144014411442144314441445144614471448144914501451145214531454145514561457145814591460146114621463146414651466146714681469147014711472147314741475147614771478147914801481148214831484148514861487148814891490149114921493149414951496149714981499150015011502150315041505150615071508150915101511151215131514151515161517151815191520152115221523152415251526152715281529153015311532153315341535153615371538153915401541154215431544154515461547154815491550155115521553155415551556155715581559156015611562156315641565156615671568156915701571157215731574157515761577157815791580158115821583158415851586158715881589159015911592159315941595159615971598159916001601160216031604160516061607160816091610161116121613161416151616161716181619162016211622162316241625162616271628162916301631163216331634163516361637163816391640164116421643164416451646164716481649165016511652165316541655165616571658165916601661166216631664166516661667166816691670167116721673167416751676167716781679168016811682168316841685168616871688168916901691169216931694169516961697169816991700170117021703170417051706170717081709171017111712171317141715171617171718171917201721172217231724172517261727172817291730173117321733173417351736173717381739174017411742174317441745174617471748174917501751175217531754175517561757175817591760176117621763176417651766176717681769177017711772177317741775177617771778177917801781178217831784178517861787178817891790179117921793179417951796179717981799180018011802180318041805180618071808180918101811181218131814181518161817181818191820182118221823182418251826182718281829183018311832183318341835183618371838183918401841184218431844184518461847184818491850185118521853185418551856185718581859186018611862186318641865186618671868186918701871187218731874187518761877187818791880188118821883188418851886188718881889189018911892189318941895189618971898189919001901190219031904190519061907190819091910191119121913191419151916191719181919192019211922192319241925192619271928192919301931193219331934193519361937193819391940194119421943194419451946194719481949195019511952195319541955195619571958195919601961196219631964196519661967196819691970197119721973197419751976197719781979198019811982198319841985198619871988198919901991199219931994199519961997199819992000200120022003200420052006200720082009201020112012201320142015201620172018201920202021202220232024202520262027202820292030203120322033203420352036203720382039204020412042204320442045204620472048204920502051205220532054205520562057205820592060206120622063206420652066206720682069207020712072207320742075207620772078207920802081208220832084208520862087208820892090209120922093209420952096209720982099210021012102210321042105210621072108210921102111211221132114211521162117211821192120212121222123212421252126212721282129213021312132213321342135213621372138213921402141214221432144214521462147214821492150215121522153215421552156215721582159216021612162216321642165216621672168216921702171217221732174217521762177217821792180218121822183218421852186218721882189219021912192219321942195219621972198219922002201220222032204220522062207220822092210221122122213221422152216221722182219222022212222222322242225222622272228222922302231223222332234223522362237223822392240224122422243224422452246224722482249225022512252225322542255225622572258225922602261226222632264226522662267226822692270227122722273227422752276227722782279228022812282228322842285228622872288228922902291229222932294229522962297229822992300230123022303230423052306230723082309231023112312231323142315231623172318231923202321232223232324232523262327232823292330233123322333233423352336233723382339234023412342234323442345234623472348234923502351235223532354235523562357235823592360236123622363236423652366236723682369237023712372237323742375237623772378237923802381238223832384238523862387238823892390239123922393239423952396239723982399240024012402240324042405240624072408240924102411241224132414241524162417241824192420242124222423242424252426242724282429243024312432243324342435243624372438243924402441244224432444244524462447244824492450245124522453245424552456245724582459246024612462246324642465246624672468246924702471247224732474247524762477247824792480248124822483248424852486248724882489249024912492249324942495249624972498249925002501250225032504250525062507250825092510251125122513251425152516251725182519252025212522252325242525252625272528252925302531253225332534253525362537253825392540254125422543254425452546254725482549255025512552255325542555255625572558255925602561256225632564256525662567256825692570257125722573257425752576257725782579258025812582258325842585258625872588258925902591259225932594259525962597259825992600260126022603260426052606260726082609261026112612261326142615261626172618261926202621262226232624262526262627262826292630263126322633263426352636263726382639264026412642264326442645264626472648264926502651265226532654265526562657265826592660266126622663266426652666266726682669267026712672267326742675267626772678267926802681268226832684268526862687268826892690269126922693269426952696269726982699270027012702270327042705270627072708270927102711271227132714271527162717271827192720272127222723272427252726272727282729273027312732273327342735273627372738273927402741274227432744274527462747274827492750275127522753275427552756275727582759276027612762276327642765276627672768276927702771277227732774277527762777277827792780278127822783278427852786278727882789279027912792279327942795279627972798279928002801280228032804280528062807280828092810281128122813281428152816281728182819282028212822282328242825282628272828282928302831283228332834283528362837283828392840284128422843284428452846284728482849285028512852285328542855285628572858285928602861286228632864286528662867286828692870287128722873287428752876287728782879288028812882288328842885288628872888288928902891289228932894289528962897289828992900290129022903290429052906290729082909291029112912291329142915291629172918291929202921292229232924292529262927292829292930293129322933293429352936293729382939294029412942294329442945294629472948294929502951295229532954295529562957295829592960296129622963296429652966296729682969297029712972297329742975297629772978297929802981298229832984298529862987298829892990299129922993299429952996299729982999300030013002300330043005300630073008300930103011301230133014301530163017301830193020302130223023302430253026302730283029303030313032303330343035303630373038303930403041304230433044304530463047304830493050305130523053305430553056305730583059306030613062306330643065306630673068306930703071307230733074307530763077307830793080308130823083308430853086308730883089309030913092309330943095309630973098309931003101310231033104310531063107310831093110311131123113311431153116311731183119312031213122312331243125312631273128312931303131313231333134313531363137313831393140314131423143314431453146314731483149315031513152315331543155315631573158315931603161316231633164316531663167316831693170317131723173317431753176317731783179318031813182318331843185318631873188318931903191319231933194319531963197319831993200320132023203320432053206320732083209321032113212321332143215321632173218321932203221322232233224322532263227322832293230323132323233323432353236323732383239324032413242324332443245324632473248324932503251325232533254325532563257325832593260326132623263326432653266326732683269327032713272327332743275327632773278327932803281328232833284328532863287328832893290329132923293329432953296329732983299330033013302330333043305330633073308330933103311331233133314331533163317331833193320332133223323332433253326332733283329333033313332333333343335333633373338333933403341334233433344334533463347334833493350335133523353335433553356335733583359336033613362336333643365336633673368336933703371337233733374337533763377337833793380338133823383338433853386338733883389339033913392339333943395339633973398339934003401340234033404340534063407340834093410341134123413341434153416341734183419342034213422342334243425342634273428342934303431343234333434343534363437343834393440344134423443344434453446344734483449345034513452345334543455345634573458345934603461346234633464346534663467346834693470347134723473347434753476347734783479348034813482348334843485348634873488348934903491349234933494349534963497349834993500350135023503350435053506350735083509351035113512351335143515351635173518351935203521352235233524352535263527352835293530353135323533353435353536353735383539354035413542354335443545354635473548354935503551355235533554355535563557355835593560356135623563356435653566356735683569357035713572357335743575357635773578357935803581358235833584358535863587358835893590359135923593359435953596359735983599360036013602360336043605360636073608360936103611361236133614361536163617361836193620362136223623362436253626362736283629363036313632363336343635363636373638363936403641364236433644364536463647364836493650365136523653365436553656365736583659366036613662366336643665366636673668366936703671367236733674367536763677367836793680368136823683368436853686368736883689369036913692369336943695369636973698369937003701370237033704370537063707370837093710371137123713371437153716371737183719372037213722372337243725372637273728372937303731373237333734373537363737373837393740374137423743374437453746374737483749375037513752375337543755375637573758375937603761376237633764376537663767376837693770377137723773377437753776377737783779378037813782378337843785378637873788378937903791379237933794379537963797379837993800380138023803380438053806380738083809381038113812381338143815
  1. package service
  2. import (
  3. "context"
  4. "errors"
  5. "fmt"
  6. "math"
  7. "strconv"
  8. "sync"
  9. "testing"
  10. "time"
  11. "github.com/Wei-Shaw/sub2api/internal/config"
  12. "github.com/stretchr/testify/require"
  13. )
  14. type openAISnapshotCacheStub struct {
  15. SchedulerCache
  16. snapshotAccounts []*Account
  17. accountsByID map[int64]*Account
  18. }
  19. type schedulerTestOpenAIAccountRepo struct {
  20. AccountRepository
  21. accounts []Account
  22. }
  23. func (r schedulerTestOpenAIAccountRepo) GetByID(ctx context.Context, id int64) (*Account, error) {
  24. for i := range r.accounts {
  25. if r.accounts[i].ID == id {
  26. return &r.accounts[i], nil
  27. }
  28. }
  29. return nil, errors.New("account not found")
  30. }
  31. func (r schedulerTestOpenAIAccountRepo) ListSchedulableByGroupIDAndPlatform(ctx context.Context, groupID int64, platform string) ([]Account, error) {
  32. var result []Account
  33. for _, acc := range r.accounts {
  34. if acc.Platform == platform {
  35. result = append(result, acc)
  36. }
  37. }
  38. return result, nil
  39. }
  40. func (r schedulerTestOpenAIAccountRepo) ListSchedulableByPlatform(ctx context.Context, platform string) ([]Account, error) {
  41. var result []Account
  42. for _, acc := range r.accounts {
  43. if acc.Platform == platform {
  44. result = append(result, acc)
  45. }
  46. }
  47. return result, nil
  48. }
  49. func (r schedulerTestOpenAIAccountRepo) ListSchedulableUngroupedByPlatform(ctx context.Context, platform string) ([]Account, error) {
  50. return r.ListSchedulableByPlatform(ctx, platform)
  51. }
  52. type schedulerGroupAwareOpenAIAccountRepo struct {
  53. schedulerTestOpenAIAccountRepo
  54. }
  55. func (r schedulerGroupAwareOpenAIAccountRepo) ListSchedulableByGroupIDAndPlatform(ctx context.Context, groupID int64, platform string) ([]Account, error) {
  56. var result []Account
  57. for _, acc := range r.accounts {
  58. if acc.Platform == platform && openAIStickyAccountMatchesGroup(&acc, &groupID) {
  59. result = append(result, acc)
  60. }
  61. }
  62. return result, nil
  63. }
  64. func (r schedulerGroupAwareOpenAIAccountRepo) ListSchedulableUngroupedByPlatform(ctx context.Context, platform string) ([]Account, error) {
  65. var result []Account
  66. for _, acc := range r.accounts {
  67. if acc.Platform == platform && openAIStickyAccountMatchesGroup(&acc, nil) {
  68. result = append(result, acc)
  69. }
  70. }
  71. return result, nil
  72. }
  73. type schedulerTestConcurrencyCache struct {
  74. ConcurrencyCache
  75. loadBatchErr error
  76. loadMap map[int64]*AccountLoadInfo
  77. acquireResults map[int64]bool
  78. waitCounts map[int64]int
  79. skipDefaultLoad bool
  80. acquiredIDs *[]int64
  81. releasedIDs *[]int64
  82. }
  83. func (c schedulerTestConcurrencyCache) AcquireAccountSlot(ctx context.Context, accountID int64, maxConcurrency int, requestID string) (bool, error) {
  84. if c.acquiredIDs != nil {
  85. *c.acquiredIDs = append(*c.acquiredIDs, accountID)
  86. }
  87. if c.acquireResults != nil {
  88. if result, ok := c.acquireResults[accountID]; ok {
  89. return result, nil
  90. }
  91. }
  92. return true, nil
  93. }
  94. func (c schedulerTestConcurrencyCache) ReleaseAccountSlot(ctx context.Context, accountID int64, requestID string) error {
  95. if c.releasedIDs != nil {
  96. *c.releasedIDs = append(*c.releasedIDs, accountID)
  97. }
  98. return nil
  99. }
  100. func (c schedulerTestConcurrencyCache) GetAccountsLoadBatch(ctx context.Context, accounts []AccountWithConcurrency) (map[int64]*AccountLoadInfo, error) {
  101. if c.loadBatchErr != nil {
  102. return nil, c.loadBatchErr
  103. }
  104. out := make(map[int64]*AccountLoadInfo, len(accounts))
  105. if c.skipDefaultLoad && c.loadMap != nil {
  106. for _, acc := range accounts {
  107. if load, ok := c.loadMap[acc.ID]; ok {
  108. out[acc.ID] = load
  109. }
  110. }
  111. return out, nil
  112. }
  113. for _, acc := range accounts {
  114. if c.loadMap != nil {
  115. if load, ok := c.loadMap[acc.ID]; ok {
  116. out[acc.ID] = load
  117. continue
  118. }
  119. }
  120. out[acc.ID] = &AccountLoadInfo{AccountID: acc.ID, LoadRate: 0}
  121. }
  122. return out, nil
  123. }
  124. func (c schedulerTestConcurrencyCache) GetAccountWaitingCount(ctx context.Context, accountID int64) (int, error) {
  125. if c.waitCounts != nil {
  126. if count, ok := c.waitCounts[accountID]; ok {
  127. return count, nil
  128. }
  129. }
  130. return 0, nil
  131. }
  132. type schedulerTestGatewayCache struct {
  133. sessionBindings map[string]int64
  134. deletedSessions map[string]int
  135. }
  136. func (c *schedulerTestGatewayCache) GetSessionAccountID(ctx context.Context, groupID int64, sessionHash string) (int64, error) {
  137. if id, ok := c.sessionBindings[sessionHash]; ok {
  138. return id, nil
  139. }
  140. return 0, ErrStickySessionNotFound
  141. }
  142. func (c *schedulerTestGatewayCache) SetSessionAccountID(ctx context.Context, groupID int64, sessionHash string, accountID int64, ttl time.Duration) error {
  143. if c.sessionBindings == nil {
  144. c.sessionBindings = make(map[string]int64)
  145. }
  146. c.sessionBindings[sessionHash] = accountID
  147. return nil
  148. }
  149. func (c *schedulerTestGatewayCache) RefreshSessionTTL(ctx context.Context, groupID int64, sessionHash string, ttl time.Duration) error {
  150. return nil
  151. }
  152. func (c *schedulerTestGatewayCache) DeleteSessionAccountID(ctx context.Context, groupID int64, sessionHash string) error {
  153. if c.sessionBindings == nil {
  154. return nil
  155. }
  156. if c.deletedSessions == nil {
  157. c.deletedSessions = make(map[string]int)
  158. }
  159. c.deletedSessions[sessionHash]++
  160. delete(c.sessionBindings, sessionHash)
  161. return nil
  162. }
  163. func (c *schedulerTestGatewayCache) SetGrokVideoPendingBilling(_ context.Context, _ string, _ []byte, _ time.Duration) error {
  164. return nil
  165. }
  166. func (c *schedulerTestGatewayCache) GetGrokVideoPendingBilling(_ context.Context, _ string) ([]byte, error) {
  167. return nil, nil
  168. }
  169. func (c *schedulerTestGatewayCache) ClaimGrokVideoBilled(_ context.Context, _ string, _ time.Duration) (bool, error) {
  170. return true, nil
  171. }
  172. func (c *schedulerTestGatewayCache) ReleaseGrokVideoBilled(_ context.Context, _ string) error {
  173. return nil
  174. }
  175. func (c *schedulerTestGatewayCache) SetReasoningContent(_ context.Context, _ string, _ string, _ time.Duration) error {
  176. return nil
  177. }
  178. func (c *schedulerTestGatewayCache) GetReasoningContent(_ context.Context, _ string) (string, error) {
  179. return "", ErrReasoningContentNotFound
  180. }
  181. func newSchedulerTestOpenAIWSV2Config() *config.Config {
  182. cfg := &config.Config{}
  183. cfg.Gateway.OpenAIWS.Enabled = true
  184. cfg.Gateway.OpenAIWS.OAuthEnabled = true
  185. cfg.Gateway.OpenAIWS.APIKeyEnabled = true
  186. cfg.Gateway.OpenAIWS.ResponsesWebsocketsV2 = true
  187. cfg.Gateway.OpenAIWS.StickyResponseIDTTLSeconds = 3600
  188. return cfg
  189. }
  190. func newSchedulerTestSubscriptionPriorityConfig() *config.Config {
  191. cfg := &config.Config{}
  192. cfg.Gateway.OpenAIWS.LBTopK = 1
  193. cfg.Gateway.OpenAIWS.SchedulerScoreWeights.Priority = 1
  194. cfg.Gateway.OpenAIWS.SchedulerScoreWeights.Load = 1
  195. cfg.Gateway.OpenAIWS.SchedulerScoreWeights.Queue = 1
  196. cfg.Gateway.OpenAIWS.SchedulerScoreWeights.ErrorRate = 0
  197. cfg.Gateway.OpenAIWS.SchedulerScoreWeights.TTFT = 0
  198. return cfg
  199. }
  200. type openAIAdvancedSchedulerSettingRepoStub struct {
  201. values map[string]string
  202. }
  203. func (s *openAIAdvancedSchedulerSettingRepoStub) Get(ctx context.Context, key string) (*Setting, error) {
  204. value, err := s.GetValue(ctx, key)
  205. if err != nil {
  206. return nil, err
  207. }
  208. return &Setting{Key: key, Value: value}, nil
  209. }
  210. func (s *openAIAdvancedSchedulerSettingRepoStub) GetValue(_ context.Context, key string) (string, error) {
  211. if s == nil || s.values == nil {
  212. return "", ErrSettingNotFound
  213. }
  214. value, ok := s.values[key]
  215. if !ok {
  216. return "", ErrSettingNotFound
  217. }
  218. return value, nil
  219. }
  220. func (s *openAIAdvancedSchedulerSettingRepoStub) Set(context.Context, string, string) error {
  221. panic("unexpected call to Set")
  222. }
  223. func (s *openAIAdvancedSchedulerSettingRepoStub) GetMultiple(_ context.Context, keys []string) (map[string]string, error) {
  224. result := make(map[string]string, len(keys))
  225. for _, key := range keys {
  226. if value, err := s.GetValue(context.Background(), key); err == nil {
  227. result[key] = value
  228. }
  229. }
  230. return result, nil
  231. }
  232. func (s *openAIAdvancedSchedulerSettingRepoStub) SetMultiple(context.Context, map[string]string) error {
  233. panic("unexpected call to SetMultiple")
  234. }
  235. func (s *openAIAdvancedSchedulerSettingRepoStub) GetAll(context.Context) (map[string]string, error) {
  236. panic("unexpected call to GetAll")
  237. }
  238. func (s *openAIAdvancedSchedulerSettingRepoStub) Delete(context.Context, string) error {
  239. panic("unexpected call to Delete")
  240. }
  241. func newOpenAIAdvancedSchedulerRateLimitService(enabled string, values ...string) *RateLimitService {
  242. resetOpenAIAdvancedSchedulerSettingCacheForTest()
  243. repo := &openAIAdvancedSchedulerSettingRepoStub{
  244. values: map[string]string{},
  245. }
  246. if enabled != "" {
  247. repo.values[openAIAdvancedSchedulerSettingKey] = enabled
  248. }
  249. if len(values) > 0 && values[0] != "" {
  250. repo.values[SettingKeyOpenAIAdvancedSchedulerStickyWeightedEnabled] = values[0]
  251. }
  252. if len(values) > 1 && values[1] != "" {
  253. repo.values[SettingKeyOpenAIAdvancedSchedulerSubscriptionPriorityEnabled] = values[1]
  254. }
  255. return &RateLimitService{
  256. settingService: NewSettingService(repo, &config.Config{}),
  257. }
  258. }
  259. func (s *openAISnapshotCacheStub) GetSnapshot(ctx context.Context, bucket SchedulerBucket) ([]*Account, bool, error) {
  260. if len(s.snapshotAccounts) == 0 {
  261. return nil, false, nil
  262. }
  263. out := make([]*Account, 0, len(s.snapshotAccounts))
  264. for _, account := range s.snapshotAccounts {
  265. if account == nil {
  266. continue
  267. }
  268. cloned := *account
  269. out = append(out, &cloned)
  270. }
  271. return out, true, nil
  272. }
  273. func (s *openAISnapshotCacheStub) GetAccount(ctx context.Context, accountID int64) (*Account, error) {
  274. if s.accountsByID == nil {
  275. return nil, nil
  276. }
  277. account := s.accountsByID[accountID]
  278. if account == nil {
  279. return nil, nil
  280. }
  281. cloned := *account
  282. return &cloned, nil
  283. }
  284. func TestOpenAIGatewayService_OpenAIAdvancedSchedulerRuntimeSettings_DBOverridesConfig(t *testing.T) {
  285. resetOpenAIAdvancedSchedulerSettingCacheForTest()
  286. defer resetOpenAIAdvancedSchedulerSettingCacheForTest()
  287. cfg := &config.Config{}
  288. cfg.Gateway.OpenAIWS.LBTopK = 11
  289. cfg.Gateway.OpenAIWS.SchedulerScoreWeights = config.GatewayOpenAIWSSchedulerScoreWeights{
  290. Priority: 1,
  291. Load: 2,
  292. Queue: 3,
  293. ErrorRate: 4,
  294. TTFT: 5,
  295. Reset: 6,
  296. QuotaHeadroom: 7,
  297. UpstreamCost: 8,
  298. PreviousResponse: 9,
  299. SessionSticky: 10,
  300. }
  301. repo := &openAIAdvancedSchedulerSettingRepoStub{
  302. values: map[string]string{
  303. openAIAdvancedSchedulerSettingKey: "true",
  304. SettingKeyOpenAIAdvancedSchedulerLBTopK: "3",
  305. SettingKeyOpenAIAdvancedSchedulerWeightPriority: "2.5",
  306. SettingKeyOpenAIAdvancedSchedulerWeightReset: "0.25",
  307. SettingKeyOpenAIAdvancedSchedulerWeightPreviousResponse: "12",
  308. },
  309. }
  310. svc := &OpenAIGatewayService{
  311. cfg: cfg,
  312. rateLimitService: &RateLimitService{settingService: NewSettingService(repo, cfg)},
  313. }
  314. ctx := context.Background()
  315. require.Equal(t, 3, svc.openAIWSLBTopKForRequest(ctx))
  316. weights := svc.openAIWSSchedulerWeightsForRequest(ctx)
  317. require.Equal(t, 2.5, weights.Priority)
  318. require.Equal(t, 2.0, weights.Load)
  319. require.Equal(t, 8.0, weights.UpstreamCost)
  320. require.Equal(t, 0.25, weights.Reset)
  321. require.Equal(t, 12.0, weights.Previous)
  322. require.Equal(t, 10.0, weights.SessionSticky)
  323. }
  324. func TestOpenAIGatewayService_OpenAIAdvancedSchedulerRuntimeSettings_InvalidWeightSumsFallBackToConfig(t *testing.T) {
  325. base := config.GatewayOpenAIWSSchedulerScoreWeights{
  326. Priority: 1, Load: 2, Queue: 3, ErrorRate: 4, TTFT: 5, Reset: 6,
  327. QuotaHeadroom: 7, UpstreamCost: 8, PreviousResponse: 9, SessionSticky: 10,
  328. }
  329. maxFloat := strconv.FormatFloat(math.MaxFloat64, 'g', -1, 64)
  330. tests := []struct {
  331. name string
  332. values map[string]string
  333. }{
  334. {
  335. name: "invalid single value",
  336. values: map[string]string{
  337. SettingKeyOpenAIAdvancedSchedulerWeightPriority: "NaN",
  338. },
  339. },
  340. {
  341. name: "base sum overflow",
  342. values: map[string]string{
  343. SettingKeyOpenAIAdvancedSchedulerWeightPriority: maxFloat,
  344. SettingKeyOpenAIAdvancedSchedulerWeightLoad: maxFloat,
  345. },
  346. },
  347. {
  348. name: "sticky total sum overflow",
  349. values: map[string]string{
  350. SettingKeyOpenAIAdvancedSchedulerWeightPriority: maxFloat,
  351. SettingKeyOpenAIAdvancedSchedulerWeightPreviousResponse: maxFloat,
  352. },
  353. },
  354. }
  355. for _, tt := range tests {
  356. t.Run(tt.name, func(t *testing.T) {
  357. resetOpenAIAdvancedSchedulerSettingCacheForTest()
  358. defer resetOpenAIAdvancedSchedulerSettingCacheForTest()
  359. tt.values[openAIAdvancedSchedulerSettingKey] = "true"
  360. cfg := &config.Config{}
  361. cfg.Gateway.OpenAIWS.SchedulerScoreWeights = base
  362. repo := &openAIAdvancedSchedulerSettingRepoStub{values: tt.values}
  363. svc := &OpenAIGatewayService{cfg: cfg, rateLimitService: &RateLimitService{settingService: NewSettingService(repo, cfg)}}
  364. require.Equal(t, (&OpenAIGatewayService{cfg: cfg}).openAIWSSchedulerWeights(), svc.openAIWSSchedulerWeightsForRequest(context.Background()))
  365. })
  366. }
  367. }
  368. func TestOpenAIGatewayService_SelectAccountWithScheduler_DefaultDisabledUsesLegacyLoadAwareness(t *testing.T) {
  369. resetOpenAIAdvancedSchedulerSettingCacheForTest()
  370. ctx := context.Background()
  371. groupID := int64(10106)
  372. accounts := []Account{
  373. {
  374. ID: 36001,
  375. Platform: PlatformOpenAI,
  376. Type: AccountTypeAPIKey,
  377. Status: StatusActive,
  378. Schedulable: true,
  379. Concurrency: 1,
  380. Priority: 5,
  381. },
  382. {
  383. ID: 36002,
  384. Platform: PlatformOpenAI,
  385. Type: AccountTypeAPIKey,
  386. Status: StatusActive,
  387. Schedulable: true,
  388. Concurrency: 1,
  389. Priority: 0,
  390. },
  391. }
  392. cfg := &config.Config{}
  393. cfg.Gateway.Scheduling.LoadBatchEnabled = false
  394. cache := &schedulerTestGatewayCache{}
  395. svc := &OpenAIGatewayService{
  396. accountRepo: schedulerTestOpenAIAccountRepo{accounts: accounts},
  397. cache: cache,
  398. cfg: cfg,
  399. concurrencyService: NewConcurrencyService(schedulerTestConcurrencyCache{}),
  400. }
  401. store := svc.getOpenAIWSStateStore()
  402. require.NoError(t, store.BindResponseAccount(ctx, groupID, "resp_disabled_001", 36001, time.Hour))
  403. require.False(t, svc.isOpenAIAdvancedSchedulerEnabled(ctx))
  404. selection, decision, err := svc.SelectAccountWithScheduler(
  405. ctx,
  406. &groupID,
  407. "resp_disabled_001",
  408. "",
  409. "gpt-5.1",
  410. nil,
  411. OpenAIUpstreamTransportAny,
  412. false,
  413. )
  414. require.NoError(t, err)
  415. require.NotNil(t, selection)
  416. require.NotNil(t, selection.Account)
  417. require.Equal(t, int64(36002), selection.Account.ID)
  418. require.Equal(t, openAIAccountScheduleLayerLoadBalance, decision.Layer)
  419. require.False(t, decision.StickyPreviousHit)
  420. }
  421. // Regression: the legacy load-batch path had two bare ErrNoAvailableAccounts
  422. // exits that bypassed the diagnostics added for both the advanced scheduler and
  423. // the non-batched legacy selector. This is the default path when load batching
  424. // is enabled, so quota auto-pause could still surface as an opaque 503.
  425. func TestOpenAIGatewayService_SelectAccountWithScheduler_DefaultDisabled_LoadBatchReportsFilterReasons(t *testing.T) {
  426. resetOpenAIAdvancedSchedulerSettingCacheForTest()
  427. ctx := withOpenAIQuotaAutoPauseSettings(context.Background(), OpsOpenAIAccountQuotaAutoPauseSettings{DefaultThreshold7d: 0.9})
  428. groupID := int64(10107)
  429. quotaPaused := Account{
  430. ID: 36003,
  431. Platform: PlatformOpenAI,
  432. Type: AccountTypeOAuth,
  433. Status: StatusActive,
  434. Schedulable: true,
  435. Concurrency: 1,
  436. Extra: map[string]any{
  437. "codex_7d_used_percent": 95.0,
  438. "codex_7d_reset_at": time.Now().Add(24 * time.Hour).Format(time.RFC3339),
  439. "codex_usage_updated_at": time.Now().Add(-time.Minute).Format(time.RFC3339),
  440. },
  441. }
  442. mappingMiss := Account{
  443. ID: 36004,
  444. Platform: PlatformOpenAI,
  445. Type: AccountTypeAPIKey,
  446. Status: StatusActive,
  447. Schedulable: true,
  448. Concurrency: 1,
  449. Credentials: map[string]any{
  450. "model_mapping": map[string]any{"gpt-4o": "gpt-4o"},
  451. },
  452. }
  453. excluded := Account{
  454. ID: 36005,
  455. Platform: PlatformOpenAI,
  456. Type: AccountTypeAPIKey,
  457. Status: StatusActive,
  458. Schedulable: true,
  459. Concurrency: 1,
  460. }
  461. cfg := &config.Config{}
  462. cfg.Gateway.Scheduling.LoadBatchEnabled = true
  463. svc := &OpenAIGatewayService{
  464. accountRepo: schedulerTestOpenAIAccountRepo{accounts: []Account{quotaPaused, mappingMiss, excluded}},
  465. cache: &schedulerTestGatewayCache{},
  466. cfg: cfg,
  467. concurrencyService: NewConcurrencyService(schedulerTestConcurrencyCache{}),
  468. }
  469. require.False(t, svc.isOpenAIAdvancedSchedulerEnabled(ctx))
  470. selection, decision, err := svc.SelectAccountWithScheduler(
  471. ctx,
  472. &groupID,
  473. "",
  474. "",
  475. "gpt-5.4-mini",
  476. map[int64]struct{}{excluded.ID: {}},
  477. OpenAIUpstreamTransportAny,
  478. false,
  479. )
  480. require.Error(t, err)
  481. require.ErrorIs(t, err, ErrNoAvailableAccounts)
  482. require.Nil(t, selection)
  483. require.Equal(t, openAIAccountScheduleLayerLoadBalance, decision.Layer)
  484. require.EqualError(t, err, "no available OpenAI accounts supporting model: gpt-5.4-mini (pool=3, filtered: excluded=1 model_not_supported=1 quota_auto_pause_7d=1)")
  485. }
  486. func TestOpenAIGatewayService_SelectAccountWithScheduler_DefaultDisabled_RequiredWSV2_SkipsHTTPOnlyAccount(t *testing.T) {
  487. resetOpenAIAdvancedSchedulerSettingCacheForTest()
  488. ctx := context.Background()
  489. groupID := int64(10108)
  490. accounts := []Account{
  491. {
  492. ID: 36011,
  493. Platform: PlatformOpenAI,
  494. Type: AccountTypeAPIKey,
  495. Status: StatusActive,
  496. Schedulable: true,
  497. Concurrency: 1,
  498. Priority: 0,
  499. },
  500. {
  501. ID: 36012,
  502. Platform: PlatformOpenAI,
  503. Type: AccountTypeAPIKey,
  504. Status: StatusActive,
  505. Schedulable: true,
  506. Concurrency: 1,
  507. Priority: 5,
  508. Extra: map[string]any{
  509. "openai_apikey_responses_websockets_v2_enabled": true,
  510. },
  511. },
  512. }
  513. cfg := newSchedulerTestOpenAIWSV2Config()
  514. cfg.Gateway.Scheduling.LoadBatchEnabled = false
  515. svc := &OpenAIGatewayService{
  516. accountRepo: schedulerTestOpenAIAccountRepo{accounts: accounts},
  517. cache: &schedulerTestGatewayCache{},
  518. cfg: cfg,
  519. concurrencyService: NewConcurrencyService(schedulerTestConcurrencyCache{}),
  520. }
  521. selection, decision, err := svc.SelectAccountWithScheduler(
  522. ctx,
  523. &groupID,
  524. "",
  525. "",
  526. "gpt-5.1",
  527. nil,
  528. OpenAIUpstreamTransportResponsesWebsocketV2,
  529. false,
  530. )
  531. require.NoError(t, err)
  532. require.NotNil(t, selection)
  533. require.NotNil(t, selection.Account)
  534. require.Equal(t, int64(36012), selection.Account.ID)
  535. require.Equal(t, openAIAccountScheduleLayerLoadBalance, decision.Layer)
  536. }
  537. func TestOpenAIGatewayService_SelectAccountWithScheduler_DefaultDisabled_RequiredWSV2_NoAvailableAccount(t *testing.T) {
  538. resetOpenAIAdvancedSchedulerSettingCacheForTest()
  539. ctx := context.Background()
  540. groupID := int64(10109)
  541. accounts := []Account{
  542. {
  543. ID: 36021,
  544. Platform: PlatformOpenAI,
  545. Type: AccountTypeAPIKey,
  546. Status: StatusActive,
  547. Schedulable: true,
  548. Concurrency: 1,
  549. Priority: 0,
  550. },
  551. }
  552. cfg := newSchedulerTestOpenAIWSV2Config()
  553. cfg.Gateway.Scheduling.LoadBatchEnabled = false
  554. svc := &OpenAIGatewayService{
  555. accountRepo: schedulerTestOpenAIAccountRepo{accounts: accounts},
  556. cache: &schedulerTestGatewayCache{},
  557. cfg: cfg,
  558. concurrencyService: NewConcurrencyService(schedulerTestConcurrencyCache{}),
  559. }
  560. selection, decision, err := svc.SelectAccountWithScheduler(
  561. ctx,
  562. &groupID,
  563. "",
  564. "",
  565. "gpt-5.1",
  566. nil,
  567. OpenAIUpstreamTransportResponsesWebsocketV2,
  568. false,
  569. )
  570. require.ErrorContains(t, err, "no available OpenAI accounts")
  571. require.Nil(t, selection)
  572. require.Equal(t, openAIAccountScheduleLayerLoadBalance, decision.Layer)
  573. }
  574. func TestOpenAIGatewayService_SelectAccountWithScheduler_DefaultDisabled_EmbeddingsSkipsChatOnlyAccount(t *testing.T) {
  575. resetOpenAIAdvancedSchedulerSettingCacheForTest()
  576. ctx := context.Background()
  577. groupID := int64(10110)
  578. accounts := []Account{
  579. {
  580. ID: 36031,
  581. Platform: PlatformOpenAI,
  582. Type: AccountTypeAPIKey,
  583. Status: StatusActive,
  584. Schedulable: true,
  585. Concurrency: 1,
  586. Priority: 0,
  587. Credentials: map[string]any{
  588. "openai_capabilities": []any{"chat_completions"},
  589. },
  590. },
  591. {
  592. ID: 36032,
  593. Platform: PlatformOpenAI,
  594. Type: AccountTypeAPIKey,
  595. Status: StatusActive,
  596. Schedulable: true,
  597. Concurrency: 1,
  598. Priority: 5,
  599. Credentials: map[string]any{
  600. "openai_capabilities": []any{"chat_completions", "embeddings"},
  601. },
  602. },
  603. }
  604. cfg := &config.Config{}
  605. cfg.Gateway.Scheduling.LoadBatchEnabled = false
  606. svc := &OpenAIGatewayService{
  607. accountRepo: schedulerTestOpenAIAccountRepo{accounts: accounts},
  608. cache: &schedulerTestGatewayCache{},
  609. cfg: cfg,
  610. concurrencyService: NewConcurrencyService(schedulerTestConcurrencyCache{}),
  611. }
  612. selection, decision, err := svc.SelectAccountWithSchedulerForCapability(
  613. ctx,
  614. &groupID,
  615. "",
  616. "",
  617. "text-embedding-3-small",
  618. nil,
  619. OpenAIUpstreamTransportHTTPSSE,
  620. OpenAIEndpointCapabilityEmbeddings,
  621. false,
  622. false,
  623. true,
  624. )
  625. require.NoError(t, err)
  626. require.NotNil(t, selection)
  627. require.NotNil(t, selection.Account)
  628. require.Equal(t, int64(36032), selection.Account.ID)
  629. require.Equal(t, openAIAccountScheduleLayerLoadBalance, decision.Layer)
  630. }
  631. func TestOpenAIGatewayService_SelectAccountForTokenCount_DoesNotAcquireGenerationSlot(t *testing.T) {
  632. ctx := context.Background()
  633. groupID := int64(10115)
  634. acquiredIDs := make([]int64, 0)
  635. accounts := []Account{
  636. {
  637. ID: 36501, Platform: PlatformOpenAI, Type: AccountTypeAPIKey,
  638. Status: StatusActive, Schedulable: true, Concurrency: 1, Priority: 0,
  639. Credentials: map[string]any{"openai_capabilities": []any{"chat_completions"}},
  640. },
  641. {
  642. ID: 36502, Platform: PlatformOpenAI, Type: AccountTypeAPIKey,
  643. Status: StatusActive, Schedulable: true, Concurrency: 1, Priority: 5,
  644. Credentials: map[string]any{"openai_capabilities": []any{"embeddings"}},
  645. },
  646. {
  647. ID: 36503, Platform: PlatformGrok, Type: AccountTypeAPIKey,
  648. Status: StatusActive, Schedulable: true, Concurrency: 1, Priority: 10,
  649. Credentials: map[string]any{"openai_capabilities": []any{"chat_completions"}},
  650. },
  651. {
  652. ID: 36504, Platform: PlatformOpenAI, Type: AccountTypeAPIKey,
  653. Status: StatusActive, Schedulable: true, Concurrency: 1, Priority: 15,
  654. Credentials: map[string]any{
  655. "openai_capabilities": []any{"chat_completions"},
  656. "model_mapping": map[string]any{"gpt-4o": "gpt-4o"},
  657. },
  658. },
  659. }
  660. svc := &OpenAIGatewayService{
  661. accountRepo: schedulerTestOpenAIAccountRepo{accounts: accounts},
  662. cache: &schedulerTestGatewayCache{},
  663. cfg: &config.Config{},
  664. concurrencyService: NewConcurrencyService(schedulerTestConcurrencyCache{
  665. acquireResults: map[int64]bool{36501: false},
  666. acquiredIDs: &acquiredIDs,
  667. }),
  668. }
  669. account, err := svc.SelectAccountForTokenCount(
  670. ctx,
  671. &groupID,
  672. "",
  673. "gpt-5.1",
  674. OpenAIEndpointCapabilityChatCompletions,
  675. PlatformOpenAI,
  676. )
  677. require.NoError(t, err)
  678. require.NotNil(t, account)
  679. require.Equal(t, int64(36501), account.ID)
  680. require.Empty(t, acquiredIDs, "token counting must not acquire a generation slot")
  681. }
  682. // 生图意图的 /v1/responses 请求要求 OpenAIEndpointCapabilityResponses:探测确认
  683. // 不支持 Responses API 的 APIKey 账号必须被排除,避免 forward 阶段降级为无法生图
  684. // 的 Chat Completions 直转(#4417)。
  685. func TestOpenAIGatewayService_SelectAccountWithScheduler_ResponsesCapabilityExcludesUnsupportedAPIKey(t *testing.T) {
  686. resetOpenAIAdvancedSchedulerSettingCacheForTest()
  687. ctx := context.Background()
  688. groupID := int64(10120)
  689. newSvc := func(accounts []Account) *OpenAIGatewayService {
  690. cfg := &config.Config{}
  691. cfg.Gateway.Scheduling.LoadBatchEnabled = false
  692. return &OpenAIGatewayService{
  693. accountRepo: schedulerTestOpenAIAccountRepo{accounts: accounts},
  694. cache: &schedulerTestGatewayCache{},
  695. cfg: cfg,
  696. concurrencyService: NewConcurrencyService(schedulerTestConcurrencyCache{}),
  697. }
  698. }
  699. supported := Account{
  700. ID: 37001, Platform: PlatformOpenAI, Type: AccountTypeAPIKey,
  701. Status: StatusActive, Schedulable: true, Concurrency: 1, Priority: 0,
  702. }
  703. // 更高优先级但探测确认不支持 Responses——若门控失效会被优先选中。
  704. unsupported := Account{
  705. ID: 37002, Platform: PlatformOpenAI, Type: AccountTypeAPIKey,
  706. Status: StatusActive, Schedulable: true, Concurrency: 1, Priority: 5,
  707. Extra: map[string]any{"openai_responses_supported": false},
  708. }
  709. t.Run("生图意图仅选中支持 responses 的账号", func(t *testing.T) {
  710. svc := newSvc([]Account{supported, unsupported})
  711. selection, _, err := svc.SelectAccountWithSchedulerForCapability(
  712. ctx, &groupID, "", "", "gpt-image-2", nil,
  713. OpenAIUpstreamTransportAny, OpenAIEndpointCapabilityResponses,
  714. false, false, false,
  715. )
  716. require.NoError(t, err)
  717. require.NotNil(t, selection)
  718. require.NotNil(t, selection.Account)
  719. require.Equal(t, int64(37001), selection.Account.ID)
  720. })
  721. t.Run("仅有不支持 responses 的账号时生图意图无可用账号", func(t *testing.T) {
  722. svc := newSvc([]Account{unsupported})
  723. selection, _, err := svc.SelectAccountWithSchedulerForCapability(
  724. ctx, &groupID, "", "", "gpt-image-2", nil,
  725. OpenAIUpstreamTransportAny, OpenAIEndpointCapabilityResponses,
  726. false, false, false,
  727. )
  728. require.Error(t, err)
  729. require.Nil(t, selection)
  730. })
  731. t.Run("非生图路径仍可选中不支持 responses 的账号", func(t *testing.T) {
  732. svc := newSvc([]Account{unsupported})
  733. selection, _, err := svc.SelectAccountWithSchedulerForCapability(
  734. ctx, &groupID, "", "", "gpt-5.1", nil,
  735. OpenAIUpstreamTransportAny, OpenAIEndpointCapabilityChatCompletions,
  736. false, false, true,
  737. )
  738. require.NoError(t, err)
  739. require.NotNil(t, selection)
  740. require.NotNil(t, selection.Account)
  741. require.Equal(t, int64(37002), selection.Account.ID)
  742. })
  743. }
  744. // alpha/search 调度必须同时放行 OAuth 与 APIKey 账号:v0.1.157 曾因 OAuth-only
  745. // 门控把 APIKey 账号从候选池剔除,纯 APIKey 分组的独立搜索请求在选号阶段就
  746. // 报无可用账号,Codex 网页搜索整体失效(转发层其实一直支持 APIKey 路径)。
  747. func TestOpenAIGatewayService_SelectAccountWithScheduler_AlphaSearchAllowsAPIKeyAccount(t *testing.T) {
  748. resetOpenAIAdvancedSchedulerSettingCacheForTest()
  749. ctx := context.Background()
  750. groupID := int64(10125)
  751. accounts := []Account{
  752. {
  753. ID: 38001,
  754. Platform: PlatformOpenAI,
  755. Type: AccountTypeAPIKey,
  756. Status: StatusActive,
  757. Schedulable: true,
  758. Concurrency: 1,
  759. Priority: 0,
  760. },
  761. }
  762. cfg := &config.Config{}
  763. cfg.Gateway.Scheduling.LoadBatchEnabled = false
  764. svc := &OpenAIGatewayService{
  765. accountRepo: schedulerTestOpenAIAccountRepo{accounts: accounts},
  766. cache: &schedulerTestGatewayCache{},
  767. cfg: cfg,
  768. concurrencyService: NewConcurrencyService(schedulerTestConcurrencyCache{}),
  769. }
  770. selection, _, err := svc.SelectAccountWithSchedulerForCapability(
  771. ctx,
  772. &groupID,
  773. "",
  774. "",
  775. "gpt-5.6-sol",
  776. nil,
  777. OpenAIUpstreamTransportHTTPSSE,
  778. OpenAIEndpointCapabilityAlphaSearch,
  779. false,
  780. false,
  781. false,
  782. )
  783. require.NoError(t, err)
  784. require.NotNil(t, selection)
  785. require.NotNil(t, selection.Account)
  786. require.Equal(t, int64(38001), selection.Account.ID)
  787. }
  788. func TestOpenAIGatewayService_SelectAccountWithScheduler_DefaultDisabled_AllowsGrokChatAccount(t *testing.T) {
  789. resetOpenAIAdvancedSchedulerSettingCacheForTest()
  790. ctx := context.Background()
  791. groupID := int64(10113)
  792. accounts := []Account{
  793. {
  794. ID: 36041,
  795. Platform: PlatformGrok,
  796. Type: AccountTypeOAuth,
  797. Status: StatusActive,
  798. Schedulable: true,
  799. Concurrency: 1,
  800. Priority: 0,
  801. },
  802. }
  803. cfg := &config.Config{}
  804. cfg.Gateway.Scheduling.LoadBatchEnabled = false
  805. svc := &OpenAIGatewayService{
  806. accountRepo: schedulerTestOpenAIAccountRepo{accounts: accounts},
  807. cache: &schedulerTestGatewayCache{},
  808. cfg: cfg,
  809. concurrencyService: NewConcurrencyService(schedulerTestConcurrencyCache{}),
  810. }
  811. selection, decision, err := svc.SelectAccountWithSchedulerForCapability(
  812. ctx,
  813. &groupID,
  814. "",
  815. "",
  816. "grok-4.3",
  817. nil,
  818. OpenAIUpstreamTransportAny,
  819. OpenAIEndpointCapabilityChatCompletions,
  820. false,
  821. false,
  822. false,
  823. PlatformGrok,
  824. )
  825. require.NoError(t, err)
  826. require.NotNil(t, selection)
  827. require.NotNil(t, selection.Account)
  828. require.Equal(t, int64(36041), selection.Account.ID)
  829. require.Equal(t, openAIAccountScheduleLayerLoadBalance, decision.Layer)
  830. }
  831. func TestOpenAIGatewayService_SelectAccountWithScheduler_GrokMediaCapabilityFiltersIneligibleAccounts(t *testing.T) {
  832. resetOpenAIAdvancedSchedulerSettingCacheForTest()
  833. ctx := context.Background()
  834. groupID := int64(10114)
  835. ineligible := Account{
  836. ID: 36051, Platform: PlatformGrok, Type: AccountTypeOAuth,
  837. Status: StatusActive, Schedulable: true, Concurrency: 1, Priority: 5,
  838. Extra: map[string]any{GrokMediaEligibleExtraKey: false},
  839. }
  840. eligible := Account{
  841. ID: 36052, Platform: PlatformGrok, Type: AccountTypeOAuth,
  842. Status: StatusActive, Schedulable: true, Concurrency: 1, Priority: 0,
  843. Extra: map[string]any{GrokMediaEligibleExtraKey: true},
  844. }
  845. newService := func(accounts []Account) *OpenAIGatewayService {
  846. cfg := &config.Config{}
  847. cfg.Gateway.Scheduling.LoadBatchEnabled = false
  848. return &OpenAIGatewayService{
  849. accountRepo: schedulerTestOpenAIAccountRepo{accounts: accounts},
  850. cache: &schedulerTestGatewayCache{},
  851. cfg: cfg,
  852. concurrencyService: NewConcurrencyService(schedulerTestConcurrencyCache{}),
  853. }
  854. }
  855. t.Run("media generation skips higher priority ineligible account", func(t *testing.T) {
  856. selection, _, err := newService([]Account{ineligible, eligible}).SelectAccountWithSchedulerForCapability(
  857. ctx, &groupID, "", "", "grok-imagine-video", nil,
  858. OpenAIUpstreamTransportHTTPSSE, OpenAIEndpointCapabilityGrokMediaGeneration,
  859. false, false, false, PlatformGrok,
  860. )
  861. require.NoError(t, err)
  862. require.NotNil(t, selection)
  863. require.NotNil(t, selection.Account)
  864. require.Equal(t, eligible.ID, selection.Account.ID)
  865. })
  866. t.Run("media generation fails closed when all accounts are ineligible", func(t *testing.T) {
  867. selection, _, err := newService([]Account{ineligible}).SelectAccountWithSchedulerForCapability(
  868. ctx, &groupID, "", "", "grok-imagine-video", nil,
  869. OpenAIUpstreamTransportHTTPSSE, OpenAIEndpointCapabilityGrokMediaGeneration,
  870. false, false, false, PlatformGrok,
  871. )
  872. require.Error(t, err)
  873. require.ErrorIs(t, err, ErrNoAvailableAccounts)
  874. require.Nil(t, selection)
  875. })
  876. t.Run("chat remains routable on media-ineligible account", func(t *testing.T) {
  877. selection, _, err := newService([]Account{ineligible}).SelectAccountWithSchedulerForCapability(
  878. ctx, &groupID, "", "", "grok-4.3", nil,
  879. OpenAIUpstreamTransportHTTPSSE, OpenAIEndpointCapabilityChatCompletions,
  880. false, false, false, PlatformGrok,
  881. )
  882. require.NoError(t, err)
  883. require.NotNil(t, selection)
  884. require.NotNil(t, selection.Account)
  885. require.Equal(t, ineligible.ID, selection.Account.ID)
  886. })
  887. }
  888. // Regression #4599: when the advanced scheduler's load-balance initial filter
  889. // drops every candidate, the resulting "no available OpenAI accounts" error must
  890. // carry per-reason exclusion counts. Previously all filter branches were silent
  891. // (debug logs at best), so a 503 with excluded_account_count=0 could not be
  892. // diagnosed from the error alone.
  893. func TestOpenAIGatewayService_SelectAccountWithScheduler_NoAvailableErrorReportsQuotaAutoPauseExclusion(t *testing.T) {
  894. resetOpenAIAdvancedSchedulerSettingCacheForTest()
  895. ctx := withOpenAIQuotaAutoPauseSettings(context.Background(), OpsOpenAIAccountQuotaAutoPauseSettings{DefaultThreshold7d: 0.9})
  896. groupID := int64(101201)
  897. accounts := []Account{
  898. {
  899. ID: 38101,
  900. Platform: PlatformOpenAI,
  901. Type: AccountTypeOAuth,
  902. Status: StatusActive,
  903. Schedulable: true,
  904. Concurrency: 1,
  905. Extra: map[string]any{
  906. "codex_7d_used_percent": 95.0,
  907. "codex_7d_reset_at": time.Now().Add(24 * time.Hour).Format(time.RFC3339),
  908. "codex_usage_updated_at": time.Now().Add(-time.Minute).Format(time.RFC3339),
  909. },
  910. },
  911. }
  912. cfg := &config.Config{}
  913. svc := &OpenAIGatewayService{
  914. accountRepo: schedulerTestOpenAIAccountRepo{accounts: accounts},
  915. cache: &schedulerTestGatewayCache{},
  916. cfg: cfg,
  917. rateLimitService: newOpenAIAdvancedSchedulerRateLimitService("true"),
  918. concurrencyService: NewConcurrencyService(schedulerTestConcurrencyCache{}),
  919. }
  920. selection, decision, err := svc.SelectAccountWithScheduler(
  921. ctx, &groupID, "", "", "gpt-5.4-mini", nil, OpenAIUpstreamTransportAny, false,
  922. )
  923. require.Error(t, err)
  924. require.ErrorIs(t, err, ErrNoAvailableAccounts)
  925. require.Nil(t, selection)
  926. require.Equal(t, openAIAccountScheduleLayerLoadBalance, decision.Layer)
  927. require.EqualError(t, err, "no available OpenAI accounts supporting model: gpt-5.4-mini (pool=1, filtered: quota_auto_pause_7d=1)")
  928. }
  929. func TestOpenAIGatewayService_SelectAccountWithScheduler_NoAvailableErrorReportsModelNotSupported(t *testing.T) {
  930. resetOpenAIAdvancedSchedulerSettingCacheForTest()
  931. ctx := context.Background()
  932. groupID := int64(101202)
  933. // OAuth account with empty model_mapping: foreign-family models (grok-*)
  934. // are not servable by Codex upstream (#3662) and are dropped by the filter.
  935. accounts := []Account{
  936. {
  937. ID: 38111,
  938. Platform: PlatformOpenAI,
  939. Type: AccountTypeOAuth,
  940. Status: StatusActive,
  941. Schedulable: true,
  942. Concurrency: 1,
  943. },
  944. }
  945. svc := &OpenAIGatewayService{
  946. accountRepo: schedulerTestOpenAIAccountRepo{accounts: accounts},
  947. cache: &schedulerTestGatewayCache{},
  948. cfg: &config.Config{},
  949. rateLimitService: newOpenAIAdvancedSchedulerRateLimitService("true"),
  950. concurrencyService: NewConcurrencyService(schedulerTestConcurrencyCache{}),
  951. }
  952. selection, _, err := svc.SelectAccountWithScheduler(
  953. ctx, &groupID, "", "", "grok-4.5", nil, OpenAIUpstreamTransportAny, false,
  954. )
  955. require.Error(t, err)
  956. require.ErrorIs(t, err, ErrNoAvailableAccounts)
  957. require.Nil(t, selection)
  958. require.EqualError(t, err, "no available OpenAI accounts supporting model: grok-4.5 (pool=1, filtered: model_not_supported=1)")
  959. }
  960. func TestOpenAIGatewayService_SelectAccountWithScheduler_NoAvailableErrorAggregatesReasonsDeterministically(t *testing.T) {
  961. resetOpenAIAdvancedSchedulerSettingCacheForTest()
  962. ctx := withOpenAIQuotaAutoPauseSettings(context.Background(), OpsOpenAIAccountQuotaAutoPauseSettings{DefaultThreshold7d: 0.9})
  963. groupID := int64(101203)
  964. quotaPaused := Account{
  965. ID: 38121,
  966. Platform: PlatformOpenAI,
  967. Type: AccountTypeOAuth,
  968. Status: StatusActive,
  969. Schedulable: true,
  970. Concurrency: 1,
  971. Extra: map[string]any{
  972. "codex_7d_used_percent": 95.0,
  973. "codex_7d_reset_at": time.Now().Add(24 * time.Hour).Format(time.RFC3339),
  974. "codex_usage_updated_at": time.Now().Add(-time.Minute).Format(time.RFC3339),
  975. },
  976. }
  977. mappingMiss := Account{
  978. ID: 38122,
  979. Platform: PlatformOpenAI,
  980. Type: AccountTypeAPIKey,
  981. Status: StatusActive,
  982. Schedulable: true,
  983. Concurrency: 1,
  984. Credentials: map[string]any{
  985. "model_mapping": map[string]any{"gpt-4o": "gpt-4o"},
  986. },
  987. }
  988. excluded := Account{
  989. ID: 38123,
  990. Platform: PlatformOpenAI,
  991. Type: AccountTypeAPIKey,
  992. Status: StatusActive,
  993. Schedulable: true,
  994. Concurrency: 1,
  995. }
  996. svc := &OpenAIGatewayService{
  997. accountRepo: schedulerTestOpenAIAccountRepo{accounts: []Account{quotaPaused, mappingMiss, excluded}},
  998. cache: &schedulerTestGatewayCache{},
  999. cfg: &config.Config{},
  1000. rateLimitService: newOpenAIAdvancedSchedulerRateLimitService("true"),
  1001. concurrencyService: NewConcurrencyService(schedulerTestConcurrencyCache{}),
  1002. }
  1003. selection, _, err := svc.SelectAccountWithScheduler(
  1004. ctx, &groupID, "", "", "gpt-5.4-mini", map[int64]struct{}{38123: {}}, OpenAIUpstreamTransportAny, false,
  1005. )
  1006. require.Error(t, err)
  1007. require.ErrorIs(t, err, ErrNoAvailableAccounts)
  1008. require.Nil(t, selection)
  1009. // Reasons are sorted lexicographically, so the message is deterministic.
  1010. require.EqualError(t, err, "no available OpenAI accounts supporting model: gpt-5.4-mini (pool=3, filtered: excluded=1 model_not_supported=1 quota_auto_pause_7d=1)")
  1011. }
  1012. func TestOpenAIGatewayService_SelectAccountWithScheduler_NoAvailableErrorReportsEmptyPool(t *testing.T) {
  1013. resetOpenAIAdvancedSchedulerSettingCacheForTest()
  1014. ctx := context.Background()
  1015. groupID := int64(101204)
  1016. svc := &OpenAIGatewayService{
  1017. accountRepo: schedulerTestOpenAIAccountRepo{},
  1018. cache: &schedulerTestGatewayCache{},
  1019. cfg: &config.Config{},
  1020. rateLimitService: newOpenAIAdvancedSchedulerRateLimitService("true"),
  1021. concurrencyService: NewConcurrencyService(schedulerTestConcurrencyCache{}),
  1022. }
  1023. selection, _, err := svc.SelectAccountWithScheduler(
  1024. ctx, &groupID, "", "", "gpt-5.1", nil, OpenAIUpstreamTransportAny, false,
  1025. )
  1026. require.Error(t, err)
  1027. require.ErrorIs(t, err, ErrNoAvailableAccounts)
  1028. require.Nil(t, selection)
  1029. require.EqualError(t, err, "no available OpenAI accounts supporting model: gpt-5.1 (pool=0)")
  1030. }
  1031. func TestOpenAIGatewayService_SelectAccountWithScheduler_EnabledUsesAdvancedPreviousResponseRouting(t *testing.T) {
  1032. resetOpenAIAdvancedSchedulerSettingCacheForTest()
  1033. ctx := context.Background()
  1034. groupID := int64(10107)
  1035. accounts := []Account{
  1036. {
  1037. ID: 37001,
  1038. Platform: PlatformOpenAI,
  1039. Type: AccountTypeAPIKey,
  1040. Status: StatusActive,
  1041. Schedulable: true,
  1042. Concurrency: 1,
  1043. Priority: 5,
  1044. Extra: map[string]any{
  1045. "openai_apikey_responses_websockets_v2_enabled": true,
  1046. },
  1047. },
  1048. {
  1049. ID: 37002,
  1050. Platform: PlatformOpenAI,
  1051. Type: AccountTypeAPIKey,
  1052. Status: StatusActive,
  1053. Schedulable: true,
  1054. Concurrency: 1,
  1055. Priority: 0,
  1056. },
  1057. }
  1058. cfg := &config.Config{}
  1059. cfg.Gateway.Scheduling.LoadBatchEnabled = false
  1060. cfg.Gateway.OpenAIWS.Enabled = true
  1061. cfg.Gateway.OpenAIWS.OAuthEnabled = true
  1062. cfg.Gateway.OpenAIWS.APIKeyEnabled = true
  1063. cfg.Gateway.OpenAIWS.ResponsesWebsocketsV2 = true
  1064. cfg.Gateway.OpenAIWS.StickyResponseIDTTLSeconds = 3600
  1065. svc := &OpenAIGatewayService{
  1066. accountRepo: schedulerTestOpenAIAccountRepo{accounts: accounts},
  1067. cache: &schedulerTestGatewayCache{},
  1068. cfg: cfg,
  1069. rateLimitService: newOpenAIAdvancedSchedulerRateLimitService("true"),
  1070. concurrencyService: NewConcurrencyService(schedulerTestConcurrencyCache{}),
  1071. }
  1072. store := svc.getOpenAIWSStateStore()
  1073. require.NoError(t, store.BindResponseAccount(ctx, groupID, "resp_enabled_001", 37001, time.Hour))
  1074. require.True(t, svc.isOpenAIAdvancedSchedulerEnabled(ctx))
  1075. selection, decision, err := svc.SelectAccountWithScheduler(
  1076. ctx,
  1077. &groupID,
  1078. "resp_enabled_001",
  1079. "",
  1080. "gpt-5.1",
  1081. nil,
  1082. OpenAIUpstreamTransportAny,
  1083. false,
  1084. )
  1085. require.NoError(t, err)
  1086. require.NotNil(t, selection)
  1087. require.NotNil(t, selection.Account)
  1088. require.Equal(t, int64(37001), selection.Account.ID)
  1089. require.Equal(t, openAIAccountScheduleLayerPreviousResponse, decision.Layer)
  1090. require.True(t, decision.StickyPreviousHit)
  1091. }
  1092. func TestOpenAIGatewayService_SelectAccountWithScheduler_StickyWeightedSessionInTopKUsesStickyFirst(t *testing.T) {
  1093. resetOpenAIAdvancedSchedulerSettingCacheForTest()
  1094. ctx := context.Background()
  1095. groupID := int64(101071)
  1096. accounts := []Account{
  1097. {
  1098. ID: 37101,
  1099. Platform: PlatformOpenAI,
  1100. Type: AccountTypeAPIKey,
  1101. Status: StatusActive,
  1102. Schedulable: true,
  1103. Concurrency: 1,
  1104. Priority: 100,
  1105. GroupIDs: []int64{groupID},
  1106. },
  1107. {
  1108. ID: 37102,
  1109. Platform: PlatformOpenAI,
  1110. Type: AccountTypeAPIKey,
  1111. Status: StatusActive,
  1112. Schedulable: true,
  1113. Concurrency: 1,
  1114. Priority: 0,
  1115. GroupIDs: []int64{groupID},
  1116. },
  1117. }
  1118. cfg := &config.Config{}
  1119. cfg.Gateway.OpenAIWS.LBTopK = 2
  1120. cfg.Gateway.OpenAIWS.SchedulerScoreWeights.Priority = 1
  1121. cfg.Gateway.OpenAIWS.SchedulerScoreWeights.Load = 1
  1122. cfg.Gateway.OpenAIWS.SchedulerScoreWeights.Queue = 0.7
  1123. cfg.Gateway.OpenAIWS.SchedulerScoreWeights.ErrorRate = 0.8
  1124. cfg.Gateway.OpenAIWS.SchedulerScoreWeights.TTFT = 0.5
  1125. cfg.Gateway.OpenAIWS.SchedulerScoreWeights.SessionSticky = 3
  1126. cache := &schedulerTestGatewayCache{sessionBindings: map[string]int64{
  1127. "openai:session_hash_weighted_topk": 37101,
  1128. }}
  1129. svc := &OpenAIGatewayService{
  1130. accountRepo: schedulerTestOpenAIAccountRepo{accounts: accounts},
  1131. cache: cache,
  1132. cfg: cfg,
  1133. rateLimitService: newOpenAIAdvancedSchedulerRateLimitService("true", "true"),
  1134. concurrencyService: NewConcurrencyService(schedulerTestConcurrencyCache{}),
  1135. }
  1136. selection, decision, err := svc.SelectAccountWithScheduler(
  1137. ctx,
  1138. &groupID,
  1139. "",
  1140. "session_hash_weighted_topk",
  1141. "gpt-5.1",
  1142. nil,
  1143. OpenAIUpstreamTransportAny,
  1144. false,
  1145. )
  1146. require.NoError(t, err)
  1147. require.NotNil(t, selection)
  1148. require.NotNil(t, selection.Account)
  1149. require.Equal(t, int64(37101), selection.Account.ID)
  1150. require.Equal(t, openAIAccountScheduleLayerLoadBalance, decision.Layer)
  1151. require.True(t, decision.StickySessionHit)
  1152. require.Equal(t, 2, decision.TopK)
  1153. if selection.ReleaseFunc != nil {
  1154. selection.ReleaseFunc()
  1155. }
  1156. }
  1157. func TestOpenAIGatewayService_SelectAccountWithScheduler_StickyWeightedPreviousRequiresMovableContext(t *testing.T) {
  1158. resetOpenAIAdvancedSchedulerSettingCacheForTest()
  1159. ctx := context.Background()
  1160. groupID := int64(101072)
  1161. accounts := []Account{
  1162. {
  1163. ID: 37111,
  1164. Platform: PlatformOpenAI,
  1165. Type: AccountTypeAPIKey,
  1166. Status: StatusActive,
  1167. Schedulable: true,
  1168. Concurrency: 1,
  1169. Priority: 100,
  1170. GroupIDs: []int64{groupID},
  1171. Extra: map[string]any{
  1172. "openai_apikey_responses_websockets_v2_enabled": true,
  1173. },
  1174. },
  1175. {
  1176. ID: 37112,
  1177. Platform: PlatformOpenAI,
  1178. Type: AccountTypeAPIKey,
  1179. Status: StatusActive,
  1180. Schedulable: true,
  1181. Concurrency: 1,
  1182. Priority: 0,
  1183. GroupIDs: []int64{groupID},
  1184. Extra: map[string]any{
  1185. "openai_apikey_responses_websockets_v2_enabled": true,
  1186. },
  1187. },
  1188. }
  1189. cfg := newSchedulerTestOpenAIWSV2Config()
  1190. cfg.Gateway.OpenAIWS.LBTopK = 1
  1191. cfg.Gateway.OpenAIWS.SchedulerScoreWeights.Priority = 1
  1192. cfg.Gateway.OpenAIWS.SchedulerScoreWeights.Load = 1
  1193. cfg.Gateway.OpenAIWS.SchedulerScoreWeights.Queue = 0.7
  1194. cfg.Gateway.OpenAIWS.SchedulerScoreWeights.ErrorRate = 0.8
  1195. cfg.Gateway.OpenAIWS.SchedulerScoreWeights.TTFT = 0.5
  1196. svc := &OpenAIGatewayService{
  1197. accountRepo: schedulerTestOpenAIAccountRepo{accounts: accounts},
  1198. cache: &schedulerTestGatewayCache{},
  1199. cfg: cfg,
  1200. rateLimitService: newOpenAIAdvancedSchedulerRateLimitService("true", "true"),
  1201. concurrencyService: NewConcurrencyService(schedulerTestConcurrencyCache{}),
  1202. }
  1203. store := svc.getOpenAIWSStateStore()
  1204. require.NoError(t, store.BindResponseAccount(ctx, groupID, "resp_weighted_unmovable", 37111, time.Hour))
  1205. selection, decision, err := svc.SelectAccountWithSchedulerForCapability(
  1206. ctx,
  1207. &groupID,
  1208. "resp_weighted_unmovable",
  1209. "",
  1210. "gpt-5.1",
  1211. nil,
  1212. OpenAIUpstreamTransportAny,
  1213. OpenAIEndpointCapabilityChatCompletions,
  1214. false,
  1215. false,
  1216. true,
  1217. PlatformOpenAI,
  1218. )
  1219. require.NoError(t, err)
  1220. require.NotNil(t, selection)
  1221. require.NotNil(t, selection.Account)
  1222. require.Equal(t, int64(37111), selection.Account.ID)
  1223. require.Equal(t, openAIAccountScheduleLayerPreviousResponse, decision.Layer)
  1224. require.True(t, decision.StickyPreviousHit)
  1225. if selection.ReleaseFunc != nil {
  1226. selection.ReleaseFunc()
  1227. }
  1228. selection, decision, err = svc.SelectAccountWithSchedulerForCapability(
  1229. ctx,
  1230. &groupID,
  1231. "resp_weighted_unmovable",
  1232. "",
  1233. "gpt-5.1",
  1234. nil,
  1235. OpenAIUpstreamTransportAny,
  1236. OpenAIEndpointCapabilityChatCompletions,
  1237. false,
  1238. true,
  1239. true,
  1240. PlatformOpenAI,
  1241. )
  1242. require.NoError(t, err)
  1243. require.NotNil(t, selection)
  1244. require.NotNil(t, selection.Account)
  1245. require.Equal(t, int64(37112), selection.Account.ID)
  1246. require.Equal(t, openAIAccountScheduleLayerLoadBalance, decision.Layer)
  1247. require.False(t, decision.StickyPreviousHit)
  1248. if selection.ReleaseFunc != nil {
  1249. selection.ReleaseFunc()
  1250. }
  1251. }
  1252. func TestOpenAIGatewayService_SelectAccountWithScheduler_PreviousResponseCompactUnsupportedDeletesBinding(t *testing.T) {
  1253. resetOpenAIAdvancedSchedulerSettingCacheForTest()
  1254. ctx := context.Background()
  1255. groupID := int64(101073)
  1256. accounts := []Account{
  1257. {
  1258. ID: 37121,
  1259. Platform: PlatformOpenAI,
  1260. Type: AccountTypeAPIKey,
  1261. Status: StatusActive,
  1262. Schedulable: true,
  1263. Concurrency: 1,
  1264. Priority: 0,
  1265. GroupIDs: []int64{groupID},
  1266. Extra: map[string]any{
  1267. "openai_apikey_responses_websockets_v2_enabled": true,
  1268. "openai_compact_mode": OpenAICompactModeForceOff,
  1269. },
  1270. },
  1271. {
  1272. ID: 37122,
  1273. Platform: PlatformOpenAI,
  1274. Type: AccountTypeAPIKey,
  1275. Status: StatusActive,
  1276. Schedulable: true,
  1277. Concurrency: 1,
  1278. Priority: 10,
  1279. GroupIDs: []int64{groupID},
  1280. Extra: map[string]any{
  1281. "openai_apikey_responses_websockets_v2_enabled": true,
  1282. "openai_compact_mode": OpenAICompactModeForceOn,
  1283. },
  1284. },
  1285. }
  1286. cfg := newSchedulerTestOpenAIWSV2Config()
  1287. cfg.Gateway.OpenAIWS.LBTopK = 2
  1288. cfg.Gateway.OpenAIWS.SchedulerScoreWeights.Priority = 1
  1289. cfg.Gateway.OpenAIWS.SchedulerScoreWeights.Load = 1
  1290. cfg.Gateway.OpenAIWS.SchedulerScoreWeights.Queue = 0.7
  1291. cfg.Gateway.OpenAIWS.SchedulerScoreWeights.ErrorRate = 0.8
  1292. cfg.Gateway.OpenAIWS.SchedulerScoreWeights.TTFT = 0.5
  1293. svc := &OpenAIGatewayService{
  1294. accountRepo: schedulerTestOpenAIAccountRepo{accounts: accounts},
  1295. cache: &schedulerTestGatewayCache{},
  1296. cfg: cfg,
  1297. rateLimitService: newOpenAIAdvancedSchedulerRateLimitService("true"),
  1298. concurrencyService: NewConcurrencyService(schedulerTestConcurrencyCache{}),
  1299. }
  1300. store := svc.getOpenAIWSStateStore()
  1301. require.NoError(t, store.BindResponseAccount(ctx, groupID, "resp_compact_unsupported", 37121, time.Hour))
  1302. selection, decision, err := svc.SelectAccountWithScheduler(
  1303. ctx,
  1304. &groupID,
  1305. "resp_compact_unsupported",
  1306. "",
  1307. "gpt-5.1",
  1308. nil,
  1309. OpenAIUpstreamTransportAny,
  1310. true,
  1311. )
  1312. require.NoError(t, err)
  1313. require.NotNil(t, selection)
  1314. require.NotNil(t, selection.Account)
  1315. require.Equal(t, int64(37122), selection.Account.ID)
  1316. require.Equal(t, openAIAccountScheduleLayerLoadBalance, decision.Layer)
  1317. require.False(t, decision.StickyPreviousHit)
  1318. if selection.ReleaseFunc != nil {
  1319. selection.ReleaseFunc()
  1320. }
  1321. accountID, err := store.GetResponseAccount(ctx, groupID, "resp_compact_unsupported")
  1322. require.NoError(t, err)
  1323. require.Zero(t, accountID)
  1324. }
  1325. func TestOpenAIGatewayService_SelectAccountWithScheduler_Enabled_EmbeddingsSkipsChatOnlyAccount(t *testing.T) {
  1326. resetOpenAIAdvancedSchedulerSettingCacheForTest()
  1327. ctx := context.Background()
  1328. groupID := int64(10111)
  1329. accounts := []Account{
  1330. {
  1331. ID: 37011,
  1332. Platform: PlatformOpenAI,
  1333. Type: AccountTypeAPIKey,
  1334. Status: StatusActive,
  1335. Schedulable: true,
  1336. Concurrency: 1,
  1337. Priority: 0,
  1338. Credentials: map[string]any{
  1339. "openai_capabilities": []any{"chat_completions"},
  1340. },
  1341. },
  1342. {
  1343. ID: 37012,
  1344. Platform: PlatformOpenAI,
  1345. Type: AccountTypeAPIKey,
  1346. Status: StatusActive,
  1347. Schedulable: true,
  1348. Concurrency: 1,
  1349. Priority: 5,
  1350. Credentials: map[string]any{
  1351. "openai_capabilities": []any{"chat_completions", "embeddings"},
  1352. },
  1353. },
  1354. }
  1355. cfg := &config.Config{}
  1356. cfg.Gateway.Scheduling.LoadBatchEnabled = false
  1357. svc := &OpenAIGatewayService{
  1358. accountRepo: schedulerTestOpenAIAccountRepo{accounts: accounts},
  1359. cache: &schedulerTestGatewayCache{},
  1360. cfg: cfg,
  1361. rateLimitService: newOpenAIAdvancedSchedulerRateLimitService("true"),
  1362. concurrencyService: NewConcurrencyService(schedulerTestConcurrencyCache{}),
  1363. }
  1364. selection, decision, err := svc.SelectAccountWithSchedulerForCapability(
  1365. ctx,
  1366. &groupID,
  1367. "",
  1368. "",
  1369. "text-embedding-3-small",
  1370. nil,
  1371. OpenAIUpstreamTransportHTTPSSE,
  1372. OpenAIEndpointCapabilityEmbeddings,
  1373. false,
  1374. false,
  1375. true,
  1376. )
  1377. require.NoError(t, err)
  1378. require.NotNil(t, selection)
  1379. require.NotNil(t, selection.Account)
  1380. require.Equal(t, int64(37012), selection.Account.ID)
  1381. require.Equal(t, openAIAccountScheduleLayerLoadBalance, decision.Layer)
  1382. require.Equal(t, 1, decision.CandidateCount)
  1383. }
  1384. func TestOpenAIGatewayService_SelectAccountWithScheduler_Enabled_EmbeddingsSkipsChatOnlyStickyBindings(t *testing.T) {
  1385. resetOpenAIAdvancedSchedulerSettingCacheForTest()
  1386. ctx := context.Background()
  1387. groupID := int64(10112)
  1388. accounts := []Account{
  1389. {
  1390. ID: 37021,
  1391. Platform: PlatformOpenAI,
  1392. Type: AccountTypeAPIKey,
  1393. Status: StatusActive,
  1394. Schedulable: true,
  1395. Concurrency: 1,
  1396. Priority: 0,
  1397. Credentials: map[string]any{
  1398. "openai_capabilities": []any{"chat_completions"},
  1399. },
  1400. Extra: map[string]any{
  1401. "openai_apikey_responses_websockets_v2_enabled": true,
  1402. },
  1403. },
  1404. {
  1405. ID: 37022,
  1406. Platform: PlatformOpenAI,
  1407. Type: AccountTypeAPIKey,
  1408. Status: StatusActive,
  1409. Schedulable: true,
  1410. Concurrency: 1,
  1411. Priority: 5,
  1412. Credentials: map[string]any{
  1413. "openai_capabilities": []any{"chat_completions", "embeddings"},
  1414. },
  1415. Extra: map[string]any{
  1416. "openai_apikey_responses_websockets_v2_enabled": true,
  1417. },
  1418. },
  1419. }
  1420. cfg := newSchedulerTestOpenAIWSV2Config()
  1421. cfg.Gateway.Scheduling.LoadBatchEnabled = false
  1422. cache := &schedulerTestGatewayCache{
  1423. sessionBindings: map[string]int64{
  1424. "openai:session_hash_embeddings": 37021,
  1425. },
  1426. }
  1427. svc := &OpenAIGatewayService{
  1428. accountRepo: schedulerTestOpenAIAccountRepo{accounts: accounts},
  1429. cache: cache,
  1430. cfg: cfg,
  1431. rateLimitService: newOpenAIAdvancedSchedulerRateLimitService("true"),
  1432. concurrencyService: NewConcurrencyService(schedulerTestConcurrencyCache{}),
  1433. }
  1434. store := svc.getOpenAIWSStateStore()
  1435. require.NoError(t, store.BindResponseAccount(ctx, groupID, "resp_embeddings_chat_only", 37021, time.Hour))
  1436. selection, decision, err := svc.SelectAccountWithSchedulerForCapability(
  1437. ctx,
  1438. &groupID,
  1439. "resp_embeddings_chat_only",
  1440. "session_hash_embeddings",
  1441. "text-embedding-3-small",
  1442. nil,
  1443. OpenAIUpstreamTransportHTTPSSE,
  1444. OpenAIEndpointCapabilityEmbeddings,
  1445. false,
  1446. false,
  1447. true,
  1448. )
  1449. require.NoError(t, err)
  1450. require.NotNil(t, selection)
  1451. require.NotNil(t, selection.Account)
  1452. require.Equal(t, int64(37022), selection.Account.ID)
  1453. require.Equal(t, openAIAccountScheduleLayerLoadBalance, decision.Layer)
  1454. require.False(t, decision.StickyPreviousHit)
  1455. require.False(t, decision.StickySessionHit)
  1456. require.Equal(t, int64(37022), cache.sessionBindings["openai:session_hash_embeddings"])
  1457. }
  1458. func TestOpenAIGatewayService_OpenAIAccountSchedulerMetrics_DisabledNoOp(t *testing.T) {
  1459. resetOpenAIAdvancedSchedulerSettingCacheForTest()
  1460. svc := &OpenAIGatewayService{}
  1461. ttft := 120
  1462. svc.ReportOpenAIAccountScheduleResult(&Account{ID: 10}, "", true, &ttft)
  1463. svc.RecordOpenAIAccountSwitch()
  1464. snapshot := svc.SnapshotOpenAIAccountSchedulerMetrics()
  1465. require.Equal(t, OpenAIAccountSchedulerMetricsSnapshot{}, snapshot)
  1466. }
  1467. func TestOpenAIGatewayService_SelectAccountWithScheduler_SkipsQuarantinedSharedProxy(t *testing.T) {
  1468. resetOpenAIAdvancedSchedulerSettingCacheForTest()
  1469. proxyA := int64(4698)
  1470. proxyB := int64(4699)
  1471. accounts := []Account{
  1472. {ID: 469801, Platform: PlatformOpenAI, Type: AccountTypeOAuth, Status: StatusActive, Schedulable: true, Concurrency: 1, Priority: 0, ProxyID: &proxyA},
  1473. {ID: 469802, Platform: PlatformOpenAI, Type: AccountTypeOAuth, Status: StatusActive, Schedulable: true, Concurrency: 1, Priority: 0, ProxyID: &proxyA},
  1474. {ID: 469803, Platform: PlatformOpenAI, Type: AccountTypeOAuth, Status: StatusActive, Schedulable: true, Concurrency: 1, Priority: 5, ProxyID: &proxyB},
  1475. }
  1476. cfg := &config.Config{}
  1477. cfg.Gateway.Scheduling.LoadBatchEnabled = false
  1478. svc := &OpenAIGatewayService{
  1479. accountRepo: schedulerTestOpenAIAccountRepo{accounts: accounts},
  1480. cfg: cfg,
  1481. concurrencyService: NewConcurrencyService(schedulerTestConcurrencyCache{}),
  1482. openaiProxyStreamCircuit: newOpenAIProxyStreamCircuit(openAIProxyStreamCircuitSettings{
  1483. failureThreshold: 1,
  1484. failureWindow: time.Minute,
  1485. quarantineTTL: 10 * time.Minute,
  1486. maxEntries: 16,
  1487. }),
  1488. }
  1489. svc.openaiProxyStreamCircuit.recordFailure(proxyA, time.Now())
  1490. selection, _, err := svc.SelectAccountWithScheduler(
  1491. context.Background(), nil, "", "", "gpt-5.6-sol", nil, OpenAIUpstreamTransportAny, false,
  1492. )
  1493. require.NoError(t, err)
  1494. require.NotNil(t, selection)
  1495. require.NotNil(t, selection.Account)
  1496. require.Equal(t, int64(469803), selection.Account.ID)
  1497. }
  1498. func TestOpenAIGatewayService_SelectAccountWithScheduler_FailsOpenWhenAllProxiesQuarantined(t *testing.T) {
  1499. resetOpenAIAdvancedSchedulerSettingCacheForTest()
  1500. // Every schedulable account shares the quarantined proxy: the circuit must
  1501. // degrade to a preference instead of zeroing out capacity (#5056).
  1502. proxyA := int64(5056)
  1503. accounts := []Account{
  1504. {ID: 505601, Platform: PlatformOpenAI, Type: AccountTypeOAuth, Status: StatusActive, Schedulable: true, Concurrency: 1, Priority: 0, ProxyID: &proxyA},
  1505. {ID: 505602, Platform: PlatformOpenAI, Type: AccountTypeOAuth, Status: StatusActive, Schedulable: true, Concurrency: 1, Priority: 5, ProxyID: &proxyA},
  1506. }
  1507. cfg := &config.Config{}
  1508. cfg.Gateway.Scheduling.LoadBatchEnabled = false
  1509. svc := &OpenAIGatewayService{
  1510. accountRepo: schedulerTestOpenAIAccountRepo{accounts: accounts},
  1511. cfg: cfg,
  1512. concurrencyService: NewConcurrencyService(schedulerTestConcurrencyCache{}),
  1513. openaiProxyStreamCircuit: newOpenAIProxyStreamCircuit(openAIProxyStreamCircuitSettings{
  1514. failureThreshold: 1,
  1515. failureWindow: time.Minute,
  1516. quarantineTTL: 10 * time.Minute,
  1517. maxEntries: 16,
  1518. }),
  1519. }
  1520. tripped, _ := svc.openaiProxyStreamCircuit.recordFailure(proxyA, time.Now())
  1521. require.True(t, tripped)
  1522. selection, _, err := svc.SelectAccountWithScheduler(
  1523. context.Background(), nil, "", "", "gpt-5.6-sol", nil, OpenAIUpstreamTransportAny, false,
  1524. )
  1525. require.NoError(t, err, "quarantine must fail open instead of returning no available accounts")
  1526. require.NotNil(t, selection)
  1527. require.NotNil(t, selection.Account)
  1528. require.NotNil(t, selection.Account.ProxyID)
  1529. require.Equal(t, proxyA, *selection.Account.ProxyID)
  1530. require.True(t, svc.openaiProxyStreamCircuit.isBlocked(proxyA, time.Now()),
  1531. "fail-open must not clear the quarantine; only a completed stream or TTL expiry does")
  1532. }
  1533. func TestOpenAIGatewayService_SelectAccountWithScheduler_SessionStickyRateLimitedAccountFallsBackToFreshCandidate(t *testing.T) {
  1534. ctx := context.Background()
  1535. groupID := int64(10101)
  1536. rateLimitedUntil := time.Now().Add(30 * time.Minute)
  1537. staleSticky := &Account{ID: 31001, Platform: PlatformOpenAI, Type: AccountTypeOAuth, Status: StatusActive, Schedulable: true, Concurrency: 1, Priority: 0, GroupIDs: []int64{groupID}}
  1538. staleBackup := &Account{ID: 31002, Platform: PlatformOpenAI, Type: AccountTypeOAuth, Status: StatusActive, Schedulable: true, Concurrency: 1, Priority: 5, GroupIDs: []int64{groupID}}
  1539. freshSticky := &Account{ID: 31001, Platform: PlatformOpenAI, Type: AccountTypeOAuth, Status: StatusActive, Schedulable: true, Concurrency: 1, Priority: 0, GroupIDs: []int64{groupID}, RateLimitResetAt: &rateLimitedUntil}
  1540. freshBackup := &Account{ID: 31002, Platform: PlatformOpenAI, Type: AccountTypeOAuth, Status: StatusActive, Schedulable: true, Concurrency: 1, Priority: 5, GroupIDs: []int64{groupID}}
  1541. cache := &schedulerTestGatewayCache{sessionBindings: map[string]int64{"openai:session_hash_rate_limited": 31001}}
  1542. snapshotCache := &openAISnapshotCacheStub{snapshotAccounts: []*Account{staleSticky, staleBackup}, accountsByID: map[int64]*Account{31001: freshSticky, 31002: freshBackup}}
  1543. snapshotService := &SchedulerSnapshotService{cache: snapshotCache}
  1544. svc := &OpenAIGatewayService{
  1545. accountRepo: schedulerTestOpenAIAccountRepo{accounts: []Account{*freshSticky, *freshBackup}},
  1546. cache: cache,
  1547. cfg: &config.Config{},
  1548. rateLimitService: newOpenAIAdvancedSchedulerRateLimitService("true"),
  1549. schedulerSnapshot: snapshotService,
  1550. concurrencyService: NewConcurrencyService(schedulerTestConcurrencyCache{}),
  1551. }
  1552. selection, decision, err := svc.SelectAccountWithScheduler(ctx, &groupID, "", "session_hash_rate_limited", "gpt-5.1", nil, OpenAIUpstreamTransportAny, false)
  1553. require.NoError(t, err)
  1554. require.NotNil(t, selection)
  1555. require.NotNil(t, selection.Account)
  1556. require.Equal(t, int64(31002), selection.Account.ID)
  1557. require.Equal(t, openAIAccountScheduleLayerLoadBalance, decision.Layer)
  1558. }
  1559. func TestOpenAIGatewayService_SelectAccountForModelWithExclusions_AutoPauseBy5hThreshold(t *testing.T) {
  1560. ctx := context.Background()
  1561. primary := Account{
  1562. ID: 35001,
  1563. Platform: PlatformOpenAI,
  1564. Type: AccountTypeAPIKey,
  1565. Status: StatusActive,
  1566. Schedulable: true,
  1567. Concurrency: 1,
  1568. Priority: 0,
  1569. Extra: map[string]any{
  1570. "codex_5h_used_percent": 95.0,
  1571. "auto_pause_5h_threshold": 0.95,
  1572. },
  1573. }
  1574. secondary := Account{ID: 35002, Platform: PlatformOpenAI, Type: AccountTypeAPIKey, Status: StatusActive, Schedulable: true, Concurrency: 1, Priority: 5}
  1575. svc := &OpenAIGatewayService{accountRepo: schedulerTestOpenAIAccountRepo{accounts: []Account{primary, secondary}}, cfg: &config.Config{}}
  1576. account, err := svc.SelectAccountForModelWithExclusions(ctx, nil, "", "gpt-5.1", nil)
  1577. require.NoError(t, err)
  1578. require.NotNil(t, account)
  1579. require.Equal(t, int64(35002), account.ID)
  1580. }
  1581. func TestOpenAIGatewayService_SelectAccountForModelWithExclusions_AllowsBelow5hThreshold(t *testing.T) {
  1582. ctx := context.Background()
  1583. primary := Account{
  1584. ID: 35101,
  1585. Platform: PlatformOpenAI,
  1586. Type: AccountTypeAPIKey,
  1587. Status: StatusActive,
  1588. Schedulable: true,
  1589. Concurrency: 1,
  1590. Priority: 0,
  1591. Extra: map[string]any{
  1592. "codex_5h_used_percent": 80.0,
  1593. "auto_pause_5h_threshold": 0.95,
  1594. },
  1595. }
  1596. secondary := Account{ID: 35102, Platform: PlatformOpenAI, Type: AccountTypeAPIKey, Status: StatusActive, Schedulable: true, Concurrency: 1, Priority: 5}
  1597. svc := &OpenAIGatewayService{accountRepo: schedulerTestOpenAIAccountRepo{accounts: []Account{primary, secondary}}, cfg: &config.Config{}}
  1598. account, err := svc.SelectAccountForModelWithExclusions(ctx, nil, "", "gpt-5.1", nil)
  1599. require.NoError(t, err)
  1600. require.NotNil(t, account)
  1601. require.Equal(t, int64(35101), account.ID)
  1602. }
  1603. func TestOpenAIGatewayService_SelectAccountForModelWithExclusions_AutoPauseBy7dThreshold(t *testing.T) {
  1604. ctx := context.Background()
  1605. primary := Account{
  1606. ID: 35201,
  1607. Platform: PlatformOpenAI,
  1608. Type: AccountTypeAPIKey,
  1609. Status: StatusActive,
  1610. Schedulable: true,
  1611. Concurrency: 1,
  1612. Priority: 0,
  1613. Extra: map[string]any{
  1614. "codex_7d_used_percent": 95.0,
  1615. "auto_pause_7d_threshold": 0.95,
  1616. },
  1617. }
  1618. secondary := Account{ID: 35202, Platform: PlatformOpenAI, Type: AccountTypeAPIKey, Status: StatusActive, Schedulable: true, Concurrency: 1, Priority: 5}
  1619. svc := &OpenAIGatewayService{accountRepo: schedulerTestOpenAIAccountRepo{accounts: []Account{primary, secondary}}, cfg: &config.Config{}}
  1620. account, err := svc.SelectAccountForModelWithExclusions(ctx, nil, "", "gpt-5.1", nil)
  1621. require.NoError(t, err)
  1622. require.NotNil(t, account)
  1623. require.Equal(t, int64(35202), account.ID)
  1624. }
  1625. func TestOpenAIGatewayService_SelectAccountForModelWithExclusions_UnconfiguredThresholdKeepsLegacyBehavior(t *testing.T) {
  1626. ctx := context.Background()
  1627. primary := Account{ID: 35301, Platform: PlatformOpenAI, Type: AccountTypeAPIKey, Status: StatusActive, Schedulable: true, Concurrency: 1, Priority: 0, Extra: map[string]any{"codex_5h_used_percent": 99.0, "codex_7d_used_percent": 99.0}}
  1628. secondary := Account{ID: 35302, Platform: PlatformOpenAI, Type: AccountTypeAPIKey, Status: StatusActive, Schedulable: true, Concurrency: 1, Priority: 5}
  1629. svc := &OpenAIGatewayService{accountRepo: schedulerTestOpenAIAccountRepo{accounts: []Account{primary, secondary}}, cfg: &config.Config{}}
  1630. account, err := svc.SelectAccountForModelWithExclusions(ctx, nil, "", "gpt-5.1", nil)
  1631. require.NoError(t, err)
  1632. require.NotNil(t, account)
  1633. require.Equal(t, int64(35301), account.ID)
  1634. }
  1635. func TestOpenAIGatewayService_SelectAccountForModelWithExclusions_UsesGlobalDefaultThreshold(t *testing.T) {
  1636. ctx := withOpenAIQuotaAutoPauseSettings(context.Background(), OpsOpenAIAccountQuotaAutoPauseSettings{DefaultThreshold5h: 0.95})
  1637. primary := Account{
  1638. ID: 35401,
  1639. Platform: PlatformOpenAI,
  1640. Type: AccountTypeAPIKey,
  1641. Status: StatusActive,
  1642. Schedulable: true,
  1643. Concurrency: 1,
  1644. Priority: 0,
  1645. Extra: map[string]any{
  1646. "codex_5h_used_percent": 95.0,
  1647. },
  1648. }
  1649. secondary := Account{ID: 35402, Platform: PlatformOpenAI, Type: AccountTypeAPIKey, Status: StatusActive, Schedulable: true, Concurrency: 1, Priority: 5}
  1650. svc := &OpenAIGatewayService{accountRepo: schedulerTestOpenAIAccountRepo{accounts: []Account{primary, secondary}}, cfg: &config.Config{}}
  1651. account, err := svc.SelectAccountForModelWithExclusions(ctx, nil, "", "gpt-5.1", nil)
  1652. require.NoError(t, err)
  1653. require.NotNil(t, account)
  1654. require.Equal(t, int64(35402), account.ID)
  1655. }
  1656. // Regression: a per-account explicit-disable flag exempts the account from auto-pause
  1657. // even when a global default threshold is set. Without this, "leave threshold blank"
  1658. // silently falls back to global default and admins have no way to whitelist a single
  1659. // account.
  1660. func TestOpenAIGatewayService_SelectAccountForModelWithExclusions_PerAccountDisableOverridesGlobalDefault(t *testing.T) {
  1661. ctx := withOpenAIQuotaAutoPauseSettings(context.Background(), OpsOpenAIAccountQuotaAutoPauseSettings{DefaultThreshold5h: 0.95})
  1662. // Account has high usage AND no per-account threshold (would normally fall back to
  1663. // the global default and get paused), but the explicit disable flag is set.
  1664. primary := Account{
  1665. ID: 35701,
  1666. Platform: PlatformOpenAI,
  1667. Type: AccountTypeAPIKey,
  1668. Status: StatusActive,
  1669. Schedulable: true,
  1670. Concurrency: 1,
  1671. Priority: 0,
  1672. Extra: map[string]any{
  1673. "codex_5h_used_percent": 99.0,
  1674. "auto_pause_5h_disabled": true,
  1675. },
  1676. }
  1677. secondary := Account{ID: 35702, Platform: PlatformOpenAI, Type: AccountTypeAPIKey, Status: StatusActive, Schedulable: true, Concurrency: 1, Priority: 5}
  1678. svc := &OpenAIGatewayService{accountRepo: schedulerTestOpenAIAccountRepo{accounts: []Account{primary, secondary}}, cfg: &config.Config{}}
  1679. account, err := svc.SelectAccountForModelWithExclusions(ctx, nil, "", "gpt-5.1", nil)
  1680. require.NoError(t, err)
  1681. require.NotNil(t, account)
  1682. require.Equal(t, int64(35701), account.ID)
  1683. }
  1684. // Disable is per-window: disabling only 5h must still allow 7d auto-pause to fire.
  1685. func TestOpenAIGatewayService_SelectAccountForModelWithExclusions_PerWindowDisableScoped(t *testing.T) {
  1686. ctx := context.Background()
  1687. primary := Account{
  1688. ID: 35801,
  1689. Platform: PlatformOpenAI,
  1690. Type: AccountTypeAPIKey,
  1691. Status: StatusActive,
  1692. Schedulable: true,
  1693. Concurrency: 1,
  1694. Priority: 0,
  1695. Extra: map[string]any{
  1696. "codex_5h_used_percent": 99.0,
  1697. "codex_7d_used_percent": 99.0,
  1698. "auto_pause_5h_disabled": true,
  1699. "auto_pause_7d_threshold": 0.95,
  1700. },
  1701. }
  1702. secondary := Account{ID: 35802, Platform: PlatformOpenAI, Type: AccountTypeAPIKey, Status: StatusActive, Schedulable: true, Concurrency: 1, Priority: 5}
  1703. svc := &OpenAIGatewayService{accountRepo: schedulerTestOpenAIAccountRepo{accounts: []Account{primary, secondary}}, cfg: &config.Config{}}
  1704. account, err := svc.SelectAccountForModelWithExclusions(ctx, nil, "", "gpt-5.1", nil)
  1705. require.NoError(t, err)
  1706. require.NotNil(t, account)
  1707. require.Equal(t, int64(35802), account.ID, "7d auto-pause must still fire even though 5h is disabled")
  1708. }
  1709. func TestOpenAIGatewayService_SelectAccountForModelWithExclusions_StaleUsageWindowResetSkipsPause(t *testing.T) {
  1710. ctx := context.Background()
  1711. // Usage is over threshold but the window's reset time has already passed, so the
  1712. // cached percentage is stale (the real window rolled over) and the account must NOT
  1713. // stay paused — otherwise it could be skipped forever with no traffic to refresh it.
  1714. primary := Account{
  1715. ID: 35501,
  1716. Platform: PlatformOpenAI,
  1717. Type: AccountTypeAPIKey,
  1718. Status: StatusActive,
  1719. Schedulable: true,
  1720. Concurrency: 1,
  1721. Priority: 0,
  1722. Extra: map[string]any{
  1723. "codex_5h_used_percent": 99.0,
  1724. "auto_pause_5h_threshold": 0.95,
  1725. "codex_5h_reset_at": time.Now().Add(-time.Minute).Format(time.RFC3339),
  1726. },
  1727. }
  1728. secondary := Account{ID: 35502, Platform: PlatformOpenAI, Type: AccountTypeAPIKey, Status: StatusActive, Schedulable: true, Concurrency: 1, Priority: 5}
  1729. svc := &OpenAIGatewayService{accountRepo: schedulerTestOpenAIAccountRepo{accounts: []Account{primary, secondary}}, cfg: &config.Config{}}
  1730. account, err := svc.SelectAccountForModelWithExclusions(ctx, nil, "", "gpt-5.1", nil)
  1731. require.NoError(t, err)
  1732. require.NotNil(t, account)
  1733. require.Equal(t, int64(35501), account.ID)
  1734. }
  1735. func TestOpenAIGatewayService_SelectAccountForModelWithExclusions_FreshUsageWindowStillPauses(t *testing.T) {
  1736. ctx := context.Background()
  1737. // Same as above but the window has not reset yet, so the account stays paused.
  1738. primary := Account{
  1739. ID: 35601,
  1740. Platform: PlatformOpenAI,
  1741. Type: AccountTypeAPIKey,
  1742. Status: StatusActive,
  1743. Schedulable: true,
  1744. Concurrency: 1,
  1745. Priority: 0,
  1746. Extra: map[string]any{
  1747. "codex_5h_used_percent": 99.0,
  1748. "auto_pause_5h_threshold": 0.95,
  1749. "codex_5h_reset_at": time.Now().Add(time.Hour).Format(time.RFC3339),
  1750. },
  1751. }
  1752. secondary := Account{ID: 35602, Platform: PlatformOpenAI, Type: AccountTypeAPIKey, Status: StatusActive, Schedulable: true, Concurrency: 1, Priority: 5}
  1753. svc := &OpenAIGatewayService{accountRepo: schedulerTestOpenAIAccountRepo{accounts: []Account{primary, secondary}}, cfg: &config.Config{}}
  1754. account, err := svc.SelectAccountForModelWithExclusions(ctx, nil, "", "gpt-5.1", nil)
  1755. require.NoError(t, err)
  1756. require.NotNil(t, account)
  1757. require.Equal(t, int64(35602), account.ID)
  1758. }
  1759. // Issue #2994: an account poisoned with an inflated used% (e.g. from the reverted #2918
  1760. // inversion) gets excluded from scheduling, and a paused account never receives traffic to
  1761. // refresh its snapshot. When the snapshot is stale (codex_usage_updated_at older than the
  1762. // staleness bound) the account must be allowed a request so it can self-heal from the real
  1763. // response headers — independent of the window's reset time.
  1764. func TestOpenAIGatewayService_SelectAccountForModelWithExclusions_StaleUsageSnapshotSkipsPause_Issue2994(t *testing.T) {
  1765. ctx := context.Background()
  1766. primary := Account{
  1767. ID: 35701,
  1768. Platform: PlatformOpenAI,
  1769. Type: AccountTypeAPIKey,
  1770. Status: StatusActive,
  1771. Schedulable: true,
  1772. Concurrency: 1,
  1773. Priority: 0,
  1774. Extra: map[string]any{
  1775. "codex_5h_used_percent": 99.0,
  1776. "auto_pause_5h_threshold": 0.95,
  1777. // Window has NOT reset yet, so the reset guard stays inactive.
  1778. "codex_5h_reset_at": time.Now().Add(time.Hour).Format(time.RFC3339),
  1779. // Snapshot is stale: older than openAICodexAutoPauseStaleAfter (2h).
  1780. "codex_usage_updated_at": time.Now().Add(-3 * time.Hour).Format(time.RFC3339),
  1781. },
  1782. }
  1783. secondary := Account{ID: 35702, Platform: PlatformOpenAI, Type: AccountTypeAPIKey, Status: StatusActive, Schedulable: true, Concurrency: 1, Priority: 5}
  1784. svc := &OpenAIGatewayService{accountRepo: schedulerTestOpenAIAccountRepo{accounts: []Account{primary, secondary}}, cfg: &config.Config{}}
  1785. account, err := svc.SelectAccountForModelWithExclusions(ctx, nil, "", "gpt-5.1", nil)
  1786. require.NoError(t, err)
  1787. require.NotNil(t, account)
  1788. require.Equal(t, int64(35701), account.ID)
  1789. }
  1790. // Issue #2994 guardrail: a genuinely-exhausted account whose snapshot was refreshed recently
  1791. // (codex_usage_updated_at fresh) must STILL be auto-paused. The stale self-heal must not let a
  1792. // real 99%-used account escape pause.
  1793. func TestOpenAIGatewayService_SelectAccountForModelWithExclusions_FreshExhaustedSnapshotStillPauses_Issue2994(t *testing.T) {
  1794. ctx := context.Background()
  1795. primary := Account{
  1796. ID: 35801,
  1797. Platform: PlatformOpenAI,
  1798. Type: AccountTypeAPIKey,
  1799. Status: StatusActive,
  1800. Schedulable: true,
  1801. Concurrency: 1,
  1802. Priority: 0,
  1803. Extra: map[string]any{
  1804. "codex_5h_used_percent": 99.0,
  1805. "auto_pause_5h_threshold": 0.95,
  1806. "codex_5h_reset_at": time.Now().Add(time.Hour).Format(time.RFC3339),
  1807. // Snapshot refreshed 1 minute ago: not stale, so the account stays paused.
  1808. "codex_usage_updated_at": time.Now().Add(-time.Minute).Format(time.RFC3339),
  1809. },
  1810. }
  1811. secondary := Account{ID: 35802, Platform: PlatformOpenAI, Type: AccountTypeAPIKey, Status: StatusActive, Schedulable: true, Concurrency: 1, Priority: 5}
  1812. svc := &OpenAIGatewayService{accountRepo: schedulerTestOpenAIAccountRepo{accounts: []Account{primary, secondary}}, cfg: &config.Config{}}
  1813. account, err := svc.SelectAccountForModelWithExclusions(ctx, nil, "", "gpt-5.1", nil)
  1814. require.NoError(t, err)
  1815. require.NotNil(t, account)
  1816. require.Equal(t, int64(35802), account.ID)
  1817. }
  1818. func TestOpenAIGatewayService_SelectAccountForModelWithExclusions_SkipsFreshlyRateLimitedSnapshotCandidate(t *testing.T) {
  1819. ctx := context.Background()
  1820. groupID := int64(10102)
  1821. rateLimitedUntil := time.Now().Add(30 * time.Minute)
  1822. stalePrimary := &Account{ID: 32001, Platform: PlatformOpenAI, Type: AccountTypeOAuth, Status: StatusActive, Schedulable: true, Concurrency: 1, Priority: 0, GroupIDs: []int64{groupID}}
  1823. staleSecondary := &Account{ID: 32002, Platform: PlatformOpenAI, Type: AccountTypeOAuth, Status: StatusActive, Schedulable: true, Concurrency: 1, Priority: 5, GroupIDs: []int64{groupID}}
  1824. freshPrimary := &Account{ID: 32001, Platform: PlatformOpenAI, Type: AccountTypeOAuth, Status: StatusActive, Schedulable: true, Concurrency: 1, Priority: 0, GroupIDs: []int64{groupID}, RateLimitResetAt: &rateLimitedUntil}
  1825. freshSecondary := &Account{ID: 32002, Platform: PlatformOpenAI, Type: AccountTypeOAuth, Status: StatusActive, Schedulable: true, Concurrency: 1, Priority: 5, GroupIDs: []int64{groupID}}
  1826. snapshotCache := &openAISnapshotCacheStub{snapshotAccounts: []*Account{stalePrimary, staleSecondary}, accountsByID: map[int64]*Account{32001: freshPrimary, 32002: freshSecondary}}
  1827. snapshotService := &SchedulerSnapshotService{cache: snapshotCache}
  1828. svc := &OpenAIGatewayService{
  1829. accountRepo: schedulerTestOpenAIAccountRepo{accounts: []Account{*freshPrimary, *freshSecondary}},
  1830. cfg: &config.Config{},
  1831. rateLimitService: newOpenAIAdvancedSchedulerRateLimitService("true"),
  1832. schedulerSnapshot: snapshotService,
  1833. }
  1834. account, err := svc.SelectAccountForModelWithExclusions(ctx, &groupID, "", "gpt-5.1", nil)
  1835. require.NoError(t, err)
  1836. require.NotNil(t, account)
  1837. require.Equal(t, int64(32002), account.ID)
  1838. }
  1839. func TestOpenAIGatewayService_SelectAccountForModelWithExclusions_ModelRateLimitOnlySkipsThatModel(t *testing.T) {
  1840. ctx := context.Background()
  1841. resetAt := time.Now().Add(30 * time.Minute).Format(time.RFC3339)
  1842. primary := Account{
  1843. ID: 32101,
  1844. Platform: PlatformOpenAI,
  1845. Type: AccountTypeAPIKey,
  1846. Status: StatusActive,
  1847. Schedulable: true,
  1848. Concurrency: 1,
  1849. Priority: 0,
  1850. Extra: map[string]any{
  1851. modelRateLimitsKey: map[string]any{
  1852. "gpt-5.4": map[string]any{
  1853. "rate_limit_reset_at": resetAt,
  1854. },
  1855. },
  1856. },
  1857. }
  1858. secondary := Account{
  1859. ID: 32102,
  1860. Platform: PlatformOpenAI,
  1861. Type: AccountTypeAPIKey,
  1862. Status: StatusActive,
  1863. Schedulable: true,
  1864. Concurrency: 1,
  1865. Priority: 5,
  1866. }
  1867. svc := &OpenAIGatewayService{
  1868. accountRepo: schedulerTestOpenAIAccountRepo{accounts: []Account{primary, secondary}},
  1869. cfg: &config.Config{},
  1870. }
  1871. account, err := svc.SelectAccountForModelWithExclusions(ctx, nil, "", "gpt-5.4", nil)
  1872. require.NoError(t, err)
  1873. require.NotNil(t, account)
  1874. require.Equal(t, int64(32102), account.ID)
  1875. account, err = svc.SelectAccountForModelWithExclusions(ctx, nil, "", "gpt-5.3", nil)
  1876. require.NoError(t, err)
  1877. require.NotNil(t, account)
  1878. require.Equal(t, int64(32101), account.ID)
  1879. }
  1880. func TestOpenAIGatewayService_SelectAccountWithScheduler_SessionStickyDBRuntimeRecheckSkipsStaleCachedAccount(t *testing.T) {
  1881. ctx := context.Background()
  1882. groupID := int64(10103)
  1883. rateLimitedUntil := time.Now().Add(30 * time.Minute)
  1884. staleSticky := &Account{ID: 33001, Platform: PlatformOpenAI, Type: AccountTypeOAuth, Status: StatusActive, Schedulable: true, Concurrency: 1, Priority: 0, GroupIDs: []int64{groupID}}
  1885. staleBackup := &Account{ID: 33002, Platform: PlatformOpenAI, Type: AccountTypeOAuth, Status: StatusActive, Schedulable: true, Concurrency: 1, Priority: 5, GroupIDs: []int64{groupID}}
  1886. dbSticky := Account{ID: 33001, Platform: PlatformOpenAI, Type: AccountTypeOAuth, Status: StatusActive, Schedulable: true, Concurrency: 1, Priority: 0, GroupIDs: []int64{groupID}, RateLimitResetAt: &rateLimitedUntil}
  1887. dbBackup := Account{ID: 33002, Platform: PlatformOpenAI, Type: AccountTypeOAuth, Status: StatusActive, Schedulable: true, Concurrency: 1, Priority: 5, GroupIDs: []int64{groupID}}
  1888. cache := &schedulerTestGatewayCache{sessionBindings: map[string]int64{"openai:session_hash_db_runtime_recheck": 33001}}
  1889. snapshotCache := &openAISnapshotCacheStub{
  1890. snapshotAccounts: []*Account{staleSticky, staleBackup},
  1891. accountsByID: map[int64]*Account{33001: staleSticky, 33002: staleBackup},
  1892. }
  1893. snapshotService := &SchedulerSnapshotService{cache: snapshotCache}
  1894. svc := &OpenAIGatewayService{
  1895. accountRepo: schedulerTestOpenAIAccountRepo{accounts: []Account{dbSticky, dbBackup}},
  1896. cache: cache,
  1897. cfg: &config.Config{},
  1898. rateLimitService: newOpenAIAdvancedSchedulerRateLimitService("true"),
  1899. schedulerSnapshot: snapshotService,
  1900. concurrencyService: NewConcurrencyService(schedulerTestConcurrencyCache{}),
  1901. }
  1902. selection, decision, err := svc.SelectAccountWithScheduler(ctx, &groupID, "", "session_hash_db_runtime_recheck", "gpt-5.1", nil, OpenAIUpstreamTransportAny, false)
  1903. require.NoError(t, err)
  1904. require.NotNil(t, selection)
  1905. require.NotNil(t, selection.Account)
  1906. require.Equal(t, int64(33002), selection.Account.ID)
  1907. require.Equal(t, openAIAccountScheduleLayerLoadBalance, decision.Layer)
  1908. }
  1909. func TestOpenAIGatewayService_SelectAccountForModelWithExclusions_DBRuntimeRecheckSkipsStaleCachedCandidate(t *testing.T) {
  1910. ctx := context.Background()
  1911. groupID := int64(10104)
  1912. rateLimitedUntil := time.Now().Add(30 * time.Minute)
  1913. stalePrimary := &Account{ID: 34001, Platform: PlatformOpenAI, Type: AccountTypeOAuth, Status: StatusActive, Schedulable: true, Concurrency: 1, Priority: 0, GroupIDs: []int64{groupID}}
  1914. staleSecondary := &Account{ID: 34002, Platform: PlatformOpenAI, Type: AccountTypeOAuth, Status: StatusActive, Schedulable: true, Concurrency: 1, Priority: 5, GroupIDs: []int64{groupID}}
  1915. dbPrimary := Account{ID: 34001, Platform: PlatformOpenAI, Type: AccountTypeOAuth, Status: StatusActive, Schedulable: true, Concurrency: 1, Priority: 0, GroupIDs: []int64{groupID}, RateLimitResetAt: &rateLimitedUntil}
  1916. dbSecondary := Account{ID: 34002, Platform: PlatformOpenAI, Type: AccountTypeOAuth, Status: StatusActive, Schedulable: true, Concurrency: 1, Priority: 5, GroupIDs: []int64{groupID}}
  1917. snapshotCache := &openAISnapshotCacheStub{
  1918. snapshotAccounts: []*Account{stalePrimary, staleSecondary},
  1919. accountsByID: map[int64]*Account{34001: stalePrimary, 34002: staleSecondary},
  1920. }
  1921. snapshotService := &SchedulerSnapshotService{cache: snapshotCache}
  1922. svc := &OpenAIGatewayService{
  1923. accountRepo: schedulerTestOpenAIAccountRepo{accounts: []Account{dbPrimary, dbSecondary}},
  1924. cfg: &config.Config{},
  1925. rateLimitService: newOpenAIAdvancedSchedulerRateLimitService("true"),
  1926. schedulerSnapshot: snapshotService,
  1927. }
  1928. account, err := svc.SelectAccountForModelWithExclusions(ctx, &groupID, "", "gpt-5.1", nil)
  1929. require.NoError(t, err)
  1930. require.NotNil(t, account)
  1931. require.Equal(t, int64(34002), account.ID)
  1932. }
  1933. func TestOpenAIGatewayService_SelectAccountWithScheduler_DBFreshGroupRecheckReleasesMovedAccount(t *testing.T) {
  1934. ctx := context.Background()
  1935. groupID, otherGroupID := int64(10105), int64(10106)
  1936. stalePrimary := &Account{ID: 34101, Platform: PlatformOpenAI, Type: AccountTypeAPIKey, Status: StatusActive, Schedulable: true, Concurrency: 1, Priority: 0, GroupIDs: []int64{groupID}}
  1937. staleBackup := &Account{ID: 34102, Platform: PlatformOpenAI, Type: AccountTypeAPIKey, Status: StatusActive, Schedulable: true, Concurrency: 1, Priority: 10, GroupIDs: []int64{groupID}}
  1938. dbPrimary := *stalePrimary
  1939. dbPrimary.GroupIDs = []int64{otherGroupID}
  1940. dbBackup := *staleBackup
  1941. snapshotCache := &openAISnapshotCacheStub{
  1942. snapshotAccounts: []*Account{stalePrimary, staleBackup},
  1943. accountsByID: map[int64]*Account{stalePrimary.ID: stalePrimary, staleBackup.ID: staleBackup},
  1944. }
  1945. acquiredIDs, releasedIDs := []int64{}, []int64{}
  1946. cfg := &config.Config{RunMode: config.RunModeStandard}
  1947. svc := &OpenAIGatewayService{
  1948. accountRepo: schedulerTestOpenAIAccountRepo{accounts: []Account{dbPrimary, dbBackup}},
  1949. cfg: cfg,
  1950. schedulerSnapshot: &SchedulerSnapshotService{cache: snapshotCache},
  1951. concurrencyService: NewConcurrencyService(schedulerTestConcurrencyCache{
  1952. acquiredIDs: &acquiredIDs,
  1953. releasedIDs: &releasedIDs,
  1954. }),
  1955. }
  1956. scheduler := &defaultOpenAIAccountScheduler{service: svc}
  1957. selection, _, err := scheduler.tryAcquireOpenAISelectionOrder(ctx, OpenAIAccountScheduleRequest{
  1958. GroupID: &groupID, Platform: PlatformOpenAI, RequestedModel: "gpt-5.1",
  1959. }, []openAIAccountCandidateScore{
  1960. {account: stalePrimary, loadInfo: &AccountLoadInfo{AccountID: stalePrimary.ID}},
  1961. {account: staleBackup, loadInfo: &AccountLoadInfo{AccountID: staleBackup.ID}},
  1962. })
  1963. require.NoError(t, err)
  1964. require.Equal(t, staleBackup.ID, selection.Account.ID)
  1965. require.Equal(t, []int64{stalePrimary.ID, staleBackup.ID}, acquiredIDs)
  1966. require.Equal(t, []int64{stalePrimary.ID}, releasedIDs)
  1967. selection.ReleaseFunc()
  1968. }
  1969. func TestOpenAIGatewayService_SelectAccountWithLoadAwareness_DBFreshGroupRecheckWaitsOnValidAccount(t *testing.T) {
  1970. ctx := context.Background()
  1971. groupID, otherGroupID := int64(10107), int64(10108)
  1972. stalePrimary := &Account{ID: 34201, Platform: PlatformOpenAI, Type: AccountTypeAPIKey, Status: StatusActive, Schedulable: true, Concurrency: 1, Priority: 0, GroupIDs: []int64{groupID}}
  1973. staleBackup := &Account{ID: 34202, Platform: PlatformOpenAI, Type: AccountTypeAPIKey, Status: StatusActive, Schedulable: true, Concurrency: 1, Priority: 10, GroupIDs: []int64{groupID}}
  1974. dbPrimary := *stalePrimary
  1975. dbPrimary.GroupIDs = []int64{otherGroupID}
  1976. dbBackup := *staleBackup
  1977. snapshotCache := &openAISnapshotCacheStub{
  1978. snapshotAccounts: []*Account{stalePrimary, staleBackup},
  1979. accountsByID: map[int64]*Account{stalePrimary.ID: stalePrimary, staleBackup.ID: staleBackup},
  1980. }
  1981. cfg := &config.Config{RunMode: config.RunModeStandard}
  1982. cfg.Gateway.Scheduling.LoadBatchEnabled = true
  1983. svc := &OpenAIGatewayService{
  1984. accountRepo: schedulerTestOpenAIAccountRepo{accounts: []Account{dbPrimary, dbBackup}},
  1985. cfg: cfg,
  1986. schedulerSnapshot: &SchedulerSnapshotService{cache: snapshotCache},
  1987. concurrencyService: NewConcurrencyService(schedulerTestConcurrencyCache{
  1988. acquireResults: map[int64]bool{staleBackup.ID: false},
  1989. }),
  1990. }
  1991. selection, err := svc.SelectAccountWithLoadAwareness(ctx, &groupID, "", "gpt-5.1", nil)
  1992. require.NoError(t, err)
  1993. require.NotNil(t, selection.WaitPlan)
  1994. require.Equal(t, staleBackup.ID, selection.Account.ID)
  1995. require.Equal(t, staleBackup.ID, selection.WaitPlan.AccountID)
  1996. }
  1997. func TestOpenAIGatewayService_RecheckSelectedOpenAIAccountFromDB_SimpleModeUsesFullPool(t *testing.T) {
  1998. grouped := Account{ID: 34301, Platform: PlatformOpenAI, Type: AccountTypeAPIKey, Status: StatusActive, Schedulable: true, Concurrency: 1, GroupIDs: []int64{99}}
  1999. svc := &OpenAIGatewayService{
  2000. accountRepo: schedulerTestOpenAIAccountRepo{accounts: []Account{grouped}},
  2001. cfg: &config.Config{RunMode: config.RunModeSimple},
  2002. schedulerSnapshot: &SchedulerSnapshotService{cache: &openAISnapshotCacheStub{}},
  2003. }
  2004. requestedGroupID := int64(100)
  2005. for _, groupID := range []*int64{nil, &requestedGroupID} {
  2006. fresh := svc.recheckSelectedOpenAIAccountFromDB(context.Background(), &grouped, groupID, PlatformOpenAI, "gpt-5.1", false, "")
  2007. require.NotNil(t, fresh)
  2008. require.Equal(t, grouped.ID, fresh.ID)
  2009. }
  2010. ungrouped := grouped
  2011. ungrouped.ID++
  2012. ungrouped.GroupIDs = nil
  2013. standardSvc := &OpenAIGatewayService{
  2014. accountRepo: schedulerTestOpenAIAccountRepo{accounts: []Account{grouped, ungrouped}},
  2015. cfg: &config.Config{RunMode: config.RunModeStandard},
  2016. schedulerSnapshot: &SchedulerSnapshotService{cache: &openAISnapshotCacheStub{}},
  2017. }
  2018. require.Nil(t, standardSvc.recheckSelectedOpenAIAccountFromDB(context.Background(), &grouped, nil, PlatformOpenAI, "gpt-5.1", false, ""))
  2019. require.NotNil(t, standardSvc.recheckSelectedOpenAIAccountFromDB(context.Background(), &ungrouped, nil, PlatformOpenAI, "gpt-5.1", false, ""))
  2020. }
  2021. func TestOpenAIGatewayService_SelectAccountWithScheduler_PreviousResponseSticky(t *testing.T) {
  2022. ctx := context.Background()
  2023. groupID := int64(9)
  2024. account := Account{
  2025. ID: 1001,
  2026. Platform: PlatformOpenAI,
  2027. Type: AccountTypeAPIKey,
  2028. Status: StatusActive,
  2029. Schedulable: true,
  2030. Concurrency: 2,
  2031. Extra: map[string]any{
  2032. "openai_apikey_responses_websockets_v2_enabled": true,
  2033. },
  2034. }
  2035. cache := &schedulerTestGatewayCache{}
  2036. cfg := &config.Config{}
  2037. cfg.Gateway.OpenAIWS.Enabled = true
  2038. cfg.Gateway.OpenAIWS.OAuthEnabled = true
  2039. cfg.Gateway.OpenAIWS.APIKeyEnabled = true
  2040. cfg.Gateway.OpenAIWS.ResponsesWebsocketsV2 = true
  2041. cfg.Gateway.OpenAIWS.StickySessionTTLSeconds = 1800
  2042. cfg.Gateway.OpenAIWS.StickyResponseIDTTLSeconds = 3600
  2043. svc := &OpenAIGatewayService{
  2044. accountRepo: schedulerTestOpenAIAccountRepo{accounts: []Account{account}},
  2045. cache: cache,
  2046. cfg: cfg,
  2047. rateLimitService: newOpenAIAdvancedSchedulerRateLimitService("true"),
  2048. concurrencyService: NewConcurrencyService(schedulerTestConcurrencyCache{}),
  2049. }
  2050. store := svc.getOpenAIWSStateStore()
  2051. require.NoError(t, store.BindResponseAccount(ctx, groupID, "resp_prev_001", account.ID, time.Hour))
  2052. selection, decision, err := svc.SelectAccountWithScheduler(
  2053. ctx,
  2054. &groupID,
  2055. "resp_prev_001",
  2056. "session_hash_001",
  2057. "gpt-5.1",
  2058. nil,
  2059. OpenAIUpstreamTransportAny,
  2060. false,
  2061. )
  2062. require.NoError(t, err)
  2063. require.NotNil(t, selection)
  2064. require.NotNil(t, selection.Account)
  2065. require.Equal(t, account.ID, selection.Account.ID)
  2066. require.Equal(t, openAIAccountScheduleLayerPreviousResponse, decision.Layer)
  2067. require.True(t, decision.StickyPreviousHit)
  2068. require.Equal(t, account.ID, cache.sessionBindings["openai:session_hash_001"])
  2069. if selection.ReleaseFunc != nil {
  2070. selection.ReleaseFunc()
  2071. }
  2072. }
  2073. func TestOpenAIGatewayService_SelectAccountWithScheduler_SessionSticky(t *testing.T) {
  2074. ctx := context.Background()
  2075. groupID := int64(10)
  2076. account := Account{
  2077. ID: 2001,
  2078. Platform: PlatformOpenAI,
  2079. Type: AccountTypeOAuth,
  2080. Status: StatusActive,
  2081. Schedulable: true,
  2082. Concurrency: 1,
  2083. GroupIDs: []int64{groupID},
  2084. }
  2085. cache := &schedulerTestGatewayCache{
  2086. sessionBindings: map[string]int64{
  2087. "openai:session_hash_abc": account.ID,
  2088. },
  2089. }
  2090. svc := &OpenAIGatewayService{
  2091. accountRepo: schedulerTestOpenAIAccountRepo{accounts: []Account{account}},
  2092. cache: cache,
  2093. cfg: &config.Config{},
  2094. rateLimitService: newOpenAIAdvancedSchedulerRateLimitService("true"),
  2095. concurrencyService: NewConcurrencyService(schedulerTestConcurrencyCache{}),
  2096. }
  2097. selection, decision, err := svc.SelectAccountWithScheduler(
  2098. ctx,
  2099. &groupID,
  2100. "",
  2101. "session_hash_abc",
  2102. "gpt-5.1",
  2103. nil,
  2104. OpenAIUpstreamTransportAny,
  2105. false,
  2106. )
  2107. require.NoError(t, err)
  2108. require.NotNil(t, selection)
  2109. require.NotNil(t, selection.Account)
  2110. require.Equal(t, account.ID, selection.Account.ID)
  2111. require.Equal(t, openAIAccountScheduleLayerSessionSticky, decision.Layer)
  2112. require.True(t, decision.StickySessionHit)
  2113. if selection.ReleaseFunc != nil {
  2114. selection.ReleaseFunc()
  2115. }
  2116. }
  2117. func TestOpenAIGatewayService_SelectAccountWithScheduler_SessionStickyBusyKeepsSticky(t *testing.T) {
  2118. ctx := context.Background()
  2119. groupID := int64(10100)
  2120. accounts := []Account{
  2121. {
  2122. ID: 21001,
  2123. Platform: PlatformOpenAI,
  2124. Type: AccountTypeAPIKey,
  2125. Status: StatusActive,
  2126. Schedulable: true,
  2127. Concurrency: 1,
  2128. Priority: 0,
  2129. GroupIDs: []int64{groupID},
  2130. },
  2131. {
  2132. ID: 21002,
  2133. Platform: PlatformOpenAI,
  2134. Type: AccountTypeAPIKey,
  2135. Status: StatusActive,
  2136. Schedulable: true,
  2137. Concurrency: 1,
  2138. Priority: 9,
  2139. GroupIDs: []int64{groupID},
  2140. },
  2141. }
  2142. cache := &schedulerTestGatewayCache{
  2143. sessionBindings: map[string]int64{
  2144. "openai:session_hash_sticky_busy": 21001,
  2145. },
  2146. }
  2147. cfg := &config.Config{}
  2148. cfg.Gateway.Scheduling.StickySessionMaxWaiting = 2
  2149. cfg.Gateway.Scheduling.StickySessionWaitTimeout = 45 * time.Second
  2150. cfg.Gateway.OpenAIScheduler.StickyEscapeEnabled = false
  2151. cfg.Gateway.OpenAIScheduler.StickyEscapeTTFTMs = 15000
  2152. cfg.Gateway.OpenAIScheduler.StickyEscapeErrorRate = 0.5
  2153. cfg.Gateway.OpenAIWS.Enabled = true
  2154. cfg.Gateway.OpenAIWS.APIKeyEnabled = true
  2155. cfg.Gateway.OpenAIWS.OAuthEnabled = true
  2156. cfg.Gateway.OpenAIWS.ResponsesWebsocketsV2 = true
  2157. concurrencyCache := schedulerTestConcurrencyCache{
  2158. acquireResults: map[int64]bool{
  2159. 21001: false, // sticky 账号已满
  2160. 21002: true, // 若回退负载均衡会命中该账号(本测试要求不能切换)
  2161. },
  2162. waitCounts: map[int64]int{
  2163. 21001: 999,
  2164. },
  2165. loadMap: map[int64]*AccountLoadInfo{
  2166. 21001: {AccountID: 21001, LoadRate: 90, WaitingCount: 9},
  2167. 21002: {AccountID: 21002, LoadRate: 1, WaitingCount: 0},
  2168. },
  2169. }
  2170. svc := &OpenAIGatewayService{
  2171. accountRepo: schedulerTestOpenAIAccountRepo{accounts: accounts},
  2172. cache: cache,
  2173. cfg: cfg,
  2174. rateLimitService: newOpenAIAdvancedSchedulerRateLimitService("true"),
  2175. concurrencyService: NewConcurrencyService(concurrencyCache),
  2176. }
  2177. selection, decision, err := svc.SelectAccountWithScheduler(
  2178. ctx,
  2179. &groupID,
  2180. "",
  2181. "session_hash_sticky_busy",
  2182. "gpt-5.1",
  2183. nil,
  2184. OpenAIUpstreamTransportAny,
  2185. false,
  2186. )
  2187. require.NoError(t, err)
  2188. require.NotNil(t, selection)
  2189. require.NotNil(t, selection.Account)
  2190. require.Equal(t, int64(21001), selection.Account.ID, "busy sticky account should remain selected")
  2191. require.False(t, selection.Acquired)
  2192. require.NotNil(t, selection.WaitPlan)
  2193. require.Equal(t, int64(21001), selection.WaitPlan.AccountID)
  2194. require.Equal(t, openAIAccountScheduleLayerSessionSticky, decision.Layer)
  2195. require.True(t, decision.StickySessionHit)
  2196. }
  2197. func TestOpenAIGatewayService_SelectAccountWithScheduler_SessionStickyEscapeByTTFT(t *testing.T) {
  2198. ctx := context.Background()
  2199. groupID := int64(10101)
  2200. accounts := []Account{
  2201. {
  2202. ID: 21101,
  2203. Platform: PlatformOpenAI,
  2204. Type: AccountTypeAPIKey,
  2205. Status: StatusActive,
  2206. Schedulable: true,
  2207. Concurrency: 1,
  2208. Priority: 0,
  2209. GroupIDs: []int64{groupID},
  2210. },
  2211. {
  2212. ID: 21102,
  2213. Platform: PlatformOpenAI,
  2214. Type: AccountTypeAPIKey,
  2215. Status: StatusActive,
  2216. Schedulable: true,
  2217. Concurrency: 1,
  2218. Priority: 1,
  2219. GroupIDs: []int64{groupID},
  2220. },
  2221. }
  2222. cache := &schedulerTestGatewayCache{sessionBindings: map[string]int64{"openai:session_hash_sticky_ttft": 21101}}
  2223. cfg := &config.Config{}
  2224. cfg.Gateway.OpenAIScheduler.StickyEscapeEnabled = true
  2225. cfg.Gateway.OpenAIScheduler.StickyEscapeTTFTMs = 15000
  2226. cfg.Gateway.OpenAIScheduler.StickyEscapeErrorRate = 0.5
  2227. concurrencyCache := schedulerTestConcurrencyCache{acquireResults: map[int64]bool{21102: true}}
  2228. svc := &OpenAIGatewayService{
  2229. accountRepo: schedulerTestOpenAIAccountRepo{accounts: accounts},
  2230. cache: cache,
  2231. cfg: cfg,
  2232. rateLimitService: newOpenAIAdvancedSchedulerRateLimitService("true"),
  2233. concurrencyService: NewConcurrencyService(concurrencyCache),
  2234. openaiAccountStats: newOpenAIAccountRuntimeStats(),
  2235. }
  2236. fastTTFT := 14999
  2237. svc.openaiAccountStats.report(21101, true, &fastTTFT)
  2238. stableTTFT := 14999
  2239. svc.openaiAccountStats.report(21101, true, &stableTTFT)
  2240. selection, decision, err := svc.SelectAccountWithScheduler(ctx, &groupID, "", "session_hash_sticky_ttft", "gpt-5.1", nil, OpenAIUpstreamTransportAny, false)
  2241. require.NoError(t, err)
  2242. require.NotNil(t, selection)
  2243. require.NotNil(t, selection.Account)
  2244. require.Equal(t, int64(21101), selection.Account.ID)
  2245. require.Equal(t, openAIAccountScheduleLayerSessionSticky, decision.Layer)
  2246. require.True(t, decision.StickySessionHit)
  2247. if selection.ReleaseFunc != nil {
  2248. selection.ReleaseFunc()
  2249. }
  2250. slowTTFT := 20000
  2251. for i := 0; i < 3; i++ {
  2252. svc.openaiAccountStats.report(21101, true, &slowTTFT)
  2253. }
  2254. selection, decision, err = svc.SelectAccountWithScheduler(ctx, &groupID, "", "session_hash_sticky_ttft", "gpt-5.1", nil, OpenAIUpstreamTransportAny, false)
  2255. require.NoError(t, err)
  2256. require.NotNil(t, selection)
  2257. require.NotNil(t, selection.Account)
  2258. require.Equal(t, int64(21102), selection.Account.ID)
  2259. require.Equal(t, openAIAccountScheduleLayerLoadBalance, decision.Layer)
  2260. require.False(t, decision.StickySessionHit)
  2261. require.Equal(t, int64(21101), cache.sessionBindings["openai:session_hash_sticky_ttft"])
  2262. if selection.ReleaseFunc != nil {
  2263. selection.ReleaseFunc()
  2264. }
  2265. }
  2266. func TestOpenAIGatewayService_SelectAccountWithScheduler_SessionStickyEscapeByErrorRate(t *testing.T) {
  2267. ctx := context.Background()
  2268. groupID := int64(10102)
  2269. accounts := []Account{
  2270. {ID: 21201, Platform: PlatformOpenAI, Type: AccountTypeAPIKey, Status: StatusActive, Schedulable: true, Concurrency: 1, Priority: 0, GroupIDs: []int64{groupID}},
  2271. {ID: 21202, Platform: PlatformOpenAI, Type: AccountTypeAPIKey, Status: StatusActive, Schedulable: true, Concurrency: 1, Priority: 1, GroupIDs: []int64{groupID}},
  2272. }
  2273. cache := &schedulerTestGatewayCache{sessionBindings: map[string]int64{"openai:session_hash_sticky_error_rate": 21201}}
  2274. cfg := &config.Config{}
  2275. cfg.Gateway.OpenAIScheduler.StickyEscapeEnabled = true
  2276. cfg.Gateway.OpenAIScheduler.StickyEscapeTTFTMs = 15000
  2277. cfg.Gateway.OpenAIScheduler.StickyEscapeErrorRate = 0.5
  2278. svc := &OpenAIGatewayService{
  2279. accountRepo: schedulerTestOpenAIAccountRepo{accounts: accounts},
  2280. cache: cache,
  2281. cfg: cfg,
  2282. rateLimitService: newOpenAIAdvancedSchedulerRateLimitService("true"),
  2283. concurrencyService: NewConcurrencyService(schedulerTestConcurrencyCache{acquireResults: map[int64]bool{21202: true}}),
  2284. openaiAccountStats: newOpenAIAccountRuntimeStats(),
  2285. }
  2286. for i := 0; i < 3; i++ {
  2287. svc.openaiAccountStats.report(21201, false, nil)
  2288. }
  2289. selection, decision, err := svc.SelectAccountWithScheduler(ctx, &groupID, "", "session_hash_sticky_error_rate", "gpt-5.1", nil, OpenAIUpstreamTransportAny, false)
  2290. require.NoError(t, err)
  2291. require.NotNil(t, selection)
  2292. require.NotNil(t, selection.Account)
  2293. require.Equal(t, int64(21201), selection.Account.ID)
  2294. require.Equal(t, openAIAccountScheduleLayerSessionSticky, decision.Layer)
  2295. require.True(t, decision.StickySessionHit)
  2296. if selection.ReleaseFunc != nil {
  2297. selection.ReleaseFunc()
  2298. }
  2299. for i := 0; i < 2; i++ {
  2300. svc.openaiAccountStats.report(21201, false, nil)
  2301. }
  2302. selection, decision, err = svc.SelectAccountWithScheduler(ctx, &groupID, "", "session_hash_sticky_error_rate", "gpt-5.1", nil, OpenAIUpstreamTransportAny, false)
  2303. require.NoError(t, err)
  2304. require.NotNil(t, selection)
  2305. require.NotNil(t, selection.Account)
  2306. require.Equal(t, int64(21202), selection.Account.ID)
  2307. require.Equal(t, openAIAccountScheduleLayerLoadBalance, decision.Layer)
  2308. require.False(t, decision.StickySessionHit)
  2309. require.Equal(t, int64(21201), cache.sessionBindings["openai:session_hash_sticky_error_rate"])
  2310. if selection.ReleaseFunc != nil {
  2311. selection.ReleaseFunc()
  2312. }
  2313. }
  2314. func TestOpenAIGatewayService_SelectAccountWithScheduler_SessionStickyBusyEscapes(t *testing.T) {
  2315. ctx := context.Background()
  2316. groupID := int64(10103)
  2317. accounts := []Account{
  2318. {ID: 21301, Platform: PlatformOpenAI, Type: AccountTypeAPIKey, Status: StatusActive, Schedulable: true, Concurrency: 1, Priority: 0, GroupIDs: []int64{groupID}},
  2319. {ID: 21302, Platform: PlatformOpenAI, Type: AccountTypeAPIKey, Status: StatusActive, Schedulable: true, Concurrency: 1, Priority: 1, GroupIDs: []int64{groupID}},
  2320. }
  2321. cache := &schedulerTestGatewayCache{sessionBindings: map[string]int64{"openai:session_hash_sticky_busy_escape": 21301}}
  2322. cfg := &config.Config{}
  2323. cfg.Gateway.OpenAIScheduler.StickyEscapeEnabled = true
  2324. cfg.Gateway.OpenAIScheduler.StickyEscapeTTFTMs = 15000
  2325. cfg.Gateway.OpenAIScheduler.StickyEscapeErrorRate = 0.5
  2326. cfg.Gateway.Scheduling.StickySessionMaxWaiting = 2
  2327. cfg.Gateway.Scheduling.StickySessionWaitTimeout = 45 * time.Second
  2328. concurrencyCache := schedulerTestConcurrencyCache{
  2329. acquireResults: map[int64]bool{21301: false, 21302: true},
  2330. waitCounts: map[int64]int{21301: 999},
  2331. loadMap: map[int64]*AccountLoadInfo{
  2332. 21301: {AccountID: 21301, LoadRate: 95, WaitingCount: 9},
  2333. 21302: {AccountID: 21302, LoadRate: 1, WaitingCount: 0},
  2334. },
  2335. }
  2336. svc := &OpenAIGatewayService{
  2337. accountRepo: schedulerTestOpenAIAccountRepo{accounts: accounts},
  2338. cache: cache,
  2339. cfg: cfg,
  2340. rateLimitService: newOpenAIAdvancedSchedulerRateLimitService("true"),
  2341. concurrencyService: NewConcurrencyService(concurrencyCache),
  2342. }
  2343. selection, decision, err := svc.SelectAccountWithScheduler(ctx, &groupID, "", "session_hash_sticky_busy_escape", "gpt-5.1", nil, OpenAIUpstreamTransportAny, false)
  2344. require.NoError(t, err)
  2345. require.NotNil(t, selection)
  2346. require.NotNil(t, selection.Account)
  2347. require.Equal(t, int64(21302), selection.Account.ID)
  2348. require.Nil(t, selection.WaitPlan)
  2349. require.Equal(t, openAIAccountScheduleLayerLoadBalance, decision.Layer)
  2350. require.False(t, decision.StickySessionHit)
  2351. if selection.ReleaseFunc != nil {
  2352. selection.ReleaseFunc()
  2353. }
  2354. }
  2355. func TestOpenAIGatewayService_SelectAccountWithScheduler_SessionStickyEscapeDisabledKeepsLegacyBehavior(t *testing.T) {
  2356. ctx := context.Background()
  2357. groupID := int64(10104)
  2358. accounts := []Account{
  2359. {ID: 21401, Platform: PlatformOpenAI, Type: AccountTypeAPIKey, Status: StatusActive, Schedulable: true, Concurrency: 1, Priority: 0, GroupIDs: []int64{groupID}},
  2360. {ID: 21402, Platform: PlatformOpenAI, Type: AccountTypeAPIKey, Status: StatusActive, Schedulable: true, Concurrency: 1, Priority: 1, GroupIDs: []int64{groupID}},
  2361. }
  2362. cache := &schedulerTestGatewayCache{sessionBindings: map[string]int64{"openai:session_hash_sticky_disabled": 21401}}
  2363. cfg := &config.Config{}
  2364. cfg.Gateway.OpenAIScheduler.StickyEscapeEnabled = false
  2365. cfg.Gateway.OpenAIScheduler.StickyEscapeTTFTMs = 15000
  2366. cfg.Gateway.OpenAIScheduler.StickyEscapeErrorRate = 0.5
  2367. cfg.Gateway.Scheduling.StickySessionMaxWaiting = 2
  2368. cfg.Gateway.Scheduling.StickySessionWaitTimeout = 45 * time.Second
  2369. concurrencyCache := schedulerTestConcurrencyCache{
  2370. acquireResults: map[int64]bool{21401: false, 21402: true},
  2371. waitCounts: map[int64]int{21401: 999},
  2372. }
  2373. svc := &OpenAIGatewayService{
  2374. accountRepo: schedulerTestOpenAIAccountRepo{accounts: accounts},
  2375. cache: cache,
  2376. cfg: cfg,
  2377. rateLimitService: newOpenAIAdvancedSchedulerRateLimitService("true"),
  2378. concurrencyService: NewConcurrencyService(concurrencyCache),
  2379. openaiAccountStats: newOpenAIAccountRuntimeStats(),
  2380. }
  2381. slowTTFT := 20000
  2382. svc.openaiAccountStats.report(21401, true, &slowTTFT)
  2383. for i := 0; i < 5; i++ {
  2384. svc.openaiAccountStats.report(21401, false, nil)
  2385. }
  2386. selection, decision, err := svc.SelectAccountWithScheduler(ctx, &groupID, "", "session_hash_sticky_disabled", "gpt-5.1", nil, OpenAIUpstreamTransportAny, false)
  2387. require.NoError(t, err)
  2388. require.NotNil(t, selection)
  2389. require.NotNil(t, selection.Account)
  2390. require.Equal(t, int64(21401), selection.Account.ID)
  2391. require.NotNil(t, selection.WaitPlan)
  2392. require.Equal(t, int64(21401), selection.WaitPlan.AccountID)
  2393. require.Equal(t, openAIAccountScheduleLayerSessionSticky, decision.Layer)
  2394. require.True(t, decision.StickySessionHit)
  2395. }
  2396. func TestOpenAIGatewayService_SelectAccountWithScheduler_SubscriptionPriorityChoosesSubscriptionPoolFirst(t *testing.T) {
  2397. ctx := context.Background()
  2398. groupID := int64(10120)
  2399. cheapAPIKey := upstreamCostTestAccount(21602, UpstreamBillingProbeStatusOK, 0.01, time.Now().Add(-time.Minute), 30*time.Minute)
  2400. cheapAPIKey.Status = StatusActive
  2401. cheapAPIKey.Schedulable = true
  2402. cheapAPIKey.Concurrency = 1
  2403. cheapAPIKey.Priority = 0
  2404. cheapAPIKey.GroupIDs = []int64{groupID}
  2405. accounts := []Account{
  2406. {
  2407. ID: 21601,
  2408. Platform: PlatformOpenAI,
  2409. Type: AccountTypeOAuth,
  2410. Status: StatusActive,
  2411. Schedulable: true,
  2412. Concurrency: 1,
  2413. Priority: 10,
  2414. GroupIDs: []int64{groupID},
  2415. Credentials: map[string]any{"plan_type": "plus"},
  2416. },
  2417. *cheapAPIKey,
  2418. }
  2419. concurrencyCache := schedulerTestConcurrencyCache{
  2420. acquireResults: map[int64]bool{21601: true, 21602: true},
  2421. loadMap: map[int64]*AccountLoadInfo{
  2422. 21601: {AccountID: 21601, LoadRate: 90, WaitingCount: 1},
  2423. 21602: {AccountID: 21602, LoadRate: 0, WaitingCount: 0},
  2424. },
  2425. }
  2426. cfg := newSchedulerTestSubscriptionPriorityConfig()
  2427. cfg.Gateway.OpenAIWS.SchedulerScoreWeights.UpstreamCost = 100
  2428. svc := &OpenAIGatewayService{
  2429. accountRepo: schedulerTestOpenAIAccountRepo{accounts: accounts},
  2430. cache: &schedulerTestGatewayCache{},
  2431. cfg: cfg,
  2432. rateLimitService: newOpenAIAdvancedSchedulerRateLimitService("true", "", "true"),
  2433. concurrencyService: NewConcurrencyService(concurrencyCache),
  2434. }
  2435. selection, decision, err := svc.SelectAccountWithScheduler(ctx, &groupID, "", "session_subscription_first", "gpt-5.1", nil, OpenAIUpstreamTransportAny, false)
  2436. require.NoError(t, err)
  2437. require.NotNil(t, selection)
  2438. require.NotNil(t, selection.Account)
  2439. require.Equal(t, int64(21601), selection.Account.ID)
  2440. require.Equal(t, openAIAccountScheduleLayerLoadBalance, decision.Layer)
  2441. require.Equal(t, 1, decision.TopK)
  2442. if selection.ReleaseFunc != nil {
  2443. selection.ReleaseFunc()
  2444. }
  2445. }
  2446. func TestOpenAIGatewayService_SelectAccountWithScheduler_SubscriptionPriorityFallsBackWhenSubscriptionFull(t *testing.T) {
  2447. ctx := context.Background()
  2448. groupID := int64(10121)
  2449. accounts := []Account{
  2450. {
  2451. ID: 21611,
  2452. Platform: PlatformOpenAI,
  2453. Type: AccountTypeOAuth,
  2454. Status: StatusActive,
  2455. Schedulable: true,
  2456. Concurrency: 1,
  2457. Priority: 0,
  2458. GroupIDs: []int64{groupID},
  2459. Credentials: map[string]any{"plan_type": "team"},
  2460. },
  2461. {
  2462. ID: 21612,
  2463. Platform: PlatformOpenAI,
  2464. Type: AccountTypeAPIKey,
  2465. Status: StatusActive,
  2466. Schedulable: true,
  2467. Concurrency: 1,
  2468. Priority: 9,
  2469. GroupIDs: []int64{groupID},
  2470. },
  2471. }
  2472. concurrencyCache := schedulerTestConcurrencyCache{
  2473. acquireResults: map[int64]bool{21611: false, 21612: true},
  2474. loadMap: map[int64]*AccountLoadInfo{
  2475. 21611: {AccountID: 21611, LoadRate: 0, WaitingCount: 0},
  2476. 21612: {AccountID: 21612, LoadRate: 90, WaitingCount: 1},
  2477. },
  2478. }
  2479. svc := &OpenAIGatewayService{
  2480. accountRepo: schedulerTestOpenAIAccountRepo{accounts: accounts},
  2481. cache: &schedulerTestGatewayCache{},
  2482. cfg: newSchedulerTestSubscriptionPriorityConfig(),
  2483. rateLimitService: newOpenAIAdvancedSchedulerRateLimitService("true", "", "true"),
  2484. concurrencyService: NewConcurrencyService(concurrencyCache),
  2485. }
  2486. selection, decision, err := svc.SelectAccountWithScheduler(ctx, &groupID, "", "session_subscription_fallback", "gpt-5.1", nil, OpenAIUpstreamTransportAny, false)
  2487. require.NoError(t, err)
  2488. require.NotNil(t, selection)
  2489. require.NotNil(t, selection.Account)
  2490. require.Equal(t, int64(21612), selection.Account.ID)
  2491. require.Equal(t, openAIAccountScheduleLayerLoadBalance, decision.Layer)
  2492. require.True(t, selection.Acquired)
  2493. if selection.ReleaseFunc != nil {
  2494. selection.ReleaseFunc()
  2495. }
  2496. }
  2497. func TestOpenAIGatewayService_SelectAccountWithScheduler_SubscriptionPriorityDisabledUsesScore(t *testing.T) {
  2498. ctx := context.Background()
  2499. groupID := int64(10122)
  2500. accounts := []Account{
  2501. {
  2502. ID: 21621,
  2503. Platform: PlatformOpenAI,
  2504. Type: AccountTypeOAuth,
  2505. Status: StatusActive,
  2506. Schedulable: true,
  2507. Concurrency: 1,
  2508. Priority: 10,
  2509. GroupIDs: []int64{groupID},
  2510. Credentials: map[string]any{"plan_type": "pro"},
  2511. },
  2512. {
  2513. ID: 21622,
  2514. Platform: PlatformOpenAI,
  2515. Type: AccountTypeAPIKey,
  2516. Status: StatusActive,
  2517. Schedulable: true,
  2518. Concurrency: 1,
  2519. Priority: 0,
  2520. GroupIDs: []int64{groupID},
  2521. },
  2522. }
  2523. concurrencyCache := schedulerTestConcurrencyCache{
  2524. acquireResults: map[int64]bool{21621: true, 21622: true},
  2525. loadMap: map[int64]*AccountLoadInfo{
  2526. 21621: {AccountID: 21621, LoadRate: 90, WaitingCount: 1},
  2527. 21622: {AccountID: 21622, LoadRate: 0, WaitingCount: 0},
  2528. },
  2529. }
  2530. svc := &OpenAIGatewayService{
  2531. accountRepo: schedulerTestOpenAIAccountRepo{accounts: accounts},
  2532. cache: &schedulerTestGatewayCache{},
  2533. cfg: newSchedulerTestSubscriptionPriorityConfig(),
  2534. rateLimitService: newOpenAIAdvancedSchedulerRateLimitService("true", "", "false"),
  2535. concurrencyService: NewConcurrencyService(concurrencyCache),
  2536. }
  2537. selection, decision, err := svc.SelectAccountWithScheduler(ctx, &groupID, "", "session_subscription_disabled", "gpt-5.1", nil, OpenAIUpstreamTransportAny, false)
  2538. require.NoError(t, err)
  2539. require.NotNil(t, selection)
  2540. require.NotNil(t, selection.Account)
  2541. require.Equal(t, int64(21622), selection.Account.ID)
  2542. require.Equal(t, openAIAccountScheduleLayerLoadBalance, decision.Layer)
  2543. if selection.ReleaseFunc != nil {
  2544. selection.ReleaseFunc()
  2545. }
  2546. }
  2547. func TestOpenAIGatewayService_SelectAccountWithScheduler_UsesAccountPriorityWithinGroupPool(t *testing.T) {
  2548. ctx := context.Background()
  2549. groupID := int64(10123)
  2550. accounts := []Account{
  2551. {
  2552. ID: 21631,
  2553. Platform: PlatformOpenAI,
  2554. Type: AccountTypeAPIKey,
  2555. Status: StatusActive,
  2556. Schedulable: true,
  2557. Concurrency: 1,
  2558. Priority: 1,
  2559. AccountGroups: []AccountGroup{
  2560. {AccountID: 21631, GroupID: groupID, Priority: 100},
  2561. },
  2562. GroupIDs: []int64{groupID},
  2563. },
  2564. {
  2565. ID: 21632,
  2566. Platform: PlatformOpenAI,
  2567. Type: AccountTypeAPIKey,
  2568. Status: StatusActive,
  2569. Schedulable: true,
  2570. Concurrency: 1,
  2571. Priority: 100000,
  2572. AccountGroups: []AccountGroup{
  2573. {AccountID: 21632, GroupID: groupID, Priority: 1},
  2574. },
  2575. GroupIDs: []int64{groupID},
  2576. },
  2577. }
  2578. cfg := newSchedulerTestSubscriptionPriorityConfig()
  2579. cfg.Gateway.OpenAIWS.SchedulerScoreWeights.Load = 0
  2580. cfg.Gateway.OpenAIWS.SchedulerScoreWeights.Queue = 0
  2581. svc := &OpenAIGatewayService{
  2582. accountRepo: schedulerGroupAwareOpenAIAccountRepo{schedulerTestOpenAIAccountRepo{accounts: accounts}},
  2583. cache: &schedulerTestGatewayCache{},
  2584. cfg: cfg,
  2585. rateLimitService: newOpenAIAdvancedSchedulerRateLimitService("true"),
  2586. concurrencyService: NewConcurrencyService(schedulerTestConcurrencyCache{}),
  2587. }
  2588. selection, decision, err := svc.SelectAccountWithScheduler(ctx, &groupID, "", "session_group_priority", "gpt-5.1", nil, OpenAIUpstreamTransportAny, false)
  2589. require.NoError(t, err)
  2590. require.NotNil(t, selection)
  2591. require.NotNil(t, selection.Account)
  2592. require.Equal(t, int64(21631), selection.Account.ID)
  2593. require.Equal(t, openAIAccountScheduleLayerLoadBalance, decision.Layer)
  2594. if selection.ReleaseFunc != nil {
  2595. selection.ReleaseFunc()
  2596. }
  2597. }
  2598. func TestOpenAIAccountScheduler_SkipsAccountBlockedForRequestedModel(t *testing.T) {
  2599. now := time.Now()
  2600. account := &Account{ID: 21633, Platform: PlatformOpenAI, Type: AccountTypeAPIKey}
  2601. svc := &OpenAIGatewayService{openaiModelTransient: newOpenAIAccountModelTransientState(128)}
  2602. svc.openaiModelTransient.recordFailure(account.ID, "gpt-5.5", now)
  2603. svc.openaiModelTransient.recordFailure(account.ID, "gpt-5.5", now.Add(time.Millisecond))
  2604. scheduler := &defaultOpenAIAccountScheduler{service: svc}
  2605. require.False(t, scheduler.isAccountRequestCompatible(context.Background(), account, OpenAIAccountScheduleRequest{RequestedModel: "gpt-5.5"}))
  2606. require.True(t, scheduler.isAccountRequestCompatible(context.Background(), account, OpenAIAccountScheduleRequest{RequestedModel: "gpt-5.6-sol"}))
  2607. }
  2608. func TestReportOpenAIAccountScheduleResult_SuccessClearsModelTransientState(t *testing.T) {
  2609. svc := &OpenAIGatewayService{openaiModelTransient: newOpenAIAccountModelTransientState(128)}
  2610. now := time.Now()
  2611. svc.openaiModelTransient.recordFailure(21636, "gpt-5.5", now)
  2612. svc.openaiModelTransient.recordFailure(21636, "gpt-5.5", now.Add(time.Millisecond))
  2613. require.True(t, svc.openaiModelTransient.isBlocked(21636, "gpt-5.5", now.Add(2*time.Millisecond)))
  2614. svc.ReportOpenAIAccountScheduleResult(&Account{ID: 21636}, "gpt-5.5", true, nil)
  2615. require.False(t, svc.openaiModelTransient.isBlocked(21636, "gpt-5.5", now.Add(2*time.Millisecond)))
  2616. }
  2617. func TestDefaultOpenAIAccountScheduler_ShouldEscapeStickyAccount_ThresholdBoundary(t *testing.T) {
  2618. stats := newOpenAIAccountRuntimeStats()
  2619. accountID := int64(21501)
  2620. ttft := 15000
  2621. stats.report(accountID, true, &ttft)
  2622. stats.report(accountID, false, nil)
  2623. stats.report(accountID, true, nil)
  2624. scheduler := &defaultOpenAIAccountScheduler{stats: stats}
  2625. reason, errorRate, observedTTFT, shouldEscape := scheduler.shouldEscapeStickyAccount(accountID, openAIStickyEscapeConfig{
  2626. enabled: true,
  2627. ttftMs: 15000,
  2628. errorRate: 0.5,
  2629. })
  2630. require.False(t, shouldEscape)
  2631. require.Empty(t, reason)
  2632. require.InDelta(t, 0.16, errorRate, 1e-9)
  2633. require.InDelta(t, 15000, observedTTFT, 1e-9)
  2634. for i := 0; i < 4; i++ {
  2635. stats.report(accountID, false, nil)
  2636. }
  2637. reason, errorRate, _, shouldEscape = scheduler.shouldEscapeStickyAccount(accountID, openAIStickyEscapeConfig{
  2638. enabled: true,
  2639. ttftMs: 15000,
  2640. errorRate: 1,
  2641. })
  2642. require.False(t, shouldEscape)
  2643. require.Empty(t, reason)
  2644. reason, errorRate, observedTTFT, shouldEscape = scheduler.shouldEscapeStickyAccount(accountID, openAIStickyEscapeConfig{
  2645. enabled: true,
  2646. ttftMs: 15000,
  2647. errorRate: errorRate,
  2648. })
  2649. require.False(t, shouldEscape)
  2650. require.Empty(t, reason)
  2651. require.InDelta(t, 0.655936, errorRate, 1e-9)
  2652. require.InDelta(t, 15000, observedTTFT, 1e-9)
  2653. }
  2654. func TestOpenAIGatewayService_SelectAccountWithScheduler_SessionSticky_ForceHTTP(t *testing.T) {
  2655. ctx := context.Background()
  2656. groupID := int64(1010)
  2657. account := Account{
  2658. ID: 2101,
  2659. Platform: PlatformOpenAI,
  2660. Type: AccountTypeOAuth,
  2661. Status: StatusActive,
  2662. Schedulable: true,
  2663. Concurrency: 1,
  2664. GroupIDs: []int64{groupID},
  2665. Extra: map[string]any{
  2666. "openai_ws_force_http": true,
  2667. },
  2668. }
  2669. cache := &schedulerTestGatewayCache{
  2670. sessionBindings: map[string]int64{
  2671. "openai:session_hash_force_http": account.ID,
  2672. },
  2673. }
  2674. svc := &OpenAIGatewayService{
  2675. accountRepo: schedulerTestOpenAIAccountRepo{accounts: []Account{account}},
  2676. cache: cache,
  2677. cfg: &config.Config{},
  2678. rateLimitService: newOpenAIAdvancedSchedulerRateLimitService("true"),
  2679. concurrencyService: NewConcurrencyService(schedulerTestConcurrencyCache{}),
  2680. }
  2681. selection, decision, err := svc.SelectAccountWithScheduler(
  2682. ctx,
  2683. &groupID,
  2684. "",
  2685. "session_hash_force_http",
  2686. "gpt-5.1",
  2687. nil,
  2688. OpenAIUpstreamTransportAny,
  2689. false,
  2690. )
  2691. require.NoError(t, err)
  2692. require.NotNil(t, selection)
  2693. require.NotNil(t, selection.Account)
  2694. require.Equal(t, account.ID, selection.Account.ID)
  2695. require.Equal(t, openAIAccountScheduleLayerSessionSticky, decision.Layer)
  2696. require.True(t, decision.StickySessionHit)
  2697. if selection.ReleaseFunc != nil {
  2698. selection.ReleaseFunc()
  2699. }
  2700. }
  2701. func TestOpenAIGatewayService_SelectAccountWithScheduler_RequiredWSV2_SkipsStickyHTTPAccount(t *testing.T) {
  2702. ctx := context.Background()
  2703. groupID := int64(1011)
  2704. accounts := []Account{
  2705. {
  2706. ID: 2201,
  2707. Platform: PlatformOpenAI,
  2708. Type: AccountTypeAPIKey,
  2709. Status: StatusActive,
  2710. Schedulable: true,
  2711. Concurrency: 1,
  2712. Priority: 0,
  2713. GroupIDs: []int64{groupID},
  2714. },
  2715. {
  2716. ID: 2202,
  2717. Platform: PlatformOpenAI,
  2718. Type: AccountTypeAPIKey,
  2719. Status: StatusActive,
  2720. Schedulable: true,
  2721. Concurrency: 1,
  2722. Priority: 5,
  2723. GroupIDs: []int64{groupID},
  2724. Extra: map[string]any{
  2725. "openai_apikey_responses_websockets_v2_enabled": true,
  2726. },
  2727. },
  2728. }
  2729. cache := &schedulerTestGatewayCache{
  2730. sessionBindings: map[string]int64{
  2731. "openai:session_hash_ws_only": 2201,
  2732. },
  2733. }
  2734. cfg := newSchedulerTestOpenAIWSV2Config()
  2735. // 构造“HTTP-only 账号负载更低”的场景,验证 required transport 会强制过滤。
  2736. concurrencyCache := schedulerTestConcurrencyCache{
  2737. loadMap: map[int64]*AccountLoadInfo{
  2738. 2201: {AccountID: 2201, LoadRate: 0, WaitingCount: 0},
  2739. 2202: {AccountID: 2202, LoadRate: 90, WaitingCount: 5},
  2740. },
  2741. }
  2742. svc := &OpenAIGatewayService{
  2743. accountRepo: schedulerTestOpenAIAccountRepo{accounts: accounts},
  2744. cache: cache,
  2745. cfg: cfg,
  2746. rateLimitService: newOpenAIAdvancedSchedulerRateLimitService("true"),
  2747. concurrencyService: NewConcurrencyService(concurrencyCache),
  2748. }
  2749. selection, decision, err := svc.SelectAccountWithScheduler(
  2750. ctx,
  2751. &groupID,
  2752. "",
  2753. "session_hash_ws_only",
  2754. "gpt-5.1",
  2755. nil,
  2756. OpenAIUpstreamTransportResponsesWebsocketV2,
  2757. false,
  2758. )
  2759. require.NoError(t, err)
  2760. require.NotNil(t, selection)
  2761. require.NotNil(t, selection.Account)
  2762. require.Equal(t, int64(2202), selection.Account.ID)
  2763. require.Equal(t, openAIAccountScheduleLayerLoadBalance, decision.Layer)
  2764. require.False(t, decision.StickySessionHit)
  2765. require.Equal(t, 1, decision.CandidateCount)
  2766. if selection.ReleaseFunc != nil {
  2767. selection.ReleaseFunc()
  2768. }
  2769. }
  2770. func TestOpenAIGatewayService_SelectAccountWithScheduler_ClearsStickyAccountOutsideGroup(t *testing.T) {
  2771. ctx := context.Background()
  2772. groupID := int64(1013)
  2773. accounts := []Account{
  2774. {
  2775. ID: 2401,
  2776. Platform: PlatformOpenAI,
  2777. Type: AccountTypeAPIKey,
  2778. Status: StatusActive,
  2779. Schedulable: true,
  2780. Concurrency: 1,
  2781. Priority: 0,
  2782. },
  2783. {
  2784. ID: 2402,
  2785. Platform: PlatformOpenAI,
  2786. Type: AccountTypeAPIKey,
  2787. Status: StatusActive,
  2788. Schedulable: true,
  2789. Concurrency: 1,
  2790. Priority: 5,
  2791. AccountGroups: []AccountGroup{
  2792. {AccountID: 2402, GroupID: groupID},
  2793. },
  2794. },
  2795. }
  2796. cache := &schedulerTestGatewayCache{
  2797. sessionBindings: map[string]int64{
  2798. "openai:session_hash_removed_group": 2401,
  2799. },
  2800. }
  2801. svc := &OpenAIGatewayService{
  2802. accountRepo: schedulerGroupAwareOpenAIAccountRepo{schedulerTestOpenAIAccountRepo{accounts: accounts}},
  2803. cache: cache,
  2804. cfg: &config.Config{},
  2805. rateLimitService: newOpenAIAdvancedSchedulerRateLimitService("true"),
  2806. concurrencyService: NewConcurrencyService(schedulerTestConcurrencyCache{}),
  2807. }
  2808. selection, decision, err := svc.SelectAccountWithScheduler(
  2809. ctx,
  2810. &groupID,
  2811. "",
  2812. "session_hash_removed_group",
  2813. "gpt-5.1",
  2814. nil,
  2815. OpenAIUpstreamTransportAny,
  2816. false,
  2817. )
  2818. require.NoError(t, err)
  2819. require.NotNil(t, selection)
  2820. require.NotNil(t, selection.Account)
  2821. require.Equal(t, int64(2402), selection.Account.ID)
  2822. require.Equal(t, openAIAccountScheduleLayerLoadBalance, decision.Layer)
  2823. require.False(t, decision.StickySessionHit)
  2824. require.Equal(t, 1, cache.deletedSessions["openai:session_hash_removed_group"])
  2825. require.Equal(t, int64(2402), cache.sessionBindings["openai:session_hash_removed_group"])
  2826. if selection.ReleaseFunc != nil {
  2827. selection.ReleaseFunc()
  2828. }
  2829. }
  2830. func TestOpenAIGatewayService_SelectAccountWithScheduler_RequiredWSV2_NoAvailableAccount(t *testing.T) {
  2831. ctx := context.Background()
  2832. groupID := int64(1012)
  2833. accounts := []Account{
  2834. {
  2835. ID: 2301,
  2836. Platform: PlatformOpenAI,
  2837. Type: AccountTypeOAuth,
  2838. Status: StatusActive,
  2839. Schedulable: true,
  2840. Concurrency: 1,
  2841. },
  2842. }
  2843. svc := &OpenAIGatewayService{
  2844. accountRepo: schedulerTestOpenAIAccountRepo{accounts: accounts},
  2845. cache: &schedulerTestGatewayCache{},
  2846. cfg: newSchedulerTestOpenAIWSV2Config(),
  2847. rateLimitService: newOpenAIAdvancedSchedulerRateLimitService("true"),
  2848. concurrencyService: NewConcurrencyService(schedulerTestConcurrencyCache{}),
  2849. }
  2850. selection, decision, err := svc.SelectAccountWithScheduler(
  2851. ctx,
  2852. &groupID,
  2853. "",
  2854. "",
  2855. "gpt-5.1",
  2856. nil,
  2857. OpenAIUpstreamTransportResponsesWebsocketV2,
  2858. false,
  2859. )
  2860. require.Error(t, err)
  2861. require.Nil(t, selection)
  2862. require.Equal(t, openAIAccountScheduleLayerLoadBalance, decision.Layer)
  2863. require.Equal(t, 0, decision.CandidateCount)
  2864. }
  2865. func TestOpenAIGatewayService_SelectAccountWithScheduler_LoadBalanceTopKFallback(t *testing.T) {
  2866. ctx := context.Background()
  2867. groupID := int64(11)
  2868. accounts := []Account{
  2869. {
  2870. ID: 3001,
  2871. Platform: PlatformOpenAI,
  2872. Type: AccountTypeAPIKey,
  2873. Status: StatusActive,
  2874. Schedulable: true,
  2875. Concurrency: 1,
  2876. Priority: 0,
  2877. },
  2878. {
  2879. ID: 3002,
  2880. Platform: PlatformOpenAI,
  2881. Type: AccountTypeAPIKey,
  2882. Status: StatusActive,
  2883. Schedulable: true,
  2884. Concurrency: 1,
  2885. Priority: 0,
  2886. },
  2887. {
  2888. ID: 3003,
  2889. Platform: PlatformOpenAI,
  2890. Type: AccountTypeAPIKey,
  2891. Status: StatusActive,
  2892. Schedulable: true,
  2893. Concurrency: 1,
  2894. Priority: 0,
  2895. },
  2896. }
  2897. cfg := &config.Config{}
  2898. cfg.Gateway.OpenAIWS.LBTopK = 2
  2899. cfg.Gateway.OpenAIWS.SchedulerScoreWeights.Priority = 0.4
  2900. cfg.Gateway.OpenAIWS.SchedulerScoreWeights.Load = 1.0
  2901. cfg.Gateway.OpenAIWS.SchedulerScoreWeights.Queue = 1.0
  2902. cfg.Gateway.OpenAIWS.SchedulerScoreWeights.ErrorRate = 0.2
  2903. cfg.Gateway.OpenAIWS.SchedulerScoreWeights.TTFT = 0.1
  2904. concurrencyCache := schedulerTestConcurrencyCache{
  2905. loadMap: map[int64]*AccountLoadInfo{
  2906. 3001: {AccountID: 3001, LoadRate: 95, WaitingCount: 8},
  2907. 3002: {AccountID: 3002, LoadRate: 20, WaitingCount: 1},
  2908. 3003: {AccountID: 3003, LoadRate: 10, WaitingCount: 0},
  2909. },
  2910. acquireResults: map[int64]bool{
  2911. 3003: false, // top1 失败,必须回退到 top-K 的下一候选
  2912. 3002: true,
  2913. },
  2914. }
  2915. svc := &OpenAIGatewayService{
  2916. accountRepo: schedulerTestOpenAIAccountRepo{accounts: accounts},
  2917. cache: &schedulerTestGatewayCache{},
  2918. cfg: cfg,
  2919. rateLimitService: newOpenAIAdvancedSchedulerRateLimitService("true"),
  2920. concurrencyService: NewConcurrencyService(concurrencyCache),
  2921. }
  2922. selection, decision, err := svc.SelectAccountWithScheduler(
  2923. ctx,
  2924. &groupID,
  2925. "",
  2926. "",
  2927. "gpt-5.1",
  2928. nil,
  2929. OpenAIUpstreamTransportAny,
  2930. false,
  2931. )
  2932. require.NoError(t, err)
  2933. require.NotNil(t, selection)
  2934. require.NotNil(t, selection.Account)
  2935. require.Equal(t, int64(3002), selection.Account.ID)
  2936. require.Equal(t, openAIAccountScheduleLayerLoadBalance, decision.Layer)
  2937. require.Equal(t, 3, decision.CandidateCount)
  2938. require.Equal(t, 2, decision.TopK)
  2939. require.Greater(t, decision.LoadSkew, 0.0)
  2940. if selection.ReleaseFunc != nil {
  2941. selection.ReleaseFunc()
  2942. }
  2943. }
  2944. // Regression: TopK initial filter must drop quota-auto-paused accounts. Otherwise
  2945. // the candidate pool is filled with paused accounts, healthy accounts fall outside
  2946. // TopK, and the scheduler returns "no available accounts" even though healthy ones
  2947. // exist.
  2948. func TestOpenAIGatewayService_SelectAccountWithScheduler_LoadBalanceTopKExcludesQuotaPaused(t *testing.T) {
  2949. ctx := context.Background()
  2950. groupID := int64(110)
  2951. accounts := []Account{
  2952. {
  2953. ID: 37001,
  2954. Platform: PlatformOpenAI,
  2955. Type: AccountTypeAPIKey,
  2956. Status: StatusActive,
  2957. Schedulable: true,
  2958. Concurrency: 1,
  2959. Priority: 0,
  2960. Extra: map[string]any{
  2961. "codex_5h_used_percent": 96.0,
  2962. "auto_pause_5h_threshold": 0.95,
  2963. },
  2964. },
  2965. {
  2966. ID: 37002,
  2967. Platform: PlatformOpenAI,
  2968. Type: AccountTypeAPIKey,
  2969. Status: StatusActive,
  2970. Schedulable: true,
  2971. Concurrency: 1,
  2972. Priority: 5,
  2973. },
  2974. }
  2975. cfg := &config.Config{}
  2976. cfg.Gateway.OpenAIWS.LBTopK = 1 // TopK=1 makes the bug fatal: paused account would crowd out the healthy one entirely
  2977. cfg.Gateway.OpenAIWS.SchedulerScoreWeights.Priority = 0.4
  2978. cfg.Gateway.OpenAIWS.SchedulerScoreWeights.Load = 1.0
  2979. cfg.Gateway.OpenAIWS.SchedulerScoreWeights.Queue = 1.0
  2980. concurrencyCache := schedulerTestConcurrencyCache{
  2981. loadMap: map[int64]*AccountLoadInfo{
  2982. 37001: {AccountID: 37001, LoadRate: 5, WaitingCount: 0},
  2983. 37002: {AccountID: 37002, LoadRate: 5, WaitingCount: 0},
  2984. },
  2985. acquireResults: map[int64]bool{
  2986. 37002: true,
  2987. },
  2988. }
  2989. svc := &OpenAIGatewayService{
  2990. accountRepo: schedulerTestOpenAIAccountRepo{accounts: accounts},
  2991. cache: &schedulerTestGatewayCache{},
  2992. cfg: cfg,
  2993. rateLimitService: newOpenAIAdvancedSchedulerRateLimitService("true"),
  2994. concurrencyService: NewConcurrencyService(concurrencyCache),
  2995. }
  2996. selection, decision, err := svc.SelectAccountWithScheduler(
  2997. ctx,
  2998. &groupID,
  2999. "",
  3000. "",
  3001. "gpt-5.1",
  3002. nil,
  3003. OpenAIUpstreamTransportAny,
  3004. false,
  3005. )
  3006. require.NoError(t, err)
  3007. require.NotNil(t, selection)
  3008. require.NotNil(t, selection.Account)
  3009. require.Equal(t, int64(37002), selection.Account.ID)
  3010. require.Equal(t, openAIAccountScheduleLayerLoadBalance, decision.Layer)
  3011. // Only the healthy account should ever enter the candidate pool; the paused one
  3012. // must be filtered out at the initial-filter stage.
  3013. require.Equal(t, 1, decision.CandidateCount)
  3014. if selection.ReleaseFunc != nil {
  3015. selection.ReleaseFunc()
  3016. }
  3017. }
  3018. func TestOpenAIGatewayService_OpenAIAccountSchedulerMetrics(t *testing.T) {
  3019. ctx := context.Background()
  3020. groupID := int64(12)
  3021. account := Account{
  3022. ID: 4001,
  3023. Platform: PlatformOpenAI,
  3024. Type: AccountTypeAPIKey,
  3025. Status: StatusActive,
  3026. Schedulable: true,
  3027. Concurrency: 1,
  3028. GroupIDs: []int64{groupID},
  3029. }
  3030. cache := &schedulerTestGatewayCache{
  3031. sessionBindings: map[string]int64{
  3032. "openai:session_hash_metrics": account.ID,
  3033. },
  3034. }
  3035. svc := &OpenAIGatewayService{
  3036. accountRepo: schedulerTestOpenAIAccountRepo{accounts: []Account{account}},
  3037. cache: cache,
  3038. cfg: &config.Config{},
  3039. rateLimitService: newOpenAIAdvancedSchedulerRateLimitService("true"),
  3040. concurrencyService: NewConcurrencyService(schedulerTestConcurrencyCache{}),
  3041. }
  3042. selection, _, err := svc.SelectAccountWithScheduler(ctx, &groupID, "", "session_hash_metrics", "gpt-5.1", nil, OpenAIUpstreamTransportAny, false)
  3043. require.NoError(t, err)
  3044. require.NotNil(t, selection)
  3045. svc.ReportOpenAIAccountScheduleResult(&account, "", true, intPtrForTest(120))
  3046. svc.RecordOpenAIAccountSwitch()
  3047. snapshot := svc.SnapshotOpenAIAccountSchedulerMetrics()
  3048. require.GreaterOrEqual(t, snapshot.SelectTotal, int64(1))
  3049. require.GreaterOrEqual(t, snapshot.StickySessionHitTotal, int64(1))
  3050. require.GreaterOrEqual(t, snapshot.AccountSwitchTotal, int64(1))
  3051. require.GreaterOrEqual(t, snapshot.SchedulerLatencyMsAvg, float64(0))
  3052. require.GreaterOrEqual(t, snapshot.StickyHitRatio, 0.0)
  3053. require.GreaterOrEqual(t, snapshot.RuntimeStatsAccountCount, 1)
  3054. }
  3055. func intPtrForTest(v int) *int {
  3056. return &v
  3057. }
  3058. func TestOpenAIAccountRuntimeStats_ReportAndSnapshot(t *testing.T) {
  3059. stats := newOpenAIAccountRuntimeStats()
  3060. stats.report(1001, true, nil)
  3061. firstTTFT := 100
  3062. stats.report(1001, false, &firstTTFT)
  3063. secondTTFT := 200
  3064. stats.report(1001, false, &secondTTFT)
  3065. errorRate, ttft, hasTTFT := stats.snapshot(1001)
  3066. require.True(t, hasTTFT)
  3067. require.InDelta(t, 0.36, errorRate, 1e-9)
  3068. require.InDelta(t, 120.0, ttft, 1e-9)
  3069. require.Equal(t, 1, stats.size())
  3070. }
  3071. func TestOpenAIAccountRuntimeStats_ReportConcurrent(t *testing.T) {
  3072. stats := newOpenAIAccountRuntimeStats()
  3073. const (
  3074. accountCount = 4
  3075. workers = 16
  3076. iterations = 800
  3077. )
  3078. var wg sync.WaitGroup
  3079. wg.Add(workers)
  3080. for worker := 0; worker < workers; worker++ {
  3081. worker := worker
  3082. go func() {
  3083. defer wg.Done()
  3084. for i := 0; i < iterations; i++ {
  3085. accountID := int64(i%accountCount + 1)
  3086. success := (i+worker)%3 != 0
  3087. ttft := 80 + (i+worker)%40
  3088. stats.report(accountID, success, &ttft)
  3089. }
  3090. }()
  3091. }
  3092. wg.Wait()
  3093. require.Equal(t, accountCount, stats.size())
  3094. for accountID := int64(1); accountID <= accountCount; accountID++ {
  3095. errorRate, ttft, hasTTFT := stats.snapshot(accountID)
  3096. require.GreaterOrEqual(t, errorRate, 0.0)
  3097. require.LessOrEqual(t, errorRate, 1.0)
  3098. require.True(t, hasTTFT)
  3099. require.Greater(t, ttft, 0.0)
  3100. }
  3101. }
  3102. func TestSelectTopKOpenAICandidates(t *testing.T) {
  3103. candidates := []openAIAccountCandidateScore{
  3104. {
  3105. account: &Account{ID: 11, Priority: 2},
  3106. loadInfo: &AccountLoadInfo{LoadRate: 10, WaitingCount: 1},
  3107. score: 10.0,
  3108. },
  3109. {
  3110. account: &Account{ID: 12, Priority: 1},
  3111. loadInfo: &AccountLoadInfo{LoadRate: 20, WaitingCount: 1},
  3112. score: 9.5,
  3113. },
  3114. {
  3115. account: &Account{ID: 13, Priority: 1},
  3116. loadInfo: &AccountLoadInfo{LoadRate: 30, WaitingCount: 0},
  3117. score: 10.0,
  3118. },
  3119. {
  3120. account: &Account{ID: 14, Priority: 0},
  3121. loadInfo: &AccountLoadInfo{LoadRate: 40, WaitingCount: 0},
  3122. score: 8.0,
  3123. },
  3124. }
  3125. top2 := selectTopKOpenAICandidates(candidates, 2)
  3126. require.Len(t, top2, 2)
  3127. require.Equal(t, int64(13), top2[0].account.ID)
  3128. require.Equal(t, int64(11), top2[1].account.ID)
  3129. topAll := selectTopKOpenAICandidates(candidates, 8)
  3130. require.Len(t, topAll, len(candidates))
  3131. require.Equal(t, int64(13), topAll[0].account.ID)
  3132. require.Equal(t, int64(11), topAll[1].account.ID)
  3133. require.Equal(t, int64(12), topAll[2].account.ID)
  3134. require.Equal(t, int64(14), topAll[3].account.ID)
  3135. }
  3136. func TestBuildOpenAIWeightedSelectionOrder_DeterministicBySessionSeed(t *testing.T) {
  3137. candidates := []openAIAccountCandidateScore{
  3138. {
  3139. account: &Account{ID: 101},
  3140. loadInfo: &AccountLoadInfo{LoadRate: 10, WaitingCount: 0},
  3141. score: 4.2,
  3142. },
  3143. {
  3144. account: &Account{ID: 102},
  3145. loadInfo: &AccountLoadInfo{LoadRate: 30, WaitingCount: 1},
  3146. score: 3.5,
  3147. },
  3148. {
  3149. account: &Account{ID: 103},
  3150. loadInfo: &AccountLoadInfo{LoadRate: 50, WaitingCount: 2},
  3151. score: 2.1,
  3152. },
  3153. }
  3154. req := OpenAIAccountScheduleRequest{
  3155. GroupID: int64PtrForTest(99),
  3156. SessionHash: "session_seed_fixed",
  3157. RequestedModel: "gpt-5.1",
  3158. }
  3159. first := buildOpenAIWeightedSelectionOrder(candidates, req)
  3160. second := buildOpenAIWeightedSelectionOrder(candidates, req)
  3161. require.Len(t, first, len(candidates))
  3162. require.Len(t, second, len(candidates))
  3163. for i := range first {
  3164. require.Equal(t, first[i].account.ID, second[i].account.ID)
  3165. }
  3166. }
  3167. func TestOpenAIGatewayService_SelectAccountWithScheduler_LoadBalanceDistributesAcrossSessions(t *testing.T) {
  3168. ctx := context.Background()
  3169. groupID := int64(15)
  3170. accounts := []Account{
  3171. {
  3172. ID: 5101,
  3173. Platform: PlatformOpenAI,
  3174. Type: AccountTypeAPIKey,
  3175. Status: StatusActive,
  3176. Schedulable: true,
  3177. Concurrency: 3,
  3178. Priority: 0,
  3179. },
  3180. {
  3181. ID: 5102,
  3182. Platform: PlatformOpenAI,
  3183. Type: AccountTypeAPIKey,
  3184. Status: StatusActive,
  3185. Schedulable: true,
  3186. Concurrency: 3,
  3187. Priority: 0,
  3188. },
  3189. {
  3190. ID: 5103,
  3191. Platform: PlatformOpenAI,
  3192. Type: AccountTypeAPIKey,
  3193. Status: StatusActive,
  3194. Schedulable: true,
  3195. Concurrency: 3,
  3196. Priority: 0,
  3197. },
  3198. }
  3199. cfg := &config.Config{}
  3200. cfg.Gateway.OpenAIWS.LBTopK = 3
  3201. cfg.Gateway.OpenAIWS.SchedulerScoreWeights.Priority = 1
  3202. cfg.Gateway.OpenAIWS.SchedulerScoreWeights.Load = 1
  3203. cfg.Gateway.OpenAIWS.SchedulerScoreWeights.Queue = 1
  3204. cfg.Gateway.OpenAIWS.SchedulerScoreWeights.ErrorRate = 1
  3205. cfg.Gateway.OpenAIWS.SchedulerScoreWeights.TTFT = 1
  3206. concurrencyCache := schedulerTestConcurrencyCache{
  3207. loadMap: map[int64]*AccountLoadInfo{
  3208. 5101: {AccountID: 5101, LoadRate: 20, WaitingCount: 1},
  3209. 5102: {AccountID: 5102, LoadRate: 20, WaitingCount: 1},
  3210. 5103: {AccountID: 5103, LoadRate: 20, WaitingCount: 1},
  3211. },
  3212. }
  3213. svc := &OpenAIGatewayService{
  3214. accountRepo: schedulerTestOpenAIAccountRepo{accounts: accounts},
  3215. cache: &schedulerTestGatewayCache{sessionBindings: map[string]int64{}},
  3216. cfg: cfg,
  3217. rateLimitService: newOpenAIAdvancedSchedulerRateLimitService("true"),
  3218. concurrencyService: NewConcurrencyService(concurrencyCache),
  3219. }
  3220. selected := make(map[int64]int, len(accounts))
  3221. for i := 0; i < 60; i++ {
  3222. sessionHash := fmt.Sprintf("session_hash_lb_%d", i)
  3223. selection, decision, err := svc.SelectAccountWithScheduler(
  3224. ctx,
  3225. &groupID,
  3226. "",
  3227. sessionHash,
  3228. "gpt-5.1",
  3229. nil,
  3230. OpenAIUpstreamTransportAny,
  3231. false,
  3232. )
  3233. require.NoError(t, err)
  3234. require.NotNil(t, selection)
  3235. require.NotNil(t, selection.Account)
  3236. require.Equal(t, openAIAccountScheduleLayerLoadBalance, decision.Layer)
  3237. selected[selection.Account.ID]++
  3238. if selection.ReleaseFunc != nil {
  3239. selection.ReleaseFunc()
  3240. }
  3241. }
  3242. // 多 session 应该能打散到多个账号,避免“恒定单账号命中”。
  3243. require.GreaterOrEqual(t, len(selected), 2)
  3244. }
  3245. func TestDeriveOpenAISelectionSeed_NoAffinityAddsEntropy(t *testing.T) {
  3246. req := OpenAIAccountScheduleRequest{
  3247. RequestedModel: "gpt-5.1",
  3248. }
  3249. seed1 := deriveOpenAISelectionSeed(req)
  3250. time.Sleep(1 * time.Millisecond)
  3251. seed2 := deriveOpenAISelectionSeed(req)
  3252. require.NotZero(t, seed1)
  3253. require.NotZero(t, seed2)
  3254. require.NotEqual(t, seed1, seed2)
  3255. }
  3256. func TestBuildOpenAIWeightedSelectionOrder_HandlesInvalidScores(t *testing.T) {
  3257. candidates := []openAIAccountCandidateScore{
  3258. {
  3259. account: &Account{ID: 901},
  3260. loadInfo: &AccountLoadInfo{LoadRate: 5, WaitingCount: 0},
  3261. score: math.NaN(),
  3262. },
  3263. {
  3264. account: &Account{ID: 902},
  3265. loadInfo: &AccountLoadInfo{LoadRate: 5, WaitingCount: 0},
  3266. score: math.Inf(1),
  3267. },
  3268. {
  3269. account: &Account{ID: 903},
  3270. loadInfo: &AccountLoadInfo{LoadRate: 5, WaitingCount: 0},
  3271. score: -1,
  3272. },
  3273. }
  3274. req := OpenAIAccountScheduleRequest{
  3275. SessionHash: "seed_invalid_scores",
  3276. }
  3277. order := buildOpenAIWeightedSelectionOrder(candidates, req)
  3278. require.Len(t, order, len(candidates))
  3279. seen := map[int64]struct{}{}
  3280. for _, item := range order {
  3281. seen[item.account.ID] = struct{}{}
  3282. }
  3283. require.Len(t, seen, len(candidates))
  3284. }
  3285. func TestOpenAISelectionRNG_SeedZeroStillWorks(t *testing.T) {
  3286. rng := newOpenAISelectionRNG(0)
  3287. v1 := rng.nextUint64()
  3288. v2 := rng.nextUint64()
  3289. require.NotEqual(t, v1, v2)
  3290. require.GreaterOrEqual(t, rng.nextFloat64(), 0.0)
  3291. require.Less(t, rng.nextFloat64(), 1.0)
  3292. }
  3293. func TestOpenAIAccountCandidateHeap_PushPopAndInvalidType(t *testing.T) {
  3294. h := openAIAccountCandidateHeap{}
  3295. h.Push(openAIAccountCandidateScore{
  3296. account: &Account{ID: 7001},
  3297. loadInfo: &AccountLoadInfo{LoadRate: 0, WaitingCount: 0},
  3298. score: 1.0,
  3299. })
  3300. require.Equal(t, 1, h.Len())
  3301. popped, ok := h.Pop().(openAIAccountCandidateScore)
  3302. require.True(t, ok)
  3303. require.Equal(t, int64(7001), popped.account.ID)
  3304. require.Equal(t, 0, h.Len())
  3305. require.Panics(t, func() {
  3306. h.Push("bad_element_type")
  3307. })
  3308. }
  3309. func TestClamp01_AllBranches(t *testing.T) {
  3310. require.Equal(t, 0.0, clamp01(-0.2))
  3311. require.Equal(t, 1.0, clamp01(1.3))
  3312. require.Equal(t, 0.5, clamp01(0.5))
  3313. }
  3314. func TestCalcLoadSkewByMoments_Branches(t *testing.T) {
  3315. require.Equal(t, 0.0, calcLoadSkewByMoments(1, 1, 1))
  3316. // variance < 0 分支:sumSquares/count - mean^2 为负值时应钳制为 0。
  3317. require.Equal(t, 0.0, calcLoadSkewByMoments(1, 0, 2))
  3318. require.GreaterOrEqual(t, calcLoadSkewByMoments(6, 20, 3), 0.0)
  3319. }
  3320. func TestDefaultOpenAIAccountScheduler_ReportSwitchAndSnapshot(t *testing.T) {
  3321. schedulerAny := newDefaultOpenAIAccountScheduler(&OpenAIGatewayService{}, nil)
  3322. scheduler, ok := schedulerAny.(*defaultOpenAIAccountScheduler)
  3323. require.True(t, ok)
  3324. ttft := 100
  3325. scheduler.ReportResult(1001, true, &ttft)
  3326. scheduler.ReportSwitch()
  3327. scheduler.metrics.recordSelect(OpenAIAccountScheduleDecision{
  3328. Layer: openAIAccountScheduleLayerLoadBalance,
  3329. LatencyMs: 8,
  3330. LoadSkew: 0.5,
  3331. StickyPreviousHit: true,
  3332. })
  3333. scheduler.metrics.recordSelect(OpenAIAccountScheduleDecision{
  3334. Layer: openAIAccountScheduleLayerSessionSticky,
  3335. LatencyMs: 6,
  3336. LoadSkew: 0.2,
  3337. StickySessionHit: true,
  3338. })
  3339. snapshot := scheduler.SnapshotMetrics()
  3340. require.Equal(t, int64(2), snapshot.SelectTotal)
  3341. require.Equal(t, int64(1), snapshot.StickyPreviousHitTotal)
  3342. require.Equal(t, int64(1), snapshot.StickySessionHitTotal)
  3343. require.Equal(t, int64(1), snapshot.LoadBalanceSelectTotal)
  3344. require.Equal(t, int64(1), snapshot.AccountSwitchTotal)
  3345. require.Greater(t, snapshot.SchedulerLatencyMsAvg, 0.0)
  3346. require.Greater(t, snapshot.StickyHitRatio, 0.0)
  3347. require.Greater(t, snapshot.LoadSkewAvg, 0.0)
  3348. }
  3349. func TestOpenAIGatewayService_SchedulerWrappersAndDefaults(t *testing.T) {
  3350. resetOpenAIAdvancedSchedulerSettingCacheForTest()
  3351. svc := &OpenAIGatewayService{}
  3352. ttft := 120
  3353. svc.ReportOpenAIAccountScheduleResult(&Account{ID: 10}, "", true, &ttft)
  3354. svc.RecordOpenAIAccountSwitch()
  3355. snapshot := svc.SnapshotOpenAIAccountSchedulerMetrics()
  3356. require.Equal(t, OpenAIAccountSchedulerMetricsSnapshot{}, snapshot)
  3357. require.Equal(t, 7, svc.openAIWSLBTopK())
  3358. require.Equal(t, openaiStickySessionTTL, svc.openAIWSSessionStickyTTL())
  3359. defaultWeights := svc.openAIWSSchedulerWeights()
  3360. require.Equal(t, 1.0, defaultWeights.Priority)
  3361. require.Equal(t, 1.0, defaultWeights.Load)
  3362. require.Equal(t, 0.7, defaultWeights.Queue)
  3363. require.Equal(t, 0.8, defaultWeights.ErrorRate)
  3364. require.Equal(t, 0.5, defaultWeights.TTFT)
  3365. cfg := &config.Config{}
  3366. cfg.Gateway.OpenAIWS.LBTopK = 9
  3367. cfg.Gateway.OpenAIWS.StickySessionTTLSeconds = 180
  3368. cfg.Gateway.OpenAIWS.SchedulerScoreWeights.Priority = 0.2
  3369. cfg.Gateway.OpenAIWS.SchedulerScoreWeights.Load = 0.3
  3370. cfg.Gateway.OpenAIWS.SchedulerScoreWeights.Queue = 0.4
  3371. cfg.Gateway.OpenAIWS.SchedulerScoreWeights.ErrorRate = 0.5
  3372. cfg.Gateway.OpenAIWS.SchedulerScoreWeights.TTFT = 0.6
  3373. svcWithCfg := &OpenAIGatewayService{cfg: cfg}
  3374. require.Equal(t, 9, svcWithCfg.openAIWSLBTopK())
  3375. require.Equal(t, 180*time.Second, svcWithCfg.openAIWSSessionStickyTTL())
  3376. customWeights := svcWithCfg.openAIWSSchedulerWeights()
  3377. require.Equal(t, 0.2, customWeights.Priority)
  3378. require.Equal(t, 0.3, customWeights.Load)
  3379. require.Equal(t, 0.4, customWeights.Queue)
  3380. require.Equal(t, 0.5, customWeights.ErrorRate)
  3381. require.Equal(t, 0.6, customWeights.TTFT)
  3382. }
  3383. func TestDefaultOpenAIAccountScheduler_IsAccountTransportCompatible_Branches(t *testing.T) {
  3384. scheduler := &defaultOpenAIAccountScheduler{}
  3385. require.True(t, scheduler.isAccountTransportCompatible(nil, OpenAIUpstreamTransportAny))
  3386. require.True(t, scheduler.isAccountTransportCompatible(nil, OpenAIUpstreamTransportHTTPSSE))
  3387. require.False(t, scheduler.isAccountTransportCompatible(nil, OpenAIUpstreamTransportResponsesWebsocketV2))
  3388. require.False(t, scheduler.isAccountTransportCompatible(nil, OpenAIUpstreamTransportResponsesWebsocketV2Ingress))
  3389. cfg := newSchedulerTestOpenAIWSV2Config()
  3390. scheduler.service = &OpenAIGatewayService{cfg: cfg}
  3391. account := &Account{
  3392. ID: 8801,
  3393. Platform: PlatformOpenAI,
  3394. Type: AccountTypeAPIKey,
  3395. Status: StatusActive,
  3396. Schedulable: true,
  3397. Concurrency: 1,
  3398. Extra: map[string]any{
  3399. "openai_apikey_responses_websockets_v2_enabled": true,
  3400. },
  3401. }
  3402. require.True(t, scheduler.isAccountTransportCompatible(account, OpenAIUpstreamTransportResponsesWebsocketV2))
  3403. require.True(t, scheduler.isAccountTransportCompatible(account, OpenAIUpstreamTransportResponsesWebsocketV2Ingress))
  3404. cfg.Gateway.OpenAIWS.ModeRouterV2Enabled = true
  3405. account.Extra["openai_apikey_responses_websockets_v2_mode"] = OpenAIWSIngressModeHTTPBridge
  3406. require.False(t, scheduler.isAccountTransportCompatible(account, OpenAIUpstreamTransportResponsesWebsocketV2))
  3407. require.True(t, scheduler.isAccountTransportCompatible(account, OpenAIUpstreamTransportResponsesWebsocketV2Ingress))
  3408. account.Extra["openai_apikey_responses_websockets_v2_mode"] = OpenAIWSIngressModeOff
  3409. require.False(t, scheduler.isAccountTransportCompatible(account, OpenAIUpstreamTransportResponsesWebsocketV2Ingress))
  3410. }
  3411. func int64PtrForTest(v int64) *int64 {
  3412. return &v
  3413. }
  3414. func TestOpenAIGatewayService_SelectAccountWithScheduler_StickyWeightedFallbackSkipsOutOfGroupStickyAccount(t *testing.T) {
  3415. resetOpenAIAdvancedSchedulerSettingCacheForTest()
  3416. ctx := context.Background()
  3417. groupID := int64(101081)
  3418. otherGroupID := int64(101082)
  3419. accounts := []Account{
  3420. {
  3421. ID: 38001,
  3422. Platform: PlatformOpenAI,
  3423. Type: AccountTypeAPIKey,
  3424. Status: StatusActive,
  3425. Schedulable: true,
  3426. Concurrency: 1,
  3427. Priority: 10,
  3428. GroupIDs: []int64{groupID},
  3429. },
  3430. {
  3431. // 会话粘连绑定指向的账号已被移出请求分组(绑定 TTL 内账号改组的场景)。
  3432. ID: 38002,
  3433. Platform: PlatformOpenAI,
  3434. Type: AccountTypeAPIKey,
  3435. Status: StatusActive,
  3436. Schedulable: true,
  3437. Concurrency: 1,
  3438. Priority: 0,
  3439. GroupIDs: []int64{otherGroupID},
  3440. },
  3441. }
  3442. cfg := &config.Config{}
  3443. cfg.Gateway.OpenAIWS.LBTopK = 2
  3444. cfg.Gateway.OpenAIWS.SchedulerScoreWeights.Priority = 1
  3445. cfg.Gateway.OpenAIWS.SchedulerScoreWeights.Load = 1
  3446. cfg.Gateway.OpenAIWS.SchedulerScoreWeights.Queue = 0.7
  3447. cfg.Gateway.OpenAIWS.SchedulerScoreWeights.ErrorRate = 0.8
  3448. cfg.Gateway.OpenAIWS.SchedulerScoreWeights.TTFT = 0.5
  3449. cfg.Gateway.OpenAIWS.SchedulerScoreWeights.SessionSticky = 3
  3450. cache := &schedulerTestGatewayCache{sessionBindings: map[string]int64{
  3451. "openai:session_weighted_out_of_group": 38002,
  3452. }}
  3453. concurrencyCache := schedulerTestConcurrencyCache{
  3454. acquireResults: map[int64]bool{38001: false, 38002: true},
  3455. }
  3456. svc := &OpenAIGatewayService{
  3457. accountRepo: schedulerGroupAwareOpenAIAccountRepo{schedulerTestOpenAIAccountRepo{accounts: accounts}},
  3458. cache: cache,
  3459. cfg: cfg,
  3460. rateLimitService: newOpenAIAdvancedSchedulerRateLimitService("true", "true"),
  3461. concurrencyService: NewConcurrencyService(concurrencyCache),
  3462. }
  3463. selection, decision, err := svc.SelectAccountWithScheduler(
  3464. ctx,
  3465. &groupID,
  3466. "",
  3467. "session_weighted_out_of_group",
  3468. "gpt-5.1",
  3469. nil,
  3470. OpenAIUpstreamTransportAny,
  3471. false,
  3472. )
  3473. require.NoError(t, err)
  3474. require.NotNil(t, selection)
  3475. require.NotNil(t, selection.Account)
  3476. // 组内唯一候选 38001 满并发:必须返回其等待计划,绝不能把请求泄漏到组外的粘连账号 38002。
  3477. require.Equal(t, int64(38001), selection.Account.ID)
  3478. require.False(t, selection.Acquired)
  3479. require.NotNil(t, selection.WaitPlan)
  3480. require.Equal(t, int64(38001), selection.WaitPlan.AccountID)
  3481. require.Equal(t, openAIAccountScheduleLayerLoadBalance, decision.Layer)
  3482. // 失效的粘连绑定应被清理,避免后续请求反复走同一条泄漏路径。
  3483. require.Positive(t, cache.deletedSessions["openai:session_weighted_out_of_group"])
  3484. }
  3485. func TestOpenAIGatewayService_SelectAccountWithScheduler_SubscriptionPriorityWaitsOnBusySubscriptionWhenRegularUnusable(t *testing.T) {
  3486. resetOpenAIAdvancedSchedulerSettingCacheForTest()
  3487. ctx := context.Background()
  3488. groupID := int64(101091)
  3489. accounts := []Account{
  3490. {
  3491. // 订阅账号:支持 compact,但并发已满(busy-but-waitable)。
  3492. ID: 38011,
  3493. Platform: PlatformOpenAI,
  3494. Type: AccountTypeOAuth,
  3495. Status: StatusActive,
  3496. Schedulable: true,
  3497. Concurrency: 1,
  3498. Priority: 0,
  3499. GroupIDs: []int64{groupID},
  3500. Credentials: map[string]any{"plan_type": "team"},
  3501. Extra: map[string]any{"openai_compact_supported": true},
  3502. },
  3503. {
  3504. // 常规账号:明确不支持 compact,无法服务本次请求。
  3505. ID: 38012,
  3506. Platform: PlatformOpenAI,
  3507. Type: AccountTypeAPIKey,
  3508. Status: StatusActive,
  3509. Schedulable: true,
  3510. Concurrency: 1,
  3511. Priority: 9,
  3512. GroupIDs: []int64{groupID},
  3513. Extra: map[string]any{"openai_compact_supported": false},
  3514. },
  3515. }
  3516. concurrencyCache := schedulerTestConcurrencyCache{
  3517. acquireResults: map[int64]bool{38011: false, 38012: true},
  3518. }
  3519. svc := &OpenAIGatewayService{
  3520. accountRepo: schedulerTestOpenAIAccountRepo{accounts: accounts},
  3521. cache: &schedulerTestGatewayCache{},
  3522. cfg: newSchedulerTestSubscriptionPriorityConfig(),
  3523. rateLimitService: newOpenAIAdvancedSchedulerRateLimitService("true", "", "true"),
  3524. concurrencyService: NewConcurrencyService(concurrencyCache),
  3525. }
  3526. selection, decision, err := svc.SelectAccountWithScheduler(
  3527. ctx,
  3528. &groupID,
  3529. "",
  3530. "session_subscription_wait",
  3531. "gpt-5.1",
  3532. nil,
  3533. OpenAIUpstreamTransportAny,
  3534. true,
  3535. )
  3536. // 常规池无可用候选时,忙碌的订阅账号应产生等待计划,而不是直接返回 no available accounts。
  3537. require.NoError(t, err)
  3538. require.NotNil(t, selection)
  3539. require.NotNil(t, selection.Account)
  3540. require.Equal(t, int64(38011), selection.Account.ID)
  3541. require.False(t, selection.Acquired)
  3542. require.NotNil(t, selection.WaitPlan)
  3543. require.Equal(t, int64(38011), selection.WaitPlan.AccountID)
  3544. require.Equal(t, openAIAccountScheduleLayerLoadBalance, decision.Layer)
  3545. }