| 12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667686970717273747576777879808182838485868788899091929394959697989910010110210310410510610710810911011111211311411511611711811912012112212312412512612712812913013113213313413513613713813914014114214314414514614714814915015115215315415515615715815916016116216316416516616716816917017117217317417517617717817918018118218318418518618718818919019119219319419519619719819920020120220320420520620720820921021121221321421521621721821922022122222322422522622722822923023123223323423523623723823924024124224324424524624724824925025125225325425525625725825926026126226326426526626726826927027127227327427527627727827928028128228328428528628728828929029129229329429529629729829930030130230330430530630730830931031131231331431531631731831932032132232332432532632732832933033133233333433533633733833934034134234334434534634734834935035135235335435535635735835936036136236336436536636736836937037137237337437537637737837938038138238338438538638738838939039139239339439539639739839940040140240340440540640740840941041141241341441541641741841942042142242342442542642742842943043143243343443543643743843944044144244344444544644744844945045145245345445545645745845946046146246346446546646746846947047147247347447547647747847948048148248348448548648748848949049149249349449549649749849950050150250350450550650750850951051151251351451551651751851952052152252352452552652752852953053153253353453553653753853954054154254354454554654754854955055155255355455555655755855956056156256356456556656756856957057157257357457557657757857958058158258358458558658758858959059159259359459559659759859960060160260360460560660760860961061161261361461561661761861962062162262362462562662762862963063163263363463563663763863964064164264364464564664764864965065165265365465565665765865966066166266366466566666766866967067167267367467567667767867968068168268368468568668768868969069169269369469569669769869970070170270370470570670770870971071171271371471571671771871972072172272372472572672772872973073173273373473573673773873974074174274374474574674774874975075175275375475575675775875976076176276376476576676776876977077177277377477577677777877978078178278378478578678778878979079179279379479579679779879980080180280380480580680780880981081181281381481581681781881982082182282382482582682782882983083183283383483583683783883984084184284384484584684784884985085185285385485585685785885986086186286386486586686786886987087187287387487587687787887988088188288388488588688788888989089189289389489589689789889990090190290390490590690790890991091191291391491591691791891992092192292392492592692792892993093193293393493593693793893994094194294394494594694794894995095195295395495595695795895996096196296396496596696796896997097197297397497597697797897998098198298398498598698798898999099199299399499599699799899910001001100210031004100510061007100810091010101110121013101410151016101710181019102010211022102310241025102610271028102910301031103210331034103510361037103810391040104110421043104410451046104710481049105010511052105310541055105610571058105910601061106210631064106510661067106810691070107110721073107410751076107710781079108010811082108310841085108610871088108910901091109210931094109510961097109810991100110111021103110411051106110711081109111011111112111311141115111611171118111911201121112211231124112511261127112811291130113111321133113411351136113711381139114011411142114311441145114611471148114911501151115211531154115511561157115811591160116111621163116411651166116711681169117011711172117311741175117611771178117911801181118211831184118511861187118811891190119111921193119411951196119711981199120012011202120312041205120612071208120912101211121212131214121512161217121812191220122112221223122412251226122712281229123012311232123312341235123612371238123912401241124212431244124512461247124812491250125112521253125412551256125712581259126012611262126312641265126612671268126912701271127212731274127512761277127812791280128112821283128412851286128712881289129012911292129312941295129612971298129913001301130213031304130513061307130813091310131113121313131413151316131713181319132013211322132313241325132613271328132913301331133213331334133513361337133813391340134113421343134413451346134713481349135013511352135313541355135613571358135913601361136213631364136513661367136813691370137113721373137413751376137713781379138013811382138313841385138613871388138913901391139213931394139513961397139813991400140114021403140414051406140714081409141014111412141314141415141614171418141914201421142214231424142514261427142814291430143114321433143414351436143714381439144014411442144314441445144614471448144914501451145214531454145514561457145814591460146114621463146414651466146714681469147014711472147314741475147614771478147914801481148214831484148514861487148814891490149114921493149414951496149714981499150015011502150315041505150615071508150915101511151215131514151515161517151815191520152115221523152415251526152715281529153015311532153315341535153615371538153915401541154215431544154515461547154815491550155115521553155415551556155715581559156015611562156315641565156615671568156915701571157215731574157515761577157815791580158115821583158415851586158715881589159015911592159315941595159615971598159916001601160216031604160516061607160816091610161116121613161416151616161716181619162016211622162316241625162616271628162916301631163216331634163516361637163816391640164116421643164416451646164716481649165016511652165316541655165616571658165916601661166216631664166516661667166816691670167116721673167416751676167716781679168016811682168316841685168616871688168916901691169216931694169516961697169816991700170117021703170417051706170717081709171017111712171317141715171617171718171917201721172217231724172517261727172817291730173117321733173417351736173717381739174017411742174317441745174617471748174917501751175217531754175517561757175817591760176117621763176417651766176717681769177017711772177317741775177617771778177917801781178217831784178517861787178817891790179117921793179417951796179717981799180018011802180318041805180618071808180918101811181218131814181518161817181818191820182118221823182418251826182718281829183018311832183318341835183618371838183918401841184218431844184518461847184818491850185118521853185418551856185718581859186018611862186318641865186618671868186918701871187218731874187518761877187818791880188118821883188418851886188718881889189018911892189318941895189618971898189919001901190219031904190519061907190819091910191119121913191419151916191719181919192019211922192319241925192619271928192919301931193219331934193519361937193819391940194119421943194419451946194719481949195019511952195319541955195619571958195919601961196219631964196519661967196819691970197119721973197419751976197719781979198019811982198319841985198619871988198919901991199219931994199519961997199819992000200120022003200420052006200720082009201020112012201320142015201620172018201920202021202220232024202520262027202820292030203120322033203420352036203720382039204020412042204320442045204620472048204920502051205220532054205520562057205820592060206120622063206420652066206720682069207020712072207320742075207620772078207920802081208220832084208520862087208820892090209120922093209420952096209720982099210021012102210321042105210621072108210921102111211221132114211521162117211821192120212121222123212421252126212721282129213021312132213321342135213621372138213921402141214221432144214521462147214821492150215121522153215421552156215721582159216021612162216321642165216621672168216921702171217221732174217521762177217821792180218121822183218421852186218721882189219021912192219321942195219621972198219922002201220222032204220522062207220822092210221122122213221422152216221722182219222022212222222322242225222622272228222922302231223222332234223522362237223822392240224122422243224422452246224722482249225022512252225322542255225622572258225922602261226222632264226522662267226822692270227122722273227422752276227722782279228022812282228322842285228622872288228922902291229222932294229522962297229822992300230123022303230423052306230723082309231023112312231323142315231623172318231923202321232223232324232523262327232823292330233123322333233423352336233723382339234023412342234323442345234623472348234923502351235223532354235523562357235823592360236123622363236423652366236723682369237023712372237323742375237623772378237923802381238223832384238523862387238823892390239123922393239423952396239723982399240024012402240324042405240624072408240924102411241224132414241524162417241824192420242124222423242424252426242724282429243024312432243324342435243624372438243924402441244224432444244524462447244824492450245124522453245424552456245724582459246024612462246324642465246624672468246924702471247224732474247524762477247824792480248124822483248424852486248724882489249024912492249324942495249624972498249925002501250225032504250525062507250825092510251125122513251425152516251725182519252025212522252325242525252625272528252925302531253225332534253525362537253825392540254125422543254425452546254725482549255025512552255325542555255625572558255925602561256225632564256525662567256825692570257125722573257425752576257725782579258025812582258325842585258625872588258925902591259225932594259525962597259825992600260126022603260426052606260726082609261026112612261326142615261626172618261926202621262226232624262526262627262826292630263126322633263426352636263726382639264026412642264326442645264626472648264926502651265226532654265526562657265826592660266126622663266426652666266726682669267026712672267326742675267626772678267926802681268226832684268526862687268826892690269126922693269426952696269726982699270027012702270327042705270627072708270927102711271227132714271527162717271827192720272127222723272427252726272727282729273027312732273327342735273627372738273927402741274227432744274527462747274827492750275127522753275427552756275727582759276027612762276327642765276627672768276927702771277227732774277527762777277827792780278127822783278427852786278727882789279027912792279327942795279627972798279928002801280228032804280528062807280828092810281128122813281428152816281728182819282028212822282328242825282628272828282928302831283228332834283528362837283828392840284128422843284428452846284728482849285028512852285328542855285628572858285928602861286228632864286528662867286828692870287128722873287428752876287728782879288028812882288328842885288628872888288928902891289228932894289528962897289828992900290129022903290429052906290729082909291029112912291329142915291629172918291929202921292229232924292529262927292829292930293129322933293429352936293729382939294029412942294329442945294629472948294929502951295229532954295529562957295829592960296129622963296429652966296729682969297029712972297329742975297629772978297929802981298229832984298529862987298829892990299129922993299429952996299729982999300030013002300330043005300630073008300930103011301230133014301530163017301830193020302130223023302430253026302730283029303030313032303330343035303630373038303930403041304230433044304530463047304830493050305130523053305430553056305730583059306030613062306330643065306630673068306930703071307230733074307530763077307830793080308130823083308430853086308730883089309030913092309330943095309630973098309931003101310231033104310531063107310831093110311131123113311431153116311731183119312031213122312331243125312631273128312931303131313231333134313531363137313831393140314131423143314431453146314731483149315031513152315331543155315631573158315931603161316231633164316531663167316831693170317131723173317431753176317731783179318031813182318331843185318631873188318931903191319231933194319531963197319831993200320132023203320432053206320732083209321032113212321332143215321632173218321932203221322232233224322532263227322832293230323132323233323432353236323732383239324032413242324332443245324632473248324932503251325232533254325532563257325832593260326132623263326432653266326732683269327032713272327332743275327632773278327932803281328232833284328532863287328832893290329132923293329432953296329732983299330033013302330333043305330633073308330933103311331233133314331533163317331833193320332133223323332433253326332733283329333033313332333333343335333633373338333933403341334233433344334533463347334833493350335133523353335433553356335733583359336033613362336333643365336633673368336933703371337233733374337533763377337833793380338133823383338433853386338733883389339033913392339333943395339633973398339934003401340234033404340534063407340834093410341134123413341434153416341734183419342034213422342334243425342634273428342934303431343234333434343534363437343834393440344134423443344434453446344734483449345034513452345334543455345634573458345934603461346234633464346534663467346834693470347134723473347434753476347734783479348034813482348334843485348634873488348934903491349234933494349534963497349834993500350135023503350435053506350735083509351035113512351335143515351635173518351935203521352235233524352535263527352835293530353135323533353435353536353735383539354035413542354335443545354635473548354935503551355235533554355535563557355835593560356135623563356435653566356735683569357035713572357335743575357635773578357935803581358235833584358535863587358835893590359135923593359435953596359735983599360036013602360336043605360636073608360936103611361236133614361536163617361836193620362136223623362436253626362736283629363036313632363336343635363636373638363936403641364236433644364536463647364836493650365136523653365436553656365736583659366036613662366336643665366636673668366936703671367236733674367536763677367836793680368136823683368436853686368736883689369036913692369336943695369636973698369937003701370237033704370537063707370837093710371137123713371437153716371737183719372037213722372337243725372637273728372937303731373237333734373537363737373837393740374137423743374437453746374737483749375037513752375337543755375637573758375937603761376237633764376537663767376837693770377137723773377437753776377737783779378037813782378337843785378637873788378937903791379237933794379537963797379837993800380138023803380438053806380738083809381038113812381338143815 |
- package service
- import (
- "context"
- "errors"
- "fmt"
- "math"
- "strconv"
- "sync"
- "testing"
- "time"
- "github.com/Wei-Shaw/sub2api/internal/config"
- "github.com/stretchr/testify/require"
- )
- type openAISnapshotCacheStub struct {
- SchedulerCache
- snapshotAccounts []*Account
- accountsByID map[int64]*Account
- }
- type schedulerTestOpenAIAccountRepo struct {
- AccountRepository
- accounts []Account
- }
- func (r schedulerTestOpenAIAccountRepo) GetByID(ctx context.Context, id int64) (*Account, error) {
- for i := range r.accounts {
- if r.accounts[i].ID == id {
- return &r.accounts[i], nil
- }
- }
- return nil, errors.New("account not found")
- }
- func (r schedulerTestOpenAIAccountRepo) ListSchedulableByGroupIDAndPlatform(ctx context.Context, groupID int64, platform string) ([]Account, error) {
- var result []Account
- for _, acc := range r.accounts {
- if acc.Platform == platform {
- result = append(result, acc)
- }
- }
- return result, nil
- }
- func (r schedulerTestOpenAIAccountRepo) ListSchedulableByPlatform(ctx context.Context, platform string) ([]Account, error) {
- var result []Account
- for _, acc := range r.accounts {
- if acc.Platform == platform {
- result = append(result, acc)
- }
- }
- return result, nil
- }
- func (r schedulerTestOpenAIAccountRepo) ListSchedulableUngroupedByPlatform(ctx context.Context, platform string) ([]Account, error) {
- return r.ListSchedulableByPlatform(ctx, platform)
- }
- type schedulerGroupAwareOpenAIAccountRepo struct {
- schedulerTestOpenAIAccountRepo
- }
- func (r schedulerGroupAwareOpenAIAccountRepo) ListSchedulableByGroupIDAndPlatform(ctx context.Context, groupID int64, platform string) ([]Account, error) {
- var result []Account
- for _, acc := range r.accounts {
- if acc.Platform == platform && openAIStickyAccountMatchesGroup(&acc, &groupID) {
- result = append(result, acc)
- }
- }
- return result, nil
- }
- func (r schedulerGroupAwareOpenAIAccountRepo) ListSchedulableUngroupedByPlatform(ctx context.Context, platform string) ([]Account, error) {
- var result []Account
- for _, acc := range r.accounts {
- if acc.Platform == platform && openAIStickyAccountMatchesGroup(&acc, nil) {
- result = append(result, acc)
- }
- }
- return result, nil
- }
- type schedulerTestConcurrencyCache struct {
- ConcurrencyCache
- loadBatchErr error
- loadMap map[int64]*AccountLoadInfo
- acquireResults map[int64]bool
- waitCounts map[int64]int
- skipDefaultLoad bool
- acquiredIDs *[]int64
- releasedIDs *[]int64
- }
- func (c schedulerTestConcurrencyCache) AcquireAccountSlot(ctx context.Context, accountID int64, maxConcurrency int, requestID string) (bool, error) {
- if c.acquiredIDs != nil {
- *c.acquiredIDs = append(*c.acquiredIDs, accountID)
- }
- if c.acquireResults != nil {
- if result, ok := c.acquireResults[accountID]; ok {
- return result, nil
- }
- }
- return true, nil
- }
- func (c schedulerTestConcurrencyCache) ReleaseAccountSlot(ctx context.Context, accountID int64, requestID string) error {
- if c.releasedIDs != nil {
- *c.releasedIDs = append(*c.releasedIDs, accountID)
- }
- return nil
- }
- func (c schedulerTestConcurrencyCache) GetAccountsLoadBatch(ctx context.Context, accounts []AccountWithConcurrency) (map[int64]*AccountLoadInfo, error) {
- if c.loadBatchErr != nil {
- return nil, c.loadBatchErr
- }
- out := make(map[int64]*AccountLoadInfo, len(accounts))
- if c.skipDefaultLoad && c.loadMap != nil {
- for _, acc := range accounts {
- if load, ok := c.loadMap[acc.ID]; ok {
- out[acc.ID] = load
- }
- }
- return out, nil
- }
- for _, acc := range accounts {
- if c.loadMap != nil {
- if load, ok := c.loadMap[acc.ID]; ok {
- out[acc.ID] = load
- continue
- }
- }
- out[acc.ID] = &AccountLoadInfo{AccountID: acc.ID, LoadRate: 0}
- }
- return out, nil
- }
- func (c schedulerTestConcurrencyCache) GetAccountWaitingCount(ctx context.Context, accountID int64) (int, error) {
- if c.waitCounts != nil {
- if count, ok := c.waitCounts[accountID]; ok {
- return count, nil
- }
- }
- return 0, nil
- }
- type schedulerTestGatewayCache struct {
- sessionBindings map[string]int64
- deletedSessions map[string]int
- }
- func (c *schedulerTestGatewayCache) GetSessionAccountID(ctx context.Context, groupID int64, sessionHash string) (int64, error) {
- if id, ok := c.sessionBindings[sessionHash]; ok {
- return id, nil
- }
- return 0, ErrStickySessionNotFound
- }
- func (c *schedulerTestGatewayCache) SetSessionAccountID(ctx context.Context, groupID int64, sessionHash string, accountID int64, ttl time.Duration) error {
- if c.sessionBindings == nil {
- c.sessionBindings = make(map[string]int64)
- }
- c.sessionBindings[sessionHash] = accountID
- return nil
- }
- func (c *schedulerTestGatewayCache) RefreshSessionTTL(ctx context.Context, groupID int64, sessionHash string, ttl time.Duration) error {
- return nil
- }
- func (c *schedulerTestGatewayCache) DeleteSessionAccountID(ctx context.Context, groupID int64, sessionHash string) error {
- if c.sessionBindings == nil {
- return nil
- }
- if c.deletedSessions == nil {
- c.deletedSessions = make(map[string]int)
- }
- c.deletedSessions[sessionHash]++
- delete(c.sessionBindings, sessionHash)
- return nil
- }
- func (c *schedulerTestGatewayCache) SetGrokVideoPendingBilling(_ context.Context, _ string, _ []byte, _ time.Duration) error {
- return nil
- }
- func (c *schedulerTestGatewayCache) GetGrokVideoPendingBilling(_ context.Context, _ string) ([]byte, error) {
- return nil, nil
- }
- func (c *schedulerTestGatewayCache) ClaimGrokVideoBilled(_ context.Context, _ string, _ time.Duration) (bool, error) {
- return true, nil
- }
- func (c *schedulerTestGatewayCache) ReleaseGrokVideoBilled(_ context.Context, _ string) error {
- return nil
- }
- func (c *schedulerTestGatewayCache) SetReasoningContent(_ context.Context, _ string, _ string, _ time.Duration) error {
- return nil
- }
- func (c *schedulerTestGatewayCache) GetReasoningContent(_ context.Context, _ string) (string, error) {
- return "", ErrReasoningContentNotFound
- }
- func newSchedulerTestOpenAIWSV2Config() *config.Config {
- cfg := &config.Config{}
- cfg.Gateway.OpenAIWS.Enabled = true
- cfg.Gateway.OpenAIWS.OAuthEnabled = true
- cfg.Gateway.OpenAIWS.APIKeyEnabled = true
- cfg.Gateway.OpenAIWS.ResponsesWebsocketsV2 = true
- cfg.Gateway.OpenAIWS.StickyResponseIDTTLSeconds = 3600
- return cfg
- }
- func newSchedulerTestSubscriptionPriorityConfig() *config.Config {
- cfg := &config.Config{}
- cfg.Gateway.OpenAIWS.LBTopK = 1
- cfg.Gateway.OpenAIWS.SchedulerScoreWeights.Priority = 1
- cfg.Gateway.OpenAIWS.SchedulerScoreWeights.Load = 1
- cfg.Gateway.OpenAIWS.SchedulerScoreWeights.Queue = 1
- cfg.Gateway.OpenAIWS.SchedulerScoreWeights.ErrorRate = 0
- cfg.Gateway.OpenAIWS.SchedulerScoreWeights.TTFT = 0
- return cfg
- }
- type openAIAdvancedSchedulerSettingRepoStub struct {
- values map[string]string
- }
- func (s *openAIAdvancedSchedulerSettingRepoStub) Get(ctx context.Context, key string) (*Setting, error) {
- value, err := s.GetValue(ctx, key)
- if err != nil {
- return nil, err
- }
- return &Setting{Key: key, Value: value}, nil
- }
- func (s *openAIAdvancedSchedulerSettingRepoStub) GetValue(_ context.Context, key string) (string, error) {
- if s == nil || s.values == nil {
- return "", ErrSettingNotFound
- }
- value, ok := s.values[key]
- if !ok {
- return "", ErrSettingNotFound
- }
- return value, nil
- }
- func (s *openAIAdvancedSchedulerSettingRepoStub) Set(context.Context, string, string) error {
- panic("unexpected call to Set")
- }
- func (s *openAIAdvancedSchedulerSettingRepoStub) GetMultiple(_ context.Context, keys []string) (map[string]string, error) {
- result := make(map[string]string, len(keys))
- for _, key := range keys {
- if value, err := s.GetValue(context.Background(), key); err == nil {
- result[key] = value
- }
- }
- return result, nil
- }
- func (s *openAIAdvancedSchedulerSettingRepoStub) SetMultiple(context.Context, map[string]string) error {
- panic("unexpected call to SetMultiple")
- }
- func (s *openAIAdvancedSchedulerSettingRepoStub) GetAll(context.Context) (map[string]string, error) {
- panic("unexpected call to GetAll")
- }
- func (s *openAIAdvancedSchedulerSettingRepoStub) Delete(context.Context, string) error {
- panic("unexpected call to Delete")
- }
- func newOpenAIAdvancedSchedulerRateLimitService(enabled string, values ...string) *RateLimitService {
- resetOpenAIAdvancedSchedulerSettingCacheForTest()
- repo := &openAIAdvancedSchedulerSettingRepoStub{
- values: map[string]string{},
- }
- if enabled != "" {
- repo.values[openAIAdvancedSchedulerSettingKey] = enabled
- }
- if len(values) > 0 && values[0] != "" {
- repo.values[SettingKeyOpenAIAdvancedSchedulerStickyWeightedEnabled] = values[0]
- }
- if len(values) > 1 && values[1] != "" {
- repo.values[SettingKeyOpenAIAdvancedSchedulerSubscriptionPriorityEnabled] = values[1]
- }
- return &RateLimitService{
- settingService: NewSettingService(repo, &config.Config{}),
- }
- }
- func (s *openAISnapshotCacheStub) GetSnapshot(ctx context.Context, bucket SchedulerBucket) ([]*Account, bool, error) {
- if len(s.snapshotAccounts) == 0 {
- return nil, false, nil
- }
- out := make([]*Account, 0, len(s.snapshotAccounts))
- for _, account := range s.snapshotAccounts {
- if account == nil {
- continue
- }
- cloned := *account
- out = append(out, &cloned)
- }
- return out, true, nil
- }
- func (s *openAISnapshotCacheStub) GetAccount(ctx context.Context, accountID int64) (*Account, error) {
- if s.accountsByID == nil {
- return nil, nil
- }
- account := s.accountsByID[accountID]
- if account == nil {
- return nil, nil
- }
- cloned := *account
- return &cloned, nil
- }
- func TestOpenAIGatewayService_OpenAIAdvancedSchedulerRuntimeSettings_DBOverridesConfig(t *testing.T) {
- resetOpenAIAdvancedSchedulerSettingCacheForTest()
- defer resetOpenAIAdvancedSchedulerSettingCacheForTest()
- cfg := &config.Config{}
- cfg.Gateway.OpenAIWS.LBTopK = 11
- cfg.Gateway.OpenAIWS.SchedulerScoreWeights = config.GatewayOpenAIWSSchedulerScoreWeights{
- Priority: 1,
- Load: 2,
- Queue: 3,
- ErrorRate: 4,
- TTFT: 5,
- Reset: 6,
- QuotaHeadroom: 7,
- UpstreamCost: 8,
- PreviousResponse: 9,
- SessionSticky: 10,
- }
- repo := &openAIAdvancedSchedulerSettingRepoStub{
- values: map[string]string{
- openAIAdvancedSchedulerSettingKey: "true",
- SettingKeyOpenAIAdvancedSchedulerLBTopK: "3",
- SettingKeyOpenAIAdvancedSchedulerWeightPriority: "2.5",
- SettingKeyOpenAIAdvancedSchedulerWeightReset: "0.25",
- SettingKeyOpenAIAdvancedSchedulerWeightPreviousResponse: "12",
- },
- }
- svc := &OpenAIGatewayService{
- cfg: cfg,
- rateLimitService: &RateLimitService{settingService: NewSettingService(repo, cfg)},
- }
- ctx := context.Background()
- require.Equal(t, 3, svc.openAIWSLBTopKForRequest(ctx))
- weights := svc.openAIWSSchedulerWeightsForRequest(ctx)
- require.Equal(t, 2.5, weights.Priority)
- require.Equal(t, 2.0, weights.Load)
- require.Equal(t, 8.0, weights.UpstreamCost)
- require.Equal(t, 0.25, weights.Reset)
- require.Equal(t, 12.0, weights.Previous)
- require.Equal(t, 10.0, weights.SessionSticky)
- }
- func TestOpenAIGatewayService_OpenAIAdvancedSchedulerRuntimeSettings_InvalidWeightSumsFallBackToConfig(t *testing.T) {
- base := config.GatewayOpenAIWSSchedulerScoreWeights{
- Priority: 1, Load: 2, Queue: 3, ErrorRate: 4, TTFT: 5, Reset: 6,
- QuotaHeadroom: 7, UpstreamCost: 8, PreviousResponse: 9, SessionSticky: 10,
- }
- maxFloat := strconv.FormatFloat(math.MaxFloat64, 'g', -1, 64)
- tests := []struct {
- name string
- values map[string]string
- }{
- {
- name: "invalid single value",
- values: map[string]string{
- SettingKeyOpenAIAdvancedSchedulerWeightPriority: "NaN",
- },
- },
- {
- name: "base sum overflow",
- values: map[string]string{
- SettingKeyOpenAIAdvancedSchedulerWeightPriority: maxFloat,
- SettingKeyOpenAIAdvancedSchedulerWeightLoad: maxFloat,
- },
- },
- {
- name: "sticky total sum overflow",
- values: map[string]string{
- SettingKeyOpenAIAdvancedSchedulerWeightPriority: maxFloat,
- SettingKeyOpenAIAdvancedSchedulerWeightPreviousResponse: maxFloat,
- },
- },
- }
- for _, tt := range tests {
- t.Run(tt.name, func(t *testing.T) {
- resetOpenAIAdvancedSchedulerSettingCacheForTest()
- defer resetOpenAIAdvancedSchedulerSettingCacheForTest()
- tt.values[openAIAdvancedSchedulerSettingKey] = "true"
- cfg := &config.Config{}
- cfg.Gateway.OpenAIWS.SchedulerScoreWeights = base
- repo := &openAIAdvancedSchedulerSettingRepoStub{values: tt.values}
- svc := &OpenAIGatewayService{cfg: cfg, rateLimitService: &RateLimitService{settingService: NewSettingService(repo, cfg)}}
- require.Equal(t, (&OpenAIGatewayService{cfg: cfg}).openAIWSSchedulerWeights(), svc.openAIWSSchedulerWeightsForRequest(context.Background()))
- })
- }
- }
- func TestOpenAIGatewayService_SelectAccountWithScheduler_DefaultDisabledUsesLegacyLoadAwareness(t *testing.T) {
- resetOpenAIAdvancedSchedulerSettingCacheForTest()
- ctx := context.Background()
- groupID := int64(10106)
- accounts := []Account{
- {
- ID: 36001,
- Platform: PlatformOpenAI,
- Type: AccountTypeAPIKey,
- Status: StatusActive,
- Schedulable: true,
- Concurrency: 1,
- Priority: 5,
- },
- {
- ID: 36002,
- Platform: PlatformOpenAI,
- Type: AccountTypeAPIKey,
- Status: StatusActive,
- Schedulable: true,
- Concurrency: 1,
- Priority: 0,
- },
- }
- cfg := &config.Config{}
- cfg.Gateway.Scheduling.LoadBatchEnabled = false
- cache := &schedulerTestGatewayCache{}
- svc := &OpenAIGatewayService{
- accountRepo: schedulerTestOpenAIAccountRepo{accounts: accounts},
- cache: cache,
- cfg: cfg,
- concurrencyService: NewConcurrencyService(schedulerTestConcurrencyCache{}),
- }
- store := svc.getOpenAIWSStateStore()
- require.NoError(t, store.BindResponseAccount(ctx, groupID, "resp_disabled_001", 36001, time.Hour))
- require.False(t, svc.isOpenAIAdvancedSchedulerEnabled(ctx))
- selection, decision, err := svc.SelectAccountWithScheduler(
- ctx,
- &groupID,
- "resp_disabled_001",
- "",
- "gpt-5.1",
- nil,
- OpenAIUpstreamTransportAny,
- false,
- )
- require.NoError(t, err)
- require.NotNil(t, selection)
- require.NotNil(t, selection.Account)
- require.Equal(t, int64(36002), selection.Account.ID)
- require.Equal(t, openAIAccountScheduleLayerLoadBalance, decision.Layer)
- require.False(t, decision.StickyPreviousHit)
- }
- // Regression: the legacy load-batch path had two bare ErrNoAvailableAccounts
- // exits that bypassed the diagnostics added for both the advanced scheduler and
- // the non-batched legacy selector. This is the default path when load batching
- // is enabled, so quota auto-pause could still surface as an opaque 503.
- func TestOpenAIGatewayService_SelectAccountWithScheduler_DefaultDisabled_LoadBatchReportsFilterReasons(t *testing.T) {
- resetOpenAIAdvancedSchedulerSettingCacheForTest()
- ctx := withOpenAIQuotaAutoPauseSettings(context.Background(), OpsOpenAIAccountQuotaAutoPauseSettings{DefaultThreshold7d: 0.9})
- groupID := int64(10107)
- quotaPaused := Account{
- ID: 36003,
- Platform: PlatformOpenAI,
- Type: AccountTypeOAuth,
- Status: StatusActive,
- Schedulable: true,
- Concurrency: 1,
- Extra: map[string]any{
- "codex_7d_used_percent": 95.0,
- "codex_7d_reset_at": time.Now().Add(24 * time.Hour).Format(time.RFC3339),
- "codex_usage_updated_at": time.Now().Add(-time.Minute).Format(time.RFC3339),
- },
- }
- mappingMiss := Account{
- ID: 36004,
- Platform: PlatformOpenAI,
- Type: AccountTypeAPIKey,
- Status: StatusActive,
- Schedulable: true,
- Concurrency: 1,
- Credentials: map[string]any{
- "model_mapping": map[string]any{"gpt-4o": "gpt-4o"},
- },
- }
- excluded := Account{
- ID: 36005,
- Platform: PlatformOpenAI,
- Type: AccountTypeAPIKey,
- Status: StatusActive,
- Schedulable: true,
- Concurrency: 1,
- }
- cfg := &config.Config{}
- cfg.Gateway.Scheduling.LoadBatchEnabled = true
- svc := &OpenAIGatewayService{
- accountRepo: schedulerTestOpenAIAccountRepo{accounts: []Account{quotaPaused, mappingMiss, excluded}},
- cache: &schedulerTestGatewayCache{},
- cfg: cfg,
- concurrencyService: NewConcurrencyService(schedulerTestConcurrencyCache{}),
- }
- require.False(t, svc.isOpenAIAdvancedSchedulerEnabled(ctx))
- selection, decision, err := svc.SelectAccountWithScheduler(
- ctx,
- &groupID,
- "",
- "",
- "gpt-5.4-mini",
- map[int64]struct{}{excluded.ID: {}},
- OpenAIUpstreamTransportAny,
- false,
- )
- require.Error(t, err)
- require.ErrorIs(t, err, ErrNoAvailableAccounts)
- require.Nil(t, selection)
- require.Equal(t, openAIAccountScheduleLayerLoadBalance, decision.Layer)
- 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)")
- }
- func TestOpenAIGatewayService_SelectAccountWithScheduler_DefaultDisabled_RequiredWSV2_SkipsHTTPOnlyAccount(t *testing.T) {
- resetOpenAIAdvancedSchedulerSettingCacheForTest()
- ctx := context.Background()
- groupID := int64(10108)
- accounts := []Account{
- {
- ID: 36011,
- Platform: PlatformOpenAI,
- Type: AccountTypeAPIKey,
- Status: StatusActive,
- Schedulable: true,
- Concurrency: 1,
- Priority: 0,
- },
- {
- ID: 36012,
- Platform: PlatformOpenAI,
- Type: AccountTypeAPIKey,
- Status: StatusActive,
- Schedulable: true,
- Concurrency: 1,
- Priority: 5,
- Extra: map[string]any{
- "openai_apikey_responses_websockets_v2_enabled": true,
- },
- },
- }
- cfg := newSchedulerTestOpenAIWSV2Config()
- cfg.Gateway.Scheduling.LoadBatchEnabled = false
- svc := &OpenAIGatewayService{
- accountRepo: schedulerTestOpenAIAccountRepo{accounts: accounts},
- cache: &schedulerTestGatewayCache{},
- cfg: cfg,
- concurrencyService: NewConcurrencyService(schedulerTestConcurrencyCache{}),
- }
- selection, decision, err := svc.SelectAccountWithScheduler(
- ctx,
- &groupID,
- "",
- "",
- "gpt-5.1",
- nil,
- OpenAIUpstreamTransportResponsesWebsocketV2,
- false,
- )
- require.NoError(t, err)
- require.NotNil(t, selection)
- require.NotNil(t, selection.Account)
- require.Equal(t, int64(36012), selection.Account.ID)
- require.Equal(t, openAIAccountScheduleLayerLoadBalance, decision.Layer)
- }
- func TestOpenAIGatewayService_SelectAccountWithScheduler_DefaultDisabled_RequiredWSV2_NoAvailableAccount(t *testing.T) {
- resetOpenAIAdvancedSchedulerSettingCacheForTest()
- ctx := context.Background()
- groupID := int64(10109)
- accounts := []Account{
- {
- ID: 36021,
- Platform: PlatformOpenAI,
- Type: AccountTypeAPIKey,
- Status: StatusActive,
- Schedulable: true,
- Concurrency: 1,
- Priority: 0,
- },
- }
- cfg := newSchedulerTestOpenAIWSV2Config()
- cfg.Gateway.Scheduling.LoadBatchEnabled = false
- svc := &OpenAIGatewayService{
- accountRepo: schedulerTestOpenAIAccountRepo{accounts: accounts},
- cache: &schedulerTestGatewayCache{},
- cfg: cfg,
- concurrencyService: NewConcurrencyService(schedulerTestConcurrencyCache{}),
- }
- selection, decision, err := svc.SelectAccountWithScheduler(
- ctx,
- &groupID,
- "",
- "",
- "gpt-5.1",
- nil,
- OpenAIUpstreamTransportResponsesWebsocketV2,
- false,
- )
- require.ErrorContains(t, err, "no available OpenAI accounts")
- require.Nil(t, selection)
- require.Equal(t, openAIAccountScheduleLayerLoadBalance, decision.Layer)
- }
- func TestOpenAIGatewayService_SelectAccountWithScheduler_DefaultDisabled_EmbeddingsSkipsChatOnlyAccount(t *testing.T) {
- resetOpenAIAdvancedSchedulerSettingCacheForTest()
- ctx := context.Background()
- groupID := int64(10110)
- accounts := []Account{
- {
- ID: 36031,
- Platform: PlatformOpenAI,
- Type: AccountTypeAPIKey,
- Status: StatusActive,
- Schedulable: true,
- Concurrency: 1,
- Priority: 0,
- Credentials: map[string]any{
- "openai_capabilities": []any{"chat_completions"},
- },
- },
- {
- ID: 36032,
- Platform: PlatformOpenAI,
- Type: AccountTypeAPIKey,
- Status: StatusActive,
- Schedulable: true,
- Concurrency: 1,
- Priority: 5,
- Credentials: map[string]any{
- "openai_capabilities": []any{"chat_completions", "embeddings"},
- },
- },
- }
- cfg := &config.Config{}
- cfg.Gateway.Scheduling.LoadBatchEnabled = false
- svc := &OpenAIGatewayService{
- accountRepo: schedulerTestOpenAIAccountRepo{accounts: accounts},
- cache: &schedulerTestGatewayCache{},
- cfg: cfg,
- concurrencyService: NewConcurrencyService(schedulerTestConcurrencyCache{}),
- }
- selection, decision, err := svc.SelectAccountWithSchedulerForCapability(
- ctx,
- &groupID,
- "",
- "",
- "text-embedding-3-small",
- nil,
- OpenAIUpstreamTransportHTTPSSE,
- OpenAIEndpointCapabilityEmbeddings,
- false,
- false,
- true,
- )
- require.NoError(t, err)
- require.NotNil(t, selection)
- require.NotNil(t, selection.Account)
- require.Equal(t, int64(36032), selection.Account.ID)
- require.Equal(t, openAIAccountScheduleLayerLoadBalance, decision.Layer)
- }
- func TestOpenAIGatewayService_SelectAccountForTokenCount_DoesNotAcquireGenerationSlot(t *testing.T) {
- ctx := context.Background()
- groupID := int64(10115)
- acquiredIDs := make([]int64, 0)
- accounts := []Account{
- {
- ID: 36501, Platform: PlatformOpenAI, Type: AccountTypeAPIKey,
- Status: StatusActive, Schedulable: true, Concurrency: 1, Priority: 0,
- Credentials: map[string]any{"openai_capabilities": []any{"chat_completions"}},
- },
- {
- ID: 36502, Platform: PlatformOpenAI, Type: AccountTypeAPIKey,
- Status: StatusActive, Schedulable: true, Concurrency: 1, Priority: 5,
- Credentials: map[string]any{"openai_capabilities": []any{"embeddings"}},
- },
- {
- ID: 36503, Platform: PlatformGrok, Type: AccountTypeAPIKey,
- Status: StatusActive, Schedulable: true, Concurrency: 1, Priority: 10,
- Credentials: map[string]any{"openai_capabilities": []any{"chat_completions"}},
- },
- {
- ID: 36504, Platform: PlatformOpenAI, Type: AccountTypeAPIKey,
- Status: StatusActive, Schedulable: true, Concurrency: 1, Priority: 15,
- Credentials: map[string]any{
- "openai_capabilities": []any{"chat_completions"},
- "model_mapping": map[string]any{"gpt-4o": "gpt-4o"},
- },
- },
- }
- svc := &OpenAIGatewayService{
- accountRepo: schedulerTestOpenAIAccountRepo{accounts: accounts},
- cache: &schedulerTestGatewayCache{},
- cfg: &config.Config{},
- concurrencyService: NewConcurrencyService(schedulerTestConcurrencyCache{
- acquireResults: map[int64]bool{36501: false},
- acquiredIDs: &acquiredIDs,
- }),
- }
- account, err := svc.SelectAccountForTokenCount(
- ctx,
- &groupID,
- "",
- "gpt-5.1",
- OpenAIEndpointCapabilityChatCompletions,
- PlatformOpenAI,
- )
- require.NoError(t, err)
- require.NotNil(t, account)
- require.Equal(t, int64(36501), account.ID)
- require.Empty(t, acquiredIDs, "token counting must not acquire a generation slot")
- }
- // 生图意图的 /v1/responses 请求要求 OpenAIEndpointCapabilityResponses:探测确认
- // 不支持 Responses API 的 APIKey 账号必须被排除,避免 forward 阶段降级为无法生图
- // 的 Chat Completions 直转(#4417)。
- func TestOpenAIGatewayService_SelectAccountWithScheduler_ResponsesCapabilityExcludesUnsupportedAPIKey(t *testing.T) {
- resetOpenAIAdvancedSchedulerSettingCacheForTest()
- ctx := context.Background()
- groupID := int64(10120)
- newSvc := func(accounts []Account) *OpenAIGatewayService {
- cfg := &config.Config{}
- cfg.Gateway.Scheduling.LoadBatchEnabled = false
- return &OpenAIGatewayService{
- accountRepo: schedulerTestOpenAIAccountRepo{accounts: accounts},
- cache: &schedulerTestGatewayCache{},
- cfg: cfg,
- concurrencyService: NewConcurrencyService(schedulerTestConcurrencyCache{}),
- }
- }
- supported := Account{
- ID: 37001, Platform: PlatformOpenAI, Type: AccountTypeAPIKey,
- Status: StatusActive, Schedulable: true, Concurrency: 1, Priority: 0,
- }
- // 更高优先级但探测确认不支持 Responses——若门控失效会被优先选中。
- unsupported := Account{
- ID: 37002, Platform: PlatformOpenAI, Type: AccountTypeAPIKey,
- Status: StatusActive, Schedulable: true, Concurrency: 1, Priority: 5,
- Extra: map[string]any{"openai_responses_supported": false},
- }
- t.Run("生图意图仅选中支持 responses 的账号", func(t *testing.T) {
- svc := newSvc([]Account{supported, unsupported})
- selection, _, err := svc.SelectAccountWithSchedulerForCapability(
- ctx, &groupID, "", "", "gpt-image-2", nil,
- OpenAIUpstreamTransportAny, OpenAIEndpointCapabilityResponses,
- false, false, false,
- )
- require.NoError(t, err)
- require.NotNil(t, selection)
- require.NotNil(t, selection.Account)
- require.Equal(t, int64(37001), selection.Account.ID)
- })
- t.Run("仅有不支持 responses 的账号时生图意图无可用账号", func(t *testing.T) {
- svc := newSvc([]Account{unsupported})
- selection, _, err := svc.SelectAccountWithSchedulerForCapability(
- ctx, &groupID, "", "", "gpt-image-2", nil,
- OpenAIUpstreamTransportAny, OpenAIEndpointCapabilityResponses,
- false, false, false,
- )
- require.Error(t, err)
- require.Nil(t, selection)
- })
- t.Run("非生图路径仍可选中不支持 responses 的账号", func(t *testing.T) {
- svc := newSvc([]Account{unsupported})
- selection, _, err := svc.SelectAccountWithSchedulerForCapability(
- ctx, &groupID, "", "", "gpt-5.1", nil,
- OpenAIUpstreamTransportAny, OpenAIEndpointCapabilityChatCompletions,
- false, false, true,
- )
- require.NoError(t, err)
- require.NotNil(t, selection)
- require.NotNil(t, selection.Account)
- require.Equal(t, int64(37002), selection.Account.ID)
- })
- }
- // alpha/search 调度必须同时放行 OAuth 与 APIKey 账号:v0.1.157 曾因 OAuth-only
- // 门控把 APIKey 账号从候选池剔除,纯 APIKey 分组的独立搜索请求在选号阶段就
- // 报无可用账号,Codex 网页搜索整体失效(转发层其实一直支持 APIKey 路径)。
- func TestOpenAIGatewayService_SelectAccountWithScheduler_AlphaSearchAllowsAPIKeyAccount(t *testing.T) {
- resetOpenAIAdvancedSchedulerSettingCacheForTest()
- ctx := context.Background()
- groupID := int64(10125)
- accounts := []Account{
- {
- ID: 38001,
- Platform: PlatformOpenAI,
- Type: AccountTypeAPIKey,
- Status: StatusActive,
- Schedulable: true,
- Concurrency: 1,
- Priority: 0,
- },
- }
- cfg := &config.Config{}
- cfg.Gateway.Scheduling.LoadBatchEnabled = false
- svc := &OpenAIGatewayService{
- accountRepo: schedulerTestOpenAIAccountRepo{accounts: accounts},
- cache: &schedulerTestGatewayCache{},
- cfg: cfg,
- concurrencyService: NewConcurrencyService(schedulerTestConcurrencyCache{}),
- }
- selection, _, err := svc.SelectAccountWithSchedulerForCapability(
- ctx,
- &groupID,
- "",
- "",
- "gpt-5.6-sol",
- nil,
- OpenAIUpstreamTransportHTTPSSE,
- OpenAIEndpointCapabilityAlphaSearch,
- false,
- false,
- false,
- )
- require.NoError(t, err)
- require.NotNil(t, selection)
- require.NotNil(t, selection.Account)
- require.Equal(t, int64(38001), selection.Account.ID)
- }
- func TestOpenAIGatewayService_SelectAccountWithScheduler_DefaultDisabled_AllowsGrokChatAccount(t *testing.T) {
- resetOpenAIAdvancedSchedulerSettingCacheForTest()
- ctx := context.Background()
- groupID := int64(10113)
- accounts := []Account{
- {
- ID: 36041,
- Platform: PlatformGrok,
- Type: AccountTypeOAuth,
- Status: StatusActive,
- Schedulable: true,
- Concurrency: 1,
- Priority: 0,
- },
- }
- cfg := &config.Config{}
- cfg.Gateway.Scheduling.LoadBatchEnabled = false
- svc := &OpenAIGatewayService{
- accountRepo: schedulerTestOpenAIAccountRepo{accounts: accounts},
- cache: &schedulerTestGatewayCache{},
- cfg: cfg,
- concurrencyService: NewConcurrencyService(schedulerTestConcurrencyCache{}),
- }
- selection, decision, err := svc.SelectAccountWithSchedulerForCapability(
- ctx,
- &groupID,
- "",
- "",
- "grok-4.3",
- nil,
- OpenAIUpstreamTransportAny,
- OpenAIEndpointCapabilityChatCompletions,
- false,
- false,
- false,
- PlatformGrok,
- )
- require.NoError(t, err)
- require.NotNil(t, selection)
- require.NotNil(t, selection.Account)
- require.Equal(t, int64(36041), selection.Account.ID)
- require.Equal(t, openAIAccountScheduleLayerLoadBalance, decision.Layer)
- }
- func TestOpenAIGatewayService_SelectAccountWithScheduler_GrokMediaCapabilityFiltersIneligibleAccounts(t *testing.T) {
- resetOpenAIAdvancedSchedulerSettingCacheForTest()
- ctx := context.Background()
- groupID := int64(10114)
- ineligible := Account{
- ID: 36051, Platform: PlatformGrok, Type: AccountTypeOAuth,
- Status: StatusActive, Schedulable: true, Concurrency: 1, Priority: 5,
- Extra: map[string]any{GrokMediaEligibleExtraKey: false},
- }
- eligible := Account{
- ID: 36052, Platform: PlatformGrok, Type: AccountTypeOAuth,
- Status: StatusActive, Schedulable: true, Concurrency: 1, Priority: 0,
- Extra: map[string]any{GrokMediaEligibleExtraKey: true},
- }
- newService := func(accounts []Account) *OpenAIGatewayService {
- cfg := &config.Config{}
- cfg.Gateway.Scheduling.LoadBatchEnabled = false
- return &OpenAIGatewayService{
- accountRepo: schedulerTestOpenAIAccountRepo{accounts: accounts},
- cache: &schedulerTestGatewayCache{},
- cfg: cfg,
- concurrencyService: NewConcurrencyService(schedulerTestConcurrencyCache{}),
- }
- }
- t.Run("media generation skips higher priority ineligible account", func(t *testing.T) {
- selection, _, err := newService([]Account{ineligible, eligible}).SelectAccountWithSchedulerForCapability(
- ctx, &groupID, "", "", "grok-imagine-video", nil,
- OpenAIUpstreamTransportHTTPSSE, OpenAIEndpointCapabilityGrokMediaGeneration,
- false, false, false, PlatformGrok,
- )
- require.NoError(t, err)
- require.NotNil(t, selection)
- require.NotNil(t, selection.Account)
- require.Equal(t, eligible.ID, selection.Account.ID)
- })
- t.Run("media generation fails closed when all accounts are ineligible", func(t *testing.T) {
- selection, _, err := newService([]Account{ineligible}).SelectAccountWithSchedulerForCapability(
- ctx, &groupID, "", "", "grok-imagine-video", nil,
- OpenAIUpstreamTransportHTTPSSE, OpenAIEndpointCapabilityGrokMediaGeneration,
- false, false, false, PlatformGrok,
- )
- require.Error(t, err)
- require.ErrorIs(t, err, ErrNoAvailableAccounts)
- require.Nil(t, selection)
- })
- t.Run("chat remains routable on media-ineligible account", func(t *testing.T) {
- selection, _, err := newService([]Account{ineligible}).SelectAccountWithSchedulerForCapability(
- ctx, &groupID, "", "", "grok-4.3", nil,
- OpenAIUpstreamTransportHTTPSSE, OpenAIEndpointCapabilityChatCompletions,
- false, false, false, PlatformGrok,
- )
- require.NoError(t, err)
- require.NotNil(t, selection)
- require.NotNil(t, selection.Account)
- require.Equal(t, ineligible.ID, selection.Account.ID)
- })
- }
- // Regression #4599: when the advanced scheduler's load-balance initial filter
- // drops every candidate, the resulting "no available OpenAI accounts" error must
- // carry per-reason exclusion counts. Previously all filter branches were silent
- // (debug logs at best), so a 503 with excluded_account_count=0 could not be
- // diagnosed from the error alone.
- func TestOpenAIGatewayService_SelectAccountWithScheduler_NoAvailableErrorReportsQuotaAutoPauseExclusion(t *testing.T) {
- resetOpenAIAdvancedSchedulerSettingCacheForTest()
- ctx := withOpenAIQuotaAutoPauseSettings(context.Background(), OpsOpenAIAccountQuotaAutoPauseSettings{DefaultThreshold7d: 0.9})
- groupID := int64(101201)
- accounts := []Account{
- {
- ID: 38101,
- Platform: PlatformOpenAI,
- Type: AccountTypeOAuth,
- Status: StatusActive,
- Schedulable: true,
- Concurrency: 1,
- Extra: map[string]any{
- "codex_7d_used_percent": 95.0,
- "codex_7d_reset_at": time.Now().Add(24 * time.Hour).Format(time.RFC3339),
- "codex_usage_updated_at": time.Now().Add(-time.Minute).Format(time.RFC3339),
- },
- },
- }
- cfg := &config.Config{}
- svc := &OpenAIGatewayService{
- accountRepo: schedulerTestOpenAIAccountRepo{accounts: accounts},
- cache: &schedulerTestGatewayCache{},
- cfg: cfg,
- rateLimitService: newOpenAIAdvancedSchedulerRateLimitService("true"),
- concurrencyService: NewConcurrencyService(schedulerTestConcurrencyCache{}),
- }
- selection, decision, err := svc.SelectAccountWithScheduler(
- ctx, &groupID, "", "", "gpt-5.4-mini", nil, OpenAIUpstreamTransportAny, false,
- )
- require.Error(t, err)
- require.ErrorIs(t, err, ErrNoAvailableAccounts)
- require.Nil(t, selection)
- require.Equal(t, openAIAccountScheduleLayerLoadBalance, decision.Layer)
- require.EqualError(t, err, "no available OpenAI accounts supporting model: gpt-5.4-mini (pool=1, filtered: quota_auto_pause_7d=1)")
- }
- func TestOpenAIGatewayService_SelectAccountWithScheduler_NoAvailableErrorReportsModelNotSupported(t *testing.T) {
- resetOpenAIAdvancedSchedulerSettingCacheForTest()
- ctx := context.Background()
- groupID := int64(101202)
- // OAuth account with empty model_mapping: foreign-family models (grok-*)
- // are not servable by Codex upstream (#3662) and are dropped by the filter.
- accounts := []Account{
- {
- ID: 38111,
- Platform: PlatformOpenAI,
- Type: AccountTypeOAuth,
- Status: StatusActive,
- Schedulable: true,
- Concurrency: 1,
- },
- }
- svc := &OpenAIGatewayService{
- accountRepo: schedulerTestOpenAIAccountRepo{accounts: accounts},
- cache: &schedulerTestGatewayCache{},
- cfg: &config.Config{},
- rateLimitService: newOpenAIAdvancedSchedulerRateLimitService("true"),
- concurrencyService: NewConcurrencyService(schedulerTestConcurrencyCache{}),
- }
- selection, _, err := svc.SelectAccountWithScheduler(
- ctx, &groupID, "", "", "grok-4.5", nil, OpenAIUpstreamTransportAny, false,
- )
- require.Error(t, err)
- require.ErrorIs(t, err, ErrNoAvailableAccounts)
- require.Nil(t, selection)
- require.EqualError(t, err, "no available OpenAI accounts supporting model: grok-4.5 (pool=1, filtered: model_not_supported=1)")
- }
- func TestOpenAIGatewayService_SelectAccountWithScheduler_NoAvailableErrorAggregatesReasonsDeterministically(t *testing.T) {
- resetOpenAIAdvancedSchedulerSettingCacheForTest()
- ctx := withOpenAIQuotaAutoPauseSettings(context.Background(), OpsOpenAIAccountQuotaAutoPauseSettings{DefaultThreshold7d: 0.9})
- groupID := int64(101203)
- quotaPaused := Account{
- ID: 38121,
- Platform: PlatformOpenAI,
- Type: AccountTypeOAuth,
- Status: StatusActive,
- Schedulable: true,
- Concurrency: 1,
- Extra: map[string]any{
- "codex_7d_used_percent": 95.0,
- "codex_7d_reset_at": time.Now().Add(24 * time.Hour).Format(time.RFC3339),
- "codex_usage_updated_at": time.Now().Add(-time.Minute).Format(time.RFC3339),
- },
- }
- mappingMiss := Account{
- ID: 38122,
- Platform: PlatformOpenAI,
- Type: AccountTypeAPIKey,
- Status: StatusActive,
- Schedulable: true,
- Concurrency: 1,
- Credentials: map[string]any{
- "model_mapping": map[string]any{"gpt-4o": "gpt-4o"},
- },
- }
- excluded := Account{
- ID: 38123,
- Platform: PlatformOpenAI,
- Type: AccountTypeAPIKey,
- Status: StatusActive,
- Schedulable: true,
- Concurrency: 1,
- }
- svc := &OpenAIGatewayService{
- accountRepo: schedulerTestOpenAIAccountRepo{accounts: []Account{quotaPaused, mappingMiss, excluded}},
- cache: &schedulerTestGatewayCache{},
- cfg: &config.Config{},
- rateLimitService: newOpenAIAdvancedSchedulerRateLimitService("true"),
- concurrencyService: NewConcurrencyService(schedulerTestConcurrencyCache{}),
- }
- selection, _, err := svc.SelectAccountWithScheduler(
- ctx, &groupID, "", "", "gpt-5.4-mini", map[int64]struct{}{38123: {}}, OpenAIUpstreamTransportAny, false,
- )
- require.Error(t, err)
- require.ErrorIs(t, err, ErrNoAvailableAccounts)
- require.Nil(t, selection)
- // Reasons are sorted lexicographically, so the message is deterministic.
- 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)")
- }
- func TestOpenAIGatewayService_SelectAccountWithScheduler_NoAvailableErrorReportsEmptyPool(t *testing.T) {
- resetOpenAIAdvancedSchedulerSettingCacheForTest()
- ctx := context.Background()
- groupID := int64(101204)
- svc := &OpenAIGatewayService{
- accountRepo: schedulerTestOpenAIAccountRepo{},
- cache: &schedulerTestGatewayCache{},
- cfg: &config.Config{},
- rateLimitService: newOpenAIAdvancedSchedulerRateLimitService("true"),
- concurrencyService: NewConcurrencyService(schedulerTestConcurrencyCache{}),
- }
- selection, _, err := svc.SelectAccountWithScheduler(
- ctx, &groupID, "", "", "gpt-5.1", nil, OpenAIUpstreamTransportAny, false,
- )
- require.Error(t, err)
- require.ErrorIs(t, err, ErrNoAvailableAccounts)
- require.Nil(t, selection)
- require.EqualError(t, err, "no available OpenAI accounts supporting model: gpt-5.1 (pool=0)")
- }
- func TestOpenAIGatewayService_SelectAccountWithScheduler_EnabledUsesAdvancedPreviousResponseRouting(t *testing.T) {
- resetOpenAIAdvancedSchedulerSettingCacheForTest()
- ctx := context.Background()
- groupID := int64(10107)
- accounts := []Account{
- {
- ID: 37001,
- Platform: PlatformOpenAI,
- Type: AccountTypeAPIKey,
- Status: StatusActive,
- Schedulable: true,
- Concurrency: 1,
- Priority: 5,
- Extra: map[string]any{
- "openai_apikey_responses_websockets_v2_enabled": true,
- },
- },
- {
- ID: 37002,
- Platform: PlatformOpenAI,
- Type: AccountTypeAPIKey,
- Status: StatusActive,
- Schedulable: true,
- Concurrency: 1,
- Priority: 0,
- },
- }
- cfg := &config.Config{}
- cfg.Gateway.Scheduling.LoadBatchEnabled = false
- cfg.Gateway.OpenAIWS.Enabled = true
- cfg.Gateway.OpenAIWS.OAuthEnabled = true
- cfg.Gateway.OpenAIWS.APIKeyEnabled = true
- cfg.Gateway.OpenAIWS.ResponsesWebsocketsV2 = true
- cfg.Gateway.OpenAIWS.StickyResponseIDTTLSeconds = 3600
- svc := &OpenAIGatewayService{
- accountRepo: schedulerTestOpenAIAccountRepo{accounts: accounts},
- cache: &schedulerTestGatewayCache{},
- cfg: cfg,
- rateLimitService: newOpenAIAdvancedSchedulerRateLimitService("true"),
- concurrencyService: NewConcurrencyService(schedulerTestConcurrencyCache{}),
- }
- store := svc.getOpenAIWSStateStore()
- require.NoError(t, store.BindResponseAccount(ctx, groupID, "resp_enabled_001", 37001, time.Hour))
- require.True(t, svc.isOpenAIAdvancedSchedulerEnabled(ctx))
- selection, decision, err := svc.SelectAccountWithScheduler(
- ctx,
- &groupID,
- "resp_enabled_001",
- "",
- "gpt-5.1",
- nil,
- OpenAIUpstreamTransportAny,
- false,
- )
- require.NoError(t, err)
- require.NotNil(t, selection)
- require.NotNil(t, selection.Account)
- require.Equal(t, int64(37001), selection.Account.ID)
- require.Equal(t, openAIAccountScheduleLayerPreviousResponse, decision.Layer)
- require.True(t, decision.StickyPreviousHit)
- }
- func TestOpenAIGatewayService_SelectAccountWithScheduler_StickyWeightedSessionInTopKUsesStickyFirst(t *testing.T) {
- resetOpenAIAdvancedSchedulerSettingCacheForTest()
- ctx := context.Background()
- groupID := int64(101071)
- accounts := []Account{
- {
- ID: 37101,
- Platform: PlatformOpenAI,
- Type: AccountTypeAPIKey,
- Status: StatusActive,
- Schedulable: true,
- Concurrency: 1,
- Priority: 100,
- GroupIDs: []int64{groupID},
- },
- {
- ID: 37102,
- Platform: PlatformOpenAI,
- Type: AccountTypeAPIKey,
- Status: StatusActive,
- Schedulable: true,
- Concurrency: 1,
- Priority: 0,
- GroupIDs: []int64{groupID},
- },
- }
- cfg := &config.Config{}
- cfg.Gateway.OpenAIWS.LBTopK = 2
- cfg.Gateway.OpenAIWS.SchedulerScoreWeights.Priority = 1
- cfg.Gateway.OpenAIWS.SchedulerScoreWeights.Load = 1
- cfg.Gateway.OpenAIWS.SchedulerScoreWeights.Queue = 0.7
- cfg.Gateway.OpenAIWS.SchedulerScoreWeights.ErrorRate = 0.8
- cfg.Gateway.OpenAIWS.SchedulerScoreWeights.TTFT = 0.5
- cfg.Gateway.OpenAIWS.SchedulerScoreWeights.SessionSticky = 3
- cache := &schedulerTestGatewayCache{sessionBindings: map[string]int64{
- "openai:session_hash_weighted_topk": 37101,
- }}
- svc := &OpenAIGatewayService{
- accountRepo: schedulerTestOpenAIAccountRepo{accounts: accounts},
- cache: cache,
- cfg: cfg,
- rateLimitService: newOpenAIAdvancedSchedulerRateLimitService("true", "true"),
- concurrencyService: NewConcurrencyService(schedulerTestConcurrencyCache{}),
- }
- selection, decision, err := svc.SelectAccountWithScheduler(
- ctx,
- &groupID,
- "",
- "session_hash_weighted_topk",
- "gpt-5.1",
- nil,
- OpenAIUpstreamTransportAny,
- false,
- )
- require.NoError(t, err)
- require.NotNil(t, selection)
- require.NotNil(t, selection.Account)
- require.Equal(t, int64(37101), selection.Account.ID)
- require.Equal(t, openAIAccountScheduleLayerLoadBalance, decision.Layer)
- require.True(t, decision.StickySessionHit)
- require.Equal(t, 2, decision.TopK)
- if selection.ReleaseFunc != nil {
- selection.ReleaseFunc()
- }
- }
- func TestOpenAIGatewayService_SelectAccountWithScheduler_StickyWeightedPreviousRequiresMovableContext(t *testing.T) {
- resetOpenAIAdvancedSchedulerSettingCacheForTest()
- ctx := context.Background()
- groupID := int64(101072)
- accounts := []Account{
- {
- ID: 37111,
- Platform: PlatformOpenAI,
- Type: AccountTypeAPIKey,
- Status: StatusActive,
- Schedulable: true,
- Concurrency: 1,
- Priority: 100,
- GroupIDs: []int64{groupID},
- Extra: map[string]any{
- "openai_apikey_responses_websockets_v2_enabled": true,
- },
- },
- {
- ID: 37112,
- Platform: PlatformOpenAI,
- Type: AccountTypeAPIKey,
- Status: StatusActive,
- Schedulable: true,
- Concurrency: 1,
- Priority: 0,
- GroupIDs: []int64{groupID},
- Extra: map[string]any{
- "openai_apikey_responses_websockets_v2_enabled": true,
- },
- },
- }
- cfg := newSchedulerTestOpenAIWSV2Config()
- cfg.Gateway.OpenAIWS.LBTopK = 1
- cfg.Gateway.OpenAIWS.SchedulerScoreWeights.Priority = 1
- cfg.Gateway.OpenAIWS.SchedulerScoreWeights.Load = 1
- cfg.Gateway.OpenAIWS.SchedulerScoreWeights.Queue = 0.7
- cfg.Gateway.OpenAIWS.SchedulerScoreWeights.ErrorRate = 0.8
- cfg.Gateway.OpenAIWS.SchedulerScoreWeights.TTFT = 0.5
- svc := &OpenAIGatewayService{
- accountRepo: schedulerTestOpenAIAccountRepo{accounts: accounts},
- cache: &schedulerTestGatewayCache{},
- cfg: cfg,
- rateLimitService: newOpenAIAdvancedSchedulerRateLimitService("true", "true"),
- concurrencyService: NewConcurrencyService(schedulerTestConcurrencyCache{}),
- }
- store := svc.getOpenAIWSStateStore()
- require.NoError(t, store.BindResponseAccount(ctx, groupID, "resp_weighted_unmovable", 37111, time.Hour))
- selection, decision, err := svc.SelectAccountWithSchedulerForCapability(
- ctx,
- &groupID,
- "resp_weighted_unmovable",
- "",
- "gpt-5.1",
- nil,
- OpenAIUpstreamTransportAny,
- OpenAIEndpointCapabilityChatCompletions,
- false,
- false,
- true,
- PlatformOpenAI,
- )
- require.NoError(t, err)
- require.NotNil(t, selection)
- require.NotNil(t, selection.Account)
- require.Equal(t, int64(37111), selection.Account.ID)
- require.Equal(t, openAIAccountScheduleLayerPreviousResponse, decision.Layer)
- require.True(t, decision.StickyPreviousHit)
- if selection.ReleaseFunc != nil {
- selection.ReleaseFunc()
- }
- selection, decision, err = svc.SelectAccountWithSchedulerForCapability(
- ctx,
- &groupID,
- "resp_weighted_unmovable",
- "",
- "gpt-5.1",
- nil,
- OpenAIUpstreamTransportAny,
- OpenAIEndpointCapabilityChatCompletions,
- false,
- true,
- true,
- PlatformOpenAI,
- )
- require.NoError(t, err)
- require.NotNil(t, selection)
- require.NotNil(t, selection.Account)
- require.Equal(t, int64(37112), selection.Account.ID)
- require.Equal(t, openAIAccountScheduleLayerLoadBalance, decision.Layer)
- require.False(t, decision.StickyPreviousHit)
- if selection.ReleaseFunc != nil {
- selection.ReleaseFunc()
- }
- }
- func TestOpenAIGatewayService_SelectAccountWithScheduler_PreviousResponseCompactUnsupportedDeletesBinding(t *testing.T) {
- resetOpenAIAdvancedSchedulerSettingCacheForTest()
- ctx := context.Background()
- groupID := int64(101073)
- accounts := []Account{
- {
- ID: 37121,
- Platform: PlatformOpenAI,
- Type: AccountTypeAPIKey,
- Status: StatusActive,
- Schedulable: true,
- Concurrency: 1,
- Priority: 0,
- GroupIDs: []int64{groupID},
- Extra: map[string]any{
- "openai_apikey_responses_websockets_v2_enabled": true,
- "openai_compact_mode": OpenAICompactModeForceOff,
- },
- },
- {
- ID: 37122,
- Platform: PlatformOpenAI,
- Type: AccountTypeAPIKey,
- Status: StatusActive,
- Schedulable: true,
- Concurrency: 1,
- Priority: 10,
- GroupIDs: []int64{groupID},
- Extra: map[string]any{
- "openai_apikey_responses_websockets_v2_enabled": true,
- "openai_compact_mode": OpenAICompactModeForceOn,
- },
- },
- }
- cfg := newSchedulerTestOpenAIWSV2Config()
- cfg.Gateway.OpenAIWS.LBTopK = 2
- cfg.Gateway.OpenAIWS.SchedulerScoreWeights.Priority = 1
- cfg.Gateway.OpenAIWS.SchedulerScoreWeights.Load = 1
- cfg.Gateway.OpenAIWS.SchedulerScoreWeights.Queue = 0.7
- cfg.Gateway.OpenAIWS.SchedulerScoreWeights.ErrorRate = 0.8
- cfg.Gateway.OpenAIWS.SchedulerScoreWeights.TTFT = 0.5
- svc := &OpenAIGatewayService{
- accountRepo: schedulerTestOpenAIAccountRepo{accounts: accounts},
- cache: &schedulerTestGatewayCache{},
- cfg: cfg,
- rateLimitService: newOpenAIAdvancedSchedulerRateLimitService("true"),
- concurrencyService: NewConcurrencyService(schedulerTestConcurrencyCache{}),
- }
- store := svc.getOpenAIWSStateStore()
- require.NoError(t, store.BindResponseAccount(ctx, groupID, "resp_compact_unsupported", 37121, time.Hour))
- selection, decision, err := svc.SelectAccountWithScheduler(
- ctx,
- &groupID,
- "resp_compact_unsupported",
- "",
- "gpt-5.1",
- nil,
- OpenAIUpstreamTransportAny,
- true,
- )
- require.NoError(t, err)
- require.NotNil(t, selection)
- require.NotNil(t, selection.Account)
- require.Equal(t, int64(37122), selection.Account.ID)
- require.Equal(t, openAIAccountScheduleLayerLoadBalance, decision.Layer)
- require.False(t, decision.StickyPreviousHit)
- if selection.ReleaseFunc != nil {
- selection.ReleaseFunc()
- }
- accountID, err := store.GetResponseAccount(ctx, groupID, "resp_compact_unsupported")
- require.NoError(t, err)
- require.Zero(t, accountID)
- }
- func TestOpenAIGatewayService_SelectAccountWithScheduler_Enabled_EmbeddingsSkipsChatOnlyAccount(t *testing.T) {
- resetOpenAIAdvancedSchedulerSettingCacheForTest()
- ctx := context.Background()
- groupID := int64(10111)
- accounts := []Account{
- {
- ID: 37011,
- Platform: PlatformOpenAI,
- Type: AccountTypeAPIKey,
- Status: StatusActive,
- Schedulable: true,
- Concurrency: 1,
- Priority: 0,
- Credentials: map[string]any{
- "openai_capabilities": []any{"chat_completions"},
- },
- },
- {
- ID: 37012,
- Platform: PlatformOpenAI,
- Type: AccountTypeAPIKey,
- Status: StatusActive,
- Schedulable: true,
- Concurrency: 1,
- Priority: 5,
- Credentials: map[string]any{
- "openai_capabilities": []any{"chat_completions", "embeddings"},
- },
- },
- }
- cfg := &config.Config{}
- cfg.Gateway.Scheduling.LoadBatchEnabled = false
- svc := &OpenAIGatewayService{
- accountRepo: schedulerTestOpenAIAccountRepo{accounts: accounts},
- cache: &schedulerTestGatewayCache{},
- cfg: cfg,
- rateLimitService: newOpenAIAdvancedSchedulerRateLimitService("true"),
- concurrencyService: NewConcurrencyService(schedulerTestConcurrencyCache{}),
- }
- selection, decision, err := svc.SelectAccountWithSchedulerForCapability(
- ctx,
- &groupID,
- "",
- "",
- "text-embedding-3-small",
- nil,
- OpenAIUpstreamTransportHTTPSSE,
- OpenAIEndpointCapabilityEmbeddings,
- false,
- false,
- true,
- )
- require.NoError(t, err)
- require.NotNil(t, selection)
- require.NotNil(t, selection.Account)
- require.Equal(t, int64(37012), selection.Account.ID)
- require.Equal(t, openAIAccountScheduleLayerLoadBalance, decision.Layer)
- require.Equal(t, 1, decision.CandidateCount)
- }
- func TestOpenAIGatewayService_SelectAccountWithScheduler_Enabled_EmbeddingsSkipsChatOnlyStickyBindings(t *testing.T) {
- resetOpenAIAdvancedSchedulerSettingCacheForTest()
- ctx := context.Background()
- groupID := int64(10112)
- accounts := []Account{
- {
- ID: 37021,
- Platform: PlatformOpenAI,
- Type: AccountTypeAPIKey,
- Status: StatusActive,
- Schedulable: true,
- Concurrency: 1,
- Priority: 0,
- Credentials: map[string]any{
- "openai_capabilities": []any{"chat_completions"},
- },
- Extra: map[string]any{
- "openai_apikey_responses_websockets_v2_enabled": true,
- },
- },
- {
- ID: 37022,
- Platform: PlatformOpenAI,
- Type: AccountTypeAPIKey,
- Status: StatusActive,
- Schedulable: true,
- Concurrency: 1,
- Priority: 5,
- Credentials: map[string]any{
- "openai_capabilities": []any{"chat_completions", "embeddings"},
- },
- Extra: map[string]any{
- "openai_apikey_responses_websockets_v2_enabled": true,
- },
- },
- }
- cfg := newSchedulerTestOpenAIWSV2Config()
- cfg.Gateway.Scheduling.LoadBatchEnabled = false
- cache := &schedulerTestGatewayCache{
- sessionBindings: map[string]int64{
- "openai:session_hash_embeddings": 37021,
- },
- }
- svc := &OpenAIGatewayService{
- accountRepo: schedulerTestOpenAIAccountRepo{accounts: accounts},
- cache: cache,
- cfg: cfg,
- rateLimitService: newOpenAIAdvancedSchedulerRateLimitService("true"),
- concurrencyService: NewConcurrencyService(schedulerTestConcurrencyCache{}),
- }
- store := svc.getOpenAIWSStateStore()
- require.NoError(t, store.BindResponseAccount(ctx, groupID, "resp_embeddings_chat_only", 37021, time.Hour))
- selection, decision, err := svc.SelectAccountWithSchedulerForCapability(
- ctx,
- &groupID,
- "resp_embeddings_chat_only",
- "session_hash_embeddings",
- "text-embedding-3-small",
- nil,
- OpenAIUpstreamTransportHTTPSSE,
- OpenAIEndpointCapabilityEmbeddings,
- false,
- false,
- true,
- )
- require.NoError(t, err)
- require.NotNil(t, selection)
- require.NotNil(t, selection.Account)
- require.Equal(t, int64(37022), selection.Account.ID)
- require.Equal(t, openAIAccountScheduleLayerLoadBalance, decision.Layer)
- require.False(t, decision.StickyPreviousHit)
- require.False(t, decision.StickySessionHit)
- require.Equal(t, int64(37022), cache.sessionBindings["openai:session_hash_embeddings"])
- }
- func TestOpenAIGatewayService_OpenAIAccountSchedulerMetrics_DisabledNoOp(t *testing.T) {
- resetOpenAIAdvancedSchedulerSettingCacheForTest()
- svc := &OpenAIGatewayService{}
- ttft := 120
- svc.ReportOpenAIAccountScheduleResult(&Account{ID: 10}, "", true, &ttft)
- svc.RecordOpenAIAccountSwitch()
- snapshot := svc.SnapshotOpenAIAccountSchedulerMetrics()
- require.Equal(t, OpenAIAccountSchedulerMetricsSnapshot{}, snapshot)
- }
- func TestOpenAIGatewayService_SelectAccountWithScheduler_SkipsQuarantinedSharedProxy(t *testing.T) {
- resetOpenAIAdvancedSchedulerSettingCacheForTest()
- proxyA := int64(4698)
- proxyB := int64(4699)
- accounts := []Account{
- {ID: 469801, Platform: PlatformOpenAI, Type: AccountTypeOAuth, Status: StatusActive, Schedulable: true, Concurrency: 1, Priority: 0, ProxyID: &proxyA},
- {ID: 469802, Platform: PlatformOpenAI, Type: AccountTypeOAuth, Status: StatusActive, Schedulable: true, Concurrency: 1, Priority: 0, ProxyID: &proxyA},
- {ID: 469803, Platform: PlatformOpenAI, Type: AccountTypeOAuth, Status: StatusActive, Schedulable: true, Concurrency: 1, Priority: 5, ProxyID: &proxyB},
- }
- cfg := &config.Config{}
- cfg.Gateway.Scheduling.LoadBatchEnabled = false
- svc := &OpenAIGatewayService{
- accountRepo: schedulerTestOpenAIAccountRepo{accounts: accounts},
- cfg: cfg,
- concurrencyService: NewConcurrencyService(schedulerTestConcurrencyCache{}),
- openaiProxyStreamCircuit: newOpenAIProxyStreamCircuit(openAIProxyStreamCircuitSettings{
- failureThreshold: 1,
- failureWindow: time.Minute,
- quarantineTTL: 10 * time.Minute,
- maxEntries: 16,
- }),
- }
- svc.openaiProxyStreamCircuit.recordFailure(proxyA, time.Now())
- selection, _, err := svc.SelectAccountWithScheduler(
- context.Background(), nil, "", "", "gpt-5.6-sol", nil, OpenAIUpstreamTransportAny, false,
- )
- require.NoError(t, err)
- require.NotNil(t, selection)
- require.NotNil(t, selection.Account)
- require.Equal(t, int64(469803), selection.Account.ID)
- }
- func TestOpenAIGatewayService_SelectAccountWithScheduler_FailsOpenWhenAllProxiesQuarantined(t *testing.T) {
- resetOpenAIAdvancedSchedulerSettingCacheForTest()
- // Every schedulable account shares the quarantined proxy: the circuit must
- // degrade to a preference instead of zeroing out capacity (#5056).
- proxyA := int64(5056)
- accounts := []Account{
- {ID: 505601, Platform: PlatformOpenAI, Type: AccountTypeOAuth, Status: StatusActive, Schedulable: true, Concurrency: 1, Priority: 0, ProxyID: &proxyA},
- {ID: 505602, Platform: PlatformOpenAI, Type: AccountTypeOAuth, Status: StatusActive, Schedulable: true, Concurrency: 1, Priority: 5, ProxyID: &proxyA},
- }
- cfg := &config.Config{}
- cfg.Gateway.Scheduling.LoadBatchEnabled = false
- svc := &OpenAIGatewayService{
- accountRepo: schedulerTestOpenAIAccountRepo{accounts: accounts},
- cfg: cfg,
- concurrencyService: NewConcurrencyService(schedulerTestConcurrencyCache{}),
- openaiProxyStreamCircuit: newOpenAIProxyStreamCircuit(openAIProxyStreamCircuitSettings{
- failureThreshold: 1,
- failureWindow: time.Minute,
- quarantineTTL: 10 * time.Minute,
- maxEntries: 16,
- }),
- }
- tripped, _ := svc.openaiProxyStreamCircuit.recordFailure(proxyA, time.Now())
- require.True(t, tripped)
- selection, _, err := svc.SelectAccountWithScheduler(
- context.Background(), nil, "", "", "gpt-5.6-sol", nil, OpenAIUpstreamTransportAny, false,
- )
- require.NoError(t, err, "quarantine must fail open instead of returning no available accounts")
- require.NotNil(t, selection)
- require.NotNil(t, selection.Account)
- require.NotNil(t, selection.Account.ProxyID)
- require.Equal(t, proxyA, *selection.Account.ProxyID)
- require.True(t, svc.openaiProxyStreamCircuit.isBlocked(proxyA, time.Now()),
- "fail-open must not clear the quarantine; only a completed stream or TTL expiry does")
- }
- func TestOpenAIGatewayService_SelectAccountWithScheduler_SessionStickyRateLimitedAccountFallsBackToFreshCandidate(t *testing.T) {
- ctx := context.Background()
- groupID := int64(10101)
- rateLimitedUntil := time.Now().Add(30 * time.Minute)
- staleSticky := &Account{ID: 31001, Platform: PlatformOpenAI, Type: AccountTypeOAuth, Status: StatusActive, Schedulable: true, Concurrency: 1, Priority: 0, GroupIDs: []int64{groupID}}
- staleBackup := &Account{ID: 31002, Platform: PlatformOpenAI, Type: AccountTypeOAuth, Status: StatusActive, Schedulable: true, Concurrency: 1, Priority: 5, GroupIDs: []int64{groupID}}
- freshSticky := &Account{ID: 31001, Platform: PlatformOpenAI, Type: AccountTypeOAuth, Status: StatusActive, Schedulable: true, Concurrency: 1, Priority: 0, GroupIDs: []int64{groupID}, RateLimitResetAt: &rateLimitedUntil}
- freshBackup := &Account{ID: 31002, Platform: PlatformOpenAI, Type: AccountTypeOAuth, Status: StatusActive, Schedulable: true, Concurrency: 1, Priority: 5, GroupIDs: []int64{groupID}}
- cache := &schedulerTestGatewayCache{sessionBindings: map[string]int64{"openai:session_hash_rate_limited": 31001}}
- snapshotCache := &openAISnapshotCacheStub{snapshotAccounts: []*Account{staleSticky, staleBackup}, accountsByID: map[int64]*Account{31001: freshSticky, 31002: freshBackup}}
- snapshotService := &SchedulerSnapshotService{cache: snapshotCache}
- svc := &OpenAIGatewayService{
- accountRepo: schedulerTestOpenAIAccountRepo{accounts: []Account{*freshSticky, *freshBackup}},
- cache: cache,
- cfg: &config.Config{},
- rateLimitService: newOpenAIAdvancedSchedulerRateLimitService("true"),
- schedulerSnapshot: snapshotService,
- concurrencyService: NewConcurrencyService(schedulerTestConcurrencyCache{}),
- }
- selection, decision, err := svc.SelectAccountWithScheduler(ctx, &groupID, "", "session_hash_rate_limited", "gpt-5.1", nil, OpenAIUpstreamTransportAny, false)
- require.NoError(t, err)
- require.NotNil(t, selection)
- require.NotNil(t, selection.Account)
- require.Equal(t, int64(31002), selection.Account.ID)
- require.Equal(t, openAIAccountScheduleLayerLoadBalance, decision.Layer)
- }
- func TestOpenAIGatewayService_SelectAccountForModelWithExclusions_AutoPauseBy5hThreshold(t *testing.T) {
- ctx := context.Background()
- primary := Account{
- ID: 35001,
- Platform: PlatformOpenAI,
- Type: AccountTypeAPIKey,
- Status: StatusActive,
- Schedulable: true,
- Concurrency: 1,
- Priority: 0,
- Extra: map[string]any{
- "codex_5h_used_percent": 95.0,
- "auto_pause_5h_threshold": 0.95,
- },
- }
- secondary := Account{ID: 35002, Platform: PlatformOpenAI, Type: AccountTypeAPIKey, Status: StatusActive, Schedulable: true, Concurrency: 1, Priority: 5}
- svc := &OpenAIGatewayService{accountRepo: schedulerTestOpenAIAccountRepo{accounts: []Account{primary, secondary}}, cfg: &config.Config{}}
- account, err := svc.SelectAccountForModelWithExclusions(ctx, nil, "", "gpt-5.1", nil)
- require.NoError(t, err)
- require.NotNil(t, account)
- require.Equal(t, int64(35002), account.ID)
- }
- func TestOpenAIGatewayService_SelectAccountForModelWithExclusions_AllowsBelow5hThreshold(t *testing.T) {
- ctx := context.Background()
- primary := Account{
- ID: 35101,
- Platform: PlatformOpenAI,
- Type: AccountTypeAPIKey,
- Status: StatusActive,
- Schedulable: true,
- Concurrency: 1,
- Priority: 0,
- Extra: map[string]any{
- "codex_5h_used_percent": 80.0,
- "auto_pause_5h_threshold": 0.95,
- },
- }
- secondary := Account{ID: 35102, Platform: PlatformOpenAI, Type: AccountTypeAPIKey, Status: StatusActive, Schedulable: true, Concurrency: 1, Priority: 5}
- svc := &OpenAIGatewayService{accountRepo: schedulerTestOpenAIAccountRepo{accounts: []Account{primary, secondary}}, cfg: &config.Config{}}
- account, err := svc.SelectAccountForModelWithExclusions(ctx, nil, "", "gpt-5.1", nil)
- require.NoError(t, err)
- require.NotNil(t, account)
- require.Equal(t, int64(35101), account.ID)
- }
- func TestOpenAIGatewayService_SelectAccountForModelWithExclusions_AutoPauseBy7dThreshold(t *testing.T) {
- ctx := context.Background()
- primary := Account{
- ID: 35201,
- Platform: PlatformOpenAI,
- Type: AccountTypeAPIKey,
- Status: StatusActive,
- Schedulable: true,
- Concurrency: 1,
- Priority: 0,
- Extra: map[string]any{
- "codex_7d_used_percent": 95.0,
- "auto_pause_7d_threshold": 0.95,
- },
- }
- secondary := Account{ID: 35202, Platform: PlatformOpenAI, Type: AccountTypeAPIKey, Status: StatusActive, Schedulable: true, Concurrency: 1, Priority: 5}
- svc := &OpenAIGatewayService{accountRepo: schedulerTestOpenAIAccountRepo{accounts: []Account{primary, secondary}}, cfg: &config.Config{}}
- account, err := svc.SelectAccountForModelWithExclusions(ctx, nil, "", "gpt-5.1", nil)
- require.NoError(t, err)
- require.NotNil(t, account)
- require.Equal(t, int64(35202), account.ID)
- }
- func TestOpenAIGatewayService_SelectAccountForModelWithExclusions_UnconfiguredThresholdKeepsLegacyBehavior(t *testing.T) {
- ctx := context.Background()
- 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}}
- secondary := Account{ID: 35302, Platform: PlatformOpenAI, Type: AccountTypeAPIKey, Status: StatusActive, Schedulable: true, Concurrency: 1, Priority: 5}
- svc := &OpenAIGatewayService{accountRepo: schedulerTestOpenAIAccountRepo{accounts: []Account{primary, secondary}}, cfg: &config.Config{}}
- account, err := svc.SelectAccountForModelWithExclusions(ctx, nil, "", "gpt-5.1", nil)
- require.NoError(t, err)
- require.NotNil(t, account)
- require.Equal(t, int64(35301), account.ID)
- }
- func TestOpenAIGatewayService_SelectAccountForModelWithExclusions_UsesGlobalDefaultThreshold(t *testing.T) {
- ctx := withOpenAIQuotaAutoPauseSettings(context.Background(), OpsOpenAIAccountQuotaAutoPauseSettings{DefaultThreshold5h: 0.95})
- primary := Account{
- ID: 35401,
- Platform: PlatformOpenAI,
- Type: AccountTypeAPIKey,
- Status: StatusActive,
- Schedulable: true,
- Concurrency: 1,
- Priority: 0,
- Extra: map[string]any{
- "codex_5h_used_percent": 95.0,
- },
- }
- secondary := Account{ID: 35402, Platform: PlatformOpenAI, Type: AccountTypeAPIKey, Status: StatusActive, Schedulable: true, Concurrency: 1, Priority: 5}
- svc := &OpenAIGatewayService{accountRepo: schedulerTestOpenAIAccountRepo{accounts: []Account{primary, secondary}}, cfg: &config.Config{}}
- account, err := svc.SelectAccountForModelWithExclusions(ctx, nil, "", "gpt-5.1", nil)
- require.NoError(t, err)
- require.NotNil(t, account)
- require.Equal(t, int64(35402), account.ID)
- }
- // Regression: a per-account explicit-disable flag exempts the account from auto-pause
- // even when a global default threshold is set. Without this, "leave threshold blank"
- // silently falls back to global default and admins have no way to whitelist a single
- // account.
- func TestOpenAIGatewayService_SelectAccountForModelWithExclusions_PerAccountDisableOverridesGlobalDefault(t *testing.T) {
- ctx := withOpenAIQuotaAutoPauseSettings(context.Background(), OpsOpenAIAccountQuotaAutoPauseSettings{DefaultThreshold5h: 0.95})
- // Account has high usage AND no per-account threshold (would normally fall back to
- // the global default and get paused), but the explicit disable flag is set.
- primary := Account{
- ID: 35701,
- Platform: PlatformOpenAI,
- Type: AccountTypeAPIKey,
- Status: StatusActive,
- Schedulable: true,
- Concurrency: 1,
- Priority: 0,
- Extra: map[string]any{
- "codex_5h_used_percent": 99.0,
- "auto_pause_5h_disabled": true,
- },
- }
- secondary := Account{ID: 35702, Platform: PlatformOpenAI, Type: AccountTypeAPIKey, Status: StatusActive, Schedulable: true, Concurrency: 1, Priority: 5}
- svc := &OpenAIGatewayService{accountRepo: schedulerTestOpenAIAccountRepo{accounts: []Account{primary, secondary}}, cfg: &config.Config{}}
- account, err := svc.SelectAccountForModelWithExclusions(ctx, nil, "", "gpt-5.1", nil)
- require.NoError(t, err)
- require.NotNil(t, account)
- require.Equal(t, int64(35701), account.ID)
- }
- // Disable is per-window: disabling only 5h must still allow 7d auto-pause to fire.
- func TestOpenAIGatewayService_SelectAccountForModelWithExclusions_PerWindowDisableScoped(t *testing.T) {
- ctx := context.Background()
- primary := Account{
- ID: 35801,
- 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,
- "auto_pause_5h_disabled": true,
- "auto_pause_7d_threshold": 0.95,
- },
- }
- secondary := Account{ID: 35802, Platform: PlatformOpenAI, Type: AccountTypeAPIKey, Status: StatusActive, Schedulable: true, Concurrency: 1, Priority: 5}
- svc := &OpenAIGatewayService{accountRepo: schedulerTestOpenAIAccountRepo{accounts: []Account{primary, secondary}}, cfg: &config.Config{}}
- account, err := svc.SelectAccountForModelWithExclusions(ctx, nil, "", "gpt-5.1", nil)
- require.NoError(t, err)
- require.NotNil(t, account)
- require.Equal(t, int64(35802), account.ID, "7d auto-pause must still fire even though 5h is disabled")
- }
- func TestOpenAIGatewayService_SelectAccountForModelWithExclusions_StaleUsageWindowResetSkipsPause(t *testing.T) {
- ctx := context.Background()
- // Usage is over threshold but the window's reset time has already passed, so the
- // cached percentage is stale (the real window rolled over) and the account must NOT
- // stay paused — otherwise it could be skipped forever with no traffic to refresh it.
- primary := Account{
- ID: 35501,
- Platform: PlatformOpenAI,
- Type: AccountTypeAPIKey,
- Status: StatusActive,
- Schedulable: true,
- Concurrency: 1,
- Priority: 0,
- Extra: map[string]any{
- "codex_5h_used_percent": 99.0,
- "auto_pause_5h_threshold": 0.95,
- "codex_5h_reset_at": time.Now().Add(-time.Minute).Format(time.RFC3339),
- },
- }
- secondary := Account{ID: 35502, Platform: PlatformOpenAI, Type: AccountTypeAPIKey, Status: StatusActive, Schedulable: true, Concurrency: 1, Priority: 5}
- svc := &OpenAIGatewayService{accountRepo: schedulerTestOpenAIAccountRepo{accounts: []Account{primary, secondary}}, cfg: &config.Config{}}
- account, err := svc.SelectAccountForModelWithExclusions(ctx, nil, "", "gpt-5.1", nil)
- require.NoError(t, err)
- require.NotNil(t, account)
- require.Equal(t, int64(35501), account.ID)
- }
- func TestOpenAIGatewayService_SelectAccountForModelWithExclusions_FreshUsageWindowStillPauses(t *testing.T) {
- ctx := context.Background()
- // Same as above but the window has not reset yet, so the account stays paused.
- primary := Account{
- ID: 35601,
- Platform: PlatformOpenAI,
- Type: AccountTypeAPIKey,
- Status: StatusActive,
- Schedulable: true,
- Concurrency: 1,
- Priority: 0,
- Extra: map[string]any{
- "codex_5h_used_percent": 99.0,
- "auto_pause_5h_threshold": 0.95,
- "codex_5h_reset_at": time.Now().Add(time.Hour).Format(time.RFC3339),
- },
- }
- secondary := Account{ID: 35602, Platform: PlatformOpenAI, Type: AccountTypeAPIKey, Status: StatusActive, Schedulable: true, Concurrency: 1, Priority: 5}
- svc := &OpenAIGatewayService{accountRepo: schedulerTestOpenAIAccountRepo{accounts: []Account{primary, secondary}}, cfg: &config.Config{}}
- account, err := svc.SelectAccountForModelWithExclusions(ctx, nil, "", "gpt-5.1", nil)
- require.NoError(t, err)
- require.NotNil(t, account)
- require.Equal(t, int64(35602), account.ID)
- }
- // Issue #2994: an account poisoned with an inflated used% (e.g. from the reverted #2918
- // inversion) gets excluded from scheduling, and a paused account never receives traffic to
- // refresh its snapshot. When the snapshot is stale (codex_usage_updated_at older than the
- // staleness bound) the account must be allowed a request so it can self-heal from the real
- // response headers — independent of the window's reset time.
- func TestOpenAIGatewayService_SelectAccountForModelWithExclusions_StaleUsageSnapshotSkipsPause_Issue2994(t *testing.T) {
- ctx := context.Background()
- primary := Account{
- ID: 35701,
- Platform: PlatformOpenAI,
- Type: AccountTypeAPIKey,
- Status: StatusActive,
- Schedulable: true,
- Concurrency: 1,
- Priority: 0,
- Extra: map[string]any{
- "codex_5h_used_percent": 99.0,
- "auto_pause_5h_threshold": 0.95,
- // Window has NOT reset yet, so the reset guard stays inactive.
- "codex_5h_reset_at": time.Now().Add(time.Hour).Format(time.RFC3339),
- // Snapshot is stale: older than openAICodexAutoPauseStaleAfter (2h).
- "codex_usage_updated_at": time.Now().Add(-3 * time.Hour).Format(time.RFC3339),
- },
- }
- secondary := Account{ID: 35702, Platform: PlatformOpenAI, Type: AccountTypeAPIKey, Status: StatusActive, Schedulable: true, Concurrency: 1, Priority: 5}
- svc := &OpenAIGatewayService{accountRepo: schedulerTestOpenAIAccountRepo{accounts: []Account{primary, secondary}}, cfg: &config.Config{}}
- account, err := svc.SelectAccountForModelWithExclusions(ctx, nil, "", "gpt-5.1", nil)
- require.NoError(t, err)
- require.NotNil(t, account)
- require.Equal(t, int64(35701), account.ID)
- }
- // Issue #2994 guardrail: a genuinely-exhausted account whose snapshot was refreshed recently
- // (codex_usage_updated_at fresh) must STILL be auto-paused. The stale self-heal must not let a
- // real 99%-used account escape pause.
- func TestOpenAIGatewayService_SelectAccountForModelWithExclusions_FreshExhaustedSnapshotStillPauses_Issue2994(t *testing.T) {
- ctx := context.Background()
- primary := Account{
- ID: 35801,
- Platform: PlatformOpenAI,
- Type: AccountTypeAPIKey,
- Status: StatusActive,
- Schedulable: true,
- Concurrency: 1,
- Priority: 0,
- Extra: map[string]any{
- "codex_5h_used_percent": 99.0,
- "auto_pause_5h_threshold": 0.95,
- "codex_5h_reset_at": time.Now().Add(time.Hour).Format(time.RFC3339),
- // Snapshot refreshed 1 minute ago: not stale, so the account stays paused.
- "codex_usage_updated_at": time.Now().Add(-time.Minute).Format(time.RFC3339),
- },
- }
- secondary := Account{ID: 35802, Platform: PlatformOpenAI, Type: AccountTypeAPIKey, Status: StatusActive, Schedulable: true, Concurrency: 1, Priority: 5}
- svc := &OpenAIGatewayService{accountRepo: schedulerTestOpenAIAccountRepo{accounts: []Account{primary, secondary}}, cfg: &config.Config{}}
- account, err := svc.SelectAccountForModelWithExclusions(ctx, nil, "", "gpt-5.1", nil)
- require.NoError(t, err)
- require.NotNil(t, account)
- require.Equal(t, int64(35802), account.ID)
- }
- func TestOpenAIGatewayService_SelectAccountForModelWithExclusions_SkipsFreshlyRateLimitedSnapshotCandidate(t *testing.T) {
- ctx := context.Background()
- groupID := int64(10102)
- rateLimitedUntil := time.Now().Add(30 * time.Minute)
- stalePrimary := &Account{ID: 32001, Platform: PlatformOpenAI, Type: AccountTypeOAuth, Status: StatusActive, Schedulable: true, Concurrency: 1, Priority: 0, GroupIDs: []int64{groupID}}
- staleSecondary := &Account{ID: 32002, Platform: PlatformOpenAI, Type: AccountTypeOAuth, Status: StatusActive, Schedulable: true, Concurrency: 1, Priority: 5, GroupIDs: []int64{groupID}}
- freshPrimary := &Account{ID: 32001, Platform: PlatformOpenAI, Type: AccountTypeOAuth, Status: StatusActive, Schedulable: true, Concurrency: 1, Priority: 0, GroupIDs: []int64{groupID}, RateLimitResetAt: &rateLimitedUntil}
- freshSecondary := &Account{ID: 32002, Platform: PlatformOpenAI, Type: AccountTypeOAuth, Status: StatusActive, Schedulable: true, Concurrency: 1, Priority: 5, GroupIDs: []int64{groupID}}
- snapshotCache := &openAISnapshotCacheStub{snapshotAccounts: []*Account{stalePrimary, staleSecondary}, accountsByID: map[int64]*Account{32001: freshPrimary, 32002: freshSecondary}}
- snapshotService := &SchedulerSnapshotService{cache: snapshotCache}
- svc := &OpenAIGatewayService{
- accountRepo: schedulerTestOpenAIAccountRepo{accounts: []Account{*freshPrimary, *freshSecondary}},
- cfg: &config.Config{},
- rateLimitService: newOpenAIAdvancedSchedulerRateLimitService("true"),
- schedulerSnapshot: snapshotService,
- }
- account, err := svc.SelectAccountForModelWithExclusions(ctx, &groupID, "", "gpt-5.1", nil)
- require.NoError(t, err)
- require.NotNil(t, account)
- require.Equal(t, int64(32002), account.ID)
- }
- func TestOpenAIGatewayService_SelectAccountForModelWithExclusions_ModelRateLimitOnlySkipsThatModel(t *testing.T) {
- ctx := context.Background()
- resetAt := time.Now().Add(30 * time.Minute).Format(time.RFC3339)
- primary := Account{
- ID: 32101,
- Platform: PlatformOpenAI,
- Type: AccountTypeAPIKey,
- Status: StatusActive,
- Schedulable: true,
- Concurrency: 1,
- Priority: 0,
- Extra: map[string]any{
- modelRateLimitsKey: map[string]any{
- "gpt-5.4": map[string]any{
- "rate_limit_reset_at": resetAt,
- },
- },
- },
- }
- secondary := Account{
- ID: 32102,
- Platform: PlatformOpenAI,
- Type: AccountTypeAPIKey,
- Status: StatusActive,
- Schedulable: true,
- Concurrency: 1,
- Priority: 5,
- }
- svc := &OpenAIGatewayService{
- accountRepo: schedulerTestOpenAIAccountRepo{accounts: []Account{primary, secondary}},
- cfg: &config.Config{},
- }
- account, err := svc.SelectAccountForModelWithExclusions(ctx, nil, "", "gpt-5.4", nil)
- require.NoError(t, err)
- require.NotNil(t, account)
- require.Equal(t, int64(32102), account.ID)
- account, err = svc.SelectAccountForModelWithExclusions(ctx, nil, "", "gpt-5.3", nil)
- require.NoError(t, err)
- require.NotNil(t, account)
- require.Equal(t, int64(32101), account.ID)
- }
- func TestOpenAIGatewayService_SelectAccountWithScheduler_SessionStickyDBRuntimeRecheckSkipsStaleCachedAccount(t *testing.T) {
- ctx := context.Background()
- groupID := int64(10103)
- rateLimitedUntil := time.Now().Add(30 * time.Minute)
- staleSticky := &Account{ID: 33001, Platform: PlatformOpenAI, Type: AccountTypeOAuth, Status: StatusActive, Schedulable: true, Concurrency: 1, Priority: 0, GroupIDs: []int64{groupID}}
- staleBackup := &Account{ID: 33002, Platform: PlatformOpenAI, Type: AccountTypeOAuth, Status: StatusActive, Schedulable: true, Concurrency: 1, Priority: 5, GroupIDs: []int64{groupID}}
- dbSticky := Account{ID: 33001, Platform: PlatformOpenAI, Type: AccountTypeOAuth, Status: StatusActive, Schedulable: true, Concurrency: 1, Priority: 0, GroupIDs: []int64{groupID}, RateLimitResetAt: &rateLimitedUntil}
- dbBackup := Account{ID: 33002, Platform: PlatformOpenAI, Type: AccountTypeOAuth, Status: StatusActive, Schedulable: true, Concurrency: 1, Priority: 5, GroupIDs: []int64{groupID}}
- cache := &schedulerTestGatewayCache{sessionBindings: map[string]int64{"openai:session_hash_db_runtime_recheck": 33001}}
- snapshotCache := &openAISnapshotCacheStub{
- snapshotAccounts: []*Account{staleSticky, staleBackup},
- accountsByID: map[int64]*Account{33001: staleSticky, 33002: staleBackup},
- }
- snapshotService := &SchedulerSnapshotService{cache: snapshotCache}
- svc := &OpenAIGatewayService{
- accountRepo: schedulerTestOpenAIAccountRepo{accounts: []Account{dbSticky, dbBackup}},
- cache: cache,
- cfg: &config.Config{},
- rateLimitService: newOpenAIAdvancedSchedulerRateLimitService("true"),
- schedulerSnapshot: snapshotService,
- concurrencyService: NewConcurrencyService(schedulerTestConcurrencyCache{}),
- }
- selection, decision, err := svc.SelectAccountWithScheduler(ctx, &groupID, "", "session_hash_db_runtime_recheck", "gpt-5.1", nil, OpenAIUpstreamTransportAny, false)
- require.NoError(t, err)
- require.NotNil(t, selection)
- require.NotNil(t, selection.Account)
- require.Equal(t, int64(33002), selection.Account.ID)
- require.Equal(t, openAIAccountScheduleLayerLoadBalance, decision.Layer)
- }
- func TestOpenAIGatewayService_SelectAccountForModelWithExclusions_DBRuntimeRecheckSkipsStaleCachedCandidate(t *testing.T) {
- ctx := context.Background()
- groupID := int64(10104)
- rateLimitedUntil := time.Now().Add(30 * time.Minute)
- stalePrimary := &Account{ID: 34001, Platform: PlatformOpenAI, Type: AccountTypeOAuth, Status: StatusActive, Schedulable: true, Concurrency: 1, Priority: 0, GroupIDs: []int64{groupID}}
- staleSecondary := &Account{ID: 34002, Platform: PlatformOpenAI, Type: AccountTypeOAuth, Status: StatusActive, Schedulable: true, Concurrency: 1, Priority: 5, GroupIDs: []int64{groupID}}
- dbPrimary := Account{ID: 34001, Platform: PlatformOpenAI, Type: AccountTypeOAuth, Status: StatusActive, Schedulable: true, Concurrency: 1, Priority: 0, GroupIDs: []int64{groupID}, RateLimitResetAt: &rateLimitedUntil}
- dbSecondary := Account{ID: 34002, Platform: PlatformOpenAI, Type: AccountTypeOAuth, Status: StatusActive, Schedulable: true, Concurrency: 1, Priority: 5, GroupIDs: []int64{groupID}}
- snapshotCache := &openAISnapshotCacheStub{
- snapshotAccounts: []*Account{stalePrimary, staleSecondary},
- accountsByID: map[int64]*Account{34001: stalePrimary, 34002: staleSecondary},
- }
- snapshotService := &SchedulerSnapshotService{cache: snapshotCache}
- svc := &OpenAIGatewayService{
- accountRepo: schedulerTestOpenAIAccountRepo{accounts: []Account{dbPrimary, dbSecondary}},
- cfg: &config.Config{},
- rateLimitService: newOpenAIAdvancedSchedulerRateLimitService("true"),
- schedulerSnapshot: snapshotService,
- }
- account, err := svc.SelectAccountForModelWithExclusions(ctx, &groupID, "", "gpt-5.1", nil)
- require.NoError(t, err)
- require.NotNil(t, account)
- require.Equal(t, int64(34002), account.ID)
- }
- func TestOpenAIGatewayService_SelectAccountWithScheduler_DBFreshGroupRecheckReleasesMovedAccount(t *testing.T) {
- ctx := context.Background()
- groupID, otherGroupID := int64(10105), int64(10106)
- stalePrimary := &Account{ID: 34101, Platform: PlatformOpenAI, Type: AccountTypeAPIKey, Status: StatusActive, Schedulable: true, Concurrency: 1, Priority: 0, GroupIDs: []int64{groupID}}
- staleBackup := &Account{ID: 34102, Platform: PlatformOpenAI, Type: AccountTypeAPIKey, Status: StatusActive, Schedulable: true, Concurrency: 1, Priority: 10, GroupIDs: []int64{groupID}}
- dbPrimary := *stalePrimary
- dbPrimary.GroupIDs = []int64{otherGroupID}
- dbBackup := *staleBackup
- snapshotCache := &openAISnapshotCacheStub{
- snapshotAccounts: []*Account{stalePrimary, staleBackup},
- accountsByID: map[int64]*Account{stalePrimary.ID: stalePrimary, staleBackup.ID: staleBackup},
- }
- acquiredIDs, releasedIDs := []int64{}, []int64{}
- cfg := &config.Config{RunMode: config.RunModeStandard}
- svc := &OpenAIGatewayService{
- accountRepo: schedulerTestOpenAIAccountRepo{accounts: []Account{dbPrimary, dbBackup}},
- cfg: cfg,
- schedulerSnapshot: &SchedulerSnapshotService{cache: snapshotCache},
- concurrencyService: NewConcurrencyService(schedulerTestConcurrencyCache{
- acquiredIDs: &acquiredIDs,
- releasedIDs: &releasedIDs,
- }),
- }
- scheduler := &defaultOpenAIAccountScheduler{service: svc}
- selection, _, err := scheduler.tryAcquireOpenAISelectionOrder(ctx, OpenAIAccountScheduleRequest{
- GroupID: &groupID, Platform: PlatformOpenAI, RequestedModel: "gpt-5.1",
- }, []openAIAccountCandidateScore{
- {account: stalePrimary, loadInfo: &AccountLoadInfo{AccountID: stalePrimary.ID}},
- {account: staleBackup, loadInfo: &AccountLoadInfo{AccountID: staleBackup.ID}},
- })
- require.NoError(t, err)
- require.Equal(t, staleBackup.ID, selection.Account.ID)
- require.Equal(t, []int64{stalePrimary.ID, staleBackup.ID}, acquiredIDs)
- require.Equal(t, []int64{stalePrimary.ID}, releasedIDs)
- selection.ReleaseFunc()
- }
- func TestOpenAIGatewayService_SelectAccountWithLoadAwareness_DBFreshGroupRecheckWaitsOnValidAccount(t *testing.T) {
- ctx := context.Background()
- groupID, otherGroupID := int64(10107), int64(10108)
- stalePrimary := &Account{ID: 34201, Platform: PlatformOpenAI, Type: AccountTypeAPIKey, Status: StatusActive, Schedulable: true, Concurrency: 1, Priority: 0, GroupIDs: []int64{groupID}}
- staleBackup := &Account{ID: 34202, Platform: PlatformOpenAI, Type: AccountTypeAPIKey, Status: StatusActive, Schedulable: true, Concurrency: 1, Priority: 10, GroupIDs: []int64{groupID}}
- dbPrimary := *stalePrimary
- dbPrimary.GroupIDs = []int64{otherGroupID}
- dbBackup := *staleBackup
- snapshotCache := &openAISnapshotCacheStub{
- snapshotAccounts: []*Account{stalePrimary, staleBackup},
- accountsByID: map[int64]*Account{stalePrimary.ID: stalePrimary, staleBackup.ID: staleBackup},
- }
- cfg := &config.Config{RunMode: config.RunModeStandard}
- cfg.Gateway.Scheduling.LoadBatchEnabled = true
- svc := &OpenAIGatewayService{
- accountRepo: schedulerTestOpenAIAccountRepo{accounts: []Account{dbPrimary, dbBackup}},
- cfg: cfg,
- schedulerSnapshot: &SchedulerSnapshotService{cache: snapshotCache},
- concurrencyService: NewConcurrencyService(schedulerTestConcurrencyCache{
- acquireResults: map[int64]bool{staleBackup.ID: false},
- }),
- }
- selection, err := svc.SelectAccountWithLoadAwareness(ctx, &groupID, "", "gpt-5.1", nil)
- require.NoError(t, err)
- require.NotNil(t, selection.WaitPlan)
- require.Equal(t, staleBackup.ID, selection.Account.ID)
- require.Equal(t, staleBackup.ID, selection.WaitPlan.AccountID)
- }
- func TestOpenAIGatewayService_RecheckSelectedOpenAIAccountFromDB_SimpleModeUsesFullPool(t *testing.T) {
- grouped := Account{ID: 34301, Platform: PlatformOpenAI, Type: AccountTypeAPIKey, Status: StatusActive, Schedulable: true, Concurrency: 1, GroupIDs: []int64{99}}
- svc := &OpenAIGatewayService{
- accountRepo: schedulerTestOpenAIAccountRepo{accounts: []Account{grouped}},
- cfg: &config.Config{RunMode: config.RunModeSimple},
- schedulerSnapshot: &SchedulerSnapshotService{cache: &openAISnapshotCacheStub{}},
- }
- requestedGroupID := int64(100)
- for _, groupID := range []*int64{nil, &requestedGroupID} {
- fresh := svc.recheckSelectedOpenAIAccountFromDB(context.Background(), &grouped, groupID, PlatformOpenAI, "gpt-5.1", false, "")
- require.NotNil(t, fresh)
- require.Equal(t, grouped.ID, fresh.ID)
- }
- ungrouped := grouped
- ungrouped.ID++
- ungrouped.GroupIDs = nil
- standardSvc := &OpenAIGatewayService{
- accountRepo: schedulerTestOpenAIAccountRepo{accounts: []Account{grouped, ungrouped}},
- cfg: &config.Config{RunMode: config.RunModeStandard},
- schedulerSnapshot: &SchedulerSnapshotService{cache: &openAISnapshotCacheStub{}},
- }
- require.Nil(t, standardSvc.recheckSelectedOpenAIAccountFromDB(context.Background(), &grouped, nil, PlatformOpenAI, "gpt-5.1", false, ""))
- require.NotNil(t, standardSvc.recheckSelectedOpenAIAccountFromDB(context.Background(), &ungrouped, nil, PlatformOpenAI, "gpt-5.1", false, ""))
- }
- func TestOpenAIGatewayService_SelectAccountWithScheduler_PreviousResponseSticky(t *testing.T) {
- ctx := context.Background()
- groupID := int64(9)
- account := Account{
- ID: 1001,
- Platform: PlatformOpenAI,
- Type: AccountTypeAPIKey,
- Status: StatusActive,
- Schedulable: true,
- Concurrency: 2,
- Extra: map[string]any{
- "openai_apikey_responses_websockets_v2_enabled": true,
- },
- }
- cache := &schedulerTestGatewayCache{}
- cfg := &config.Config{}
- cfg.Gateway.OpenAIWS.Enabled = true
- cfg.Gateway.OpenAIWS.OAuthEnabled = true
- cfg.Gateway.OpenAIWS.APIKeyEnabled = true
- cfg.Gateway.OpenAIWS.ResponsesWebsocketsV2 = true
- cfg.Gateway.OpenAIWS.StickySessionTTLSeconds = 1800
- cfg.Gateway.OpenAIWS.StickyResponseIDTTLSeconds = 3600
- svc := &OpenAIGatewayService{
- accountRepo: schedulerTestOpenAIAccountRepo{accounts: []Account{account}},
- cache: cache,
- cfg: cfg,
- rateLimitService: newOpenAIAdvancedSchedulerRateLimitService("true"),
- concurrencyService: NewConcurrencyService(schedulerTestConcurrencyCache{}),
- }
- store := svc.getOpenAIWSStateStore()
- require.NoError(t, store.BindResponseAccount(ctx, groupID, "resp_prev_001", account.ID, time.Hour))
- selection, decision, err := svc.SelectAccountWithScheduler(
- ctx,
- &groupID,
- "resp_prev_001",
- "session_hash_001",
- "gpt-5.1",
- nil,
- OpenAIUpstreamTransportAny,
- false,
- )
- require.NoError(t, err)
- require.NotNil(t, selection)
- require.NotNil(t, selection.Account)
- require.Equal(t, account.ID, selection.Account.ID)
- require.Equal(t, openAIAccountScheduleLayerPreviousResponse, decision.Layer)
- require.True(t, decision.StickyPreviousHit)
- require.Equal(t, account.ID, cache.sessionBindings["openai:session_hash_001"])
- if selection.ReleaseFunc != nil {
- selection.ReleaseFunc()
- }
- }
- func TestOpenAIGatewayService_SelectAccountWithScheduler_SessionSticky(t *testing.T) {
- ctx := context.Background()
- groupID := int64(10)
- account := Account{
- ID: 2001,
- Platform: PlatformOpenAI,
- Type: AccountTypeOAuth,
- Status: StatusActive,
- Schedulable: true,
- Concurrency: 1,
- GroupIDs: []int64{groupID},
- }
- cache := &schedulerTestGatewayCache{
- sessionBindings: map[string]int64{
- "openai:session_hash_abc": account.ID,
- },
- }
- svc := &OpenAIGatewayService{
- accountRepo: schedulerTestOpenAIAccountRepo{accounts: []Account{account}},
- cache: cache,
- cfg: &config.Config{},
- rateLimitService: newOpenAIAdvancedSchedulerRateLimitService("true"),
- concurrencyService: NewConcurrencyService(schedulerTestConcurrencyCache{}),
- }
- selection, decision, err := svc.SelectAccountWithScheduler(
- ctx,
- &groupID,
- "",
- "session_hash_abc",
- "gpt-5.1",
- nil,
- OpenAIUpstreamTransportAny,
- false,
- )
- require.NoError(t, err)
- require.NotNil(t, selection)
- require.NotNil(t, selection.Account)
- require.Equal(t, account.ID, selection.Account.ID)
- require.Equal(t, openAIAccountScheduleLayerSessionSticky, decision.Layer)
- require.True(t, decision.StickySessionHit)
- if selection.ReleaseFunc != nil {
- selection.ReleaseFunc()
- }
- }
- func TestOpenAIGatewayService_SelectAccountWithScheduler_SessionStickyBusyKeepsSticky(t *testing.T) {
- ctx := context.Background()
- groupID := int64(10100)
- accounts := []Account{
- {
- ID: 21001,
- Platform: PlatformOpenAI,
- Type: AccountTypeAPIKey,
- Status: StatusActive,
- Schedulable: true,
- Concurrency: 1,
- Priority: 0,
- GroupIDs: []int64{groupID},
- },
- {
- ID: 21002,
- Platform: PlatformOpenAI,
- Type: AccountTypeAPIKey,
- Status: StatusActive,
- Schedulable: true,
- Concurrency: 1,
- Priority: 9,
- GroupIDs: []int64{groupID},
- },
- }
- cache := &schedulerTestGatewayCache{
- sessionBindings: map[string]int64{
- "openai:session_hash_sticky_busy": 21001,
- },
- }
- cfg := &config.Config{}
- cfg.Gateway.Scheduling.StickySessionMaxWaiting = 2
- cfg.Gateway.Scheduling.StickySessionWaitTimeout = 45 * time.Second
- cfg.Gateway.OpenAIScheduler.StickyEscapeEnabled = false
- cfg.Gateway.OpenAIScheduler.StickyEscapeTTFTMs = 15000
- cfg.Gateway.OpenAIScheduler.StickyEscapeErrorRate = 0.5
- cfg.Gateway.OpenAIWS.Enabled = true
- cfg.Gateway.OpenAIWS.APIKeyEnabled = true
- cfg.Gateway.OpenAIWS.OAuthEnabled = true
- cfg.Gateway.OpenAIWS.ResponsesWebsocketsV2 = true
- concurrencyCache := schedulerTestConcurrencyCache{
- acquireResults: map[int64]bool{
- 21001: false, // sticky 账号已满
- 21002: true, // 若回退负载均衡会命中该账号(本测试要求不能切换)
- },
- waitCounts: map[int64]int{
- 21001: 999,
- },
- loadMap: map[int64]*AccountLoadInfo{
- 21001: {AccountID: 21001, LoadRate: 90, WaitingCount: 9},
- 21002: {AccountID: 21002, LoadRate: 1, WaitingCount: 0},
- },
- }
- svc := &OpenAIGatewayService{
- accountRepo: schedulerTestOpenAIAccountRepo{accounts: accounts},
- cache: cache,
- cfg: cfg,
- rateLimitService: newOpenAIAdvancedSchedulerRateLimitService("true"),
- concurrencyService: NewConcurrencyService(concurrencyCache),
- }
- selection, decision, err := svc.SelectAccountWithScheduler(
- ctx,
- &groupID,
- "",
- "session_hash_sticky_busy",
- "gpt-5.1",
- nil,
- OpenAIUpstreamTransportAny,
- false,
- )
- require.NoError(t, err)
- require.NotNil(t, selection)
- require.NotNil(t, selection.Account)
- require.Equal(t, int64(21001), selection.Account.ID, "busy sticky account should remain selected")
- require.False(t, selection.Acquired)
- require.NotNil(t, selection.WaitPlan)
- require.Equal(t, int64(21001), selection.WaitPlan.AccountID)
- require.Equal(t, openAIAccountScheduleLayerSessionSticky, decision.Layer)
- require.True(t, decision.StickySessionHit)
- }
- func TestOpenAIGatewayService_SelectAccountWithScheduler_SessionStickyEscapeByTTFT(t *testing.T) {
- ctx := context.Background()
- groupID := int64(10101)
- accounts := []Account{
- {
- ID: 21101,
- Platform: PlatformOpenAI,
- Type: AccountTypeAPIKey,
- Status: StatusActive,
- Schedulable: true,
- Concurrency: 1,
- Priority: 0,
- GroupIDs: []int64{groupID},
- },
- {
- ID: 21102,
- Platform: PlatformOpenAI,
- Type: AccountTypeAPIKey,
- Status: StatusActive,
- Schedulable: true,
- Concurrency: 1,
- Priority: 1,
- GroupIDs: []int64{groupID},
- },
- }
- cache := &schedulerTestGatewayCache{sessionBindings: map[string]int64{"openai:session_hash_sticky_ttft": 21101}}
- cfg := &config.Config{}
- cfg.Gateway.OpenAIScheduler.StickyEscapeEnabled = true
- cfg.Gateway.OpenAIScheduler.StickyEscapeTTFTMs = 15000
- cfg.Gateway.OpenAIScheduler.StickyEscapeErrorRate = 0.5
- concurrencyCache := schedulerTestConcurrencyCache{acquireResults: map[int64]bool{21102: true}}
- svc := &OpenAIGatewayService{
- accountRepo: schedulerTestOpenAIAccountRepo{accounts: accounts},
- cache: cache,
- cfg: cfg,
- rateLimitService: newOpenAIAdvancedSchedulerRateLimitService("true"),
- concurrencyService: NewConcurrencyService(concurrencyCache),
- openaiAccountStats: newOpenAIAccountRuntimeStats(),
- }
- fastTTFT := 14999
- svc.openaiAccountStats.report(21101, true, &fastTTFT)
- stableTTFT := 14999
- svc.openaiAccountStats.report(21101, true, &stableTTFT)
- selection, decision, err := svc.SelectAccountWithScheduler(ctx, &groupID, "", "session_hash_sticky_ttft", "gpt-5.1", nil, OpenAIUpstreamTransportAny, false)
- require.NoError(t, err)
- require.NotNil(t, selection)
- require.NotNil(t, selection.Account)
- require.Equal(t, int64(21101), selection.Account.ID)
- require.Equal(t, openAIAccountScheduleLayerSessionSticky, decision.Layer)
- require.True(t, decision.StickySessionHit)
- if selection.ReleaseFunc != nil {
- selection.ReleaseFunc()
- }
- slowTTFT := 20000
- for i := 0; i < 3; i++ {
- svc.openaiAccountStats.report(21101, true, &slowTTFT)
- }
- selection, decision, err = svc.SelectAccountWithScheduler(ctx, &groupID, "", "session_hash_sticky_ttft", "gpt-5.1", nil, OpenAIUpstreamTransportAny, false)
- require.NoError(t, err)
- require.NotNil(t, selection)
- require.NotNil(t, selection.Account)
- require.Equal(t, int64(21102), selection.Account.ID)
- require.Equal(t, openAIAccountScheduleLayerLoadBalance, decision.Layer)
- require.False(t, decision.StickySessionHit)
- require.Equal(t, int64(21101), cache.sessionBindings["openai:session_hash_sticky_ttft"])
- if selection.ReleaseFunc != nil {
- selection.ReleaseFunc()
- }
- }
- func TestOpenAIGatewayService_SelectAccountWithScheduler_SessionStickyEscapeByErrorRate(t *testing.T) {
- ctx := context.Background()
- groupID := int64(10102)
- accounts := []Account{
- {ID: 21201, Platform: PlatformOpenAI, Type: AccountTypeAPIKey, Status: StatusActive, Schedulable: true, Concurrency: 1, Priority: 0, GroupIDs: []int64{groupID}},
- {ID: 21202, Platform: PlatformOpenAI, Type: AccountTypeAPIKey, Status: StatusActive, Schedulable: true, Concurrency: 1, Priority: 1, GroupIDs: []int64{groupID}},
- }
- cache := &schedulerTestGatewayCache{sessionBindings: map[string]int64{"openai:session_hash_sticky_error_rate": 21201}}
- cfg := &config.Config{}
- cfg.Gateway.OpenAIScheduler.StickyEscapeEnabled = true
- cfg.Gateway.OpenAIScheduler.StickyEscapeTTFTMs = 15000
- cfg.Gateway.OpenAIScheduler.StickyEscapeErrorRate = 0.5
- svc := &OpenAIGatewayService{
- accountRepo: schedulerTestOpenAIAccountRepo{accounts: accounts},
- cache: cache,
- cfg: cfg,
- rateLimitService: newOpenAIAdvancedSchedulerRateLimitService("true"),
- concurrencyService: NewConcurrencyService(schedulerTestConcurrencyCache{acquireResults: map[int64]bool{21202: true}}),
- openaiAccountStats: newOpenAIAccountRuntimeStats(),
- }
- for i := 0; i < 3; i++ {
- svc.openaiAccountStats.report(21201, false, nil)
- }
- selection, decision, err := svc.SelectAccountWithScheduler(ctx, &groupID, "", "session_hash_sticky_error_rate", "gpt-5.1", nil, OpenAIUpstreamTransportAny, false)
- require.NoError(t, err)
- require.NotNil(t, selection)
- require.NotNil(t, selection.Account)
- require.Equal(t, int64(21201), selection.Account.ID)
- require.Equal(t, openAIAccountScheduleLayerSessionSticky, decision.Layer)
- require.True(t, decision.StickySessionHit)
- if selection.ReleaseFunc != nil {
- selection.ReleaseFunc()
- }
- for i := 0; i < 2; i++ {
- svc.openaiAccountStats.report(21201, false, nil)
- }
- selection, decision, err = svc.SelectAccountWithScheduler(ctx, &groupID, "", "session_hash_sticky_error_rate", "gpt-5.1", nil, OpenAIUpstreamTransportAny, false)
- require.NoError(t, err)
- require.NotNil(t, selection)
- require.NotNil(t, selection.Account)
- require.Equal(t, int64(21202), selection.Account.ID)
- require.Equal(t, openAIAccountScheduleLayerLoadBalance, decision.Layer)
- require.False(t, decision.StickySessionHit)
- require.Equal(t, int64(21201), cache.sessionBindings["openai:session_hash_sticky_error_rate"])
- if selection.ReleaseFunc != nil {
- selection.ReleaseFunc()
- }
- }
- func TestOpenAIGatewayService_SelectAccountWithScheduler_SessionStickyBusyEscapes(t *testing.T) {
- ctx := context.Background()
- groupID := int64(10103)
- accounts := []Account{
- {ID: 21301, Platform: PlatformOpenAI, Type: AccountTypeAPIKey, Status: StatusActive, Schedulable: true, Concurrency: 1, Priority: 0, GroupIDs: []int64{groupID}},
- {ID: 21302, Platform: PlatformOpenAI, Type: AccountTypeAPIKey, Status: StatusActive, Schedulable: true, Concurrency: 1, Priority: 1, GroupIDs: []int64{groupID}},
- }
- cache := &schedulerTestGatewayCache{sessionBindings: map[string]int64{"openai:session_hash_sticky_busy_escape": 21301}}
- cfg := &config.Config{}
- cfg.Gateway.OpenAIScheduler.StickyEscapeEnabled = true
- cfg.Gateway.OpenAIScheduler.StickyEscapeTTFTMs = 15000
- cfg.Gateway.OpenAIScheduler.StickyEscapeErrorRate = 0.5
- cfg.Gateway.Scheduling.StickySessionMaxWaiting = 2
- cfg.Gateway.Scheduling.StickySessionWaitTimeout = 45 * time.Second
- concurrencyCache := schedulerTestConcurrencyCache{
- acquireResults: map[int64]bool{21301: false, 21302: true},
- waitCounts: map[int64]int{21301: 999},
- loadMap: map[int64]*AccountLoadInfo{
- 21301: {AccountID: 21301, LoadRate: 95, WaitingCount: 9},
- 21302: {AccountID: 21302, LoadRate: 1, WaitingCount: 0},
- },
- }
- svc := &OpenAIGatewayService{
- accountRepo: schedulerTestOpenAIAccountRepo{accounts: accounts},
- cache: cache,
- cfg: cfg,
- rateLimitService: newOpenAIAdvancedSchedulerRateLimitService("true"),
- concurrencyService: NewConcurrencyService(concurrencyCache),
- }
- selection, decision, err := svc.SelectAccountWithScheduler(ctx, &groupID, "", "session_hash_sticky_busy_escape", "gpt-5.1", nil, OpenAIUpstreamTransportAny, false)
- require.NoError(t, err)
- require.NotNil(t, selection)
- require.NotNil(t, selection.Account)
- require.Equal(t, int64(21302), selection.Account.ID)
- require.Nil(t, selection.WaitPlan)
- require.Equal(t, openAIAccountScheduleLayerLoadBalance, decision.Layer)
- require.False(t, decision.StickySessionHit)
- if selection.ReleaseFunc != nil {
- selection.ReleaseFunc()
- }
- }
- func TestOpenAIGatewayService_SelectAccountWithScheduler_SessionStickyEscapeDisabledKeepsLegacyBehavior(t *testing.T) {
- ctx := context.Background()
- groupID := int64(10104)
- accounts := []Account{
- {ID: 21401, Platform: PlatformOpenAI, Type: AccountTypeAPIKey, Status: StatusActive, Schedulable: true, Concurrency: 1, Priority: 0, GroupIDs: []int64{groupID}},
- {ID: 21402, Platform: PlatformOpenAI, Type: AccountTypeAPIKey, Status: StatusActive, Schedulable: true, Concurrency: 1, Priority: 1, GroupIDs: []int64{groupID}},
- }
- cache := &schedulerTestGatewayCache{sessionBindings: map[string]int64{"openai:session_hash_sticky_disabled": 21401}}
- cfg := &config.Config{}
- cfg.Gateway.OpenAIScheduler.StickyEscapeEnabled = false
- cfg.Gateway.OpenAIScheduler.StickyEscapeTTFTMs = 15000
- cfg.Gateway.OpenAIScheduler.StickyEscapeErrorRate = 0.5
- cfg.Gateway.Scheduling.StickySessionMaxWaiting = 2
- cfg.Gateway.Scheduling.StickySessionWaitTimeout = 45 * time.Second
- concurrencyCache := schedulerTestConcurrencyCache{
- acquireResults: map[int64]bool{21401: false, 21402: true},
- waitCounts: map[int64]int{21401: 999},
- }
- svc := &OpenAIGatewayService{
- accountRepo: schedulerTestOpenAIAccountRepo{accounts: accounts},
- cache: cache,
- cfg: cfg,
- rateLimitService: newOpenAIAdvancedSchedulerRateLimitService("true"),
- concurrencyService: NewConcurrencyService(concurrencyCache),
- openaiAccountStats: newOpenAIAccountRuntimeStats(),
- }
- slowTTFT := 20000
- svc.openaiAccountStats.report(21401, true, &slowTTFT)
- for i := 0; i < 5; i++ {
- svc.openaiAccountStats.report(21401, false, nil)
- }
- selection, decision, err := svc.SelectAccountWithScheduler(ctx, &groupID, "", "session_hash_sticky_disabled", "gpt-5.1", nil, OpenAIUpstreamTransportAny, false)
- require.NoError(t, err)
- require.NotNil(t, selection)
- require.NotNil(t, selection.Account)
- require.Equal(t, int64(21401), selection.Account.ID)
- require.NotNil(t, selection.WaitPlan)
- require.Equal(t, int64(21401), selection.WaitPlan.AccountID)
- require.Equal(t, openAIAccountScheduleLayerSessionSticky, decision.Layer)
- require.True(t, decision.StickySessionHit)
- }
- func TestOpenAIGatewayService_SelectAccountWithScheduler_SubscriptionPriorityChoosesSubscriptionPoolFirst(t *testing.T) {
- ctx := context.Background()
- groupID := int64(10120)
- cheapAPIKey := upstreamCostTestAccount(21602, UpstreamBillingProbeStatusOK, 0.01, time.Now().Add(-time.Minute), 30*time.Minute)
- cheapAPIKey.Status = StatusActive
- cheapAPIKey.Schedulable = true
- cheapAPIKey.Concurrency = 1
- cheapAPIKey.Priority = 0
- cheapAPIKey.GroupIDs = []int64{groupID}
- accounts := []Account{
- {
- ID: 21601,
- Platform: PlatformOpenAI,
- Type: AccountTypeOAuth,
- Status: StatusActive,
- Schedulable: true,
- Concurrency: 1,
- Priority: 10,
- GroupIDs: []int64{groupID},
- Credentials: map[string]any{"plan_type": "plus"},
- },
- *cheapAPIKey,
- }
- concurrencyCache := schedulerTestConcurrencyCache{
- acquireResults: map[int64]bool{21601: true, 21602: true},
- loadMap: map[int64]*AccountLoadInfo{
- 21601: {AccountID: 21601, LoadRate: 90, WaitingCount: 1},
- 21602: {AccountID: 21602, LoadRate: 0, WaitingCount: 0},
- },
- }
- cfg := newSchedulerTestSubscriptionPriorityConfig()
- cfg.Gateway.OpenAIWS.SchedulerScoreWeights.UpstreamCost = 100
- svc := &OpenAIGatewayService{
- accountRepo: schedulerTestOpenAIAccountRepo{accounts: accounts},
- cache: &schedulerTestGatewayCache{},
- cfg: cfg,
- rateLimitService: newOpenAIAdvancedSchedulerRateLimitService("true", "", "true"),
- concurrencyService: NewConcurrencyService(concurrencyCache),
- }
- selection, decision, err := svc.SelectAccountWithScheduler(ctx, &groupID, "", "session_subscription_first", "gpt-5.1", nil, OpenAIUpstreamTransportAny, false)
- require.NoError(t, err)
- require.NotNil(t, selection)
- require.NotNil(t, selection.Account)
- require.Equal(t, int64(21601), selection.Account.ID)
- require.Equal(t, openAIAccountScheduleLayerLoadBalance, decision.Layer)
- require.Equal(t, 1, decision.TopK)
- if selection.ReleaseFunc != nil {
- selection.ReleaseFunc()
- }
- }
- func TestOpenAIGatewayService_SelectAccountWithScheduler_SubscriptionPriorityFallsBackWhenSubscriptionFull(t *testing.T) {
- ctx := context.Background()
- groupID := int64(10121)
- accounts := []Account{
- {
- ID: 21611,
- Platform: PlatformOpenAI,
- Type: AccountTypeOAuth,
- Status: StatusActive,
- Schedulable: true,
- Concurrency: 1,
- Priority: 0,
- GroupIDs: []int64{groupID},
- Credentials: map[string]any{"plan_type": "team"},
- },
- {
- ID: 21612,
- Platform: PlatformOpenAI,
- Type: AccountTypeAPIKey,
- Status: StatusActive,
- Schedulable: true,
- Concurrency: 1,
- Priority: 9,
- GroupIDs: []int64{groupID},
- },
- }
- concurrencyCache := schedulerTestConcurrencyCache{
- acquireResults: map[int64]bool{21611: false, 21612: true},
- loadMap: map[int64]*AccountLoadInfo{
- 21611: {AccountID: 21611, LoadRate: 0, WaitingCount: 0},
- 21612: {AccountID: 21612, LoadRate: 90, WaitingCount: 1},
- },
- }
- svc := &OpenAIGatewayService{
- accountRepo: schedulerTestOpenAIAccountRepo{accounts: accounts},
- cache: &schedulerTestGatewayCache{},
- cfg: newSchedulerTestSubscriptionPriorityConfig(),
- rateLimitService: newOpenAIAdvancedSchedulerRateLimitService("true", "", "true"),
- concurrencyService: NewConcurrencyService(concurrencyCache),
- }
- selection, decision, err := svc.SelectAccountWithScheduler(ctx, &groupID, "", "session_subscription_fallback", "gpt-5.1", nil, OpenAIUpstreamTransportAny, false)
- require.NoError(t, err)
- require.NotNil(t, selection)
- require.NotNil(t, selection.Account)
- require.Equal(t, int64(21612), selection.Account.ID)
- require.Equal(t, openAIAccountScheduleLayerLoadBalance, decision.Layer)
- require.True(t, selection.Acquired)
- if selection.ReleaseFunc != nil {
- selection.ReleaseFunc()
- }
- }
- func TestOpenAIGatewayService_SelectAccountWithScheduler_SubscriptionPriorityDisabledUsesScore(t *testing.T) {
- ctx := context.Background()
- groupID := int64(10122)
- accounts := []Account{
- {
- ID: 21621,
- Platform: PlatformOpenAI,
- Type: AccountTypeOAuth,
- Status: StatusActive,
- Schedulable: true,
- Concurrency: 1,
- Priority: 10,
- GroupIDs: []int64{groupID},
- Credentials: map[string]any{"plan_type": "pro"},
- },
- {
- ID: 21622,
- Platform: PlatformOpenAI,
- Type: AccountTypeAPIKey,
- Status: StatusActive,
- Schedulable: true,
- Concurrency: 1,
- Priority: 0,
- GroupIDs: []int64{groupID},
- },
- }
- concurrencyCache := schedulerTestConcurrencyCache{
- acquireResults: map[int64]bool{21621: true, 21622: true},
- loadMap: map[int64]*AccountLoadInfo{
- 21621: {AccountID: 21621, LoadRate: 90, WaitingCount: 1},
- 21622: {AccountID: 21622, LoadRate: 0, WaitingCount: 0},
- },
- }
- svc := &OpenAIGatewayService{
- accountRepo: schedulerTestOpenAIAccountRepo{accounts: accounts},
- cache: &schedulerTestGatewayCache{},
- cfg: newSchedulerTestSubscriptionPriorityConfig(),
- rateLimitService: newOpenAIAdvancedSchedulerRateLimitService("true", "", "false"),
- concurrencyService: NewConcurrencyService(concurrencyCache),
- }
- selection, decision, err := svc.SelectAccountWithScheduler(ctx, &groupID, "", "session_subscription_disabled", "gpt-5.1", nil, OpenAIUpstreamTransportAny, false)
- require.NoError(t, err)
- require.NotNil(t, selection)
- require.NotNil(t, selection.Account)
- require.Equal(t, int64(21622), selection.Account.ID)
- require.Equal(t, openAIAccountScheduleLayerLoadBalance, decision.Layer)
- if selection.ReleaseFunc != nil {
- selection.ReleaseFunc()
- }
- }
- func TestOpenAIGatewayService_SelectAccountWithScheduler_UsesAccountPriorityWithinGroupPool(t *testing.T) {
- ctx := context.Background()
- groupID := int64(10123)
- accounts := []Account{
- {
- ID: 21631,
- Platform: PlatformOpenAI,
- Type: AccountTypeAPIKey,
- Status: StatusActive,
- Schedulable: true,
- Concurrency: 1,
- Priority: 1,
- AccountGroups: []AccountGroup{
- {AccountID: 21631, GroupID: groupID, Priority: 100},
- },
- GroupIDs: []int64{groupID},
- },
- {
- ID: 21632,
- Platform: PlatformOpenAI,
- Type: AccountTypeAPIKey,
- Status: StatusActive,
- Schedulable: true,
- Concurrency: 1,
- Priority: 100000,
- AccountGroups: []AccountGroup{
- {AccountID: 21632, GroupID: groupID, Priority: 1},
- },
- GroupIDs: []int64{groupID},
- },
- }
- cfg := newSchedulerTestSubscriptionPriorityConfig()
- cfg.Gateway.OpenAIWS.SchedulerScoreWeights.Load = 0
- cfg.Gateway.OpenAIWS.SchedulerScoreWeights.Queue = 0
- svc := &OpenAIGatewayService{
- accountRepo: schedulerGroupAwareOpenAIAccountRepo{schedulerTestOpenAIAccountRepo{accounts: accounts}},
- cache: &schedulerTestGatewayCache{},
- cfg: cfg,
- rateLimitService: newOpenAIAdvancedSchedulerRateLimitService("true"),
- concurrencyService: NewConcurrencyService(schedulerTestConcurrencyCache{}),
- }
- selection, decision, err := svc.SelectAccountWithScheduler(ctx, &groupID, "", "session_group_priority", "gpt-5.1", nil, OpenAIUpstreamTransportAny, false)
- require.NoError(t, err)
- require.NotNil(t, selection)
- require.NotNil(t, selection.Account)
- require.Equal(t, int64(21631), selection.Account.ID)
- require.Equal(t, openAIAccountScheduleLayerLoadBalance, decision.Layer)
- if selection.ReleaseFunc != nil {
- selection.ReleaseFunc()
- }
- }
- func TestOpenAIAccountScheduler_SkipsAccountBlockedForRequestedModel(t *testing.T) {
- now := time.Now()
- account := &Account{ID: 21633, Platform: PlatformOpenAI, Type: AccountTypeAPIKey}
- svc := &OpenAIGatewayService{openaiModelTransient: newOpenAIAccountModelTransientState(128)}
- svc.openaiModelTransient.recordFailure(account.ID, "gpt-5.5", now)
- svc.openaiModelTransient.recordFailure(account.ID, "gpt-5.5", now.Add(time.Millisecond))
- scheduler := &defaultOpenAIAccountScheduler{service: svc}
- require.False(t, scheduler.isAccountRequestCompatible(context.Background(), account, OpenAIAccountScheduleRequest{RequestedModel: "gpt-5.5"}))
- require.True(t, scheduler.isAccountRequestCompatible(context.Background(), account, OpenAIAccountScheduleRequest{RequestedModel: "gpt-5.6-sol"}))
- }
- func TestReportOpenAIAccountScheduleResult_SuccessClearsModelTransientState(t *testing.T) {
- svc := &OpenAIGatewayService{openaiModelTransient: newOpenAIAccountModelTransientState(128)}
- now := time.Now()
- svc.openaiModelTransient.recordFailure(21636, "gpt-5.5", now)
- svc.openaiModelTransient.recordFailure(21636, "gpt-5.5", now.Add(time.Millisecond))
- require.True(t, svc.openaiModelTransient.isBlocked(21636, "gpt-5.5", now.Add(2*time.Millisecond)))
- svc.ReportOpenAIAccountScheduleResult(&Account{ID: 21636}, "gpt-5.5", true, nil)
- require.False(t, svc.openaiModelTransient.isBlocked(21636, "gpt-5.5", now.Add(2*time.Millisecond)))
- }
- func TestDefaultOpenAIAccountScheduler_ShouldEscapeStickyAccount_ThresholdBoundary(t *testing.T) {
- stats := newOpenAIAccountRuntimeStats()
- accountID := int64(21501)
- ttft := 15000
- stats.report(accountID, true, &ttft)
- stats.report(accountID, false, nil)
- stats.report(accountID, true, nil)
- scheduler := &defaultOpenAIAccountScheduler{stats: stats}
- reason, errorRate, observedTTFT, shouldEscape := scheduler.shouldEscapeStickyAccount(accountID, openAIStickyEscapeConfig{
- enabled: true,
- ttftMs: 15000,
- errorRate: 0.5,
- })
- require.False(t, shouldEscape)
- require.Empty(t, reason)
- require.InDelta(t, 0.16, errorRate, 1e-9)
- require.InDelta(t, 15000, observedTTFT, 1e-9)
- for i := 0; i < 4; i++ {
- stats.report(accountID, false, nil)
- }
- reason, errorRate, _, shouldEscape = scheduler.shouldEscapeStickyAccount(accountID, openAIStickyEscapeConfig{
- enabled: true,
- ttftMs: 15000,
- errorRate: 1,
- })
- require.False(t, shouldEscape)
- require.Empty(t, reason)
- reason, errorRate, observedTTFT, shouldEscape = scheduler.shouldEscapeStickyAccount(accountID, openAIStickyEscapeConfig{
- enabled: true,
- ttftMs: 15000,
- errorRate: errorRate,
- })
- require.False(t, shouldEscape)
- require.Empty(t, reason)
- require.InDelta(t, 0.655936, errorRate, 1e-9)
- require.InDelta(t, 15000, observedTTFT, 1e-9)
- }
- func TestOpenAIGatewayService_SelectAccountWithScheduler_SessionSticky_ForceHTTP(t *testing.T) {
- ctx := context.Background()
- groupID := int64(1010)
- account := Account{
- ID: 2101,
- Platform: PlatformOpenAI,
- Type: AccountTypeOAuth,
- Status: StatusActive,
- Schedulable: true,
- Concurrency: 1,
- GroupIDs: []int64{groupID},
- Extra: map[string]any{
- "openai_ws_force_http": true,
- },
- }
- cache := &schedulerTestGatewayCache{
- sessionBindings: map[string]int64{
- "openai:session_hash_force_http": account.ID,
- },
- }
- svc := &OpenAIGatewayService{
- accountRepo: schedulerTestOpenAIAccountRepo{accounts: []Account{account}},
- cache: cache,
- cfg: &config.Config{},
- rateLimitService: newOpenAIAdvancedSchedulerRateLimitService("true"),
- concurrencyService: NewConcurrencyService(schedulerTestConcurrencyCache{}),
- }
- selection, decision, err := svc.SelectAccountWithScheduler(
- ctx,
- &groupID,
- "",
- "session_hash_force_http",
- "gpt-5.1",
- nil,
- OpenAIUpstreamTransportAny,
- false,
- )
- require.NoError(t, err)
- require.NotNil(t, selection)
- require.NotNil(t, selection.Account)
- require.Equal(t, account.ID, selection.Account.ID)
- require.Equal(t, openAIAccountScheduleLayerSessionSticky, decision.Layer)
- require.True(t, decision.StickySessionHit)
- if selection.ReleaseFunc != nil {
- selection.ReleaseFunc()
- }
- }
- func TestOpenAIGatewayService_SelectAccountWithScheduler_RequiredWSV2_SkipsStickyHTTPAccount(t *testing.T) {
- ctx := context.Background()
- groupID := int64(1011)
- accounts := []Account{
- {
- ID: 2201,
- Platform: PlatformOpenAI,
- Type: AccountTypeAPIKey,
- Status: StatusActive,
- Schedulable: true,
- Concurrency: 1,
- Priority: 0,
- GroupIDs: []int64{groupID},
- },
- {
- ID: 2202,
- Platform: PlatformOpenAI,
- Type: AccountTypeAPIKey,
- Status: StatusActive,
- Schedulable: true,
- Concurrency: 1,
- Priority: 5,
- GroupIDs: []int64{groupID},
- Extra: map[string]any{
- "openai_apikey_responses_websockets_v2_enabled": true,
- },
- },
- }
- cache := &schedulerTestGatewayCache{
- sessionBindings: map[string]int64{
- "openai:session_hash_ws_only": 2201,
- },
- }
- cfg := newSchedulerTestOpenAIWSV2Config()
- // 构造“HTTP-only 账号负载更低”的场景,验证 required transport 会强制过滤。
- concurrencyCache := schedulerTestConcurrencyCache{
- loadMap: map[int64]*AccountLoadInfo{
- 2201: {AccountID: 2201, LoadRate: 0, WaitingCount: 0},
- 2202: {AccountID: 2202, LoadRate: 90, WaitingCount: 5},
- },
- }
- svc := &OpenAIGatewayService{
- accountRepo: schedulerTestOpenAIAccountRepo{accounts: accounts},
- cache: cache,
- cfg: cfg,
- rateLimitService: newOpenAIAdvancedSchedulerRateLimitService("true"),
- concurrencyService: NewConcurrencyService(concurrencyCache),
- }
- selection, decision, err := svc.SelectAccountWithScheduler(
- ctx,
- &groupID,
- "",
- "session_hash_ws_only",
- "gpt-5.1",
- nil,
- OpenAIUpstreamTransportResponsesWebsocketV2,
- false,
- )
- require.NoError(t, err)
- require.NotNil(t, selection)
- require.NotNil(t, selection.Account)
- require.Equal(t, int64(2202), selection.Account.ID)
- require.Equal(t, openAIAccountScheduleLayerLoadBalance, decision.Layer)
- require.False(t, decision.StickySessionHit)
- require.Equal(t, 1, decision.CandidateCount)
- if selection.ReleaseFunc != nil {
- selection.ReleaseFunc()
- }
- }
- func TestOpenAIGatewayService_SelectAccountWithScheduler_ClearsStickyAccountOutsideGroup(t *testing.T) {
- ctx := context.Background()
- groupID := int64(1013)
- accounts := []Account{
- {
- ID: 2401,
- Platform: PlatformOpenAI,
- Type: AccountTypeAPIKey,
- Status: StatusActive,
- Schedulable: true,
- Concurrency: 1,
- Priority: 0,
- },
- {
- ID: 2402,
- Platform: PlatformOpenAI,
- Type: AccountTypeAPIKey,
- Status: StatusActive,
- Schedulable: true,
- Concurrency: 1,
- Priority: 5,
- AccountGroups: []AccountGroup{
- {AccountID: 2402, GroupID: groupID},
- },
- },
- }
- cache := &schedulerTestGatewayCache{
- sessionBindings: map[string]int64{
- "openai:session_hash_removed_group": 2401,
- },
- }
- svc := &OpenAIGatewayService{
- accountRepo: schedulerGroupAwareOpenAIAccountRepo{schedulerTestOpenAIAccountRepo{accounts: accounts}},
- cache: cache,
- cfg: &config.Config{},
- rateLimitService: newOpenAIAdvancedSchedulerRateLimitService("true"),
- concurrencyService: NewConcurrencyService(schedulerTestConcurrencyCache{}),
- }
- selection, decision, err := svc.SelectAccountWithScheduler(
- ctx,
- &groupID,
- "",
- "session_hash_removed_group",
- "gpt-5.1",
- nil,
- OpenAIUpstreamTransportAny,
- false,
- )
- require.NoError(t, err)
- require.NotNil(t, selection)
- require.NotNil(t, selection.Account)
- require.Equal(t, int64(2402), selection.Account.ID)
- require.Equal(t, openAIAccountScheduleLayerLoadBalance, decision.Layer)
- require.False(t, decision.StickySessionHit)
- require.Equal(t, 1, cache.deletedSessions["openai:session_hash_removed_group"])
- require.Equal(t, int64(2402), cache.sessionBindings["openai:session_hash_removed_group"])
- if selection.ReleaseFunc != nil {
- selection.ReleaseFunc()
- }
- }
- func TestOpenAIGatewayService_SelectAccountWithScheduler_RequiredWSV2_NoAvailableAccount(t *testing.T) {
- ctx := context.Background()
- groupID := int64(1012)
- accounts := []Account{
- {
- ID: 2301,
- Platform: PlatformOpenAI,
- Type: AccountTypeOAuth,
- Status: StatusActive,
- Schedulable: true,
- Concurrency: 1,
- },
- }
- svc := &OpenAIGatewayService{
- accountRepo: schedulerTestOpenAIAccountRepo{accounts: accounts},
- cache: &schedulerTestGatewayCache{},
- cfg: newSchedulerTestOpenAIWSV2Config(),
- rateLimitService: newOpenAIAdvancedSchedulerRateLimitService("true"),
- concurrencyService: NewConcurrencyService(schedulerTestConcurrencyCache{}),
- }
- selection, decision, err := svc.SelectAccountWithScheduler(
- ctx,
- &groupID,
- "",
- "",
- "gpt-5.1",
- nil,
- OpenAIUpstreamTransportResponsesWebsocketV2,
- false,
- )
- require.Error(t, err)
- require.Nil(t, selection)
- require.Equal(t, openAIAccountScheduleLayerLoadBalance, decision.Layer)
- require.Equal(t, 0, decision.CandidateCount)
- }
- func TestOpenAIGatewayService_SelectAccountWithScheduler_LoadBalanceTopKFallback(t *testing.T) {
- ctx := context.Background()
- groupID := int64(11)
- accounts := []Account{
- {
- ID: 3001,
- Platform: PlatformOpenAI,
- Type: AccountTypeAPIKey,
- Status: StatusActive,
- Schedulable: true,
- Concurrency: 1,
- Priority: 0,
- },
- {
- ID: 3002,
- Platform: PlatformOpenAI,
- Type: AccountTypeAPIKey,
- Status: StatusActive,
- Schedulable: true,
- Concurrency: 1,
- Priority: 0,
- },
- {
- ID: 3003,
- Platform: PlatformOpenAI,
- Type: AccountTypeAPIKey,
- Status: StatusActive,
- Schedulable: true,
- Concurrency: 1,
- Priority: 0,
- },
- }
- cfg := &config.Config{}
- cfg.Gateway.OpenAIWS.LBTopK = 2
- cfg.Gateway.OpenAIWS.SchedulerScoreWeights.Priority = 0.4
- cfg.Gateway.OpenAIWS.SchedulerScoreWeights.Load = 1.0
- cfg.Gateway.OpenAIWS.SchedulerScoreWeights.Queue = 1.0
- cfg.Gateway.OpenAIWS.SchedulerScoreWeights.ErrorRate = 0.2
- cfg.Gateway.OpenAIWS.SchedulerScoreWeights.TTFT = 0.1
- concurrencyCache := schedulerTestConcurrencyCache{
- loadMap: map[int64]*AccountLoadInfo{
- 3001: {AccountID: 3001, LoadRate: 95, WaitingCount: 8},
- 3002: {AccountID: 3002, LoadRate: 20, WaitingCount: 1},
- 3003: {AccountID: 3003, LoadRate: 10, WaitingCount: 0},
- },
- acquireResults: map[int64]bool{
- 3003: false, // top1 失败,必须回退到 top-K 的下一候选
- 3002: true,
- },
- }
- svc := &OpenAIGatewayService{
- accountRepo: schedulerTestOpenAIAccountRepo{accounts: accounts},
- cache: &schedulerTestGatewayCache{},
- cfg: cfg,
- rateLimitService: newOpenAIAdvancedSchedulerRateLimitService("true"),
- concurrencyService: NewConcurrencyService(concurrencyCache),
- }
- selection, decision, err := svc.SelectAccountWithScheduler(
- ctx,
- &groupID,
- "",
- "",
- "gpt-5.1",
- nil,
- OpenAIUpstreamTransportAny,
- false,
- )
- require.NoError(t, err)
- require.NotNil(t, selection)
- require.NotNil(t, selection.Account)
- require.Equal(t, int64(3002), selection.Account.ID)
- require.Equal(t, openAIAccountScheduleLayerLoadBalance, decision.Layer)
- require.Equal(t, 3, decision.CandidateCount)
- require.Equal(t, 2, decision.TopK)
- require.Greater(t, decision.LoadSkew, 0.0)
- if selection.ReleaseFunc != nil {
- selection.ReleaseFunc()
- }
- }
- // Regression: TopK initial filter must drop quota-auto-paused accounts. Otherwise
- // the candidate pool is filled with paused accounts, healthy accounts fall outside
- // TopK, and the scheduler returns "no available accounts" even though healthy ones
- // exist.
- func TestOpenAIGatewayService_SelectAccountWithScheduler_LoadBalanceTopKExcludesQuotaPaused(t *testing.T) {
- ctx := context.Background()
- groupID := int64(110)
- accounts := []Account{
- {
- ID: 37001,
- Platform: PlatformOpenAI,
- Type: AccountTypeAPIKey,
- Status: StatusActive,
- Schedulable: true,
- Concurrency: 1,
- Priority: 0,
- Extra: map[string]any{
- "codex_5h_used_percent": 96.0,
- "auto_pause_5h_threshold": 0.95,
- },
- },
- {
- ID: 37002,
- Platform: PlatformOpenAI,
- Type: AccountTypeAPIKey,
- Status: StatusActive,
- Schedulable: true,
- Concurrency: 1,
- Priority: 5,
- },
- }
- cfg := &config.Config{}
- cfg.Gateway.OpenAIWS.LBTopK = 1 // TopK=1 makes the bug fatal: paused account would crowd out the healthy one entirely
- cfg.Gateway.OpenAIWS.SchedulerScoreWeights.Priority = 0.4
- cfg.Gateway.OpenAIWS.SchedulerScoreWeights.Load = 1.0
- cfg.Gateway.OpenAIWS.SchedulerScoreWeights.Queue = 1.0
- concurrencyCache := schedulerTestConcurrencyCache{
- loadMap: map[int64]*AccountLoadInfo{
- 37001: {AccountID: 37001, LoadRate: 5, WaitingCount: 0},
- 37002: {AccountID: 37002, LoadRate: 5, WaitingCount: 0},
- },
- acquireResults: map[int64]bool{
- 37002: true,
- },
- }
- svc := &OpenAIGatewayService{
- accountRepo: schedulerTestOpenAIAccountRepo{accounts: accounts},
- cache: &schedulerTestGatewayCache{},
- cfg: cfg,
- rateLimitService: newOpenAIAdvancedSchedulerRateLimitService("true"),
- concurrencyService: NewConcurrencyService(concurrencyCache),
- }
- selection, decision, err := svc.SelectAccountWithScheduler(
- ctx,
- &groupID,
- "",
- "",
- "gpt-5.1",
- nil,
- OpenAIUpstreamTransportAny,
- false,
- )
- require.NoError(t, err)
- require.NotNil(t, selection)
- require.NotNil(t, selection.Account)
- require.Equal(t, int64(37002), selection.Account.ID)
- require.Equal(t, openAIAccountScheduleLayerLoadBalance, decision.Layer)
- // Only the healthy account should ever enter the candidate pool; the paused one
- // must be filtered out at the initial-filter stage.
- require.Equal(t, 1, decision.CandidateCount)
- if selection.ReleaseFunc != nil {
- selection.ReleaseFunc()
- }
- }
- func TestOpenAIGatewayService_OpenAIAccountSchedulerMetrics(t *testing.T) {
- ctx := context.Background()
- groupID := int64(12)
- account := Account{
- ID: 4001,
- Platform: PlatformOpenAI,
- Type: AccountTypeAPIKey,
- Status: StatusActive,
- Schedulable: true,
- Concurrency: 1,
- GroupIDs: []int64{groupID},
- }
- cache := &schedulerTestGatewayCache{
- sessionBindings: map[string]int64{
- "openai:session_hash_metrics": account.ID,
- },
- }
- svc := &OpenAIGatewayService{
- accountRepo: schedulerTestOpenAIAccountRepo{accounts: []Account{account}},
- cache: cache,
- cfg: &config.Config{},
- rateLimitService: newOpenAIAdvancedSchedulerRateLimitService("true"),
- concurrencyService: NewConcurrencyService(schedulerTestConcurrencyCache{}),
- }
- selection, _, err := svc.SelectAccountWithScheduler(ctx, &groupID, "", "session_hash_metrics", "gpt-5.1", nil, OpenAIUpstreamTransportAny, false)
- require.NoError(t, err)
- require.NotNil(t, selection)
- svc.ReportOpenAIAccountScheduleResult(&account, "", true, intPtrForTest(120))
- svc.RecordOpenAIAccountSwitch()
- snapshot := svc.SnapshotOpenAIAccountSchedulerMetrics()
- require.GreaterOrEqual(t, snapshot.SelectTotal, int64(1))
- require.GreaterOrEqual(t, snapshot.StickySessionHitTotal, int64(1))
- require.GreaterOrEqual(t, snapshot.AccountSwitchTotal, int64(1))
- require.GreaterOrEqual(t, snapshot.SchedulerLatencyMsAvg, float64(0))
- require.GreaterOrEqual(t, snapshot.StickyHitRatio, 0.0)
- require.GreaterOrEqual(t, snapshot.RuntimeStatsAccountCount, 1)
- }
- func intPtrForTest(v int) *int {
- return &v
- }
- func TestOpenAIAccountRuntimeStats_ReportAndSnapshot(t *testing.T) {
- stats := newOpenAIAccountRuntimeStats()
- stats.report(1001, true, nil)
- firstTTFT := 100
- stats.report(1001, false, &firstTTFT)
- secondTTFT := 200
- stats.report(1001, false, &secondTTFT)
- errorRate, ttft, hasTTFT := stats.snapshot(1001)
- require.True(t, hasTTFT)
- require.InDelta(t, 0.36, errorRate, 1e-9)
- require.InDelta(t, 120.0, ttft, 1e-9)
- require.Equal(t, 1, stats.size())
- }
- func TestOpenAIAccountRuntimeStats_ReportConcurrent(t *testing.T) {
- stats := newOpenAIAccountRuntimeStats()
- const (
- accountCount = 4
- workers = 16
- iterations = 800
- )
- var wg sync.WaitGroup
- wg.Add(workers)
- for worker := 0; worker < workers; worker++ {
- worker := worker
- go func() {
- defer wg.Done()
- for i := 0; i < iterations; i++ {
- accountID := int64(i%accountCount + 1)
- success := (i+worker)%3 != 0
- ttft := 80 + (i+worker)%40
- stats.report(accountID, success, &ttft)
- }
- }()
- }
- wg.Wait()
- require.Equal(t, accountCount, stats.size())
- for accountID := int64(1); accountID <= accountCount; accountID++ {
- errorRate, ttft, hasTTFT := stats.snapshot(accountID)
- require.GreaterOrEqual(t, errorRate, 0.0)
- require.LessOrEqual(t, errorRate, 1.0)
- require.True(t, hasTTFT)
- require.Greater(t, ttft, 0.0)
- }
- }
- func TestSelectTopKOpenAICandidates(t *testing.T) {
- candidates := []openAIAccountCandidateScore{
- {
- account: &Account{ID: 11, Priority: 2},
- loadInfo: &AccountLoadInfo{LoadRate: 10, WaitingCount: 1},
- score: 10.0,
- },
- {
- account: &Account{ID: 12, Priority: 1},
- loadInfo: &AccountLoadInfo{LoadRate: 20, WaitingCount: 1},
- score: 9.5,
- },
- {
- account: &Account{ID: 13, Priority: 1},
- loadInfo: &AccountLoadInfo{LoadRate: 30, WaitingCount: 0},
- score: 10.0,
- },
- {
- account: &Account{ID: 14, Priority: 0},
- loadInfo: &AccountLoadInfo{LoadRate: 40, WaitingCount: 0},
- score: 8.0,
- },
- }
- top2 := selectTopKOpenAICandidates(candidates, 2)
- require.Len(t, top2, 2)
- require.Equal(t, int64(13), top2[0].account.ID)
- require.Equal(t, int64(11), top2[1].account.ID)
- topAll := selectTopKOpenAICandidates(candidates, 8)
- require.Len(t, topAll, len(candidates))
- require.Equal(t, int64(13), topAll[0].account.ID)
- require.Equal(t, int64(11), topAll[1].account.ID)
- require.Equal(t, int64(12), topAll[2].account.ID)
- require.Equal(t, int64(14), topAll[3].account.ID)
- }
- func TestBuildOpenAIWeightedSelectionOrder_DeterministicBySessionSeed(t *testing.T) {
- candidates := []openAIAccountCandidateScore{
- {
- account: &Account{ID: 101},
- loadInfo: &AccountLoadInfo{LoadRate: 10, WaitingCount: 0},
- score: 4.2,
- },
- {
- account: &Account{ID: 102},
- loadInfo: &AccountLoadInfo{LoadRate: 30, WaitingCount: 1},
- score: 3.5,
- },
- {
- account: &Account{ID: 103},
- loadInfo: &AccountLoadInfo{LoadRate: 50, WaitingCount: 2},
- score: 2.1,
- },
- }
- req := OpenAIAccountScheduleRequest{
- GroupID: int64PtrForTest(99),
- SessionHash: "session_seed_fixed",
- RequestedModel: "gpt-5.1",
- }
- first := buildOpenAIWeightedSelectionOrder(candidates, req)
- second := buildOpenAIWeightedSelectionOrder(candidates, req)
- require.Len(t, first, len(candidates))
- require.Len(t, second, len(candidates))
- for i := range first {
- require.Equal(t, first[i].account.ID, second[i].account.ID)
- }
- }
- func TestOpenAIGatewayService_SelectAccountWithScheduler_LoadBalanceDistributesAcrossSessions(t *testing.T) {
- ctx := context.Background()
- groupID := int64(15)
- accounts := []Account{
- {
- ID: 5101,
- Platform: PlatformOpenAI,
- Type: AccountTypeAPIKey,
- Status: StatusActive,
- Schedulable: true,
- Concurrency: 3,
- Priority: 0,
- },
- {
- ID: 5102,
- Platform: PlatformOpenAI,
- Type: AccountTypeAPIKey,
- Status: StatusActive,
- Schedulable: true,
- Concurrency: 3,
- Priority: 0,
- },
- {
- ID: 5103,
- Platform: PlatformOpenAI,
- Type: AccountTypeAPIKey,
- Status: StatusActive,
- Schedulable: true,
- Concurrency: 3,
- Priority: 0,
- },
- }
- cfg := &config.Config{}
- cfg.Gateway.OpenAIWS.LBTopK = 3
- cfg.Gateway.OpenAIWS.SchedulerScoreWeights.Priority = 1
- cfg.Gateway.OpenAIWS.SchedulerScoreWeights.Load = 1
- cfg.Gateway.OpenAIWS.SchedulerScoreWeights.Queue = 1
- cfg.Gateway.OpenAIWS.SchedulerScoreWeights.ErrorRate = 1
- cfg.Gateway.OpenAIWS.SchedulerScoreWeights.TTFT = 1
- concurrencyCache := schedulerTestConcurrencyCache{
- loadMap: map[int64]*AccountLoadInfo{
- 5101: {AccountID: 5101, LoadRate: 20, WaitingCount: 1},
- 5102: {AccountID: 5102, LoadRate: 20, WaitingCount: 1},
- 5103: {AccountID: 5103, LoadRate: 20, WaitingCount: 1},
- },
- }
- svc := &OpenAIGatewayService{
- accountRepo: schedulerTestOpenAIAccountRepo{accounts: accounts},
- cache: &schedulerTestGatewayCache{sessionBindings: map[string]int64{}},
- cfg: cfg,
- rateLimitService: newOpenAIAdvancedSchedulerRateLimitService("true"),
- concurrencyService: NewConcurrencyService(concurrencyCache),
- }
- selected := make(map[int64]int, len(accounts))
- for i := 0; i < 60; i++ {
- sessionHash := fmt.Sprintf("session_hash_lb_%d", i)
- selection, decision, err := svc.SelectAccountWithScheduler(
- ctx,
- &groupID,
- "",
- sessionHash,
- "gpt-5.1",
- nil,
- OpenAIUpstreamTransportAny,
- false,
- )
- require.NoError(t, err)
- require.NotNil(t, selection)
- require.NotNil(t, selection.Account)
- require.Equal(t, openAIAccountScheduleLayerLoadBalance, decision.Layer)
- selected[selection.Account.ID]++
- if selection.ReleaseFunc != nil {
- selection.ReleaseFunc()
- }
- }
- // 多 session 应该能打散到多个账号,避免“恒定单账号命中”。
- require.GreaterOrEqual(t, len(selected), 2)
- }
- func TestDeriveOpenAISelectionSeed_NoAffinityAddsEntropy(t *testing.T) {
- req := OpenAIAccountScheduleRequest{
- RequestedModel: "gpt-5.1",
- }
- seed1 := deriveOpenAISelectionSeed(req)
- time.Sleep(1 * time.Millisecond)
- seed2 := deriveOpenAISelectionSeed(req)
- require.NotZero(t, seed1)
- require.NotZero(t, seed2)
- require.NotEqual(t, seed1, seed2)
- }
- func TestBuildOpenAIWeightedSelectionOrder_HandlesInvalidScores(t *testing.T) {
- candidates := []openAIAccountCandidateScore{
- {
- account: &Account{ID: 901},
- loadInfo: &AccountLoadInfo{LoadRate: 5, WaitingCount: 0},
- score: math.NaN(),
- },
- {
- account: &Account{ID: 902},
- loadInfo: &AccountLoadInfo{LoadRate: 5, WaitingCount: 0},
- score: math.Inf(1),
- },
- {
- account: &Account{ID: 903},
- loadInfo: &AccountLoadInfo{LoadRate: 5, WaitingCount: 0},
- score: -1,
- },
- }
- req := OpenAIAccountScheduleRequest{
- SessionHash: "seed_invalid_scores",
- }
- order := buildOpenAIWeightedSelectionOrder(candidates, req)
- require.Len(t, order, len(candidates))
- seen := map[int64]struct{}{}
- for _, item := range order {
- seen[item.account.ID] = struct{}{}
- }
- require.Len(t, seen, len(candidates))
- }
- func TestOpenAISelectionRNG_SeedZeroStillWorks(t *testing.T) {
- rng := newOpenAISelectionRNG(0)
- v1 := rng.nextUint64()
- v2 := rng.nextUint64()
- require.NotEqual(t, v1, v2)
- require.GreaterOrEqual(t, rng.nextFloat64(), 0.0)
- require.Less(t, rng.nextFloat64(), 1.0)
- }
- func TestOpenAIAccountCandidateHeap_PushPopAndInvalidType(t *testing.T) {
- h := openAIAccountCandidateHeap{}
- h.Push(openAIAccountCandidateScore{
- account: &Account{ID: 7001},
- loadInfo: &AccountLoadInfo{LoadRate: 0, WaitingCount: 0},
- score: 1.0,
- })
- require.Equal(t, 1, h.Len())
- popped, ok := h.Pop().(openAIAccountCandidateScore)
- require.True(t, ok)
- require.Equal(t, int64(7001), popped.account.ID)
- require.Equal(t, 0, h.Len())
- require.Panics(t, func() {
- h.Push("bad_element_type")
- })
- }
- func TestClamp01_AllBranches(t *testing.T) {
- require.Equal(t, 0.0, clamp01(-0.2))
- require.Equal(t, 1.0, clamp01(1.3))
- require.Equal(t, 0.5, clamp01(0.5))
- }
- func TestCalcLoadSkewByMoments_Branches(t *testing.T) {
- require.Equal(t, 0.0, calcLoadSkewByMoments(1, 1, 1))
- // variance < 0 分支:sumSquares/count - mean^2 为负值时应钳制为 0。
- require.Equal(t, 0.0, calcLoadSkewByMoments(1, 0, 2))
- require.GreaterOrEqual(t, calcLoadSkewByMoments(6, 20, 3), 0.0)
- }
- func TestDefaultOpenAIAccountScheduler_ReportSwitchAndSnapshot(t *testing.T) {
- schedulerAny := newDefaultOpenAIAccountScheduler(&OpenAIGatewayService{}, nil)
- scheduler, ok := schedulerAny.(*defaultOpenAIAccountScheduler)
- require.True(t, ok)
- ttft := 100
- scheduler.ReportResult(1001, true, &ttft)
- scheduler.ReportSwitch()
- scheduler.metrics.recordSelect(OpenAIAccountScheduleDecision{
- Layer: openAIAccountScheduleLayerLoadBalance,
- LatencyMs: 8,
- LoadSkew: 0.5,
- StickyPreviousHit: true,
- })
- scheduler.metrics.recordSelect(OpenAIAccountScheduleDecision{
- Layer: openAIAccountScheduleLayerSessionSticky,
- LatencyMs: 6,
- LoadSkew: 0.2,
- StickySessionHit: true,
- })
- snapshot := scheduler.SnapshotMetrics()
- require.Equal(t, int64(2), snapshot.SelectTotal)
- require.Equal(t, int64(1), snapshot.StickyPreviousHitTotal)
- require.Equal(t, int64(1), snapshot.StickySessionHitTotal)
- require.Equal(t, int64(1), snapshot.LoadBalanceSelectTotal)
- require.Equal(t, int64(1), snapshot.AccountSwitchTotal)
- require.Greater(t, snapshot.SchedulerLatencyMsAvg, 0.0)
- require.Greater(t, snapshot.StickyHitRatio, 0.0)
- require.Greater(t, snapshot.LoadSkewAvg, 0.0)
- }
- func TestOpenAIGatewayService_SchedulerWrappersAndDefaults(t *testing.T) {
- resetOpenAIAdvancedSchedulerSettingCacheForTest()
- svc := &OpenAIGatewayService{}
- ttft := 120
- svc.ReportOpenAIAccountScheduleResult(&Account{ID: 10}, "", true, &ttft)
- svc.RecordOpenAIAccountSwitch()
- snapshot := svc.SnapshotOpenAIAccountSchedulerMetrics()
- require.Equal(t, OpenAIAccountSchedulerMetricsSnapshot{}, snapshot)
- require.Equal(t, 7, svc.openAIWSLBTopK())
- require.Equal(t, openaiStickySessionTTL, svc.openAIWSSessionStickyTTL())
- defaultWeights := svc.openAIWSSchedulerWeights()
- require.Equal(t, 1.0, defaultWeights.Priority)
- require.Equal(t, 1.0, defaultWeights.Load)
- require.Equal(t, 0.7, defaultWeights.Queue)
- require.Equal(t, 0.8, defaultWeights.ErrorRate)
- require.Equal(t, 0.5, defaultWeights.TTFT)
- cfg := &config.Config{}
- cfg.Gateway.OpenAIWS.LBTopK = 9
- cfg.Gateway.OpenAIWS.StickySessionTTLSeconds = 180
- cfg.Gateway.OpenAIWS.SchedulerScoreWeights.Priority = 0.2
- cfg.Gateway.OpenAIWS.SchedulerScoreWeights.Load = 0.3
- cfg.Gateway.OpenAIWS.SchedulerScoreWeights.Queue = 0.4
- cfg.Gateway.OpenAIWS.SchedulerScoreWeights.ErrorRate = 0.5
- cfg.Gateway.OpenAIWS.SchedulerScoreWeights.TTFT = 0.6
- svcWithCfg := &OpenAIGatewayService{cfg: cfg}
- require.Equal(t, 9, svcWithCfg.openAIWSLBTopK())
- require.Equal(t, 180*time.Second, svcWithCfg.openAIWSSessionStickyTTL())
- customWeights := svcWithCfg.openAIWSSchedulerWeights()
- require.Equal(t, 0.2, customWeights.Priority)
- require.Equal(t, 0.3, customWeights.Load)
- require.Equal(t, 0.4, customWeights.Queue)
- require.Equal(t, 0.5, customWeights.ErrorRate)
- require.Equal(t, 0.6, customWeights.TTFT)
- }
- func TestDefaultOpenAIAccountScheduler_IsAccountTransportCompatible_Branches(t *testing.T) {
- scheduler := &defaultOpenAIAccountScheduler{}
- require.True(t, scheduler.isAccountTransportCompatible(nil, OpenAIUpstreamTransportAny))
- require.True(t, scheduler.isAccountTransportCompatible(nil, OpenAIUpstreamTransportHTTPSSE))
- require.False(t, scheduler.isAccountTransportCompatible(nil, OpenAIUpstreamTransportResponsesWebsocketV2))
- require.False(t, scheduler.isAccountTransportCompatible(nil, OpenAIUpstreamTransportResponsesWebsocketV2Ingress))
- cfg := newSchedulerTestOpenAIWSV2Config()
- scheduler.service = &OpenAIGatewayService{cfg: cfg}
- account := &Account{
- ID: 8801,
- Platform: PlatformOpenAI,
- Type: AccountTypeAPIKey,
- Status: StatusActive,
- Schedulable: true,
- Concurrency: 1,
- Extra: map[string]any{
- "openai_apikey_responses_websockets_v2_enabled": true,
- },
- }
- require.True(t, scheduler.isAccountTransportCompatible(account, OpenAIUpstreamTransportResponsesWebsocketV2))
- require.True(t, scheduler.isAccountTransportCompatible(account, OpenAIUpstreamTransportResponsesWebsocketV2Ingress))
- cfg.Gateway.OpenAIWS.ModeRouterV2Enabled = true
- account.Extra["openai_apikey_responses_websockets_v2_mode"] = OpenAIWSIngressModeHTTPBridge
- require.False(t, scheduler.isAccountTransportCompatible(account, OpenAIUpstreamTransportResponsesWebsocketV2))
- require.True(t, scheduler.isAccountTransportCompatible(account, OpenAIUpstreamTransportResponsesWebsocketV2Ingress))
- account.Extra["openai_apikey_responses_websockets_v2_mode"] = OpenAIWSIngressModeOff
- require.False(t, scheduler.isAccountTransportCompatible(account, OpenAIUpstreamTransportResponsesWebsocketV2Ingress))
- }
- func int64PtrForTest(v int64) *int64 {
- return &v
- }
- func TestOpenAIGatewayService_SelectAccountWithScheduler_StickyWeightedFallbackSkipsOutOfGroupStickyAccount(t *testing.T) {
- resetOpenAIAdvancedSchedulerSettingCacheForTest()
- ctx := context.Background()
- groupID := int64(101081)
- otherGroupID := int64(101082)
- accounts := []Account{
- {
- ID: 38001,
- Platform: PlatformOpenAI,
- Type: AccountTypeAPIKey,
- Status: StatusActive,
- Schedulable: true,
- Concurrency: 1,
- Priority: 10,
- GroupIDs: []int64{groupID},
- },
- {
- // 会话粘连绑定指向的账号已被移出请求分组(绑定 TTL 内账号改组的场景)。
- ID: 38002,
- Platform: PlatformOpenAI,
- Type: AccountTypeAPIKey,
- Status: StatusActive,
- Schedulable: true,
- Concurrency: 1,
- Priority: 0,
- GroupIDs: []int64{otherGroupID},
- },
- }
- cfg := &config.Config{}
- cfg.Gateway.OpenAIWS.LBTopK = 2
- cfg.Gateway.OpenAIWS.SchedulerScoreWeights.Priority = 1
- cfg.Gateway.OpenAIWS.SchedulerScoreWeights.Load = 1
- cfg.Gateway.OpenAIWS.SchedulerScoreWeights.Queue = 0.7
- cfg.Gateway.OpenAIWS.SchedulerScoreWeights.ErrorRate = 0.8
- cfg.Gateway.OpenAIWS.SchedulerScoreWeights.TTFT = 0.5
- cfg.Gateway.OpenAIWS.SchedulerScoreWeights.SessionSticky = 3
- cache := &schedulerTestGatewayCache{sessionBindings: map[string]int64{
- "openai:session_weighted_out_of_group": 38002,
- }}
- concurrencyCache := schedulerTestConcurrencyCache{
- acquireResults: map[int64]bool{38001: false, 38002: true},
- }
- svc := &OpenAIGatewayService{
- accountRepo: schedulerGroupAwareOpenAIAccountRepo{schedulerTestOpenAIAccountRepo{accounts: accounts}},
- cache: cache,
- cfg: cfg,
- rateLimitService: newOpenAIAdvancedSchedulerRateLimitService("true", "true"),
- concurrencyService: NewConcurrencyService(concurrencyCache),
- }
- selection, decision, err := svc.SelectAccountWithScheduler(
- ctx,
- &groupID,
- "",
- "session_weighted_out_of_group",
- "gpt-5.1",
- nil,
- OpenAIUpstreamTransportAny,
- false,
- )
- require.NoError(t, err)
- require.NotNil(t, selection)
- require.NotNil(t, selection.Account)
- // 组内唯一候选 38001 满并发:必须返回其等待计划,绝不能把请求泄漏到组外的粘连账号 38002。
- require.Equal(t, int64(38001), selection.Account.ID)
- require.False(t, selection.Acquired)
- require.NotNil(t, selection.WaitPlan)
- require.Equal(t, int64(38001), selection.WaitPlan.AccountID)
- require.Equal(t, openAIAccountScheduleLayerLoadBalance, decision.Layer)
- // 失效的粘连绑定应被清理,避免后续请求反复走同一条泄漏路径。
- require.Positive(t, cache.deletedSessions["openai:session_weighted_out_of_group"])
- }
- func TestOpenAIGatewayService_SelectAccountWithScheduler_SubscriptionPriorityWaitsOnBusySubscriptionWhenRegularUnusable(t *testing.T) {
- resetOpenAIAdvancedSchedulerSettingCacheForTest()
- ctx := context.Background()
- groupID := int64(101091)
- accounts := []Account{
- {
- // 订阅账号:支持 compact,但并发已满(busy-but-waitable)。
- ID: 38011,
- Platform: PlatformOpenAI,
- Type: AccountTypeOAuth,
- Status: StatusActive,
- Schedulable: true,
- Concurrency: 1,
- Priority: 0,
- GroupIDs: []int64{groupID},
- Credentials: map[string]any{"plan_type": "team"},
- Extra: map[string]any{"openai_compact_supported": true},
- },
- {
- // 常规账号:明确不支持 compact,无法服务本次请求。
- ID: 38012,
- Platform: PlatformOpenAI,
- Type: AccountTypeAPIKey,
- Status: StatusActive,
- Schedulable: true,
- Concurrency: 1,
- Priority: 9,
- GroupIDs: []int64{groupID},
- Extra: map[string]any{"openai_compact_supported": false},
- },
- }
- concurrencyCache := schedulerTestConcurrencyCache{
- acquireResults: map[int64]bool{38011: false, 38012: true},
- }
- svc := &OpenAIGatewayService{
- accountRepo: schedulerTestOpenAIAccountRepo{accounts: accounts},
- cache: &schedulerTestGatewayCache{},
- cfg: newSchedulerTestSubscriptionPriorityConfig(),
- rateLimitService: newOpenAIAdvancedSchedulerRateLimitService("true", "", "true"),
- concurrencyService: NewConcurrencyService(concurrencyCache),
- }
- selection, decision, err := svc.SelectAccountWithScheduler(
- ctx,
- &groupID,
- "",
- "session_subscription_wait",
- "gpt-5.1",
- nil,
- OpenAIUpstreamTransportAny,
- true,
- )
- // 常规池无可用候选时,忙碌的订阅账号应产生等待计划,而不是直接返回 no available accounts。
- require.NoError(t, err)
- require.NotNil(t, selection)
- require.NotNil(t, selection.Account)
- require.Equal(t, int64(38011), selection.Account.ID)
- require.False(t, selection.Acquired)
- require.NotNil(t, selection.WaitPlan)
- require.Equal(t, int64(38011), selection.WaitPlan.AccountID)
- require.Equal(t, openAIAccountScheduleLayerLoadBalance, decision.Layer)
- }
|