| 123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797798799800801802803804805806807808809810811812813814815816817818819820821822823824825826827828829830831832833834835836837838839840841842843844845846847848849850851852853854855856857858859860861862863864865866867868869870871872873874875876877878879880881882883884885886887888889890891892893894895896897898899900901902903904905906907908909910911912913914915916917918919920921922923924925926927928929930931932933934935936937938939940941942943944945946947948949950951952953954955956957958959960961962963964965966967968969970971972973974975976977978979980981982983984985986987988989990991992993994995996997998999100010011002100310041005100610071008100910101011101210131014101510161017101810191020102110221023102410251026102710281029103010311032103310341035103610371038103910401041104210431044104510461047104810491050105110521053105410551056105710581059106010611062106310641065106610671068106910701071107210731074107510761077107810791080108110821083108410851086108710881089109010911092109310941095109610971098109911001101110211031104110511061107110811091110111111121113111411151116111711181119112011211122112311241125112611271128112911301131113211331134113511361137113811391140114111421143114411451146114711481149115011511152115311541155115611571158115911601161116211631164116511661167116811691170117111721173117411751176117711781179118011811182118311841185118611871188118911901191119211931194119511961197119811991200120112021203120412051206120712081209121012111212121312141215121612171218121912201221122212231224122512261227122812291230123112321233123412351236123712381239124012411242124312441245124612471248124912501251125212531254125512561257125812591260126112621263126412651266126712681269127012711272127312741275127612771278127912801281128212831284128512861287128812891290129112921293129412951296129712981299130013011302130313041305130613071308130913101311131213131314131513161317131813191320132113221323132413251326132713281329133013311332133313341335133613371338133913401341134213431344134513461347134813491350135113521353135413551356135713581359136013611362136313641365136613671368136913701371137213731374137513761377137813791380138113821383138413851386138713881389139013911392139313941395139613971398139914001401140214031404140514061407140814091410141114121413141414151416141714181419142014211422142314241425142614271428142914301431143214331434143514361437143814391440144114421443144414451446144714481449145014511452145314541455145614571458145914601461146214631464146514661467146814691470147114721473147414751476147714781479148014811482148314841485148614871488148914901491149214931494149514961497149814991500150115021503150415051506150715081509151015111512151315141515151615171518151915201521152215231524152515261527152815291530153115321533153415351536153715381539154015411542154315441545154615471548154915501551155215531554155515561557155815591560156115621563156415651566156715681569157015711572157315741575157615771578157915801581158215831584158515861587158815891590159115921593159415951596159715981599160016011602160316041605160616071608160916101611161216131614161516161617161816191620162116221623162416251626162716281629163016311632163316341635163616371638163916401641164216431644164516461647164816491650165116521653165416551656165716581659166016611662166316641665166616671668166916701671167216731674167516761677167816791680168116821683168416851686168716881689169016911692169316941695169616971698169917001701170217031704170517061707170817091710171117121713171417151716171717181719172017211722172317241725172617271728172917301731173217331734173517361737173817391740174117421743174417451746174717481749175017511752175317541755175617571758175917601761176217631764176517661767176817691770177117721773177417751776177717781779178017811782178317841785178617871788178917901791179217931794179517961797179817991800180118021803180418051806180718081809181018111812181318141815181618171818181918201821182218231824182518261827182818291830183118321833183418351836183718381839184018411842184318441845184618471848184918501851185218531854185518561857185818591860186118621863186418651866186718681869187018711872187318741875187618771878187918801881188218831884188518861887188818891890189118921893189418951896189718981899190019011902190319041905190619071908190919101911191219131914191519161917191819191920192119221923192419251926192719281929193019311932193319341935193619371938193919401941194219431944194519461947194819491950195119521953195419551956195719581959196019611962196319641965196619671968196919701971197219731974197519761977197819791980198119821983198419851986198719881989199019911992199319941995199619971998199920002001200220032004200520062007200820092010201120122013201420152016201720182019202020212022202320242025202620272028202920302031203220332034203520362037203820392040204120422043204420452046204720482049205020512052205320542055205620572058205920602061206220632064206520662067206820692070207120722073207420752076207720782079208020812082208320842085208620872088208920902091209220932094209520962097209820992100210121022103210421052106210721082109211021112112211321142115211621172118211921202121212221232124212521262127212821292130213121322133213421352136213721382139214021412142214321442145214621472148214921502151215221532154215521562157215821592160216121622163216421652166216721682169217021712172217321742175217621772178217921802181218221832184218521862187218821892190219121922193219421952196219721982199220022012202220322042205220622072208220922102211221222132214221522162217221822192220222122222223222422252226222722282229223022312232223322342235223622372238223922402241224222432244224522462247224822492250225122522253225422552256225722582259226022612262226322642265226622672268226922702271227222732274227522762277227822792280228122822283228422852286228722882289229022912292229322942295229622972298229923002301230223032304230523062307230823092310231123122313231423152316231723182319232023212322232323242325232623272328232923302331233223332334233523362337233823392340234123422343234423452346234723482349235023512352235323542355235623572358235923602361236223632364236523662367236823692370237123722373237423752376237723782379238023812382238323842385238623872388238923902391239223932394239523962397239823992400240124022403240424052406240724082409241024112412241324142415241624172418241924202421242224232424242524262427242824292430243124322433243424352436243724382439244024412442244324442445244624472448244924502451245224532454245524562457245824592460246124622463246424652466246724682469247024712472247324742475247624772478247924802481248224832484248524862487248824892490249124922493249424952496249724982499250025012502250325042505250625072508250925102511251225132514251525162517251825192520252125222523252425252526252725282529253025312532253325342535253625372538253925402541254225432544254525462547254825492550255125522553255425552556255725582559256025612562256325642565256625672568256925702571257225732574257525762577257825792580258125822583258425852586258725882589259025912592259325942595259625972598259926002601260226032604260526062607260826092610261126122613261426152616261726182619262026212622262326242625262626272628262926302631263226332634263526362637263826392640264126422643264426452646264726482649265026512652265326542655265626572658265926602661266226632664266526662667266826692670267126722673267426752676267726782679268026812682268326842685268626872688268926902691269226932694269526962697269826992700270127022703270427052706270727082709271027112712271327142715271627172718271927202721272227232724272527262727272827292730273127322733273427352736273727382739274027412742274327442745274627472748274927502751275227532754275527562757275827592760276127622763276427652766276727682769277027712772277327742775277627772778277927802781278227832784278527862787278827892790279127922793279427952796279727982799280028012802280328042805280628072808280928102811281228132814281528162817281828192820282128222823282428252826282728282829283028312832283328342835283628372838283928402841284228432844284528462847284828492850285128522853285428552856285728582859286028612862286328642865286628672868286928702871287228732874287528762877287828792880288128822883288428852886288728882889289028912892289328942895289628972898289929002901290229032904290529062907290829092910291129122913291429152916291729182919292029212922292329242925292629272928292929302931293229332934293529362937293829392940294129422943294429452946294729482949295029512952295329542955295629572958295929602961296229632964296529662967296829692970297129722973297429752976297729782979298029812982298329842985298629872988298929902991299229932994299529962997299829993000300130023003300430053006300730083009301030113012301330143015301630173018301930203021 |
- package service
- import (
- "container/heap"
- "context"
- "errors"
- "fmt"
- "hash/fnv"
- "log/slog"
- "math"
- "sort"
- "strconv"
- "strings"
- "sync"
- "sync/atomic"
- "time"
- "github.com/Wei-Shaw/sub2api/internal/config"
- "golang.org/x/sync/singleflight"
- )
- const (
- openAIAccountScheduleLayerPreviousResponse = "previous_response_id"
- openAIAccountScheduleLayerGuardianParent = "guardian_parent"
- openAIAccountScheduleLayerSessionSticky = "session_hash"
- openAIAccountScheduleLayerLoadBalance = "load_balance"
- openAIAdvancedSchedulerSettingKey = "openai_advanced_scheduler_enabled"
- )
- const (
- openAIAdvancedSchedulerSettingCacheTTL = 5 * time.Second
- openAIAdvancedSchedulerSettingDBTimeout = 2 * time.Second
- // ponytail: cap probes added when cost ordering expands configured Top-K;
- // use bulk acquisition if a measured workload needs a higher ceiling.
- openAIAccountSelectionProbeLimit = 64
- )
- const (
- openAIQuotaHeadroomNeutralFactor = 0.5
- openAIQuotaHeadroomSecondaryLowRemain = 0.10
- openAIQuotaHeadroomSnapshotStaleAfter = 8 * time.Hour
- openAIUpstreamCostNeutralFactor = 0.5
- defaultOpenAIOAuthSchedulingRateMultiplier = 1.0
- )
- type cachedOpenAIAdvancedSchedulerSetting struct {
- lowUpstreamRatePriorityEnabled bool
- oauthSchedulingRateMultiplier float64
- enabled bool
- stickyWeightedEnabled bool
- subscriptionPriorityEnabled bool
- lbTopKOverride int
- weightOverrides map[string]float64
- expiresAt int64
- }
- type openAIAdvancedSchedulerRuntimeSettings struct {
- lowUpstreamRatePriorityEnabled bool
- oauthSchedulingRateMultiplier float64
- enabled bool
- stickyWeightedEnabled bool
- subscriptionPriorityEnabled bool
- lbTopKOverride int
- weightOverrides map[string]float64
- }
- var openAIAdvancedSchedulerSettingCache atomic.Value // *cachedOpenAIAdvancedSchedulerSetting
- var openAIAdvancedSchedulerSettingSF singleflight.Group
- type OpenAIAccountScheduleRequest struct {
- GroupID *int64
- Platform string
- SessionHash string
- StickyAccountID int64
- GuardianParentAccountID int64
- StickyPreviousAccountID int64
- StickyWeighted bool
- SubscriptionPriority bool
- PreserveStickyBinding bool
- RequirePrivacySet bool
- PreviousResponseID string
- PreviousResponseCanMove bool
- UseUpstreamTokenCost bool
- RequestedModel string
- RequiredTransport OpenAIUpstreamTransport
- RequiredCapability OpenAIEndpointCapability
- RequiredImageCapability OpenAIImagesCapability
- // RequireCompact is only for legacy /responses/compact capability filtering
- // and compact_model_mapping; native remote compaction v2 leaves it false.
- RequireCompact bool
- ExcludedIDs map[int64]struct{}
- }
- type OpenAIAccountScheduleDecision struct {
- Layer string
- StickyPreviousHit bool
- StickySessionHit bool
- CandidateCount int
- TopK int
- LatencyMs int64
- LoadSkew float64
- SelectedAccountID int64
- SelectedAccountType string
- }
- type OpenAIAccountSchedulerMetricsSnapshot struct {
- SelectTotal int64
- StickyPreviousHitTotal int64
- StickySessionHitTotal int64
- LoadBalanceSelectTotal int64
- AccountSwitchTotal int64
- SchedulerLatencyMsTotal int64
- SchedulerLatencyMsAvg float64
- StickyHitRatio float64
- AccountSwitchRate float64
- LoadSkewAvg float64
- RuntimeStatsAccountCount int
- }
- type OpenAIAccountScheduler interface {
- Select(ctx context.Context, req OpenAIAccountScheduleRequest) (*AccountSelectionResult, OpenAIAccountScheduleDecision, error)
- ReportResult(accountID int64, success bool, firstTokenMs *int)
- ReportSwitch()
- SnapshotMetrics() OpenAIAccountSchedulerMetricsSnapshot
- }
- type openAIAccountSchedulerMetrics struct {
- selectTotal atomic.Int64
- stickyPreviousHitTotal atomic.Int64
- stickySessionHitTotal atomic.Int64
- loadBalanceSelectTotal atomic.Int64
- accountSwitchTotal atomic.Int64
- latencyMsTotal atomic.Int64
- loadSkewMilliTotal atomic.Int64
- }
- type openAIAccountLoadPlan struct {
- allCandidates []openAIAccountCandidateScore
- candidates []openAIAccountCandidateScore
- staleSnapshotCompactRetry []openAIAccountCandidateScore
- selectionOrder []openAIAccountCandidateScore
- candidateCount int
- topK int
- loadSkew float64
- includeOverflowFallback bool
- }
- type openAIAccountLoadSelectionAttempt struct {
- result *AccountSelectionResult
- selectionOrder []openAIAccountCandidateScore
- candidateCount int
- topK int
- loadSkew float64
- compactBlocked bool
- noCompactCandidates bool
- err error
- }
- func (m *openAIAccountSchedulerMetrics) recordSelect(decision OpenAIAccountScheduleDecision) {
- if m == nil {
- return
- }
- m.selectTotal.Add(1)
- m.latencyMsTotal.Add(decision.LatencyMs)
- m.loadSkewMilliTotal.Add(int64(math.Round(decision.LoadSkew * 1000)))
- if decision.StickyPreviousHit {
- m.stickyPreviousHitTotal.Add(1)
- }
- if decision.StickySessionHit {
- m.stickySessionHitTotal.Add(1)
- }
- if decision.Layer == openAIAccountScheduleLayerLoadBalance {
- m.loadBalanceSelectTotal.Add(1)
- }
- }
- func (m *openAIAccountSchedulerMetrics) recordSwitch() {
- if m == nil {
- return
- }
- m.accountSwitchTotal.Add(1)
- }
- type openAIAccountRuntimeStats struct {
- accounts sync.Map
- accountCount atomic.Int64
- }
- type openAIAccountRuntimeStat struct {
- errorRateEWMABits atomic.Uint64
- ttftEWMABits atomic.Uint64
- }
- func newOpenAIAccountRuntimeStats() *openAIAccountRuntimeStats {
- return &openAIAccountRuntimeStats{}
- }
- func (s *openAIAccountRuntimeStats) loadOrCreate(accountID int64) *openAIAccountRuntimeStat {
- if value, ok := s.accounts.Load(accountID); ok {
- stat, _ := value.(*openAIAccountRuntimeStat)
- if stat != nil {
- return stat
- }
- }
- stat := &openAIAccountRuntimeStat{}
- stat.ttftEWMABits.Store(math.Float64bits(math.NaN()))
- actual, loaded := s.accounts.LoadOrStore(accountID, stat)
- if !loaded {
- s.accountCount.Add(1)
- return stat
- }
- existing, _ := actual.(*openAIAccountRuntimeStat)
- if existing != nil {
- return existing
- }
- return stat
- }
- func updateEWMAAtomic(target *atomic.Uint64, sample float64, alpha float64) {
- for {
- oldBits := target.Load()
- oldValue := math.Float64frombits(oldBits)
- newValue := alpha*sample + (1-alpha)*oldValue
- if target.CompareAndSwap(oldBits, math.Float64bits(newValue)) {
- return
- }
- }
- }
- func (s *openAIAccountRuntimeStats) report(accountID int64, success bool, firstTokenMs *int) {
- if s == nil || accountID <= 0 {
- return
- }
- const alpha = 0.2
- stat := s.loadOrCreate(accountID)
- errorSample := 1.0
- if success {
- errorSample = 0.0
- }
- updateEWMAAtomic(&stat.errorRateEWMABits, errorSample, alpha)
- if firstTokenMs != nil && *firstTokenMs > 0 {
- ttft := float64(*firstTokenMs)
- ttftBits := math.Float64bits(ttft)
- for {
- oldBits := stat.ttftEWMABits.Load()
- oldValue := math.Float64frombits(oldBits)
- if math.IsNaN(oldValue) {
- if stat.ttftEWMABits.CompareAndSwap(oldBits, ttftBits) {
- break
- }
- continue
- }
- newValue := alpha*ttft + (1-alpha)*oldValue
- if stat.ttftEWMABits.CompareAndSwap(oldBits, math.Float64bits(newValue)) {
- break
- }
- }
- }
- }
- func (s *openAIAccountRuntimeStats) snapshot(accountID int64) (errorRate float64, ttft float64, hasTTFT bool) {
- if s == nil || accountID <= 0 {
- return 0, 0, false
- }
- value, ok := s.accounts.Load(accountID)
- if !ok {
- return 0, 0, false
- }
- stat, _ := value.(*openAIAccountRuntimeStat)
- if stat == nil {
- return 0, 0, false
- }
- errorRate = clamp01(math.Float64frombits(stat.errorRateEWMABits.Load()))
- ttftValue := math.Float64frombits(stat.ttftEWMABits.Load())
- if math.IsNaN(ttftValue) {
- return errorRate, 0, false
- }
- return errorRate, ttftValue, true
- }
- func (s *openAIAccountRuntimeStats) size() int {
- if s == nil {
- return 0
- }
- return int(s.accountCount.Load())
- }
- type defaultOpenAIAccountScheduler struct {
- service *OpenAIGatewayService
- metrics openAIAccountSchedulerMetrics
- stats *openAIAccountRuntimeStats
- grokFreeQuotaGateCache sync.Map // key: int64(accountID), value: grokFreeQuotaGateCacheEntry
- }
- type openAISelectionProbeBudget struct {
- acquires int
- rechecks int
- attempted map[int64]struct{}
- limited bool
- }
- func newOpenAISelectionProbeBudget() *openAISelectionProbeBudget {
- return &openAISelectionProbeBudget{attempted: make(map[int64]struct{})}
- }
- func (b *openAISelectionProbeBudget) enableLimit() {
- if b != nil {
- b.limited = true
- }
- }
- func (b *openAISelectionProbeBudget) recordAcquire(accountID int64) bool {
- if b == nil {
- return false
- }
- if !b.limited {
- return true
- }
- if b.acquires >= openAIAccountSelectionProbeLimit {
- return false
- }
- if b.attempted == nil {
- b.attempted = make(map[int64]struct{})
- }
- b.acquires++
- b.attempted[accountID] = struct{}{}
- return true
- }
- func (b *openAISelectionProbeBudget) recordRecheck() bool {
- if b == nil {
- return false
- }
- if !b.limited {
- return true
- }
- if b.rechecks >= openAIAccountSelectionProbeLimit {
- return false
- }
- b.rechecks++
- return true
- }
- func (b *openAISelectionProbeBudget) acquireExhausted() bool {
- return b != nil && b.limited && b.acquires >= openAIAccountSelectionProbeLimit
- }
- func (b *openAISelectionProbeBudget) wasAttempted(accountID int64) bool {
- if b == nil {
- return false
- }
- _, ok := b.attempted[accountID]
- return ok
- }
- type openAIStickyEscapeConfig struct {
- enabled bool
- ttftMs float64
- errorRate float64
- }
- func newDefaultOpenAIAccountScheduler(service *OpenAIGatewayService, stats *openAIAccountRuntimeStats) OpenAIAccountScheduler {
- if stats == nil {
- stats = newOpenAIAccountRuntimeStats()
- }
- return &defaultOpenAIAccountScheduler{
- service: service,
- stats: stats,
- }
- }
- func (s *defaultOpenAIAccountScheduler) Select(
- ctx context.Context,
- req OpenAIAccountScheduleRequest,
- ) (*AccountSelectionResult, OpenAIAccountScheduleDecision, error) {
- if s != nil && s.service != nil && s.service.openAIGroupRequiresPrivacySet(ctx, req.GroupID) {
- req.RequirePrivacySet = true
- }
- decision := OpenAIAccountScheduleDecision{}
- start := time.Now()
- defer func() {
- decision.LatencyMs = time.Since(start).Milliseconds()
- s.metrics.recordSelect(decision)
- }()
- previousResponseID := strings.TrimSpace(req.PreviousResponseID)
- if previousResponseID != "" && NormalizeOpenAICompatiblePlatform(req.Platform) == PlatformOpenAI &&
- (!req.StickyWeighted || !req.PreviousResponseCanMove) {
- selection, err := s.service.selectAccountByPreviousResponseIDForCapability(
- ctx,
- req.GroupID,
- previousResponseID,
- req.RequestedModel,
- req.ExcludedIDs,
- req.RequiredCapability,
- req.RequireCompact,
- )
- if err != nil {
- return nil, decision, err
- }
- if selection != nil && selection.Account != nil {
- compatible, _ := s.isAccountRequestCompatibleReason(ctx, selection.Account, req)
- hasGroupMetadata := len(selection.Account.GroupIDs) > 0 || len(selection.Account.AccountGroups) > 0
- groupCompatible := !hasGroupMetadata || openAIStickyAccountMatchesGroup(selection.Account, req.GroupID)
- if hasGroupMetadata && s.service != nil {
- groupCompatible = s.service.openAIAccountMatchesSchedulingGroup(selection.Account, req.GroupID)
- }
- if !groupCompatible ||
- !compatible || !s.isAccountTransportCompatible(selection.Account, req.RequiredTransport) {
- if selection.ReleaseFunc != nil {
- selection.ReleaseFunc()
- }
- selection = nil
- }
- }
- if selection != nil && selection.Account != nil {
- decision.Layer = openAIAccountScheduleLayerPreviousResponse
- decision.StickyPreviousHit = true
- decision.SelectedAccountID = selection.Account.ID
- decision.SelectedAccountType = selection.Account.Type
- if req.SessionHash != "" {
- _ = s.service.bindOpenAIStickySessionDuringSelection(ctx, req.GroupID, req.SessionHash, selection.Account.ID)
- }
- return selection, decision, nil
- }
- }
- if req.GuardianParentAccountID > 0 {
- parentReq := req
- parentReq.StickyAccountID = req.GuardianParentAccountID
- parentReq.PreserveStickyBinding = true
- selection, _, err := s.selectBySessionHash(ctx, parentReq)
- if err != nil {
- return nil, decision, err
- }
- if selection != nil && selection.Account != nil {
- decision.Layer = openAIAccountScheduleLayerGuardianParent
- decision.StickySessionHit = true
- decision.SelectedAccountID = selection.Account.ID
- decision.SelectedAccountType = selection.Account.Type
- return selection, decision, nil
- }
- }
- if !req.StickyWeighted {
- selection, escapedSticky, err := s.selectBySessionHash(ctx, req)
- if err != nil {
- return nil, decision, err
- }
- if selection != nil && selection.Account != nil {
- decision.Layer = openAIAccountScheduleLayerSessionSticky
- decision.StickySessionHit = true
- decision.SelectedAccountID = selection.Account.ID
- decision.SelectedAccountType = selection.Account.Type
- return selection, decision, nil
- }
- if escapedSticky {
- req.PreserveStickyBinding = true
- }
- }
- selection, candidateCount, topK, loadSkew, err := s.selectByLoadBalance(ctx, req)
- decision.Layer = openAIAccountScheduleLayerLoadBalance
- decision.CandidateCount = candidateCount
- decision.TopK = topK
- decision.LoadSkew = loadSkew
- if err != nil {
- return nil, decision, err
- }
- if selection != nil && selection.Account != nil {
- decision.SelectedAccountID = selection.Account.ID
- decision.SelectedAccountType = selection.Account.Type
- if req.StickyWeighted {
- if req.StickyPreviousAccountID > 0 && selection.Account.ID == req.StickyPreviousAccountID {
- decision.StickyPreviousHit = true
- }
- if req.StickyAccountID > 0 && selection.Account.ID == req.StickyAccountID {
- decision.StickySessionHit = true
- }
- }
- }
- return selection, decision, nil
- }
- func (s *defaultOpenAIAccountScheduler) selectBySessionHash(
- ctx context.Context,
- req OpenAIAccountScheduleRequest,
- ) (*AccountSelectionResult, bool, error) {
- sessionHash := strings.TrimSpace(req.SessionHash)
- if sessionHash == "" || s == nil || s.service == nil || s.service.cache == nil {
- return nil, false, nil
- }
- accountID := req.StickyAccountID
- clearBinding := func() {
- if !req.PreserveStickyBinding {
- _ = s.service.deleteStickySessionAccountID(ctx, req.GroupID, sessionHash)
- }
- }
- if accountID <= 0 {
- var err error
- accountID, err = s.service.getStickySessionAccountID(ctx, req.GroupID, sessionHash)
- if err != nil || accountID <= 0 {
- return nil, false, nil
- }
- }
- if accountID <= 0 {
- return nil, false, nil
- }
- if req.ExcludedIDs != nil {
- if _, excluded := req.ExcludedIDs[accountID]; excluded {
- return nil, false, nil
- }
- }
- account, err := s.service.getSchedulableAccount(ctx, accountID)
- if err != nil || account == nil {
- clearBinding()
- return nil, false, nil
- }
- if shouldClearStickySession(account, req.RequestedModel) || account.Platform != NormalizeOpenAICompatiblePlatform(req.Platform) || !account.IsOpenAICompatible() || !account.IsSchedulable() {
- clearBinding()
- return nil, false, nil
- }
- if !s.isAccountRequestCompatible(ctx, account, req) {
- return nil, false, nil
- }
- if !s.isAccountTransportCompatible(account, req.RequiredTransport) {
- clearBinding()
- return nil, false, nil
- }
- account = s.service.recheckSelectedOpenAIAccountFromDB(ctx, account, req.GroupID, req.Platform, req.RequestedModel, req.RequireCompact, req.RequiredCapability)
- if account == nil || !s.service.openAIAccountMatchesSchedulingGroup(account, req.GroupID) || !s.isAccountRequestCompatible(ctx, account, req) || !s.isAccountTransportCompatible(account, req.RequiredTransport) {
- clearBinding()
- return nil, false, nil
- }
- // Free-tier soft gate: sticky session must not pin an over-quota free OAuth account.
- // Admin QueryQuota / import probes do not use this path.
- if account != nil && len(s.filterGrokFreeQuotaAccounts(ctx, []Account{*account})) == 0 {
- clearBinding()
- return nil, false, nil
- }
- // Team+model cool: sticky must not pin a sibling under the same team 429 window.
- now := time.Now()
- upstreamModel := canonicalOpenAIAccountSchedulingModel(account, req.RequestedModel)
- if account != nil && isGrokTeamModelRateLimited(account, upstreamModel, now) {
- clearBinding()
- return nil, false, nil
- }
- if account != nil && isGrokModelQuotaBlocked(account.ID, upstreamModel, now) {
- clearBinding()
- return nil, false, nil
- }
- escapeCfg := s.service.openAIStickyEscapeConfig()
- if reason, errorRate, ttft, shouldEscape := s.shouldEscapeStickyAccount(accountID, escapeCfg); shouldEscape {
- slog.Info("sticky_escape_triggered",
- "account_id", accountID,
- "reason", reason,
- "error_rate", errorRate,
- "ttft", ttft,
- )
- return nil, true, nil
- }
- result, acquireErr := s.service.tryAcquireAccountSlot(ctx, accountID, account.Mode1EffectiveConcurrency())
- if acquireErr == nil && result != nil && result.Acquired {
- if !req.PreserveStickyBinding {
- _ = s.service.refreshStickySessionTTL(ctx, req.GroupID, sessionHash, s.service.openAIWSSessionStickyTTL())
- }
- return attachSelectionProfitGate(ctx, &AccountSelectionResult{
- Account: account,
- Acquired: true,
- ReleaseFunc: result.ReleaseFunc,
- }), false, nil
- }
- cfg := s.service.schedulingConfig()
- // WaitPlan.MaxConcurrency 使用 Concurrency(非 EffectiveLoadFactor),因为 WaitPlan 控制的是 Redis 实际并发槽位等待。
- if s.service.concurrencyService != nil {
- if escapeCfg.enabled && acquireErr == nil && result != nil && !result.Acquired {
- errorRate, ttft, _ := s.stats.snapshot(accountID)
- slog.Info("sticky_escape_triggered",
- "account_id", accountID,
- "reason", "concurrency_full",
- "error_rate", errorRate,
- "ttft", ttft,
- )
- return nil, true, nil
- }
- return attachSelectionProfitGate(ctx, &AccountSelectionResult{
- Account: account,
- WaitPlan: &AccountWaitPlan{
- AccountID: accountID,
- MaxConcurrency: account.Mode1EffectiveConcurrency(),
- Timeout: cfg.StickySessionWaitTimeout,
- MaxWaiting: cfg.StickySessionMaxWaiting,
- },
- }), false, nil
- }
- return nil, false, nil
- }
- func openAIStickyAccountMatchesGroup(account *Account, groupID *int64) bool {
- if account == nil {
- return false
- }
- if groupID == nil {
- return len(account.AccountGroups) == 0 && len(account.GroupIDs) == 0
- }
- for _, accountGroupID := range account.GroupIDs {
- if accountGroupID == *groupID {
- return true
- }
- }
- for _, accountGroup := range account.AccountGroups {
- if accountGroup.GroupID == *groupID {
- return true
- }
- }
- return false
- }
- func openAIAccountSchedulingPriority(account *Account) int {
- if account == nil {
- return 0
- }
- return account.Priority
- }
- func (s *defaultOpenAIAccountScheduler) shouldEscapeStickyAccount(accountID int64, cfg openAIStickyEscapeConfig) (reason string, errorRate float64, ttft float64, shouldEscape bool) {
- if !cfg.enabled || s == nil || s.stats == nil || accountID <= 0 {
- return "", 0, 0, false
- }
- errorRate, ttft, hasTTFT := s.stats.snapshot(accountID)
- if hasTTFT && ttft > cfg.ttftMs {
- return "ttft", errorRate, ttft, true
- }
- if errorRate > cfg.errorRate {
- return "error_rate", errorRate, ttft, true
- }
- return "", errorRate, ttft, false
- }
- type openAIAccountCandidateScore struct {
- account *Account
- loadInfo *AccountLoadInfo
- loadKnown bool
- score float64
- priority int
- errorRate float64
- ttft float64
- hasTTFT bool
- }
- type openAIAccountCandidateHeap []openAIAccountCandidateScore
- func (h openAIAccountCandidateHeap) Len() int {
- return len(h)
- }
- func (h openAIAccountCandidateHeap) Less(i, j int) bool {
- // 最小堆根节点保存“最差”候选,便于 O(log k) 维护 topK。
- return isOpenAIAccountCandidateBetter(h[j], h[i])
- }
- func (h openAIAccountCandidateHeap) Swap(i, j int) {
- h[i], h[j] = h[j], h[i]
- }
- func (h *openAIAccountCandidateHeap) Push(x any) {
- candidate, ok := x.(openAIAccountCandidateScore)
- if !ok {
- panic("openAIAccountCandidateHeap: invalid element type")
- }
- *h = append(*h, candidate)
- }
- func (h *openAIAccountCandidateHeap) Pop() any {
- old := *h
- n := len(old)
- last := old[n-1]
- *h = old[:n-1]
- return last
- }
- func isOpenAIAccountCandidateBetter(left openAIAccountCandidateScore, right openAIAccountCandidateScore) bool {
- if left.score != right.score {
- return left.score > right.score
- }
- if left.account.Priority != right.account.Priority {
- return left.account.Priority < right.account.Priority
- }
- if left.loadInfo.LoadRate != right.loadInfo.LoadRate {
- return left.loadInfo.LoadRate < right.loadInfo.LoadRate
- }
- if left.loadInfo.WaitingCount != right.loadInfo.WaitingCount {
- return left.loadInfo.WaitingCount < right.loadInfo.WaitingCount
- }
- return left.account.ID < right.account.ID
- }
- func selectTopKOpenAICandidates(candidates []openAIAccountCandidateScore, topK int) []openAIAccountCandidateScore {
- if len(candidates) == 0 {
- return nil
- }
- if topK <= 0 {
- topK = 1
- }
- if topK >= len(candidates) {
- ranked := append([]openAIAccountCandidateScore(nil), candidates...)
- sort.Slice(ranked, func(i, j int) bool {
- return isOpenAIAccountCandidateBetter(ranked[i], ranked[j])
- })
- return ranked
- }
- best := make(openAIAccountCandidateHeap, 0, topK)
- for _, candidate := range candidates {
- if len(best) < topK {
- heap.Push(&best, candidate)
- continue
- }
- if isOpenAIAccountCandidateBetter(candidate, best[0]) {
- best[0] = candidate
- heap.Fix(&best, 0)
- }
- }
- ranked := make([]openAIAccountCandidateScore, len(best))
- copy(ranked, best)
- sort.Slice(ranked, func(i, j int) bool {
- return isOpenAIAccountCandidateBetter(ranked[i], ranked[j])
- })
- return ranked
- }
- type openAISelectionRNG struct {
- state uint64
- }
- func newOpenAISelectionRNG(seed uint64) openAISelectionRNG {
- if seed == 0 {
- seed = 0x9e3779b97f4a7c15
- }
- return openAISelectionRNG{state: seed}
- }
- func (r *openAISelectionRNG) nextUint64() uint64 {
- // xorshift64*
- x := r.state
- x ^= x >> 12
- x ^= x << 25
- x ^= x >> 27
- r.state = x
- return x * 2685821657736338717
- }
- func (r *openAISelectionRNG) nextFloat64() float64 {
- // [0,1)
- return float64(r.nextUint64()>>11) / (1 << 53)
- }
- func deriveOpenAISelectionSeed(req OpenAIAccountScheduleRequest) uint64 {
- hasher := fnv.New64a()
- writeValue := func(value string) {
- trimmed := strings.TrimSpace(value)
- if trimmed == "" {
- return
- }
- _, _ = hasher.Write([]byte(trimmed))
- _, _ = hasher.Write([]byte{0})
- }
- writeValue(req.SessionHash)
- writeValue(req.PreviousResponseID)
- writeValue(req.RequestedModel)
- if req.GroupID != nil {
- _, _ = hasher.Write([]byte(strconv.FormatInt(*req.GroupID, 10)))
- }
- seed := hasher.Sum64()
- // 对“无会话锚点”的纯负载均衡请求引入时间熵,避免固定命中同一账号。
- if strings.TrimSpace(req.SessionHash) == "" && strings.TrimSpace(req.PreviousResponseID) == "" {
- seed ^= uint64(time.Now().UnixNano())
- }
- if seed == 0 {
- seed = uint64(time.Now().UnixNano()) ^ 0x9e3779b97f4a7c15
- }
- return seed
- }
- func buildOpenAIWeightedSelectionOrder(
- candidates []openAIAccountCandidateScore,
- req OpenAIAccountScheduleRequest,
- ) []openAIAccountCandidateScore {
- if len(candidates) <= 1 {
- return append([]openAIAccountCandidateScore(nil), candidates...)
- }
- pool := append([]openAIAccountCandidateScore(nil), candidates...)
- weights := make([]float64, len(pool))
- minScore := pool[0].score
- for i := 1; i < len(pool); i++ {
- if pool[i].score < minScore {
- minScore = pool[i].score
- }
- }
- for i := range pool {
- // 将 top-K 分值平移到正区间,避免“单一最高分账号”长期垄断。
- weight := (pool[i].score - minScore) + 1.0
- if math.IsNaN(weight) || math.IsInf(weight, 0) || weight <= 0 {
- weight = 1.0
- }
- weights[i] = weight
- }
- order := make([]openAIAccountCandidateScore, 0, len(pool))
- rng := newOpenAISelectionRNG(deriveOpenAISelectionSeed(req))
- for len(pool) > 0 {
- total := 0.0
- for _, w := range weights {
- total += w
- }
- selectedIdx := 0
- if total > 0 {
- r := rng.nextFloat64() * total
- acc := 0.0
- for i, w := range weights {
- acc += w
- if r <= acc {
- selectedIdx = i
- break
- }
- }
- } else {
- selectedIdx = int(rng.nextUint64() % uint64(len(pool)))
- }
- order = append(order, pool[selectedIdx])
- pool = append(pool[:selectedIdx], pool[selectedIdx+1:]...)
- weights = append(weights[:selectedIdx], weights[selectedIdx+1:]...)
- }
- return order
- }
- func (s *defaultOpenAIAccountScheduler) buildOpenAIAccountLoadPlan(
- ctx context.Context,
- req OpenAIAccountScheduleRequest,
- filtered []*Account,
- loadMap map[int64]*AccountLoadInfo,
- ) openAIAccountLoadPlan {
- allCandidates := make([]openAIAccountCandidateScore, 0, len(filtered))
- for _, account := range filtered {
- loadInfo, loadKnown := loadMap[account.ID]
- if !loadKnown || loadInfo == nil {
- loadInfo = &AccountLoadInfo{AccountID: account.ID}
- loadKnown = false
- }
- errorRate, ttft, hasTTFT := 0.0, 0.0, false
- if s.stats != nil {
- errorRate, ttft, hasTTFT = s.stats.snapshot(account.ID)
- }
- allCandidates = append(allCandidates, openAIAccountCandidateScore{
- account: account,
- loadInfo: loadInfo,
- loadKnown: loadKnown,
- errorRate: errorRate,
- ttft: ttft,
- hasTTFT: hasTTFT,
- })
- }
- candidates := allCandidates
- staleSnapshotCompactRetry := make([]openAIAccountCandidateScore, 0, len(allCandidates))
- if req.RequireCompact {
- candidates = make([]openAIAccountCandidateScore, 0, len(allCandidates))
- for _, candidate := range allCandidates {
- if openAICompactSupportTier(candidate.account) == 0 {
- staleSnapshotCompactRetry = append(staleSnapshotCompactRetry, candidate)
- continue
- }
- candidates = append(candidates, candidate)
- }
- }
- plan := openAIAccountLoadPlan{
- allCandidates: allCandidates,
- candidates: candidates,
- staleSnapshotCompactRetry: staleSnapshotCompactRetry,
- candidateCount: len(candidates),
- }
- if len(candidates) == 0 {
- plan.selectionOrder = s.buildOpenAISelectionOrder(req, plan)
- return plan
- }
- minPriority, maxPriority := openAIAccountSchedulingPriority(candidates[0].account), openAIAccountSchedulingPriority(candidates[0].account)
- maxWaiting := 1
- loadRateSum := 0.0
- loadRateSumSquares := 0.0
- minTTFT, maxTTFT := 0.0, 0.0
- hasTTFTSample := false
- for i := range candidates {
- candidate := &candidates[i]
- candidate.priority = openAIAccountSchedulingPriority(candidate.account)
- if candidate.priority < minPriority {
- minPriority = candidate.priority
- }
- if candidate.priority > maxPriority {
- maxPriority = candidate.priority
- }
- if candidate.loadInfo.WaitingCount > maxWaiting {
- maxWaiting = candidate.loadInfo.WaitingCount
- }
- if candidate.hasTTFT && candidate.ttft > 0 {
- if !hasTTFTSample {
- minTTFT, maxTTFT = candidate.ttft, candidate.ttft
- hasTTFTSample = true
- } else {
- if candidate.ttft < minTTFT {
- minTTFT = candidate.ttft
- }
- if candidate.ttft > maxTTFT {
- maxTTFT = candidate.ttft
- }
- }
- }
- loadRate := float64(candidate.loadInfo.LoadRate)
- loadRateSum += loadRate
- loadRateSumSquares += loadRate * loadRate
- }
- plan.loadSkew = calcLoadSkewByMoments(loadRateSum, loadRateSumSquares, len(candidates))
- weights := s.service.openAIWSSchedulerWeightsForRequest(ctx)
- now := time.Now()
- upstreamCostFactors := map[int64]float64(nil)
- if req.UseUpstreamTokenCost && weights.UpstreamCost > 0 {
- accounts := make([]*Account, 0, len(candidates))
- for _, candidate := range candidates {
- accounts = append(accounts, candidate.account)
- }
- upstreamCostFactors = openAIUpstreamCostFactors(accounts, now, s.service.openAIOAuthSchedulingRateMultiplier(ctx))
- for _, factor := range upstreamCostFactors {
- if factor != openAIUpstreamCostNeutralFactor {
- plan.includeOverflowFallback = true
- break
- }
- }
- }
- // Reset 因子(use-it-or-lose-it):在拥有「未来会话窗口结束时间」的账号中,
- // 剩余时间越短 → 因子越接近 1(越早重置越优先用尽)。无活跃窗口的账号因子为 0。
- // 仅在 weights.Reset > 0 时计算,默认关闭不影响原有行为。
- minResetRemaining, maxResetRemaining := 0.0, 0.0
- hasResetSample := false
- if weights.Reset > 0 {
- for _, candidate := range candidates {
- end := candidate.account.SessionWindowEnd
- if end == nil || !now.Before(*end) {
- continue
- }
- remaining := end.Sub(now).Seconds()
- if !hasResetSample {
- minResetRemaining, maxResetRemaining = remaining, remaining
- hasResetSample = true
- continue
- }
- if remaining < minResetRemaining {
- minResetRemaining = remaining
- }
- if remaining > maxResetRemaining {
- maxResetRemaining = remaining
- }
- }
- }
- for i := range candidates {
- item := &candidates[i]
- priorityFactor := 1.0
- if maxPriority > minPriority {
- priorityFactor = 1 - float64(item.priority-minPriority)/float64(maxPriority-minPriority)
- }
- loadFactor := 1 - clamp01(float64(item.loadInfo.LoadRate)/100.0)
- queueFactor := 1 - clamp01(float64(item.loadInfo.WaitingCount)/float64(maxWaiting))
- errorFactor := 1 - clamp01(item.errorRate)
- ttftFactor := 0.5
- if item.hasTTFT && hasTTFTSample && maxTTFT > minTTFT {
- ttftFactor = 1 - clamp01((item.ttft-minTTFT)/(maxTTFT-minTTFT))
- }
- resetFactor := 0.0
- if weights.Reset > 0 && hasResetSample {
- if end := item.account.SessionWindowEnd; end != nil && now.Before(*end) {
- if maxResetRemaining > minResetRemaining {
- resetFactor = 1 - clamp01((end.Sub(now).Seconds()-minResetRemaining)/(maxResetRemaining-minResetRemaining))
- } else {
- // 所有有窗口的账号剩余时间相同:一律给满分,让其优于无窗口账号。
- resetFactor = 1
- }
- }
- }
- quotaHeadroomFactor := 0.0
- if weights.QuotaHeadroom > 0 {
- quotaHeadroomFactor = openAIQuotaHeadroomFactor(item.account, now)
- }
- upstreamCostFactor := openAIUpstreamCostNeutralFactor
- if factor, ok := upstreamCostFactors[item.account.ID]; ok {
- upstreamCostFactor = factor
- }
- item.score = weights.Priority*priorityFactor +
- weights.Load*loadFactor +
- weights.Queue*queueFactor +
- weights.ErrorRate*errorFactor +
- weights.TTFT*ttftFactor +
- weights.Reset*resetFactor +
- weights.QuotaHeadroom*quotaHeadroomFactor +
- weights.UpstreamCost*(upstreamCostFactor-openAIUpstreamCostNeutralFactor)
- if req.StickyWeighted {
- if req.PreviousResponseCanMove && req.StickyPreviousAccountID > 0 && item.account.ID == req.StickyPreviousAccountID {
- item.score += weights.Previous
- }
- if req.StickyAccountID > 0 && item.account.ID == req.StickyAccountID {
- item.score += weights.SessionSticky
- }
- }
- }
- plan.candidates = candidates
- plan.topK = s.service.openAIWSLBTopKForRequest(ctx)
- if plan.topK > len(candidates) {
- plan.topK = len(candidates)
- }
- if plan.topK <= 0 {
- plan.topK = 1
- }
- plan.selectionOrder = s.buildOpenAISelectionOrder(req, plan)
- return plan
- }
- func (s *defaultOpenAIAccountScheduler) buildOpenAISelectionOrder(
- req OpenAIAccountScheduleRequest,
- plan openAIAccountLoadPlan,
- ) []openAIAccountCandidateScore {
- buildSelectionOrder := func(pool []openAIAccountCandidateScore) []openAIAccountCandidateScore {
- if len(pool) == 0 || plan.topK <= 0 {
- return nil
- }
- groupTopK := plan.topK
- if groupTopK > len(pool) {
- groupTopK = len(pool)
- }
- ranked := selectTopKOpenAICandidates(pool, groupTopK)
- var primary []openAIAccountCandidateScore
- if req.StickyWeighted {
- for _, stickyID := range []int64{req.StickyPreviousAccountID, req.StickyAccountID} {
- if stickyID <= 0 {
- continue
- }
- for i, candidate := range ranked {
- if candidate.account != nil && candidate.account.ID == stickyID {
- primary = append([]openAIAccountCandidateScore{candidate}, ranked[:i]...)
- primary = append(primary, ranked[i+1:]...)
- break
- }
- }
- if len(primary) > 0 {
- break
- }
- }
- }
- if len(primary) == 0 {
- primary = buildOpenAIWeightedSelectionOrder(ranked, req)
- }
- if !plan.includeOverflowFallback || groupTopK >= len(pool) {
- return primary
- }
- selected := make(map[int64]struct{}, len(primary))
- for _, candidate := range primary {
- selected[candidate.account.ID] = struct{}{}
- }
- overflow := make([]openAIAccountCandidateScore, 0, len(pool)-len(primary))
- for _, candidate := range pool {
- if _, ok := selected[candidate.account.ID]; !ok {
- overflow = append(overflow, candidate)
- }
- }
- sort.Slice(overflow, func(i, j int) bool {
- return isOpenAIAccountCandidateBetter(overflow[i], overflow[j])
- })
- return append(primary, overflow...)
- }
- if req.RequireCompact {
- supported := make([]openAIAccountCandidateScore, 0, len(plan.candidates))
- unknown := make([]openAIAccountCandidateScore, 0, len(plan.candidates))
- for _, candidate := range plan.candidates {
- switch openAICompactSupportTier(candidate.account) {
- case 2:
- supported = append(supported, candidate)
- case 1:
- unknown = append(unknown, candidate)
- }
- }
- selectionOrder := make([]openAIAccountCandidateScore, 0, len(plan.allCandidates))
- selectionOrder = append(selectionOrder, buildSelectionOrder(supported)...)
- selectionOrder = append(selectionOrder, buildSelectionOrder(unknown)...)
- if len(plan.staleSnapshotCompactRetry) > 0 && s.service.schedulerSnapshot != nil {
- selectionOrder = append(selectionOrder, sortOpenAICompactRetryCandidates(plan.staleSnapshotCompactRetry)...)
- }
- return selectionOrder
- }
- return buildSelectionOrder(plan.candidates)
- }
- func sortOpenAICompactRetryCandidates(pool []openAIAccountCandidateScore) []openAIAccountCandidateScore {
- if len(pool) == 0 {
- return nil
- }
- ordered := append([]openAIAccountCandidateScore(nil), pool...)
- sort.SliceStable(ordered, func(i, j int) bool {
- a, b := ordered[i], ordered[j]
- if a.account.Priority != b.account.Priority {
- return a.account.Priority < b.account.Priority
- }
- if a.loadInfo.LoadRate != b.loadInfo.LoadRate {
- return a.loadInfo.LoadRate < b.loadInfo.LoadRate
- }
- if a.loadInfo.WaitingCount != b.loadInfo.WaitingCount {
- return a.loadInfo.WaitingCount < b.loadInfo.WaitingCount
- }
- switch {
- case a.account.LastUsedAt == nil && b.account.LastUsedAt != nil:
- return true
- case a.account.LastUsedAt != nil && b.account.LastUsedAt == nil:
- return false
- case a.account.LastUsedAt == nil && b.account.LastUsedAt == nil:
- return false
- default:
- return a.account.LastUsedAt.Before(*b.account.LastUsedAt)
- }
- })
- return ordered
- }
- func (s *defaultOpenAIAccountScheduler) tryAcquireOpenAISelectionOrder(
- ctx context.Context,
- req OpenAIAccountScheduleRequest,
- selectionOrder []openAIAccountCandidateScore,
- ) (*AccountSelectionResult, bool, error) {
- budget := newOpenAISelectionProbeBudget()
- budget.enableLimit()
- return s.tryAcquireOpenAISelectionOrderWithBudget(ctx, req, selectionOrder, budget)
- }
- func (s *defaultOpenAIAccountScheduler) tryAcquireOpenAISelectionOrderWithBudget(
- ctx context.Context,
- req OpenAIAccountScheduleRequest,
- selectionOrder []openAIAccountCandidateScore,
- budget *openAISelectionProbeBudget,
- ) (*AccountSelectionResult, bool, error) {
- compactBlocked := false
- release := func(result *AcquireResult) {
- if result != nil && result.ReleaseFunc != nil {
- result.ReleaseFunc()
- }
- }
- for i := 0; i < len(selectionOrder); i++ {
- candidate := selectionOrder[i]
- if candidate.account == nil {
- continue
- }
- if candidate.loadKnown && candidate.account.Mode1EffectiveConcurrency() > 0 &&
- candidate.loadInfo.CurrentConcurrency >= candidate.account.Mode1EffectiveConcurrency() {
- continue
- }
- result, attempted, acquireErr := s.tryAcquireOpenAIAccountSlot(ctx, candidate.account.ID, candidate.account.Mode1EffectiveConcurrency(), budget)
- if !attempted {
- break
- }
- if acquireErr != nil {
- return nil, compactBlocked, acquireErr
- }
- if result == nil || !result.Acquired {
- continue
- }
- fresh := s.service.resolveFreshSchedulableOpenAIAccount(ctx, candidate.account, req.Platform, req.RequestedModel, false, req.RequiredCapability)
- if fresh == nil || !s.isAccountTransportCompatible(fresh, req.RequiredTransport) || !s.isAccountRequestCompatible(ctx, fresh, req) {
- release(result)
- continue
- }
- if !s.consumeOpenAISelectionDBRecheck(budget) {
- release(result)
- break
- }
- fresh = s.service.recheckSelectedOpenAIAccountFromDB(ctx, fresh, req.GroupID, req.Platform, req.RequestedModel, false, req.RequiredCapability)
- if fresh == nil || !s.isAccountTransportCompatible(fresh, req.RequiredTransport) || !s.isAccountRequestCompatible(ctx, fresh, req) {
- release(result)
- continue
- }
- if req.RequireCompact && openAICompactSupportTier(fresh) == 0 {
- compactBlocked = true
- release(result)
- continue
- }
- if fresh.Mode1EffectiveConcurrency() != candidate.account.Mode1EffectiveConcurrency() {
- release(result)
- result, attempted, acquireErr = s.tryAcquireOpenAIAccountSlot(ctx, fresh.ID, fresh.Mode1EffectiveConcurrency(), budget)
- if !attempted {
- continue
- }
- if acquireErr != nil {
- return nil, compactBlocked, acquireErr
- }
- if result == nil || !result.Acquired {
- continue
- }
- }
- if req.SessionHash != "" && !req.PreserveStickyBinding {
- _ = s.service.bindOpenAIStickySessionDuringSelection(ctx, req.GroupID, req.SessionHash, fresh.ID)
- }
- return attachSelectionProfitGate(ctx, &AccountSelectionResult{
- Account: fresh,
- Acquired: true,
- ReleaseFunc: result.ReleaseFunc,
- }), compactBlocked, nil
- }
- return nil, compactBlocked, nil
- }
- func (s *defaultOpenAIAccountScheduler) tryAcquireOpenAIAccountSlot(
- ctx context.Context,
- accountID int64,
- maxConcurrency int,
- budget *openAISelectionProbeBudget,
- ) (*AcquireResult, bool, error) {
- if s.service.concurrencyService != nil && maxConcurrency > 0 && !budget.recordAcquire(accountID) {
- return nil, false, nil
- }
- result, err := s.service.tryAcquireAccountSlot(ctx, accountID, maxConcurrency)
- return result, true, err
- }
- func (s *defaultOpenAIAccountScheduler) consumeOpenAISelectionDBRecheck(budget *openAISelectionProbeBudget) bool {
- if s.service.schedulerSnapshot == nil || s.service.accountRepo == nil {
- return true
- }
- return budget.recordRecheck()
- }
- func (s *defaultOpenAIAccountScheduler) tryFallbackToWeightedSticky(
- ctx context.Context,
- req OpenAIAccountScheduleRequest,
- ) (*AccountSelectionResult, error) {
- if !req.StickyWeighted {
- return nil, nil
- }
- for _, accountID := range []int64{req.StickyPreviousAccountID, req.StickyAccountID} {
- if accountID <= 0 {
- continue
- }
- if req.ExcludedIDs != nil {
- if _, excluded := req.ExcludedIDs[accountID]; excluded {
- continue
- }
- }
- account, err := s.service.getSchedulableAccount(ctx, accountID)
- if err != nil || account == nil {
- continue
- }
- if !s.isAccountRequestCompatible(ctx, account, req) || !s.isAccountTransportCompatible(account, req.RequiredTransport) {
- continue
- }
- account = s.service.recheckSelectedOpenAIAccountFromDB(ctx, account, req.GroupID, req.Platform, req.RequestedModel, req.RequireCompact, req.RequiredCapability)
- if account == nil {
- if accountID == req.StickyAccountID && strings.TrimSpace(req.SessionHash) != "" {
- _ = s.service.deleteStickySessionAccountID(ctx, req.GroupID, req.SessionHash)
- }
- continue
- }
- if !s.service.openAIAccountMatchesSchedulingGroup(account, req.GroupID) {
- if accountID == req.StickyAccountID && strings.TrimSpace(req.SessionHash) != "" {
- _ = s.service.deleteStickySessionAccountID(ctx, req.GroupID, req.SessionHash)
- }
- continue
- }
- if !s.isAccountRequestCompatible(ctx, account, req) || !s.isAccountTransportCompatible(account, req.RequiredTransport) {
- continue
- }
- if req.RequireCompact && openAICompactSupportTier(account) == 0 {
- continue
- }
- // Keep weighted sticky fallback subject to the same free-tier gate as the
- // normal and sticky selection paths. Otherwise an over-quota free account
- // could be reintroduced after the primary candidate pass.
- if len(s.filterGrokFreeQuotaAccounts(ctx, []Account{*account})) == 0 {
- continue
- }
- upstreamModel := canonicalOpenAIAccountSchedulingModel(account, req.RequestedModel)
- now := time.Now()
- if isGrokTeamModelRateLimited(account, upstreamModel, now) ||
- isGrokModelQuotaBlocked(account.ID, upstreamModel, now) {
- continue
- }
- result, acquireErr := s.service.tryAcquireAccountSlot(ctx, account.ID, account.Mode1EffectiveConcurrency())
- if acquireErr != nil {
- return nil, acquireErr
- }
- if result != nil && result.Acquired {
- if req.SessionHash != "" && !req.PreserveStickyBinding {
- _ = s.service.bindOpenAIStickySessionDuringSelection(ctx, req.GroupID, req.SessionHash, account.ID)
- }
- return attachSelectionProfitGate(ctx, &AccountSelectionResult{
- Account: account,
- Acquired: true,
- ReleaseFunc: result.ReleaseFunc,
- }), nil
- }
- if s.service.concurrencyService != nil {
- cfg := s.service.schedulingConfig()
- return attachSelectionProfitGate(ctx, &AccountSelectionResult{
- Account: account,
- WaitPlan: &AccountWaitPlan{
- AccountID: account.ID,
- MaxConcurrency: account.Mode1EffectiveConcurrency(),
- Timeout: cfg.StickySessionWaitTimeout,
- MaxWaiting: cfg.StickySessionMaxWaiting,
- },
- }), nil
- }
- }
- return nil, nil
- }
- // openAISelectionFilterStats counts why candidates were dropped by the
- // selectByLoadBalance initial filter. Historically these exclusions were
- // silent (debug logs at best), so a "no available accounts" failure with
- // excluded_account_count=0 was undiagnosable from the error alone (#4599).
- // The reasons map is lazily allocated: on the happy path (nothing filtered
- // out, or an account is eventually selected) no extra allocation happens.
- type openAISelectionFilterStats struct {
- pool int
- reasons map[string]int
- }
- func (s *openAISelectionFilterStats) exclude(reason string) {
- if s.reasons == nil {
- s.reasons = make(map[string]int, 4)
- }
- s.reasons[reason]++
- }
- // summary renders deterministic exclusion statistics for scheduling error
- // messages, e.g. "pool=3, filtered: model_not_supported=2 quota_auto_pause_7d=1".
- // Reasons are sorted lexicographically so the output is stable for tests and
- // log aggregation. extra, when non-empty, is appended as a trailing marker.
- func (s openAISelectionFilterStats) summary(extra string) string {
- var b strings.Builder
- _, _ = b.WriteString("pool=")
- _, _ = b.WriteString(strconv.Itoa(s.pool))
- if len(s.reasons) > 0 {
- reasons := make([]string, 0, len(s.reasons))
- for reason := range s.reasons {
- reasons = append(reasons, reason)
- }
- sort.Strings(reasons)
- _, _ = b.WriteString(", filtered:")
- for _, reason := range reasons {
- _, _ = b.WriteString(" ")
- _, _ = b.WriteString(reason)
- _, _ = b.WriteString("=")
- _, _ = b.WriteString(strconv.Itoa(s.reasons[reason]))
- }
- }
- if extra != "" {
- _, _ = b.WriteString(", ")
- _, _ = b.WriteString(extra)
- }
- return b.String()
- }
- func (s *defaultOpenAIAccountScheduler) selectByLoadBalance(
- ctx context.Context,
- req OpenAIAccountScheduleRequest,
- ) (*AccountSelectionResult, int, int, float64, error) {
- budget := newOpenAISelectionProbeBudget()
- accounts, err := s.service.listSchedulableAccounts(ctx, req.GroupID, req.Platform)
- if err != nil {
- return nil, 0, 0, 0, err
- }
- if len(accounts) == 0 {
- return nil, 0, 0, 0, noAvailableOpenAISelectionError(req.RequestedModel, false, openAISelectionFilterStats{}.summary(""))
- }
- // Local free-tier soft gate on the Grok scheduling path only (not admin probe).
- accounts = s.filterGrokFreeQuotaAccounts(ctx, accounts)
- if len(accounts) == 0 {
- return nil, 0, 0, 0, noAvailableOpenAISelectionError(req.RequestedModel, false, openAISelectionFilterStats{}.summary("grok_free_quota_soft_gate"))
- }
- // Team+model rate-limit cool: siblings of a 429'd team skip the hot model.
- if req.Platform == PlatformGrok {
- now := time.Now()
- filtered := filterGrokTeamModelRateLimitedAccounts(accounts, req.RequestedModel, now)
- if len(filtered) == 0 && len(accounts) > 0 {
- return nil, 0, 0, 0, noAvailableOpenAISelectionError(req.RequestedModel, false, openAISelectionFilterStats{}.summary("grok_team_model_rate_limit"))
- }
- if filtered != nil {
- accounts = filtered
- }
- // Per-account model free-usage soft-block (other models stay eligible).
- modelFiltered := filterGrokModelQuotaBlockedAccounts(accounts, req.RequestedModel, now)
- if len(modelFiltered) == 0 && len(accounts) > 0 {
- return nil, 0, 0, 0, noAvailableOpenAISelectionError(req.RequestedModel, false, openAISelectionFilterStats{}.summary("grok_model_quota_block"))
- }
- accounts = modelFiltered
- }
- // require_privacy_set: 获取分组信息
- var schedGroup *Group
- if req.GroupID != nil && s.service.schedulerSnapshot != nil {
- schedGroup, _ = s.service.schedulerSnapshot.GetGroupByID(ctx, *req.GroupID)
- }
- filterStats := openAISelectionFilterStats{pool: len(accounts)}
- filtered := make([]*Account, 0, len(accounts))
- loadReq := make([]AccountWithConcurrency, 0, len(accounts))
- for i := range accounts {
- account := &accounts[i]
- if req.ExcludedIDs != nil {
- if _, excluded := req.ExcludedIDs[account.ID]; excluded {
- filterStats.exclude("excluded")
- continue
- }
- }
- if !account.IsSchedulable() {
- filterStats.exclude("not_schedulable")
- continue
- }
- if account.Platform != NormalizeOpenAICompatiblePlatform(req.Platform) || !account.IsOpenAICompatible() {
- filterStats.exclude("platform_mismatch")
- continue
- }
- if s.service.isOpenAIAccountRequestRuntimeBlocked(account, req.RequestedModel) {
- filterStats.exclude("runtime_blocked")
- continue
- }
- // require_privacy_set is a group-scoped eligibility gate. Do not mutate the
- // shared account: another group may intentionally allow accounts whose
- // upstream privacy setting has not been confirmed.
- if schedGroup != nil && schedGroup.RequirePrivacySet && !account.IsPrivacySet() {
- filterStats.exclude("privacy_not_set")
- continue
- }
- if compatible, reason := s.isAccountRequestCompatibleReason(ctx, account, req); !compatible {
- filterStats.exclude(reason)
- continue
- }
- if !s.isAccountTransportCompatible(account, req.RequiredTransport) {
- filterStats.exclude("transport_incompatible")
- continue
- }
- filtered = append(filtered, account)
- loadReq = append(loadReq, AccountWithConcurrency{
- ID: account.ID,
- MaxConcurrency: account.EffectiveLoadFactor(),
- })
- }
- if len(filtered) == 0 {
- return nil, 0, 0, 0, noAvailableOpenAISelectionError(req.RequestedModel, false, filterStats.summary(""))
- }
- loadMap := map[int64]*AccountLoadInfo{}
- if s.service.concurrencyService != nil {
- if batchLoad, loadErr := s.service.concurrencyService.GetAccountsLoadBatch(ctx, loadReq); loadErr == nil {
- loadMap = batchLoad
- }
- }
- if req.SubscriptionPriority {
- subscriptionAccounts, regularAccounts := partitionOpenAIChatGPTSubscriptionAccounts(filtered)
- if len(subscriptionAccounts) > 0 {
- attempt := s.trySelectByLoadBalancePool(ctx, req, subscriptionAccounts, loadMap, budget)
- if attempt.err != nil && (!attempt.noCompactCandidates || len(regularAccounts) <= 0) {
- return nil, attempt.candidateCount, attempt.topK, attempt.loadSkew, attempt.err
- }
- if attempt.result != nil {
- return attempt.result, attempt.candidateCount, attempt.topK, attempt.loadSkew, nil
- }
- if len(regularAccounts) > 0 {
- regularAttempt := s.trySelectByLoadBalancePool(ctx, req, regularAccounts, loadMap, budget)
- if regularAttempt.err != nil && !regularAttempt.noCompactCandidates {
- return nil, regularAttempt.candidateCount, regularAttempt.topK, regularAttempt.loadSkew, regularAttempt.err
- }
- if regularAttempt.result != nil {
- return regularAttempt.result, regularAttempt.candidateCount, regularAttempt.topK, regularAttempt.loadSkew, nil
- }
- var result *AccountSelectionResult
- candidateCount, topK, loadSkew := regularAttempt.candidateCount, regularAttempt.topK, regularAttempt.loadSkew
- fallbackErr := regularAttempt.err
- if regularAttempt.err == nil {
- result, candidateCount, topK, loadSkew, fallbackErr = s.finishLoadBalanceSelectionFallback(ctx, req, regularAttempt, budget, filterStats)
- if fallbackErr == nil && result != nil {
- return result, candidateCount, topK, loadSkew, nil
- }
- }
- // 常规池既无法获取也无法排队(含仅剩不支持 compact 的候选)时,
- // 回退到订阅池的等待计划:busy-but-waitable 的订阅账号不应因常规池存在
- // 而被丢弃,否则开启订阅优先反而让本可排队成功的请求硬失败。
- subResult, subCandidateCount, subTopK, subLoadSkew, subErr := s.finishLoadBalanceSelectionFallback(ctx, req, attempt, budget, filterStats)
- if subErr == nil && subResult != nil {
- return subResult, subCandidateCount, subTopK, subLoadSkew, nil
- }
- return result, candidateCount, topK, loadSkew, fallbackErr
- }
- return s.finishLoadBalanceSelectionFallback(ctx, req, attempt, budget, filterStats)
- }
- }
- attempt := s.trySelectByLoadBalancePool(ctx, req, filtered, loadMap, budget)
- if attempt.err != nil {
- return nil, attempt.candidateCount, attempt.topK, attempt.loadSkew, attempt.err
- }
- if attempt.result != nil {
- return attempt.result, attempt.candidateCount, attempt.topK, attempt.loadSkew, nil
- }
- return s.finishLoadBalanceSelectionFallback(ctx, req, attempt, budget, filterStats)
- }
- func partitionOpenAIChatGPTSubscriptionAccounts(accounts []*Account) ([]*Account, []*Account) {
- subscriptionAccounts := make([]*Account, 0, len(accounts))
- regularAccounts := make([]*Account, 0, len(accounts))
- for _, account := range accounts {
- if account != nil && account.IsOpenAIChatGPTSubscription() {
- subscriptionAccounts = append(subscriptionAccounts, account)
- continue
- }
- regularAccounts = append(regularAccounts, account)
- }
- return subscriptionAccounts, regularAccounts
- }
- func (s *defaultOpenAIAccountScheduler) trySelectByLoadBalancePool(
- ctx context.Context,
- req OpenAIAccountScheduleRequest,
- filtered []*Account,
- loadMap map[int64]*AccountLoadInfo,
- budget *openAISelectionProbeBudget,
- ) openAIAccountLoadSelectionAttempt {
- plan := s.buildOpenAIAccountLoadPlan(ctx, req, filtered, loadMap)
- if openAICostOverflowExpanded(req, plan) {
- budget.enableLimit()
- }
- attempt := openAIAccountLoadSelectionAttempt{
- selectionOrder: plan.selectionOrder,
- candidateCount: plan.candidateCount,
- topK: plan.topK,
- loadSkew: plan.loadSkew,
- }
- if req.RequireCompact && len(plan.candidates) == 0 && len(plan.staleSnapshotCompactRetry) == 0 {
- attempt.noCompactCandidates = true
- attempt.err = ErrNoAvailableCompactAccounts
- return attempt
- }
- if req.RequireCompact && len(attempt.selectionOrder) == 0 && s.service.schedulerSnapshot == nil {
- attempt.noCompactCandidates = true
- attempt.err = ErrNoAvailableCompactAccounts
- return attempt
- }
- if len(attempt.selectionOrder) == 0 {
- attempt.compactBlocked = req.RequireCompact && len(plan.allCandidates) > 0
- return attempt
- }
- result, compactBlocked, acquireErr := s.tryAcquireOpenAISelectionOrderWithBudget(ctx, req, attempt.selectionOrder, budget)
- attempt.compactBlocked = compactBlocked
- if acquireErr != nil {
- attempt.err = acquireErr
- return attempt
- }
- if result != nil {
- attempt.result = result
- return attempt
- }
- if s.service.concurrencyService != nil && !budget.acquireExhausted() {
- loadReq := buildOpenAIAccountLoadRequest(filtered)
- if freshLoadMap, loadErr := s.service.concurrencyService.GetAccountsLoadBatchFresh(ctx, loadReq); loadErr == nil {
- freshPlan := s.buildOpenAIAccountLoadPlan(ctx, req, filtered, freshLoadMap)
- if openAICostOverflowExpanded(req, freshPlan) {
- budget.enableLimit()
- }
- if len(freshPlan.selectionOrder) > 0 {
- freshResult, freshCompactBlocked, freshAcquireErr := s.tryAcquireOpenAISelectionOrderWithBudget(ctx, req, freshPlan.selectionOrder, budget)
- if freshAcquireErr != nil {
- attempt.err = freshAcquireErr
- return attempt
- }
- if freshResult != nil {
- attempt.result = freshResult
- attempt.selectionOrder = freshPlan.selectionOrder
- attempt.candidateCount = freshPlan.candidateCount
- attempt.topK = freshPlan.topK
- attempt.loadSkew = freshPlan.loadSkew
- return attempt
- }
- attempt.compactBlocked = attempt.compactBlocked || freshCompactBlocked
- attempt.selectionOrder = freshPlan.selectionOrder
- attempt.candidateCount = freshPlan.candidateCount
- attempt.topK = freshPlan.topK
- attempt.loadSkew = freshPlan.loadSkew
- }
- }
- }
- return attempt
- }
- func openAICostOverflowExpanded(req OpenAIAccountScheduleRequest, plan openAIAccountLoadPlan) bool {
- if !plan.includeOverflowFallback || plan.topK <= 0 {
- return false
- }
- if !req.RequireCompact {
- return len(plan.candidates) > plan.topK
- }
- supported, unknown := 0, 0
- for _, candidate := range plan.candidates {
- switch openAICompactSupportTier(candidate.account) {
- case 2:
- supported++
- case 1:
- unknown++
- }
- }
- return supported > plan.topK || unknown > plan.topK
- }
- func buildOpenAIAccountLoadRequest(accounts []*Account) []AccountWithConcurrency {
- loadReq := make([]AccountWithConcurrency, 0, len(accounts))
- for _, account := range accounts {
- if account == nil {
- continue
- }
- loadReq = append(loadReq, AccountWithConcurrency{
- ID: account.ID,
- MaxConcurrency: account.EffectiveLoadFactor(),
- })
- }
- return loadReq
- }
- func (s *defaultOpenAIAccountScheduler) finishLoadBalanceSelectionFallback(
- ctx context.Context,
- req OpenAIAccountScheduleRequest,
- attempt openAIAccountLoadSelectionAttempt,
- budget *openAISelectionProbeBudget,
- filterStats openAISelectionFilterStats,
- ) (*AccountSelectionResult, int, int, float64, error) {
- candidateCount := attempt.candidateCount
- topK := attempt.topK
- loadSkew := attempt.loadSkew
- if len(attempt.selectionOrder) == 0 {
- return nil, candidateCount, topK, loadSkew, noAvailableOpenAISelectionError(req.RequestedModel, attempt.compactBlocked, filterStats.summary("selection_order_empty"))
- }
- if stickyFallback, stickyErr := s.tryFallbackToWeightedSticky(ctx, req); stickyErr != nil {
- return nil, candidateCount, topK, loadSkew, stickyErr
- } else if stickyFallback != nil {
- return stickyFallback, candidateCount, topK, loadSkew, nil
- }
- cfg := s.service.schedulingConfig()
- compactBlocked := attempt.compactBlocked
- // WaitPlan.MaxConcurrency 使用 Concurrency(非 EffectiveLoadFactor),因为 WaitPlan 控制的是 Redis 实际并发槽位等待。
- passes := 1
- if budget != nil && budget.limited {
- passes = 4
- }
- for pass := 0; pass < passes; pass++ {
- wantAttempted := pass == 1 || pass == 3
- wantKnownFull := pass >= 2
- for _, candidate := range attempt.selectionOrder {
- if candidate.account == nil {
- continue
- }
- if budget != nil && budget.limited {
- knownFull := candidate.loadKnown && candidate.account.Mode1EffectiveConcurrency() > 0 &&
- candidate.loadInfo.CurrentConcurrency >= candidate.account.Mode1EffectiveConcurrency()
- if budget.wasAttempted(candidate.account.ID) != wantAttempted || knownFull != wantKnownFull {
- continue
- }
- }
- fresh := s.service.resolveFreshSchedulableOpenAIAccount(ctx, candidate.account, req.Platform, req.RequestedModel, false, req.RequiredCapability)
- if fresh == nil || !s.isAccountTransportCompatible(fresh, req.RequiredTransport) || !s.isAccountRequestCompatible(ctx, fresh, req) {
- continue
- }
- if !s.consumeOpenAISelectionDBRecheck(budget) {
- return nil, candidateCount, topK, loadSkew, noAvailableOpenAISelectionError(req.RequestedModel, compactBlocked, filterStats.summary("selection_order_exhausted"))
- }
- fresh = s.service.recheckSelectedOpenAIAccountFromDB(ctx, fresh, req.GroupID, req.Platform, req.RequestedModel, false, req.RequiredCapability)
- if fresh == nil || !s.isAccountTransportCompatible(fresh, req.RequiredTransport) || !s.isAccountRequestCompatible(ctx, fresh, req) {
- continue
- }
- if req.RequireCompact && openAICompactSupportTier(fresh) == 0 {
- compactBlocked = true
- continue
- }
- return attachSelectionProfitGate(ctx, &AccountSelectionResult{
- Account: fresh,
- WaitPlan: &AccountWaitPlan{
- AccountID: fresh.ID,
- MaxConcurrency: fresh.Mode1EffectiveConcurrency(),
- Timeout: cfg.FallbackWaitTimeout,
- MaxWaiting: cfg.FallbackMaxWaiting,
- },
- }), candidateCount, topK, loadSkew, nil
- }
- }
- return nil, candidateCount, topK, loadSkew, noAvailableOpenAISelectionError(req.RequestedModel, compactBlocked, filterStats.summary("selection_order_exhausted"))
- }
- func (s *defaultOpenAIAccountScheduler) isAccountTransportCompatible(account *Account, requiredTransport OpenAIUpstreamTransport) bool {
- if requiredTransport == OpenAIUpstreamTransportAny || requiredTransport == OpenAIUpstreamTransportHTTPSSE {
- return true
- }
- if s == nil || s.service == nil {
- return false
- }
- return s.service.isOpenAIAccountTransportCompatible(account, requiredTransport)
- }
- func (s *defaultOpenAIAccountScheduler) lookupShadowParentAccount(ctx context.Context, id int64) *Account {
- if s == nil || s.service == nil {
- return nil
- }
- if s.service.schedulerSnapshot != nil {
- if account, err := s.service.schedulerSnapshot.GetAccount(ctx, id); err == nil && account != nil {
- return account
- }
- }
- if s.service.accountRepo == nil {
- return nil
- }
- account, _ := s.service.accountRepo.GetByID(ctx, id)
- return account
- }
- func (s *defaultOpenAIAccountScheduler) isAccountRequestCompatible(ctx context.Context, account *Account, req OpenAIAccountScheduleRequest) bool {
- compatible, _ := s.isAccountRequestCompatibleReason(ctx, account, req)
- return compatible
- }
- // isAccountRequestCompatibleReason reports whether the account can serve the
- // request, and when it cannot, names the veto point. The reason feeds
- // openAISelectionFilterStats so that "no available accounts" errors state why
- // each candidate was dropped instead of failing silently (#4599).
- func (s *defaultOpenAIAccountScheduler) isAccountRequestCompatibleReason(ctx context.Context, account *Account, req OpenAIAccountScheduleRequest) (bool, string) {
- if account == nil {
- return false, "account_nil"
- }
- if req.RequirePrivacySet && !account.IsPrivacySet() {
- return false, "privacy_not_set"
- }
- if s != nil && s.service != nil && s.service.isOpenAIAccountRequestRuntimeBlocked(account, req.RequestedModel) {
- return false, "runtime_blocked"
- }
- if s != nil && s.service != nil && s.service.isOpenAIProxyStreamQuarantined(ctx, account) {
- return false, "proxy_stream_quarantined"
- }
- // Quota auto-pause must be evaluated during the initial filter too. Without it the
- // TopK candidate pool can be filled with paused accounts and the later fresh/DB
- // rechecks won't reach healthy accounts that fell outside TopK — manifesting as
- // "no available accounts" even though healthy ones exist.
- if paused, decision := shouldAutoPauseOpenAIAccountByQuota(ctx, account); paused {
- if decision.reason != "" {
- return false, decision.reason
- }
- reason := "quota_auto_pause"
- if decision.window != "" {
- reason += "_" + decision.window
- }
- return false, reason
- }
- // 母账号健康联动:影子账号的凭据来自母账号,母账号不可调度时影子也不应被选中。
- // Parent-health gate: shadow borrows the parent's credentials; an unschedulable
- // parent must block the shadow across all scheduler paths.
- if !parentHealthyForShadow(account, func(id int64) *Account {
- return s.lookupShadowParentAccount(ctx, id)
- }) {
- return false, "shadow_parent_unhealthy"
- }
- if req.RequestedModel != "" && !account.IsModelSupported(req.RequestedModel) {
- return false, "model_not_supported"
- }
- if req.GroupID != nil && s != nil && s.service != nil &&
- s.service.needsUpstreamChannelRestrictionCheck(ctx, req.GroupID) &&
- s.service.isUpstreamModelRestrictedByChannel(ctx, *req.GroupID, account, req.RequestedModel, req.RequireCompact) {
- return false, "channel_upstream_restricted"
- }
- if !accountSupportsOpenAICapabilities(account, req.RequiredCapability, req.RequiredImageCapability) {
- return false, "capability_mismatch"
- }
- // 分组利润控制:不合格账号在候选过滤与抢槽后终检阶段即被排除,
- // 排序/评分/粘性/熔断只在合格账号之间工作;named reason 进入 filter stats。
- if vetoed, reason := openAIProfitControlVetoReason(ctx, account); vetoed {
- return false, reason
- }
- return true, ""
- }
- func (s *defaultOpenAIAccountScheduler) ReportResult(accountID int64, success bool, firstTokenMs *int) {
- if s == nil || s.stats == nil {
- return
- }
- s.stats.report(accountID, success, firstTokenMs)
- }
- func (s *defaultOpenAIAccountScheduler) ReportSwitch() {
- if s == nil {
- return
- }
- s.metrics.recordSwitch()
- }
- func (s *defaultOpenAIAccountScheduler) SnapshotMetrics() OpenAIAccountSchedulerMetricsSnapshot {
- if s == nil {
- return OpenAIAccountSchedulerMetricsSnapshot{}
- }
- selectTotal := s.metrics.selectTotal.Load()
- prevHit := s.metrics.stickyPreviousHitTotal.Load()
- sessionHit := s.metrics.stickySessionHitTotal.Load()
- switchTotal := s.metrics.accountSwitchTotal.Load()
- latencyTotal := s.metrics.latencyMsTotal.Load()
- loadSkewTotal := s.metrics.loadSkewMilliTotal.Load()
- snapshot := OpenAIAccountSchedulerMetricsSnapshot{
- SelectTotal: selectTotal,
- StickyPreviousHitTotal: prevHit,
- StickySessionHitTotal: sessionHit,
- LoadBalanceSelectTotal: s.metrics.loadBalanceSelectTotal.Load(),
- AccountSwitchTotal: switchTotal,
- SchedulerLatencyMsTotal: latencyTotal,
- RuntimeStatsAccountCount: s.stats.size(),
- }
- if selectTotal > 0 {
- snapshot.SchedulerLatencyMsAvg = float64(latencyTotal) / float64(selectTotal)
- snapshot.StickyHitRatio = float64(prevHit+sessionHit) / float64(selectTotal)
- snapshot.AccountSwitchRate = float64(switchTotal) / float64(selectTotal)
- snapshot.LoadSkewAvg = float64(loadSkewTotal) / 1000 / float64(selectTotal)
- }
- return snapshot
- }
- func (s *OpenAIGatewayService) openAIAdvancedSchedulerSettingRepo() SettingRepository {
- if s == nil || s.rateLimitService == nil || s.rateLimitService.settingService == nil {
- return nil
- }
- return s.rateLimitService.settingService.settingRepo
- }
- func (s *OpenAIGatewayService) openAIAdvancedSchedulerRuntimeSettings(ctx context.Context) openAIAdvancedSchedulerRuntimeSettings {
- if cached, ok := openAIAdvancedSchedulerSettingCache.Load().(*cachedOpenAIAdvancedSchedulerSetting); ok && cached != nil {
- if time.Now().UnixNano() < cached.expiresAt {
- return openAIAdvancedSchedulerRuntimeSettings{
- lowUpstreamRatePriorityEnabled: cached.lowUpstreamRatePriorityEnabled,
- oauthSchedulingRateMultiplier: cached.oauthSchedulingRateMultiplier,
- enabled: cached.enabled,
- stickyWeightedEnabled: cached.stickyWeightedEnabled,
- subscriptionPriorityEnabled: cached.subscriptionPriorityEnabled,
- lbTopKOverride: cached.lbTopKOverride,
- weightOverrides: cloneOpenAIAdvancedSchedulerWeightOverrides(cached.weightOverrides),
- }
- }
- }
- result, _, _ := openAIAdvancedSchedulerSettingSF.Do(openAIAdvancedSchedulerSettingKey, func() (any, error) {
- if cached, ok := openAIAdvancedSchedulerSettingCache.Load().(*cachedOpenAIAdvancedSchedulerSetting); ok && cached != nil {
- if time.Now().UnixNano() < cached.expiresAt {
- return openAIAdvancedSchedulerRuntimeSettings{
- lowUpstreamRatePriorityEnabled: cached.lowUpstreamRatePriorityEnabled,
- oauthSchedulingRateMultiplier: cached.oauthSchedulingRateMultiplier,
- enabled: cached.enabled,
- stickyWeightedEnabled: cached.stickyWeightedEnabled,
- subscriptionPriorityEnabled: cached.subscriptionPriorityEnabled,
- lbTopKOverride: cached.lbTopKOverride,
- weightOverrides: cloneOpenAIAdvancedSchedulerWeightOverrides(cached.weightOverrides),
- }, nil
- }
- }
- lowUpstreamRatePriorityEnabled := false
- oauthSchedulingRateMultiplier := defaultOpenAIOAuthSchedulingRateMultiplier
- enabled := false
- stickyWeightedEnabled := false
- subscriptionPriorityEnabled := false
- lbTopKOverride := 0
- weightOverrides := map[string]float64{}
- if repo := s.openAIAdvancedSchedulerSettingRepo(); repo != nil {
- dbCtx, cancel := context.WithTimeout(context.WithoutCancel(ctx), openAIAdvancedSchedulerSettingDBTimeout)
- defer cancel()
- if values, err := repo.GetMultiple(dbCtx, openAIAdvancedSchedulerRuntimeSettingKeys()); err == nil {
- lowUpstreamRatePriorityEnabled = strings.EqualFold(strings.TrimSpace(values[SettingKeyOpenAILowUpstreamRatePriorityEnabled]), "true")
- oauthSchedulingRateMultiplier = parseOpenAIOAuthSchedulingRateMultiplier(values[SettingKeyOpenAIOAuthSchedulingRateMultiplier])
- enabled = strings.EqualFold(strings.TrimSpace(values[openAIAdvancedSchedulerSettingKey]), "true")
- stickyWeightedEnabled = strings.EqualFold(strings.TrimSpace(values[SettingKeyOpenAIAdvancedSchedulerStickyWeightedEnabled]), "true")
- subscriptionPriorityEnabled = strings.EqualFold(strings.TrimSpace(values[SettingKeyOpenAIAdvancedSchedulerSubscriptionPriorityEnabled]), "true")
- lbTopKOverride = parsePositiveIntOverride(values[SettingKeyOpenAIAdvancedSchedulerLBTopK])
- weightOverrides = parseOpenAIAdvancedSchedulerWeightOverrides(values)
- } else {
- // 批量读取失败时逐键降级,覆盖全部键(含 TopK/权重),避免只加载布尔开关
- // 而静默丢弃管理员配置的覆盖值;降级状态会被缓存一个 TTL,必须留痕。
- slog.Warn("openai_advanced_scheduler_settings_batch_load_failed", "error", err)
- fallbackValues := make(map[string]string)
- for _, key := range openAIAdvancedSchedulerRuntimeSettingKeys() {
- if value, valueErr := repo.GetValue(dbCtx, key); valueErr == nil {
- fallbackValues[key] = value
- }
- }
- lowUpstreamRatePriorityEnabled = strings.EqualFold(strings.TrimSpace(fallbackValues[SettingKeyOpenAILowUpstreamRatePriorityEnabled]), "true")
- oauthSchedulingRateMultiplier = parseOpenAIOAuthSchedulingRateMultiplier(fallbackValues[SettingKeyOpenAIOAuthSchedulingRateMultiplier])
- enabled = strings.EqualFold(strings.TrimSpace(fallbackValues[openAIAdvancedSchedulerSettingKey]), "true")
- stickyWeightedEnabled = strings.EqualFold(strings.TrimSpace(fallbackValues[SettingKeyOpenAIAdvancedSchedulerStickyWeightedEnabled]), "true")
- subscriptionPriorityEnabled = strings.EqualFold(strings.TrimSpace(fallbackValues[SettingKeyOpenAIAdvancedSchedulerSubscriptionPriorityEnabled]), "true")
- lbTopKOverride = parsePositiveIntOverride(fallbackValues[SettingKeyOpenAIAdvancedSchedulerLBTopK])
- weightOverrides = parseOpenAIAdvancedSchedulerWeightOverrides(fallbackValues)
- }
- }
- openAIAdvancedSchedulerSettingCache.Store(&cachedOpenAIAdvancedSchedulerSetting{
- lowUpstreamRatePriorityEnabled: lowUpstreamRatePriorityEnabled,
- oauthSchedulingRateMultiplier: oauthSchedulingRateMultiplier,
- enabled: enabled,
- stickyWeightedEnabled: stickyWeightedEnabled,
- subscriptionPriorityEnabled: subscriptionPriorityEnabled,
- lbTopKOverride: lbTopKOverride,
- weightOverrides: cloneOpenAIAdvancedSchedulerWeightOverrides(weightOverrides),
- expiresAt: time.Now().Add(openAIAdvancedSchedulerSettingCacheTTL).UnixNano(),
- })
- return openAIAdvancedSchedulerRuntimeSettings{
- lowUpstreamRatePriorityEnabled: lowUpstreamRatePriorityEnabled,
- oauthSchedulingRateMultiplier: oauthSchedulingRateMultiplier,
- enabled: enabled,
- stickyWeightedEnabled: stickyWeightedEnabled,
- subscriptionPriorityEnabled: subscriptionPriorityEnabled,
- lbTopKOverride: lbTopKOverride,
- weightOverrides: weightOverrides,
- }, nil
- })
- settings, _ := result.(openAIAdvancedSchedulerRuntimeSettings)
- return settings
- }
- func (s *OpenAIGatewayService) isOpenAIAdvancedSchedulerEnabled(ctx context.Context) bool {
- return s.openAIAdvancedSchedulerRuntimeSettings(ctx).enabled
- }
- func (s *OpenAIGatewayService) isOpenAILowUpstreamRatePriorityEnabled(ctx context.Context) bool {
- settings := s.openAIAdvancedSchedulerRuntimeSettings(ctx)
- return !settings.enabled && settings.lowUpstreamRatePriorityEnabled
- }
- func (s *OpenAIGatewayService) openAIOAuthSchedulingRateMultiplier(ctx context.Context) float64 {
- return s.openAIAdvancedSchedulerRuntimeSettings(ctx).oauthSchedulingRateMultiplier
- }
- func (s *OpenAIGatewayService) isOpenAIAdvancedSchedulerStickyWeightedEnabled(ctx context.Context) bool {
- settings := s.openAIAdvancedSchedulerRuntimeSettings(ctx)
- return settings.enabled && settings.stickyWeightedEnabled
- }
- func (s *OpenAIGatewayService) isOpenAIAdvancedSchedulerSubscriptionPriorityEnabled(ctx context.Context) bool {
- settings := s.openAIAdvancedSchedulerRuntimeSettings(ctx)
- return settings.enabled && settings.subscriptionPriorityEnabled
- }
- func openAIAdvancedSchedulerRuntimeSettingKeys() []string {
- keys := []string{
- SettingKeyOpenAILowUpstreamRatePriorityEnabled,
- SettingKeyOpenAIOAuthSchedulingRateMultiplier,
- openAIAdvancedSchedulerSettingKey,
- SettingKeyOpenAIAdvancedSchedulerStickyWeightedEnabled,
- SettingKeyOpenAIAdvancedSchedulerSubscriptionPriorityEnabled,
- SettingKeyOpenAIAdvancedSchedulerLBTopK,
- }
- for _, spec := range openAIAdvancedSchedulerWeightOverrideSpecs() {
- keys = append(keys, spec.key)
- }
- return keys
- }
- type openAIAdvancedSchedulerWeightOverrideSpec struct {
- key string
- name string
- }
- func openAIAdvancedSchedulerWeightOverrideSpecs() []openAIAdvancedSchedulerWeightOverrideSpec {
- return []openAIAdvancedSchedulerWeightOverrideSpec{
- {key: SettingKeyOpenAIAdvancedSchedulerWeightPriority, name: "priority"},
- {key: SettingKeyOpenAIAdvancedSchedulerWeightLoad, name: "load"},
- {key: SettingKeyOpenAIAdvancedSchedulerWeightQueue, name: "queue"},
- {key: SettingKeyOpenAIAdvancedSchedulerWeightErrorRate, name: "error_rate"},
- {key: SettingKeyOpenAIAdvancedSchedulerWeightTTFT, name: "ttft"},
- {key: SettingKeyOpenAIAdvancedSchedulerWeightReset, name: "reset"},
- {key: SettingKeyOpenAIAdvancedSchedulerWeightQuotaHeadroom, name: "quota_headroom"},
- {key: SettingKeyOpenAIAdvancedSchedulerWeightUpstreamCost, name: "upstream_cost"},
- {key: SettingKeyOpenAIAdvancedSchedulerWeightPreviousResponse, name: "previous_response"},
- {key: SettingKeyOpenAIAdvancedSchedulerWeightSessionSticky, name: "session_sticky"},
- }
- }
- func parsePositiveIntOverride(raw string) int {
- raw = strings.TrimSpace(raw)
- if raw == "" {
- return 0
- }
- value, err := strconv.Atoi(raw)
- if err != nil || value <= 0 {
- return 0
- }
- return value
- }
- func parseOpenAIAdvancedSchedulerWeightOverrides(values map[string]string) map[string]float64 {
- overrides := map[string]float64{}
- for _, spec := range openAIAdvancedSchedulerWeightOverrideSpecs() {
- raw := strings.TrimSpace(values[spec.key])
- if raw == "" {
- continue
- }
- value, err := strconv.ParseFloat(raw, 64)
- if err != nil || value < 0 || math.IsNaN(value) || math.IsInf(value, 0) {
- continue
- }
- overrides[spec.name] = value
- }
- return overrides
- }
- func cloneOpenAIAdvancedSchedulerWeightOverrides(in map[string]float64) map[string]float64 {
- if len(in) == 0 {
- return nil
- }
- out := make(map[string]float64, len(in))
- for key, value := range in {
- out[key] = value
- }
- return out
- }
- func (s *OpenAIGatewayService) getOpenAIAccountScheduler(ctx context.Context) OpenAIAccountScheduler {
- if s == nil {
- return nil
- }
- if !s.isOpenAIAdvancedSchedulerEnabled(ctx) {
- return nil
- }
- s.openaiSchedulerOnce.Do(func() {
- if s.openaiAccountStats == nil {
- s.openaiAccountStats = newOpenAIAccountRuntimeStats()
- }
- if s.openaiScheduler == nil {
- s.openaiScheduler = newDefaultOpenAIAccountScheduler(s, s.openaiAccountStats)
- }
- })
- return s.openaiScheduler
- }
- func resetOpenAIAdvancedSchedulerSettingCacheForTest() {
- openAIAdvancedSchedulerSettingCache = atomic.Value{}
- openAIAdvancedSchedulerSettingSF = singleflight.Group{}
- }
- func (s *OpenAIGatewayService) SelectAccountWithScheduler(
- ctx context.Context,
- groupID *int64,
- previousResponseID string,
- sessionHash string,
- requestedModel string,
- excludedIDs map[int64]struct{},
- requiredTransport OpenAIUpstreamTransport,
- requireCompact bool,
- ) (*AccountSelectionResult, OpenAIAccountScheduleDecision, error) {
- return s.selectAccountWithScheduler(ctx, groupID, previousResponseID, sessionHash, requestedModel, excludedIDs, requiredTransport, "", "", requireCompact, PlatformOpenAI, false, true)
- }
- // SelectAccountWithSchedulerForCapability 按能力要求调度账号。
- // previousResponseCanMove 表示首包 input 可自行重建工具续链,previous_response_id 允许跨账号迁移
- // (粘性加权模式下改为加权偏好而非硬粘连)。
- func (s *OpenAIGatewayService) SelectAccountWithSchedulerForCapability(
- ctx context.Context,
- groupID *int64,
- previousResponseID string,
- sessionHash string,
- requestedModel string,
- excludedIDs map[int64]struct{},
- requiredTransport OpenAIUpstreamTransport,
- requiredCapability OpenAIEndpointCapability,
- requireCompact bool,
- previousResponseCanMove bool,
- useUpstreamTokenCost bool,
- platformOverride ...string,
- ) (*AccountSelectionResult, OpenAIAccountScheduleDecision, error) {
- platform := PlatformOpenAI
- if len(platformOverride) > 0 {
- platform = platformOverride[0]
- }
- return s.selectAccountWithScheduler(ctx, groupID, previousResponseID, sessionHash, requestedModel, excludedIDs, requiredTransport, requiredCapability, "", requireCompact, platform, previousResponseCanMove, useUpstreamTokenCost)
- }
- func (s *OpenAIGatewayService) SelectAccountWithSchedulerForImages(
- ctx context.Context,
- groupID *int64,
- sessionHash string,
- requestedModel string,
- excludedIDs map[int64]struct{},
- requiredCapability OpenAIImagesCapability,
- ) (*AccountSelectionResult, OpenAIAccountScheduleDecision, error) {
- selection, decision, err := s.selectAccountWithScheduler(ctx, groupID, "", sessionHash, requestedModel, excludedIDs, OpenAIUpstreamTransportHTTPSSE, "", requiredCapability, false, PlatformOpenAI, false, false)
- if err == nil && selection != nil && selection.Account != nil {
- return selection, decision, nil
- }
- // 如果要求 native 能力(如指定了模型)但没有可用的 APIKey 账号,回退到 basic(OAuth 账号)
- if requiredCapability == OpenAIImagesCapabilityNative {
- return s.selectAccountWithScheduler(ctx, groupID, "", sessionHash, requestedModel, excludedIDs, OpenAIUpstreamTransportHTTPSSE, "", OpenAIImagesCapabilityBasic, false, PlatformOpenAI, false, false)
- }
- return selection, decision, err
- }
- // selectAccountWithScheduler wraps selectAccountWithSchedulerOnce with a
- // fail-open second pass for the proxy stream circuit (#5056): when the only
- // reason no account is available is that every candidate sits behind a
- // quarantined proxy, the quarantine must degrade to a preference instead of
- // zeroing out capacity. The retry re-runs the exact same selection with the
- // quarantine checks bypassed, so healthy proxies always win the first pass
- // and quarantined ones only serve when nothing else can.
- func (s *OpenAIGatewayService) selectAccountWithScheduler(
- ctx context.Context,
- groupID *int64,
- previousResponseID string,
- sessionHash string,
- requestedModel string,
- excludedIDs map[int64]struct{},
- requiredTransport OpenAIUpstreamTransport,
- requiredCapability OpenAIEndpointCapability,
- requiredImageCapability OpenAIImagesCapability,
- requireCompact bool,
- platform string,
- previousResponseCanMove bool,
- useUpstreamTokenCost bool,
- ) (*AccountSelectionResult, OpenAIAccountScheduleDecision, error) {
- selection, decision, err := s.selectAccountWithSchedulerOnce(ctx, groupID, previousResponseID, sessionHash, requestedModel, excludedIDs, requiredTransport, requiredCapability, requiredImageCapability, requireCompact, platform, previousResponseCanMove, useUpstreamTokenCost)
- if err == nil || openAIProxyStreamQuarantineBypassed(ctx) {
- return selection, decision, err
- }
- if !errors.Is(err, ErrNoAvailableAccounts) && !errors.Is(err, ErrNoAvailableCompactAccounts) {
- return selection, decision, err
- }
- // The circuit only ever quarantines PlatformOpenAI accounts.
- if NormalizeOpenAICompatiblePlatform(platform) != PlatformOpenAI {
- return selection, decision, err
- }
- blocked := s.getOpenAIProxyStreamCircuit().activeBlockCount(time.Now())
- if blocked == 0 {
- return selection, decision, err
- }
- s.logOpenAIProxyStreamQuarantineFailOpen(requestedModel, blocked)
- return s.selectAccountWithSchedulerOnce(withOpenAIProxyStreamQuarantineBypass(ctx), groupID, previousResponseID, sessionHash, requestedModel, excludedIDs, requiredTransport, requiredCapability, requiredImageCapability, requireCompact, platform, previousResponseCanMove, useUpstreamTokenCost)
- }
- type openAIGroupPrivacyRequirementContextKey struct{}
- type openAIGroupPrivacyRequirement struct {
- groupID int64
- required bool
- }
- func (s *OpenAIGatewayService) withOpenAIGroupPrivacyRequirement(ctx context.Context, groupID *int64) context.Context {
- return context.WithValue(ctx, openAIGroupPrivacyRequirementContextKey{}, openAIGroupPrivacyRequirement{
- groupID: derefGroupID(groupID),
- required: s.loadOpenAIGroupRequiresPrivacySet(ctx, groupID),
- })
- }
- func (s *OpenAIGatewayService) openAIGroupRequiresPrivacySet(ctx context.Context, groupID *int64) bool {
- if cached, ok := ctx.Value(openAIGroupPrivacyRequirementContextKey{}).(openAIGroupPrivacyRequirement); ok && cached.groupID == derefGroupID(groupID) {
- return cached.required
- }
- return s.loadOpenAIGroupRequiresPrivacySet(ctx, groupID)
- }
- func (s *OpenAIGatewayService) loadOpenAIGroupRequiresPrivacySet(ctx context.Context, groupID *int64) bool {
- if s == nil || groupID == nil || s.schedulerSnapshot == nil {
- return false
- }
- group, err := s.schedulerSnapshot.GetGroupByID(ctx, *groupID)
- if err != nil {
- return true
- }
- return group != nil && group.RequirePrivacySet
- }
- func (s *OpenAIGatewayService) selectAccountWithSchedulerOnce(
- ctx context.Context,
- groupID *int64,
- previousResponseID string,
- sessionHash string,
- requestedModel string,
- excludedIDs map[int64]struct{},
- requiredTransport OpenAIUpstreamTransport,
- requiredCapability OpenAIEndpointCapability,
- requiredImageCapability OpenAIImagesCapability,
- requireCompact bool,
- platform string,
- previousResponseCanMove bool,
- useUpstreamTokenCost bool,
- ) (*AccountSelectionResult, OpenAIAccountScheduleDecision, error) {
- ctx = s.withOpenAIQuotaAutoPauseContext(ctx)
- ctx = s.withOpenAIGroupPrivacyRequirement(ctx, groupID)
- // 分组利润控制:唯一文本调度入口的防御性装门。handler 文本
- // 入口已在请求开始经 WithOpenAIRequestPricingContext 装门并固定 pricingAt,
- // 此处对同分组门直接复用(failover 重入阈值稳定),仅为不经 handler 装配的
- // 内部调用兜底。图片/视频调度不在利润门范围:requiredImageCapability 非空的
- // Images 调度不装门;其他使用 Responses 能力的文本请求(包括原生远程压缩)
- // 仍须装门。其余媒体路径通过 WithOpenAIProfitControlSuppressed 显式跳过。
- if requiredImageCapability == "" {
- ctx = s.withOpenAIProfitControlGate(ctx, groupID)
- }
- platform = NormalizeOpenAICompatiblePlatform(platform)
- decision := OpenAIAccountScheduleDecision{}
- preserveGuardianParentBinding := preserveOpenAIGuardianParentBinding(ctx, sessionHash)
- guardianParentAccountID := int64(0)
- if strings.TrimSpace(previousResponseID) == "" {
- guardianParentAccountID = s.resolveOpenAIGuardianParentAccountID(ctx, groupID)
- }
- scheduler := s.getOpenAIAccountScheduler(ctx)
- if scheduler == nil {
- decision.Layer = openAIAccountScheduleLayerLoadBalance
- if guardianParentAccountID > 0 {
- if s.checkChannelPricingRestriction(ctx, groupID, requestedModel) {
- return nil, decision, fmt.Errorf("%w supporting model: %s (channel pricing restriction)", ErrNoAvailableAccounts, requestedModel)
- }
- fallbackScheduler := &defaultOpenAIAccountScheduler{service: s, stats: newOpenAIAccountRuntimeStats()}
- selection, _, err := fallbackScheduler.selectBySessionHash(ctx, OpenAIAccountScheduleRequest{
- GroupID: groupID,
- Platform: platform,
- SessionHash: sessionHash,
- StickyAccountID: guardianParentAccountID,
- PreserveStickyBinding: true,
- RequestedModel: requestedModel,
- RequiredTransport: requiredTransport,
- RequiredCapability: requiredCapability,
- RequiredImageCapability: requiredImageCapability,
- RequireCompact: requireCompact,
- ExcludedIDs: excludedIDs,
- RequirePrivacySet: s.openAIGroupRequiresPrivacySet(ctx, groupID),
- })
- if err != nil {
- return nil, decision, err
- }
- if selection != nil && selection.Account != nil {
- decision.Layer = openAIAccountScheduleLayerGuardianParent
- decision.StickySessionHit = true
- decision.SelectedAccountID = selection.Account.ID
- decision.SelectedAccountType = selection.Account.Type
- return selection, decision, nil
- }
- }
- legacySessionHash := sessionHash
- if preserveGuardianParentBinding {
- legacySessionHash = ""
- }
- if requiredTransport == OpenAIUpstreamTransportAny || requiredTransport == OpenAIUpstreamTransportHTTPSSE {
- effectiveExcludedIDs := cloneExcludedAccountIDs(excludedIDs)
- for {
- selection, err := s.selectAccountWithLoadAwareness(ctx, groupID, platform, legacySessionHash, requestedModel, effectiveExcludedIDs, requireCompact, requiredCapability, useUpstreamTokenCost)
- if err != nil {
- return nil, decision, err
- }
- if selection == nil || selection.Account == nil {
- return selection, decision, nil
- }
- if accountSupportsOpenAICapabilities(selection.Account, requiredCapability, requiredImageCapability) {
- return selection, decision, nil
- }
- if selection.ReleaseFunc != nil {
- selection.ReleaseFunc()
- }
- if effectiveExcludedIDs == nil {
- effectiveExcludedIDs = make(map[int64]struct{})
- }
- if _, exists := effectiveExcludedIDs[selection.Account.ID]; exists {
- return nil, decision, ErrNoAvailableAccounts
- }
- effectiveExcludedIDs[selection.Account.ID] = struct{}{}
- }
- }
- effectiveExcludedIDs := cloneExcludedAccountIDs(excludedIDs)
- for {
- selection, err := s.selectAccountWithLoadAwareness(ctx, groupID, platform, legacySessionHash, requestedModel, effectiveExcludedIDs, requireCompact, requiredCapability, useUpstreamTokenCost)
- if err != nil {
- return nil, decision, err
- }
- if selection == nil || selection.Account == nil {
- return selection, decision, nil
- }
- if s.isOpenAIAccountTransportCompatible(selection.Account, requiredTransport) &&
- accountSupportsOpenAICapabilities(selection.Account, requiredCapability, requiredImageCapability) {
- return selection, decision, nil
- }
- if selection.ReleaseFunc != nil {
- selection.ReleaseFunc()
- }
- if effectiveExcludedIDs == nil {
- effectiveExcludedIDs = make(map[int64]struct{})
- }
- if _, exists := effectiveExcludedIDs[selection.Account.ID]; exists {
- return nil, decision, ErrNoAvailableAccounts
- }
- effectiveExcludedIDs[selection.Account.ID] = struct{}{}
- }
- }
- if s.checkChannelPricingRestriction(ctx, groupID, requestedModel) {
- slog.Warn("channel pricing restriction blocked request",
- "group_id", derefGroupID(groupID),
- "model", requestedModel)
- return nil, decision, fmt.Errorf("%w supporting model: %s (channel pricing restriction)", ErrNoAvailableAccounts, requestedModel)
- }
- var stickyAccountID int64
- if sessionHash != "" && s.cache != nil {
- if accountID, err := s.getStickySessionAccountID(ctx, groupID, sessionHash); err == nil && accountID > 0 {
- stickyAccountID = accountID
- }
- }
- stickyWeighted := s.isOpenAIAdvancedSchedulerStickyWeightedEnabled(ctx)
- subscriptionPriority := s.isOpenAIAdvancedSchedulerSubscriptionPriorityEnabled(ctx)
- stickyPreviousAccountID := int64(0)
- if stickyWeighted && previousResponseCanMove && strings.TrimSpace(previousResponseID) != "" && platform == PlatformOpenAI {
- stickyPreviousAccountID = s.ResolveAccountIDByPreviousResponseIDForScheduler(ctx, groupID, previousResponseID, requestedModel, excludedIDs, requiredCapability, requireCompact)
- }
- selection, decision, err := scheduler.Select(ctx, OpenAIAccountScheduleRequest{
- GroupID: groupID,
- Platform: platform,
- SessionHash: sessionHash,
- StickyAccountID: stickyAccountID,
- GuardianParentAccountID: guardianParentAccountID,
- StickyPreviousAccountID: stickyPreviousAccountID,
- StickyWeighted: stickyWeighted,
- SubscriptionPriority: subscriptionPriority,
- PreserveStickyBinding: preserveGuardianParentBinding,
- RequirePrivacySet: s.openAIGroupRequiresPrivacySet(ctx, groupID),
- PreviousResponseID: previousResponseID,
- PreviousResponseCanMove: previousResponseCanMove,
- UseUpstreamTokenCost: useUpstreamTokenCost,
- RequestedModel: requestedModel,
- RequiredTransport: requiredTransport,
- RequiredCapability: requiredCapability,
- RequiredImageCapability: requiredImageCapability,
- RequireCompact: requireCompact,
- ExcludedIDs: excludedIDs,
- })
- if err != nil || selection == nil || selection.Account == nil {
- return selection, decision, err
- }
- // The advanced scheduler has several direct sticky/previous-response
- // branches that construct AccountSelectionResult without going through
- // newSelectionResult. Resolve only the request-time random proxy here: the
- // other branches already re-check the account in the scheduler, and a full
- // snapshot hydration for every non-random request would add an avoidable
- // Redis/DB read on this hot path.
- if selection.Account.IsRandomProxy() {
- if hydrateErr := ResolveRandomProxyFromSource(ctx, selection.Account, s.accountRepo); hydrateErr != nil {
- if selection.ReleaseFunc != nil {
- selection.ReleaseFunc()
- selection.ReleaseFunc = nil
- }
- return nil, decision, hydrateErr
- }
- }
- return selection, decision, nil
- }
- func accountSupportsOpenAICapabilities(account *Account, requiredCapability OpenAIEndpointCapability, requiredImageCapability OpenAIImagesCapability) bool {
- if account == nil {
- return false
- }
- return account.SupportsOpenAIEndpointCapability(requiredCapability) &&
- account.SupportsOpenAIImageCapability(requiredImageCapability)
- }
- func cloneExcludedAccountIDs(excludedIDs map[int64]struct{}) map[int64]struct{} {
- if len(excludedIDs) == 0 {
- return nil
- }
- cloned := make(map[int64]struct{}, len(excludedIDs))
- for id := range excludedIDs {
- cloned[id] = struct{}{}
- }
- return cloned
- }
- func (s *OpenAIGatewayService) isOpenAIAccountTransportCompatible(account *Account, requiredTransport OpenAIUpstreamTransport) bool {
- if requiredTransport == OpenAIUpstreamTransportAny || requiredTransport == OpenAIUpstreamTransportHTTPSSE {
- return true
- }
- if s == nil || account == nil {
- return false
- }
- if requiredTransport == OpenAIUpstreamTransportResponsesWebsocketV2Ingress {
- if s.cfg == nil || !s.cfg.Gateway.OpenAIWS.ModeRouterV2Enabled {
- return s.getOpenAIWSProtocolResolver().Resolve(account).Transport == OpenAIUpstreamTransportResponsesWebsocketV2
- }
- mode := account.ResolveOpenAIResponsesWebSocketV2Mode(s.cfg.Gateway.OpenAIWS.IngressModeDefault)
- switch mode {
- case OpenAIWSIngressModeCtxPool, OpenAIWSIngressModePassthrough, OpenAIWSIngressModeHTTPBridge, OpenAIWSIngressModeShared, OpenAIWSIngressModeDedicated:
- return true
- default:
- return false
- }
- }
- return s.getOpenAIWSProtocolResolver().Resolve(account).Transport == requiredTransport
- }
- func (s *OpenAIGatewayService) ReportOpenAIAccountScheduleResult(account *Account, model string, success bool, firstTokenMs *int, observedErr ...error) bool {
- if account == nil {
- return false
- }
- accountID := account.ID
- healthTripped := false
- if s != nil && s.rateLimitService != nil {
- if success {
- s.rateLimitService.ObserveOpenAIAPIKeyHealthSuccess(context.Background(), account)
- } else if len(observedErr) > 0 && observedErr[0] != nil {
- healthTripped = s.rateLimitService.ObserveOpenAIAPIKeyHealthFailure(context.Background(), account, observedErr[0])
- }
- }
- if success {
- s.openaiOAuth429RetryStartedAt.Delete(accountID)
- s.clearOpenAIAccountModelTransientState(accountID, normalizeOpenAIAccountModelTransientModel(model))
- }
- scheduler := s.getOpenAIAccountScheduler(context.Background())
- if scheduler == nil {
- return healthTripped
- }
- scheduler.ReportResult(accountID, success, firstTokenMs)
- return healthTripped
- }
- // ObserveOpenAIAccountHealthFailure records failures that cannot reach the
- // scheduler-result path, for example after semantic response bytes were sent.
- func (s *OpenAIGatewayService) ObserveOpenAIAccountHealthFailure(ctx context.Context, account *Account, observedErr error) bool {
- if s == nil || s.rateLimitService == nil || account == nil || observedErr == nil {
- return false
- }
- return s.rateLimitService.ObserveOpenAIAPIKeyHealthFailure(ctx, account, observedErr)
- }
- func (s *OpenAIGatewayService) RecordOpenAIAccountSwitch() {
- scheduler := s.getOpenAIAccountScheduler(context.Background())
- if scheduler == nil {
- return
- }
- scheduler.ReportSwitch()
- }
- func (s *OpenAIGatewayService) SnapshotOpenAIAccountSchedulerMetrics() OpenAIAccountSchedulerMetricsSnapshot {
- scheduler := s.getOpenAIAccountScheduler(context.Background())
- if scheduler == nil {
- return OpenAIAccountSchedulerMetricsSnapshot{}
- }
- return scheduler.SnapshotMetrics()
- }
- func (s *OpenAIGatewayService) openAIWSSessionStickyTTL() time.Duration {
- if s != nil && s.cfg != nil && s.cfg.Gateway.OpenAIWS.StickySessionTTLSeconds > 0 {
- return time.Duration(s.cfg.Gateway.OpenAIWS.StickySessionTTLSeconds) * time.Second
- }
- return openaiStickySessionTTL
- }
- func (s *OpenAIGatewayService) openAIWSLBTopK() int {
- if s != nil && s.cfg != nil && s.cfg.Gateway.OpenAIWS.LBTopK > 0 {
- return s.cfg.Gateway.OpenAIWS.LBTopK
- }
- return 7
- }
- func (s *OpenAIGatewayService) openAIWSLBTopKForRequest(ctx context.Context) int {
- base := s.openAIWSLBTopK()
- settings := s.openAIAdvancedSchedulerRuntimeSettings(ctx)
- // DB 覆盖值与 stickyWeighted/subscriptionPriority 一样受总开关门控:
- // 关闭高级调度器后所有调用方(含管理页分数快照)都应回到配置/默认行为。
- if !settings.enabled {
- return base
- }
- if settings.lbTopKOverride > 0 {
- return settings.lbTopKOverride
- }
- return base
- }
- func (s *OpenAIGatewayService) openAIStickyEscapeConfig() openAIStickyEscapeConfig {
- if s != nil && s.cfg != nil {
- cfg := s.cfg.Gateway.OpenAIScheduler
- enabled := cfg.StickyEscapeEnabled
- if !enabled && cfg.StickyEscapeTTFTMs == 0 && cfg.StickyEscapeErrorRate == 0 {
- enabled = true
- }
- ttftMs := float64(cfg.StickyEscapeTTFTMs)
- if ttftMs <= 0 {
- ttftMs = 15000
- }
- errorRate := cfg.StickyEscapeErrorRate
- if errorRate < 0 || errorRate > 1 {
- errorRate = 0.5
- }
- if errorRate == 0 && cfg.StickyEscapeTTFTMs == 0 && cfg.StickyEscapeErrorRate == 0 {
- errorRate = 0.5
- }
- return openAIStickyEscapeConfig{
- enabled: enabled,
- ttftMs: ttftMs,
- errorRate: errorRate,
- }
- }
- return openAIStickyEscapeConfig{
- enabled: true,
- ttftMs: 15000,
- errorRate: 0.5,
- }
- }
- func (s *OpenAIGatewayService) openAIWSSchedulerWeights() GatewayOpenAIWSSchedulerScoreWeightsView {
- if s != nil && s.cfg != nil {
- return GatewayOpenAIWSSchedulerScoreWeightsView{
- Priority: s.cfg.Gateway.OpenAIWS.SchedulerScoreWeights.Priority,
- Load: s.cfg.Gateway.OpenAIWS.SchedulerScoreWeights.Load,
- Queue: s.cfg.Gateway.OpenAIWS.SchedulerScoreWeights.Queue,
- ErrorRate: s.cfg.Gateway.OpenAIWS.SchedulerScoreWeights.ErrorRate,
- TTFT: s.cfg.Gateway.OpenAIWS.SchedulerScoreWeights.TTFT,
- Reset: s.cfg.Gateway.OpenAIWS.SchedulerScoreWeights.Reset,
- QuotaHeadroom: s.cfg.Gateway.OpenAIWS.SchedulerScoreWeights.QuotaHeadroom,
- UpstreamCost: s.cfg.Gateway.OpenAIWS.SchedulerScoreWeights.UpstreamCost,
- Previous: s.cfg.Gateway.OpenAIWS.SchedulerScoreWeights.PreviousResponse,
- SessionSticky: s.cfg.Gateway.OpenAIWS.SchedulerScoreWeights.SessionSticky,
- }
- }
- return GatewayOpenAIWSSchedulerScoreWeightsView{
- Priority: 1.0,
- Load: 1.0,
- Queue: 0.7,
- ErrorRate: 0.8,
- TTFT: 0.5,
- Reset: 0.0,
- QuotaHeadroom: 0.0,
- UpstreamCost: 0.0,
- Previous: 5.0,
- SessionSticky: 3.0,
- }
- }
- func (s *OpenAIGatewayService) openAIWSSchedulerWeightsForRequest(ctx context.Context) GatewayOpenAIWSSchedulerScoreWeightsView {
- weights := s.openAIWSSchedulerWeights()
- settings := s.openAIAdvancedSchedulerRuntimeSettings(ctx)
- // 同 openAIWSLBTopKForRequest:总开关关闭时不应用 DB 覆盖值。
- if !settings.enabled {
- return weights
- }
- overridden := applyOpenAIAdvancedSchedulerWeightOverrides(weights, settings.weightOverrides)
- if !overridden.configWeights().IsValid() {
- return weights
- }
- return overridden
- }
- func applyOpenAIAdvancedSchedulerWeightOverrides(
- weights GatewayOpenAIWSSchedulerScoreWeightsView,
- overrides map[string]float64,
- ) GatewayOpenAIWSSchedulerScoreWeightsView {
- for key, value := range overrides {
- switch key {
- case "priority":
- weights.Priority = value
- case "load":
- weights.Load = value
- case "queue":
- weights.Queue = value
- case "error_rate":
- weights.ErrorRate = value
- case "ttft":
- weights.TTFT = value
- case "reset":
- weights.Reset = value
- case "quota_headroom":
- weights.QuotaHeadroom = value
- case "upstream_cost":
- weights.UpstreamCost = value
- case "previous_response":
- weights.Previous = value
- case "session_sticky":
- weights.SessionSticky = value
- }
- }
- return weights
- }
- type GatewayOpenAIWSSchedulerScoreWeightsView struct {
- Priority float64
- Load float64
- Queue float64
- ErrorRate float64
- TTFT float64
- // Reset 倾向「会话窗口最早重置」的账号;0 表示关闭(默认)。
- Reset float64
- QuotaHeadroom float64
- UpstreamCost float64
- Previous float64
- SessionSticky float64
- }
- func (w GatewayOpenAIWSSchedulerScoreWeightsView) configWeights() config.GatewayOpenAIWSSchedulerScoreWeights {
- return config.GatewayOpenAIWSSchedulerScoreWeights{
- Priority: w.Priority,
- Load: w.Load,
- Queue: w.Queue,
- ErrorRate: w.ErrorRate,
- TTFT: w.TTFT,
- Reset: w.Reset,
- QuotaHeadroom: w.QuotaHeadroom,
- UpstreamCost: w.UpstreamCost,
- PreviousResponse: w.Previous,
- SessionSticky: w.SessionSticky,
- }
- }
- type OpenAIAccountSchedulerScoreSnapshot struct {
- BaseScore float64
- StickyScore float64
- StickyScoreInfinity bool
- StickyWeightedEnabled bool
- }
- func (s *RateLimitService) BuildOpenAIAccountSchedulerScoreSnapshot(
- ctx context.Context,
- accounts []*Account,
- loadMap map[int64]*AccountLoadInfo,
- ) map[int64]OpenAIAccountSchedulerScoreSnapshot {
- gateway := &OpenAIGatewayService{cfg: nil, rateLimitService: s}
- if s != nil {
- gateway.cfg = s.cfg
- }
- return buildOpenAIAccountSchedulerScoreSnapshot(
- accounts,
- loadMap,
- gateway.openAIWSSchedulerWeightsForRequest(ctx),
- gateway.isOpenAIAdvancedSchedulerStickyWeightedEnabled(ctx),
- gateway.openAIOAuthSchedulingRateMultiplier(ctx),
- )
- }
- func BuildOpenAIAccountSchedulerScoreSnapshot(
- accounts []*Account,
- loadMap map[int64]*AccountLoadInfo,
- ) map[int64]OpenAIAccountSchedulerScoreSnapshot {
- gateway := &OpenAIGatewayService{}
- return buildOpenAIAccountSchedulerScoreSnapshot(accounts, loadMap, gateway.openAIWSSchedulerWeights(), false, defaultOpenAIOAuthSchedulingRateMultiplier)
- }
- func buildOpenAIAccountSchedulerScoreSnapshot(
- accounts []*Account,
- loadMap map[int64]*AccountLoadInfo,
- weights GatewayOpenAIWSSchedulerScoreWeightsView,
- stickyWeightedEnabled bool,
- oauthSchedulingRateMultiplier float64,
- ) map[int64]OpenAIAccountSchedulerScoreSnapshot {
- if len(accounts) == 0 {
- return nil
- }
- candidates := make([]openAIAccountCandidateScore, 0, len(accounts))
- for _, account := range accounts {
- if account == nil {
- continue
- }
- loadInfo := loadMap[account.ID]
- if loadInfo == nil {
- loadInfo = &AccountLoadInfo{AccountID: account.ID}
- }
- candidates = append(candidates, openAIAccountCandidateScore{
- account: account,
- loadInfo: loadInfo,
- errorRate: 0,
- ttft: 0,
- hasTTFT: false,
- })
- }
- if len(candidates) == 0 {
- return nil
- }
- minPriority, maxPriority := openAIAccountSchedulingPriority(candidates[0].account), openAIAccountSchedulingPriority(candidates[0].account)
- maxWaiting := 1
- for i := range candidates {
- candidate := &candidates[i]
- candidate.priority = openAIAccountSchedulingPriority(candidate.account)
- if candidate.priority < minPriority {
- minPriority = candidate.priority
- }
- if candidate.priority > maxPriority {
- maxPriority = candidate.priority
- }
- if candidate.loadInfo.WaitingCount > maxWaiting {
- maxWaiting = candidate.loadInfo.WaitingCount
- }
- }
- minResetRemaining, maxResetRemaining := 0.0, 0.0
- hasResetSample := false
- now := time.Now()
- upstreamCostFactors := map[int64]float64(nil)
- if weights.UpstreamCost > 0 {
- accounts := make([]*Account, 0, len(candidates))
- for _, candidate := range candidates {
- accounts = append(accounts, candidate.account)
- }
- upstreamCostFactors = openAIUpstreamCostFactors(accounts, now, oauthSchedulingRateMultiplier)
- }
- if weights.Reset > 0 {
- for _, candidate := range candidates {
- end := candidate.account.SessionWindowEnd
- if end == nil || !now.Before(*end) {
- continue
- }
- remaining := end.Sub(now).Seconds()
- if !hasResetSample {
- minResetRemaining, maxResetRemaining = remaining, remaining
- hasResetSample = true
- continue
- }
- if remaining < minResetRemaining {
- minResetRemaining = remaining
- }
- if remaining > maxResetRemaining {
- maxResetRemaining = remaining
- }
- }
- }
- result := make(map[int64]OpenAIAccountSchedulerScoreSnapshot, len(candidates))
- for _, candidate := range candidates {
- priorityFactor := 1.0
- if maxPriority > minPriority {
- priorityFactor = 1 - float64(candidate.priority-minPriority)/float64(maxPriority-minPriority)
- }
- loadFactor := 1 - clamp01(float64(candidate.loadInfo.LoadRate)/100.0)
- queueFactor := 1 - clamp01(float64(candidate.loadInfo.WaitingCount)/float64(maxWaiting))
- errorFactor := 1.0
- ttftFactor := 0.5
- resetFactor := 0.0
- if weights.Reset > 0 && hasResetSample {
- if end := candidate.account.SessionWindowEnd; end != nil && now.Before(*end) {
- if maxResetRemaining > minResetRemaining {
- resetFactor = 1 - clamp01((end.Sub(now).Seconds()-minResetRemaining)/(maxResetRemaining-minResetRemaining))
- } else {
- resetFactor = 1
- }
- }
- }
- quotaHeadroomFactor := 0.0
- if weights.QuotaHeadroom > 0 {
- quotaHeadroomFactor = openAIQuotaHeadroomFactor(candidate.account, now)
- }
- upstreamCostFactor := openAIUpstreamCostNeutralFactor
- if factor, ok := upstreamCostFactors[candidate.account.ID]; ok {
- upstreamCostFactor = factor
- }
- baseScore := weights.Priority*priorityFactor +
- weights.Load*loadFactor +
- weights.Queue*queueFactor +
- weights.ErrorRate*errorFactor +
- weights.TTFT*ttftFactor +
- weights.Reset*resetFactor +
- weights.QuotaHeadroom*quotaHeadroomFactor +
- weights.UpstreamCost*(upstreamCostFactor-openAIUpstreamCostNeutralFactor)
- score := OpenAIAccountSchedulerScoreSnapshot{
- BaseScore: baseScore,
- StickyWeightedEnabled: stickyWeightedEnabled,
- StickyScoreInfinity: !stickyWeightedEnabled,
- }
- if stickyWeightedEnabled {
- score.StickyScore = baseScore + weights.Previous + weights.SessionSticky
- }
- result[candidate.account.ID] = score
- }
- return result
- }
- func openAIUpstreamCostFactors(accounts []*Account, now time.Time, oauthSchedulingRateMultiplier float64) map[int64]float64 {
- type rateSample struct {
- accountID int64
- rate float64
- }
- factors := make(map[int64]float64, len(accounts))
- samples := make([]rateSample, 0, len(accounts))
- eligibleCount := 0
- for _, account := range accounts {
- if account == nil {
- continue
- }
- factors[account.ID] = openAIUpstreamCostNeutralFactor
- if !account.IsOpenAIApiKey() && !account.IsOpenAIOAuthLike() {
- continue
- }
- eligibleCount++
- if rate, ok := openAISchedulingRate(account, now, oauthSchedulingRateMultiplier); ok {
- samples = append(samples, rateSample{accountID: account.ID, rate: rate})
- }
- }
- if len(samples) < 2 || eligibleCount == 0 {
- return factors
- }
- allEqual := true
- positiveLogs := make([]float64, 0, len(samples))
- for i, sample := range samples {
- if i > 0 && sample.rate != samples[0].rate {
- allEqual = false
- }
- if sample.rate > 0 {
- positiveLogs = append(positiveLogs, math.Log(sample.rate))
- }
- }
- if allEqual || len(positiveLogs) == 0 {
- return factors
- }
- sort.Float64s(positiveLogs)
- middle := len(positiveLogs) / 2
- medianLog := positiveLogs[middle]
- if len(positiveLogs)%2 == 0 {
- medianLog = (positiveLogs[middle-1] + positiveLogs[middle]) / 2
- }
- center := math.Exp(medianLog)
- if center <= 0 || math.IsNaN(center) || math.IsInf(center, 0) {
- return factors
- }
- coverage := float64(len(samples)) / float64(eligibleCount)
- for _, sample := range samples {
- rawFactor := 1.0
- if sample.rate > 0 {
- rawFactor = 1 / (1 + sample.rate/center)
- }
- factors[sample.accountID] = clamp01(openAIUpstreamCostNeutralFactor + coverage*(rawFactor-openAIUpstreamCostNeutralFactor))
- }
- return factors
- }
- type openAILegacyUpstreamRateOrder struct {
- enabled bool
- rates map[int64]float64
- }
- func newOpenAILegacyUpstreamRateOrder(accounts []*Account, now time.Time, oauthSchedulingRateMultiplier float64) openAILegacyUpstreamRateOrder {
- rates := make(map[int64]float64, len(accounts))
- var first float64
- distinct := false
- for _, account := range accounts {
- if account == nil {
- continue
- }
- // 与 openAIUpstreamCostFactors 使用同一道平台门控:只有 OpenAI 平台账号
- // 的倍率参与 legacy 低倍率优先排序。上游自报倍率来自中转方,不能让它对
- // 其他平台的调度产生影响——否则自报低价即可吸走流量,而实际结算走本地倍率。
- if !account.IsOpenAIApiKey() && !account.IsOpenAIOAuthLike() {
- continue
- }
- rate, ok := openAISchedulingRate(account, now, oauthSchedulingRateMultiplier)
- if !ok {
- continue
- }
- if len(rates) == 0 {
- first = rate
- } else if rate != first {
- distinct = true
- }
- rates[account.ID] = rate
- }
- return openAILegacyUpstreamRateOrder{enabled: len(rates) >= 2 && distinct, rates: rates}
- }
- func openAISchedulingRate(account *Account, now time.Time, oauthSchedulingRateMultiplier float64) (float64, bool) {
- if account != nil && account.IsOpenAIOAuthLike() {
- return oauthSchedulingRateMultiplier, true
- }
- return openAIFreshUpstreamBillingRate(account, now)
- }
- // compare returns -1 when a should be selected before b, 1 when b should be
- // selected first, and 0 when the rate signal does not distinguish them.
- func (o openAILegacyUpstreamRateOrder) compare(a, b *Account) int {
- if !o.enabled || a == nil || b == nil {
- return 0
- }
- aRate, aKnown := o.rates[a.ID]
- bRate, bKnown := o.rates[b.ID]
- if aKnown != bKnown {
- if aKnown {
- return -1
- }
- return 1
- }
- if !aKnown || aRate == bRate {
- return 0
- }
- if aRate < bRate {
- return -1
- }
- return 1
- }
- func openAIFreshUpstreamBillingRate(account *Account, now time.Time) (float64, bool) {
- if !isUpstreamBillingProbeAccount(account) {
- return 0, false
- }
- snapshot := decodeUpstreamBillingProbeSnapshot(account.Extra)
- if snapshot == nil || (snapshot.Status != UpstreamBillingProbeStatusOK && snapshot.Status != UpstreamBillingProbeStatusFailed) ||
- snapshot.ReceivedAt == nil || snapshot.ReceivedAt.IsZero() {
- return 0, false
- }
- receivedAt := *snapshot.ReceivedAt
- freshUntil := snapshot.FreshUntil
- if freshUntil == nil && snapshot.Status == UpstreamBillingProbeStatusOK {
- interval := snapshot.NextProbeAt.Sub(receivedAt)
- if interval > 0 {
- freshUntil = probeTimePtr(receivedAt.Add(2 * interval))
- }
- }
- if freshUntil == nil || !freshUntil.After(receivedAt) || now.Before(receivedAt) || now.After(*freshUntil) {
- return 0, false
- }
- return upstreamBillingRateAt(snapshot.Data, now)
- }
- func openAIQuotaHeadroomFactor(account *Account, now time.Time) float64 {
- if account == nil || len(account.Extra) == 0 || openAIQuotaHeadroomSnapshotStale(account.Extra, now) {
- return openAIQuotaHeadroomNeutralFactor
- }
- primaryUsedPercent, ok := resolveAccountExtraNumber(account.Extra, "codex_primary_used_percent", "codex_7d_used_percent")
- if !ok || openAIQuotaWindowResetAny(account.Extra, now, "primary", "7d") {
- return openAIQuotaHeadroomNeutralFactor
- }
- factor := 1 - clamp01(primaryUsedPercent/100)
- if secondaryUsedPercent, ok := resolveAccountExtraNumber(account.Extra, "codex_secondary_used_percent", "codex_5h_used_percent"); ok &&
- !openAIQuotaWindowResetAny(account.Extra, now, "secondary", "5h") {
- secondaryRemaining := 1 - clamp01(secondaryUsedPercent/100)
- if secondaryRemaining < openAIQuotaHeadroomSecondaryLowRemain {
- factor *= openAIQuotaHeadroomNeutralFactor
- }
- }
- return factor
- }
- func openAIQuotaHeadroomSnapshotStale(extra map[string]any, now time.Time) bool {
- updatedRaw, ok := extra["codex_usage_updated_at"]
- if !ok {
- return true
- }
- updatedAt, err := parseTime(fmt.Sprint(updatedRaw))
- if err != nil {
- return true
- }
- return now.Sub(updatedAt) >= openAIQuotaHeadroomSnapshotStaleAfter
- }
- func openAIQuotaWindowResetAny(extra map[string]any, now time.Time, windows ...string) bool {
- for _, window := range windows {
- if openAIQuotaWindowReset(extra, window, now) {
- return true
- }
- }
- return false
- }
- func clamp01(value float64) float64 {
- switch {
- case value < 0:
- return 0
- case value > 1:
- return 1
- default:
- return value
- }
- }
- func calcLoadSkewByMoments(sum float64, sumSquares float64, count int) float64 {
- if count <= 1 {
- return 0
- }
- mean := sum / float64(count)
- variance := sumSquares/float64(count) - mean*mean
- if variance < 0 {
- variance = 0
- }
- return math.Sqrt(variance)
- }
|