openai_gateway_scheduling.go 66 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768697071727374757677787980818283848586878889909192939495969798991001011021031041051061071081091101111121131141151161171181191201211221231241251261271281291301311321331341351361371381391401411421431441451461471481491501511521531541551561571581591601611621631641651661671681691701711721731741751761771781791801811821831841851861871881891901911921931941951961971981992002012022032042052062072082092102112122132142152162172182192202212222232242252262272282292302312322332342352362372382392402412422432442452462472482492502512522532542552562572582592602612622632642652662672682692702712722732742752762772782792802812822832842852862872882892902912922932942952962972982993003013023033043053063073083093103113123133143153163173183193203213223233243253263273283293303313323333343353363373383393403413423433443453463473483493503513523533543553563573583593603613623633643653663673683693703713723733743753763773783793803813823833843853863873883893903913923933943953963973983994004014024034044054064074084094104114124134144154164174184194204214224234244254264274284294304314324334344354364374384394404414424434444454464474484494504514524534544554564574584594604614624634644654664674684694704714724734744754764774784794804814824834844854864874884894904914924934944954964974984995005015025035045055065075085095105115125135145155165175185195205215225235245255265275285295305315325335345355365375385395405415425435445455465475485495505515525535545555565575585595605615625635645655665675685695705715725735745755765775785795805815825835845855865875885895905915925935945955965975985996006016026036046056066076086096106116126136146156166176186196206216226236246256266276286296306316326336346356366376386396406416426436446456466476486496506516526536546556566576586596606616626636646656666676686696706716726736746756766776786796806816826836846856866876886896906916926936946956966976986997007017027037047057067077087097107117127137147157167177187197207217227237247257267277287297307317327337347357367377387397407417427437447457467477487497507517527537547557567577587597607617627637647657667677687697707717727737747757767777787797807817827837847857867877887897907917927937947957967977987998008018028038048058068078088098108118128138148158168178188198208218228238248258268278288298308318328338348358368378388398408418428438448458468478488498508518528538548558568578588598608618628638648658668678688698708718728738748758768778788798808818828838848858868878888898908918928938948958968978988999009019029039049059069079089099109119129139149159169179189199209219229239249259269279289299309319329339349359369379389399409419429439449459469479489499509519529539549559569579589599609619629639649659669679689699709719729739749759769779789799809819829839849859869879889899909919929939949959969979989991000100110021003100410051006100710081009101010111012101310141015101610171018101910201021102210231024102510261027102810291030103110321033103410351036103710381039104010411042104310441045104610471048104910501051105210531054105510561057105810591060106110621063106410651066106710681069107010711072107310741075107610771078107910801081108210831084108510861087108810891090109110921093109410951096109710981099110011011102110311041105110611071108110911101111111211131114111511161117111811191120112111221123112411251126112711281129113011311132113311341135113611371138113911401141114211431144114511461147114811491150115111521153115411551156115711581159116011611162116311641165116611671168116911701171117211731174117511761177117811791180118111821183118411851186118711881189119011911192119311941195119611971198119912001201120212031204120512061207120812091210121112121213121412151216121712181219122012211222122312241225122612271228122912301231123212331234123512361237123812391240124112421243124412451246124712481249125012511252125312541255125612571258125912601261126212631264126512661267126812691270127112721273127412751276127712781279128012811282128312841285128612871288128912901291129212931294129512961297129812991300130113021303130413051306130713081309131013111312131313141315131613171318131913201321132213231324132513261327132813291330133113321333133413351336133713381339134013411342134313441345134613471348134913501351135213531354135513561357135813591360136113621363136413651366136713681369137013711372137313741375137613771378137913801381138213831384138513861387138813891390139113921393139413951396139713981399140014011402140314041405140614071408140914101411141214131414141514161417141814191420142114221423142414251426142714281429143014311432143314341435143614371438143914401441144214431444144514461447144814491450145114521453145414551456145714581459146014611462146314641465146614671468146914701471147214731474147514761477147814791480148114821483148414851486148714881489149014911492149314941495149614971498149915001501150215031504150515061507150815091510151115121513151415151516151715181519152015211522152315241525152615271528152915301531153215331534153515361537153815391540154115421543154415451546154715481549155015511552155315541555155615571558155915601561156215631564156515661567156815691570157115721573157415751576157715781579158015811582158315841585158615871588158915901591159215931594159515961597159815991600160116021603160416051606160716081609161016111612161316141615161616171618161916201621162216231624162516261627162816291630163116321633163416351636163716381639164016411642164316441645164616471648164916501651165216531654165516561657165816591660166116621663166416651666166716681669167016711672167316741675167616771678167916801681168216831684168516861687168816891690169116921693169416951696169716981699170017011702170317041705170617071708170917101711171217131714171517161717171817191720172117221723172417251726172717281729173017311732173317341735173617371738173917401741174217431744174517461747174817491750175117521753
  1. package service
  2. // 本文件由 openai_gateway_service.go 纯移动拆分而来:粘性会话哈希、账号选择与
  3. // 负载感知调度、配额自动暂停判定、并发槽位获取。仅做代码搬迁,无任何行为变更。
  4. import (
  5. "context"
  6. "encoding/json"
  7. "fmt"
  8. "log/slog"
  9. "sort"
  10. "strconv"
  11. "strings"
  12. "time"
  13. "github.com/Wei-Shaw/sub2api/internal/config"
  14. "github.com/Wei-Shaw/sub2api/internal/pkg/openai"
  15. "github.com/Wei-Shaw/sub2api/internal/pkg/xai"
  16. "github.com/gin-gonic/gin"
  17. "github.com/tidwall/gjson"
  18. )
  19. const (
  20. openCodeSessionAffinityHeader = "X-Session-Affinity"
  21. openCodeSessionIDHeader = "X-Session-Id"
  22. openCodeNativeSessionHeader = "X-OpenCode-Session"
  23. codeBuddyConversationHeader = "X-Conversation-ID"
  24. )
  25. var explicitOpenAIHeaderSessionNames = []string{
  26. "session-id",
  27. "session_id",
  28. "conversation_id",
  29. openCodeSessionAffinityHeader,
  30. openCodeSessionIDHeader,
  31. openCodeNativeSessionHeader,
  32. codeBuddyConversationHeader,
  33. }
  34. // explicitOpenAIHeaderSessionID resolves stable conversation identifiers sent
  35. // by OpenAI-compatible clients. Keep this list limited to session-scoped
  36. // fields: request/message IDs rotate every turn and would defeat sticky routing
  37. // and upstream prompt caching.
  38. func explicitOpenAIHeaderSessionID(c *gin.Context) string {
  39. if c == nil {
  40. return ""
  41. }
  42. for _, header := range explicitOpenAIHeaderSessionNames {
  43. if sessionID := strings.TrimSpace(c.GetHeader(header)); sessionID != "" {
  44. return sessionID
  45. }
  46. }
  47. return ""
  48. }
  49. // ExtractSessionID extracts the raw session ID from headers or body without hashing.
  50. // Used by ForwardAsAnthropic to pass as prompt_cache_key for upstream cache.
  51. func (s *OpenAIGatewayService) ExtractSessionID(c *gin.Context, body []byte) string {
  52. return explicitOpenAIRequestSessionID(c, body)
  53. }
  54. func explicitOpenAISessionID(c *gin.Context, body []byte) string {
  55. if c == nil {
  56. return ""
  57. }
  58. sessionID := explicitOpenAIHeaderSessionID(c)
  59. if sessionID == "" && len(body) > 0 {
  60. sessionID = strings.TrimSpace(openAIRequestPayloadView(body).Get("prompt_cache_key").String())
  61. }
  62. return sessionID
  63. }
  64. // openAIRequestPayloadView unwraps Responses WebSocket event envelopes while
  65. // leaving ordinary HTTP objects untouched even when they contain a response
  66. // field for another purpose.
  67. func openAIRequestPayloadView(body []byte) gjson.Result {
  68. root := parseRawJSONView(body)
  69. eventType := strings.ToLower(strings.TrimSpace(root.Get("type").String()))
  70. if strings.HasPrefix(eventType, "response.") {
  71. if response := root.Get("response"); response.Exists() && response.IsObject() {
  72. return response
  73. }
  74. }
  75. return root
  76. }
  77. // explicitOpenAIRequestSessionID extends the common OpenAI session signals
  78. // with Grok's native conversation header only for requests authenticated to a
  79. // Grok group. This keeps an unrelated x-grok-conv-id header from changing
  80. // scheduling or upstream session behavior for non-Grok groups.
  81. //
  82. // For Grok groups only, previous_response_id is a last-resort sticky seed so
  83. // multi-turn Responses chains stay on the same OAuth account when no explicit
  84. // session/conversation/prompt_cache_key is present. Non-Grok groups omit this
  85. // so HTTP OpenAI paths that delete previous_response_id before upstream are
  86. // unchanged.
  87. func explicitOpenAIRequestSessionID(c *gin.Context, body []byte) string {
  88. if c == nil {
  89. return ""
  90. }
  91. sessionID := explicitOpenAIHeaderSessionID(c)
  92. if sessionID == "" && isGrokRequestContext(c) {
  93. sessionID = strings.TrimSpace(c.GetHeader(grokConversationIDHeader))
  94. }
  95. if sessionID == "" && len(body) > 0 {
  96. sessionID = strings.TrimSpace(openAIRequestPayloadView(body).Get("prompt_cache_key").String())
  97. }
  98. if sessionID == "" && isGrokRequestContext(c) && len(body) > 0 {
  99. sessionID = grokPreviousResponseSessionSeed(body)
  100. }
  101. return sessionID
  102. }
  103. // grokPreviousResponseSessionSeed returns a stable sticky seed from a Responses
  104. // previous_response_id. Only resp_* response ids are accepted; message ids and
  105. // unknown shapes must not pin sticky routing or prompt-cache identity.
  106. func grokPreviousResponseSessionSeed(body []byte) string {
  107. id := strings.TrimSpace(gjson.GetBytes(body, "previous_response_id").String())
  108. if id == "" {
  109. return ""
  110. }
  111. if ClassifyOpenAIPreviousResponseIDKind(id) != OpenAIPreviousResponseIDKindResponseID {
  112. return ""
  113. }
  114. // Namespace so content-derived seeds never collide with response ids.
  115. return "grok-prev-resp:" + id
  116. }
  117. // GenerateExplicitSessionHash generates a sticky-session hash only from explicit
  118. // client session signals. It intentionally skips content-derived fallback and is
  119. // used by stateless endpoints such as /v1/images.
  120. func (s *OpenAIGatewayService) GenerateExplicitSessionHash(c *gin.Context, body []byte) string {
  121. sessionID := explicitOpenAIRequestSessionID(c, body)
  122. if sessionID == "" {
  123. return ""
  124. }
  125. currentHash, legacyHash := deriveOpenAISessionHashes(sessionID)
  126. attachOpenAILegacySessionHashToGin(c, legacyHash)
  127. return currentHash
  128. }
  129. // GenerateSessionHash generates a sticky-session hash for OpenAI requests.
  130. //
  131. // Priority:
  132. // 1. Header: session-id / session_id
  133. // 2. Header: conversation_id
  134. // 3. Header: x-session-affinity / x-session-id / x-opencode-session (OpenCode)
  135. // 4. Header: x-conversation-id (CodeBuddy)
  136. // 5. Header: x-grok-conv-id (Grok groups only)
  137. // 6. Body: prompt_cache_key
  138. // 7. Body: content-based fallback (model + system + tools + first user message)
  139. //
  140. // Grok sticky affinity is intentionally separate from the upstream
  141. // prompt_cache_key identity (resolveGrokCacheIdentity): sticky pins an OAuth
  142. // account for multi-turn routing, while the cache identity is tenant+model
  143. // isolated for xAI server-side prompt cache. For Grok groups we scope the
  144. // sticky seed with the client-requested model so switching models does not
  145. // inherit a stale account binding (grok2api affinityKey pattern).
  146. func (s *OpenAIGatewayService) GenerateSessionHash(c *gin.Context, body []byte) string {
  147. if c == nil {
  148. return ""
  149. }
  150. sessionID := explicitOpenAIRequestSessionID(c, body)
  151. if sessionID == "" && len(body) > 0 {
  152. sessionID = deriveOpenAIContentSessionSeed(body)
  153. }
  154. if sessionID == "" {
  155. return ""
  156. }
  157. if isGrokRequestContext(c) {
  158. sessionID = grokStickyAffinitySeed(sessionID, body)
  159. }
  160. currentHash, legacyHash := deriveOpenAISessionHashes(sessionID)
  161. attachOpenAILegacySessionHashToGin(c, legacyHash)
  162. return currentHash
  163. }
  164. // grokStickyAffinitySeed scopes sticky routing by model without changing the
  165. // upstream prompt_cache_key written by applyGrokResponsesCacheIdentity.
  166. func grokStickyAffinitySeed(sessionID string, body []byte) string {
  167. sessionID = strings.TrimSpace(sessionID)
  168. if sessionID == "" {
  169. return ""
  170. }
  171. model := ""
  172. if len(body) > 0 {
  173. model = strings.ToLower(strings.TrimSpace(gjson.GetBytes(body, "model").String()))
  174. }
  175. if model == "" {
  176. return "grok-affinity:v1:" + sessionID
  177. }
  178. return "grok-affinity:v1:" + model + ":" + sessionID
  179. }
  180. // GenerateSessionHashWithFallback 先按常规信号生成会话哈希;
  181. // 当未携带 session_id/conversation_id/prompt_cache_key 时,使用 fallbackSeed 生成稳定哈希。
  182. // 该方法用于 WS ingress,避免会话信号缺失时发生跨账号漂移。
  183. func (s *OpenAIGatewayService) GenerateSessionHashWithFallback(c *gin.Context, body []byte, fallbackSeed string) string {
  184. sessionHash := s.GenerateSessionHash(c, body)
  185. if sessionHash != "" {
  186. return sessionHash
  187. }
  188. seed := strings.TrimSpace(fallbackSeed)
  189. if seed == "" {
  190. return ""
  191. }
  192. currentHash, legacyHash := deriveOpenAISessionHashes(seed)
  193. attachOpenAILegacySessionHashToGin(c, legacyHash)
  194. return currentHash
  195. }
  196. func resolveOpenAIUpstreamOriginator(c *gin.Context, isOfficialClient bool) string {
  197. if c != nil {
  198. if originator := strings.TrimSpace(c.GetHeader("originator")); originator != "" {
  199. return originator
  200. }
  201. }
  202. if isOfficialClient {
  203. return openai.CodexDefaultOriginator
  204. }
  205. return "opencode"
  206. }
  207. // BindStickySession sets session -> account binding with standard TTL.
  208. func (s *OpenAIGatewayService) BindStickySession(ctx context.Context, groupID *int64, sessionHash string, accountID int64) error {
  209. if sessionHash == "" || accountID <= 0 {
  210. return nil
  211. }
  212. ttl := openaiStickySessionTTL
  213. if s != nil && s.cfg != nil && s.cfg.Gateway.OpenAIWS.StickySessionTTLSeconds > 0 {
  214. ttl = time.Duration(s.cfg.Gateway.OpenAIWS.StickySessionTTLSeconds) * time.Second
  215. }
  216. return s.setStickySessionAccountID(ctx, groupID, sessionHash, accountID, ttl)
  217. }
  218. // SelectAccount selects an OpenAI account with sticky session support
  219. func (s *OpenAIGatewayService) SelectAccount(ctx context.Context, groupID *int64, sessionHash string) (*Account, error) {
  220. return s.SelectAccountForModel(ctx, groupID, sessionHash, "")
  221. }
  222. // SelectAccountForModel selects an account supporting the requested model
  223. func (s *OpenAIGatewayService) SelectAccountForModel(ctx context.Context, groupID *int64, sessionHash string, requestedModel string) (*Account, error) {
  224. return s.SelectAccountForModelWithExclusions(ctx, groupID, sessionHash, requestedModel, nil)
  225. }
  226. // SelectAccountForModelWithExclusions selects an account supporting the requested model while excluding specified accounts.
  227. // SelectAccountForModelWithExclusions 选择支持指定模型的账号,同时排除指定的账号。
  228. func (s *OpenAIGatewayService) SelectAccountForModelWithExclusions(ctx context.Context, groupID *int64, sessionHash string, requestedModel string, excludedIDs map[int64]struct{}) (*Account, error) {
  229. ctx = s.withOpenAIQuotaAutoPauseContext(ctx)
  230. account, err := s.selectAccountForModelWithExclusions(ctx, groupID, PlatformOpenAI, sessionHash, requestedModel, excludedIDs, false, 0, "", false)
  231. if err != nil {
  232. return nil, err
  233. }
  234. return s.hydrateSelectedAccount(ctx, account)
  235. }
  236. // SelectAccountForTokenCount selects an account for a non-billable token-count
  237. // request. It applies the normal platform, model, capability, and runtime
  238. // eligibility checks without acquiring or waiting for a generation slot.
  239. func (s *OpenAIGatewayService) SelectAccountForTokenCount(
  240. ctx context.Context,
  241. groupID *int64,
  242. sessionHash string,
  243. requestedModel string,
  244. requiredCapability OpenAIEndpointCapability,
  245. platform string,
  246. ) (*Account, error) {
  247. ctx = WithOpenAIProfitControlSuppressed(ctx)
  248. ctx = s.withOpenAIQuotaAutoPauseContext(ctx)
  249. account, err := s.selectAccountForModelWithExclusions(
  250. ctx,
  251. groupID,
  252. platform,
  253. sessionHash,
  254. requestedModel,
  255. nil,
  256. false,
  257. 0,
  258. requiredCapability,
  259. false,
  260. )
  261. if err != nil {
  262. return nil, err
  263. }
  264. return s.hydrateSelectedAccount(ctx, account)
  265. }
  266. // NormalizeOpenAICompatiblePlatform 保留 grok 与国产 OpenAI 兼容供应商(kimi/zhipu/
  267. // deepseek)的原值,其他值一律归一为 openai。调度器据此对账号与请求做精确平台匹配:
  268. // kimi 分组请求只命中 kimi 账号,语义与 openai/grok 一致。
  269. // (upstream 曾将本函数改为未导出 normalizeOpenAICompatiblePlatform,本分支的
  270. // handler 调度入口仍需导出,保持导出名。)
  271. func NormalizeOpenAICompatiblePlatform(platform string) string {
  272. switch platform {
  273. case PlatformGrok, PlatformKimi, PlatformZhipu, PlatformDeepseek, PlatformMiniMax:
  274. return platform
  275. default:
  276. return PlatformOpenAI
  277. }
  278. }
  279. // noAvailableOpenAISelectionError builds the standard "no account available" error
  280. // while preserving the legacy /responses/compact error when applicable.
  281. // details carries an optional machine-parseable exclusion summary (e.g.
  282. // "pool=2, filtered: quota_auto_pause_7d=1 runtime_blocked=1") appended in
  283. // parentheses. It is for server-side logs / ops diagnostics only: handlers
  284. // never forward this error text to OpenAI-platform clients (they respond with
  285. // the generic classification message). Callers that must preserve the legacy
  286. // message pass "".
  287. func noAvailableOpenAISelectionError(requestedModel string, compactBlocked bool, details string) error {
  288. if compactBlocked {
  289. return ErrNoAvailableCompactAccounts
  290. }
  291. message := "no available OpenAI accounts"
  292. if requestedModel != "" {
  293. message = fmt.Sprintf("no available OpenAI accounts supporting model: %s", requestedModel)
  294. }
  295. if details != "" {
  296. message += " (" + details + ")"
  297. }
  298. return openAINoAvailableSelectionError{message: message}
  299. }
  300. type openAINoAvailableSelectionError struct {
  301. message string
  302. }
  303. func (e openAINoAvailableSelectionError) Error() string {
  304. return e.message
  305. }
  306. func (e openAINoAvailableSelectionError) Unwrap() error {
  307. return ErrNoAvailableAccounts
  308. }
  309. // openAICompactSupportTier classifies an OpenAI-compatible account by compact capability.
  310. // 0 = explicitly unsupported, 1 = unknown / not yet probed, 2 = explicitly supported.
  311. func openAICompactSupportTier(account *Account) int {
  312. if account == nil {
  313. return 0
  314. }
  315. if account.IsGrok() {
  316. return 2
  317. }
  318. if !account.IsOpenAI() {
  319. return 0
  320. }
  321. supported, known := account.OpenAICompactSupportKnown()
  322. if !known {
  323. return 1
  324. }
  325. if supported {
  326. return 2
  327. }
  328. return 0
  329. }
  330. // isOpenAICompatibleAccountEligibleForRequest 判断 OpenAI 兼容账号是否满足本次请求的调度条件。
  331. // 检查内容包括:平台匹配、账号可用性、quota 自动暂停、spark 路由限制、模型支持及端点能力。
  332. //
  333. // 注意:对 spark 影子账号,调用方还须额外调用 parentHealthyForShadow(account, lookup)
  334. // 检查母账号凭据可用性;该检查未内置于本函数,以避免注入 DB 依赖。
  335. func isOpenAICompatibleAccountEligibleForRequest(ctx context.Context, account *Account, platform string, requestedModel string, requireCompact bool, requiredCapability OpenAIEndpointCapability) bool {
  336. return openAICompatibleAccountEligibilityFailureReason(ctx, account, platform, requestedModel, requireCompact, requiredCapability) == ""
  337. }
  338. // openAICompatibleAccountEligibilityFailureReason mirrors the legacy boolean
  339. // eligibility check while naming its first veto point. Load-batch selection uses
  340. // the reason only for server-side no-account diagnostics; the admission behavior
  341. // remains unchanged.
  342. func openAICompatibleAccountEligibilityFailureReason(ctx context.Context, account *Account, platform string, requestedModel string, requireCompact bool, requiredCapability OpenAIEndpointCapability) string {
  343. if reason := openAICompatibleAccountEligibilityFailureReasonBeforeProfit(ctx, account, platform, requestedModel, requireCompact, requiredCapability); reason != "" {
  344. return reason
  345. }
  346. // 分组利润控制:legacy 引擎的粘性/候选循环与 DB recheck 共用
  347. // 本判定,任何 fallback 都不能把利润不合格账号重新放回候选。
  348. if vetoed, reason := openAIProfitControlVetoReason(ctx, account); vetoed {
  349. return reason
  350. }
  351. return ""
  352. }
  353. // isOpenAICompatibleAccountEligibleForRequestBeforeProfit applies every
  354. // ordinary scheduling gate. Legacy selection uses it before classifying the
  355. // profit veto so earlier failures retain their actual reason.
  356. func isOpenAICompatibleAccountEligibleForRequestBeforeProfit(ctx context.Context, account *Account, platform string, requestedModel string, requireCompact bool, requiredCapability OpenAIEndpointCapability) bool {
  357. return openAICompatibleAccountEligibilityFailureReasonBeforeProfit(ctx, account, platform, requestedModel, requireCompact, requiredCapability) == ""
  358. }
  359. func openAICompatibleAccountEligibilityFailureReasonBeforeProfit(ctx context.Context, account *Account, platform string, requestedModel string, requireCompact bool, requiredCapability OpenAIEndpointCapability) string {
  360. platform = NormalizeOpenAICompatiblePlatform(platform)
  361. if account == nil {
  362. return "account_nil"
  363. }
  364. if account.Platform != platform || !account.IsOpenAICompatible() {
  365. return "platform_mismatch"
  366. }
  367. if !account.IsSchedulableForModelWithContext(ctx, requestedModel) {
  368. if account.IsSchedulable() {
  369. return "model_rate_limited"
  370. }
  371. return "not_schedulable"
  372. }
  373. if account.IsOpenAI() {
  374. if paused, reason := shouldAutoPauseOpenAIAccountByQuota(ctx, account); paused {
  375. // Debug level: this fires per-candidate on the scheduling hot path, so Info
  376. // would amplify into log spam once several accounts cross the threshold.
  377. slog.Debug("account_auto_paused_by_quota",
  378. "account_id", account.ID,
  379. "window", reason.window,
  380. "threshold", reason.threshold,
  381. "utilization", reason.utilization,
  382. )
  383. if reason.window != "" {
  384. return "quota_auto_pause_" + reason.window
  385. }
  386. return "quota_auto_pause"
  387. }
  388. }
  389. if account.IsGrok() {
  390. if paused, reason := shouldAutoPauseGrokAccountByQuota(account); paused {
  391. slog.Debug("grok_account_auto_paused_by_quota",
  392. "account_id", account.ID,
  393. "window", reason.window,
  394. "threshold", reason.threshold,
  395. "utilization", reason.utilization,
  396. )
  397. if reason.window != "" {
  398. return "quota_auto_pause_" + reason.window
  399. }
  400. return "quota_auto_pause"
  401. }
  402. }
  403. if requestedModel != "" && !account.IsModelSupported(requestedModel) {
  404. return "model_not_supported"
  405. }
  406. if !account.SupportsOpenAIEndpointCapability(requiredCapability) {
  407. if account.IsGrok() && requiredCapability == OpenAIEndpointCapabilityGrokMediaGeneration {
  408. _, reason := account.GrokMediaGenerationEligibility()
  409. slog.Debug("grok_media_account_ineligible", "account_id", account.ID, "reason", reason)
  410. }
  411. return "capability_mismatch"
  412. }
  413. if requireCompact && openAICompactSupportTier(account) == 0 {
  414. return "compact_unsupported"
  415. }
  416. return ""
  417. }
  418. type openAIQuotaAutoPauseDecision struct {
  419. window string
  420. threshold float64
  421. utilization float64
  422. reason string
  423. }
  424. func shouldAutoPauseGrokAccountByQuota(account *Account) (bool, openAIQuotaAutoPauseDecision) {
  425. if account == nil || !account.IsGrok() || account.Type != AccountTypeOAuth {
  426. return false, openAIQuotaAutoPauseDecision{}
  427. }
  428. snapshot, err := grokQuotaSnapshotFromExtra(account.Extra)
  429. if err != nil || snapshot == nil {
  430. return false, openAIQuotaAutoPauseDecision{}
  431. }
  432. now := time.Now()
  433. if grokQuotaSnapshotStaleForPause(snapshot, now) {
  434. return false, openAIQuotaAutoPauseDecision{}
  435. }
  436. if grokQuotaRetryAfterActive(snapshot, now) {
  437. return true, openAIQuotaAutoPauseDecision{window: "retry_after", threshold: 1, utilization: 1}
  438. }
  439. if paused, decision := shouldAutoPauseGrokQuotaWindow("requests", snapshot.Requests, now); paused {
  440. return true, decision
  441. }
  442. if paused, decision := shouldAutoPauseGrokQuotaWindow("tokens", snapshot.Tokens, now); paused {
  443. return true, decision
  444. }
  445. return false, openAIQuotaAutoPauseDecision{}
  446. }
  447. func grokQuotaRetryAfterActive(snapshot *xai.QuotaSnapshot, now time.Time) bool {
  448. if snapshot == nil || snapshot.RetryAfterSeconds == nil || *snapshot.RetryAfterSeconds <= 0 {
  449. return false
  450. }
  451. if strings.TrimSpace(snapshot.UpdatedAt) == "" {
  452. return true
  453. }
  454. updatedAt, err := parseTime(snapshot.UpdatedAt)
  455. if err != nil {
  456. return true
  457. }
  458. retryAfterUntil := updatedAt.Add(time.Duration(*snapshot.RetryAfterSeconds) * time.Second)
  459. return now.Before(retryAfterUntil)
  460. }
  461. func shouldAutoPauseGrokQuotaWindow(name string, window *xai.QuotaWindow, now time.Time) (bool, openAIQuotaAutoPauseDecision) {
  462. if window == nil || window.Limit == nil || window.Remaining == nil || *window.Limit <= 0 {
  463. return false, openAIQuotaAutoPauseDecision{}
  464. }
  465. if window.ResetUnix != nil && *window.ResetUnix > 0 && !now.Before(time.Unix(*window.ResetUnix, 0)) {
  466. return false, openAIQuotaAutoPauseDecision{}
  467. }
  468. utilization := float64(*window.Limit-*window.Remaining) / float64(*window.Limit)
  469. if *window.Remaining <= 0 || utilization >= 1 {
  470. return true, openAIQuotaAutoPauseDecision{window: name, threshold: 1, utilization: utilization}
  471. }
  472. return false, openAIQuotaAutoPauseDecision{}
  473. }
  474. func grokQuotaSnapshotStaleForPause(snapshot *xai.QuotaSnapshot, now time.Time) bool {
  475. if snapshot == nil || strings.TrimSpace(snapshot.UpdatedAt) == "" {
  476. return false
  477. }
  478. updatedAt, err := parseTime(snapshot.UpdatedAt)
  479. if err != nil {
  480. return false
  481. }
  482. return now.Sub(updatedAt) >= openAICodexAutoPauseStaleAfter
  483. }
  484. func shouldAutoPauseOpenAIAccountByQuota(ctx context.Context, account *Account) (bool, openAIQuotaAutoPauseDecision) {
  485. if account == nil || !account.IsOpenAI() {
  486. return false, openAIQuotaAutoPauseDecision{}
  487. }
  488. // 自动用卡有独立阈值:达到消费阈值时必须先退出调度;仅达到普通暂停阈值时,
  489. // 只有新鲜状态明确存在可用卡才继续放行到消费阈值。
  490. if config := ResolveOpenAIAutoResetCreditConfig(account); config.Enabled {
  491. now := time.Now()
  492. utilization5h, has5h := resolveOpenAIQuotaUtilization(account.Extra, "5h", now)
  493. utilization7d, has7d := resolveOpenAIQuotaUtilization(account.Extra, "7d", now)
  494. if has5h && utilization5h >= config.Threshold5h {
  495. notifyOpenAIAutoReset(account.ID)
  496. return true, openAIQuotaAutoPauseDecision{window: "5h", threshold: config.Threshold5h, utilization: utilization5h, reason: "quota_auto_reset_pending_5h"}
  497. }
  498. if has7d && utilization7d >= config.Threshold7d {
  499. notifyOpenAIAutoReset(account.ID)
  500. return true, openAIQuotaAutoPauseDecision{window: "7d", threshold: config.Threshold7d, utilization: utilization7d, reason: "quota_auto_reset_pending_7d"}
  501. }
  502. disabled5h := resolveAccountExtraBool(account.Extra, "auto_pause_5h_disabled")
  503. disabled7d := resolveAccountExtraBool(account.Extra, "auto_pause_7d_disabled")
  504. pause5h, pause7d := resolveOpenAIQuotaAutoPauseThresholds(ctx, account)
  505. pauseReached5h := !disabled5h && pause5h > 0 && has5h && utilization5h >= pause5h
  506. pauseReached7d := !disabled7d && pause7d > 0 && has7d && utilization7d >= pause7d
  507. if pauseReached5h || pauseReached7d {
  508. state := openAIAutoResetStateFromExtra(account.Extra)
  509. if state != nil && state.Status == OpenAIAutoResetStatusAvailable && state.AvailableCount > 0 && !openAIAutoResetStateStale(state, now) {
  510. return false, openAIQuotaAutoPauseDecision{}
  511. }
  512. notifyOpenAIAutoReset(account.ID)
  513. if pauseReached5h {
  514. return true, openAIQuotaAutoPauseDecision{window: "5h", threshold: pause5h, utilization: utilization5h, reason: "quota_auto_reset_credit_check_5h"}
  515. }
  516. return true, openAIQuotaAutoPauseDecision{window: "7d", threshold: pause7d, utilization: utilization7d, reason: "quota_auto_reset_credit_check_7d"}
  517. }
  518. }
  519. // Per-account explicit-disable flags must take precedence over the global default.
  520. // Without these, leaving the account threshold blank means "use global default",
  521. // so an admin has no way to exempt a single account from auto-pause once a global
  522. // default exists. The disable flag is per-window so an account can opt out of
  523. // only 5h or only 7d auto-pause.
  524. disabled5h := resolveAccountExtraBool(account.Extra, "auto_pause_5h_disabled")
  525. disabled7d := resolveAccountExtraBool(account.Extra, "auto_pause_7d_disabled")
  526. threshold5h, threshold7d := resolveOpenAIQuotaAutoPauseThresholds(ctx, account)
  527. now := time.Now()
  528. if !disabled5h && threshold5h > 0 {
  529. if utilization, ok := resolveOpenAIQuotaUtilization(account.Extra, "5h", now); ok && utilization >= threshold5h {
  530. return true, openAIQuotaAutoPauseDecision{window: "5h", threshold: threshold5h, utilization: utilization}
  531. }
  532. }
  533. if !disabled7d && threshold7d > 0 {
  534. if utilization, ok := resolveOpenAIQuotaUtilization(account.Extra, "7d", now); ok && utilization >= threshold7d {
  535. return true, openAIQuotaAutoPauseDecision{window: "7d", threshold: threshold7d, utilization: utilization}
  536. }
  537. }
  538. return false, openAIQuotaAutoPauseDecision{}
  539. }
  540. // resolveAccountExtraBool reads a bool-like value from account extra, tolerating
  541. // the few shapes JSON unmarshalling may produce (real bool, "true"/"false"
  542. // strings, 0/1 numbers).
  543. func resolveAccountExtraBool(extra map[string]any, key string) bool {
  544. if len(extra) == 0 {
  545. return false
  546. }
  547. value, ok := extra[key]
  548. if !ok || value == nil {
  549. return false
  550. }
  551. switch v := value.(type) {
  552. case bool:
  553. return v
  554. case string:
  555. parsed, err := strconv.ParseBool(strings.TrimSpace(v))
  556. return err == nil && parsed
  557. case float64:
  558. return v != 0
  559. case float32:
  560. return v != 0
  561. case int:
  562. return v != 0
  563. case int64:
  564. return v != 0
  565. case json.Number:
  566. if i, err := v.Int64(); err == nil {
  567. return i != 0
  568. }
  569. }
  570. return false
  571. }
  572. func resolveOpenAIQuotaAutoPauseThresholds(ctx context.Context, account *Account) (float64, float64) {
  573. threshold5h, _ := resolveAccountExtraNumber(account.Extra, "auto_pause_5h_threshold")
  574. threshold7d, _ := resolveAccountExtraNumber(account.Extra, "auto_pause_7d_threshold")
  575. threshold5h = clamp01(threshold5h)
  576. threshold7d = clamp01(threshold7d)
  577. if threshold5h > 0 && threshold7d > 0 {
  578. return threshold5h, threshold7d
  579. }
  580. settings := openAIQuotaAutoPauseSettingsFromContext(ctx)
  581. if threshold5h <= 0 {
  582. threshold5h = clamp01(settings.DefaultThreshold5h)
  583. }
  584. if threshold7d <= 0 {
  585. threshold7d = clamp01(settings.DefaultThreshold7d)
  586. }
  587. return threshold5h, threshold7d
  588. }
  589. func resolveAccountExtraNumber(extra map[string]any, keys ...string) (float64, bool) {
  590. if len(extra) == 0 {
  591. return 0, false
  592. }
  593. for _, key := range keys {
  594. value, ok := extra[key]
  595. if !ok || value == nil {
  596. continue
  597. }
  598. switch v := value.(type) {
  599. case float64:
  600. return v, true
  601. case float32:
  602. return float64(v), true
  603. case int:
  604. return float64(v), true
  605. case int64:
  606. return float64(v), true
  607. case json.Number:
  608. parsed, err := v.Float64()
  609. if err == nil {
  610. return parsed, true
  611. }
  612. case string:
  613. parsed, err := strconv.ParseFloat(strings.TrimSpace(v), 64)
  614. if err == nil {
  615. return parsed, true
  616. }
  617. }
  618. }
  619. return 0, false
  620. }
  621. // resolveOpenAIQuotaUtilization returns the current utilization ratio (0..1) for the
  622. // given Codex usage window. ok=false means there is no usable signal to pause on:
  623. // either no snapshot exists, or the window has already rolled over so the cached
  624. // percentage is stale. The stale guard matters because a paused account stops
  625. // receiving requests, so its snapshot is never refreshed from upstream headers —
  626. // without this check an old used_percent would keep the account paused forever even
  627. // after the real window reset.
  628. func resolveOpenAIQuotaUtilization(extra map[string]any, window string, now time.Time) (float64, bool) {
  629. usedPercent := readOpenAIQuotaUsedPercent(extra, window)
  630. if usedPercent <= 0 {
  631. return 0, false
  632. }
  633. if openAIQuotaWindowReset(extra, window, now) {
  634. return 0, false
  635. }
  636. // 快照过于陈旧(账号长期未收到流量刷新)时,不再据此暂停。放行后下一次响应头
  637. // 会刷新快照实现自愈,避免账号在错误/过期的 used% 上被永久跳过(issue #2994)。
  638. if openAICodexSnapshotStaleForPause(extra, now) {
  639. return 0, false
  640. }
  641. return usedPercent / 100, true
  642. }
  643. // openAICodexSnapshotStaleForPause reports whether the Codex usage snapshot is stale
  644. // enough that it should no longer keep an account auto-paused. It anchors on
  645. // codex_usage_updated_at (always written by buildCodexUsageExtraUpdates). A missing or
  646. // unparseable timestamp returns false (treated as fresh, so the account stays paused) —
  647. // this is deliberate: it prevents any snapshot without a write time from silently escaping
  648. // auto-pause, and a genuinely-exhausted account that is actively served refreshes the
  649. // timestamp on every response so it never crosses the staleness bound.
  650. func openAICodexSnapshotStaleForPause(extra map[string]any, now time.Time) bool {
  651. if len(extra) == 0 {
  652. return false
  653. }
  654. updatedRaw, ok := extra["codex_usage_updated_at"]
  655. if !ok {
  656. return false
  657. }
  658. updatedAt, err := parseTime(fmt.Sprint(updatedRaw))
  659. if err != nil {
  660. return false
  661. }
  662. return now.Sub(updatedAt) >= openAICodexAutoPauseStaleAfter
  663. }
  664. // openAIQuotaWindowReset reports whether the Codex usage window's reset time has
  665. // already passed relative to now. It prefers the absolute codex_<window>_reset_at
  666. // timestamp and falls back to codex_<window>_reset_after_seconds anchored at
  667. // codex_usage_updated_at, mirroring AccountUsageService's window-progress logic.
  668. func openAIQuotaWindowReset(extra map[string]any, window string, now time.Time) bool {
  669. if len(extra) == 0 {
  670. return false
  671. }
  672. if resetAtRaw, ok := extra["codex_"+window+"_reset_at"]; ok {
  673. if resetAt, err := parseTime(fmt.Sprint(resetAtRaw)); err == nil {
  674. return !now.Before(resetAt)
  675. }
  676. }
  677. resetAfter := parseExtraInt(extra["codex_"+window+"_reset_after_seconds"])
  678. if resetAfter <= 0 {
  679. return false
  680. }
  681. base := now
  682. if updatedRaw, ok := extra["codex_usage_updated_at"]; ok {
  683. if updatedAt, err := parseTime(fmt.Sprint(updatedRaw)); err == nil {
  684. base = updatedAt
  685. }
  686. }
  687. resetAt := base.Add(time.Duration(resetAfter) * time.Second)
  688. return !now.Before(resetAt)
  689. }
  690. func readOpenAIQuotaUsedPercent(extra map[string]any, window string) float64 {
  691. if len(extra) == 0 {
  692. return 0
  693. }
  694. if value, ok := resolveAccountExtraNumber(extra, "codex_"+window+"_used_percent"); ok {
  695. return value
  696. }
  697. return 0
  698. }
  699. type openAIQuotaAutoPauseCtxKey struct{}
  700. func withOpenAIQuotaAutoPauseSettings(ctx context.Context, settings OpsOpenAIAccountQuotaAutoPauseSettings) context.Context {
  701. if ctx == nil {
  702. ctx = context.Background()
  703. }
  704. return context.WithValue(ctx, openAIQuotaAutoPauseCtxKey{}, settings)
  705. }
  706. func openAIQuotaAutoPauseSettingsFromContext(ctx context.Context) OpsOpenAIAccountQuotaAutoPauseSettings {
  707. if ctx == nil {
  708. return OpsOpenAIAccountQuotaAutoPauseSettings{}
  709. }
  710. settings, _ := ctx.Value(openAIQuotaAutoPauseCtxKey{}).(OpsOpenAIAccountQuotaAutoPauseSettings)
  711. return settings
  712. }
  713. func (s *OpenAIGatewayService) withOpenAIQuotaAutoPauseContext(ctx context.Context) context.Context {
  714. if s == nil || s.settingService == nil {
  715. return ctx
  716. }
  717. return withOpenAIQuotaAutoPauseSettings(ctx, s.settingService.GetOpenAIQuotaAutoPauseSettings(ctx))
  718. }
  719. // prioritizeOpenAICompactAccounts re-orders a slice so that accounts with known
  720. // compact support are tried first, followed by unknown, then explicitly unsupported.
  721. // The relative order within each tier is preserved.
  722. func prioritizeOpenAICompactAccounts(accounts []*Account) []*Account {
  723. if len(accounts) == 0 {
  724. return nil
  725. }
  726. supported := make([]*Account, 0, len(accounts))
  727. unknown := make([]*Account, 0, len(accounts))
  728. unsupported := make([]*Account, 0, len(accounts))
  729. for _, account := range accounts {
  730. switch openAICompactSupportTier(account) {
  731. case 2:
  732. supported = append(supported, account)
  733. case 1:
  734. unknown = append(unknown, account)
  735. default:
  736. unsupported = append(unsupported, account)
  737. }
  738. }
  739. out := make([]*Account, 0, len(accounts))
  740. out = append(out, supported...)
  741. out = append(out, unknown...)
  742. out = append(out, unsupported...)
  743. return out
  744. }
  745. // resolveOpenAIAccountUpstreamModelForRequest resolves the upstream model that
  746. // would be sent for a given request, honoring the legacy compact-only mapping
  747. // when the caller is on the /responses/compact path.
  748. func resolveOpenAIAccountUpstreamModelForRequest(account *Account, requestedModel string, requireCompact bool) string {
  749. // Forward checks the raw Chat Completions fallback before passthrough.
  750. // These API-key accounts therefore apply normal account model_mapping and
  751. // upstream normalization, but never compact_model_mapping.
  752. if shouldForwardOpenAIResponsesViaRawChatCompletions(account) {
  753. upstreamModel := resolveOpenAIForwardModel(account, requestedModel, "")
  754. return normalizeOpenAIModelForUpstream(account, upstreamModel)
  755. }
  756. // Passthrough accounts only replace authentication. Their Forward path
  757. // keeps the channel-mapped model in the request body and does not apply the
  758. // account's normal model_mapping. Legacy /responses/compact is the one
  759. // exception: forwardOpenAIPassthrough applies compact_model_mapping
  760. // directly to that channel-mapped model.
  761. if account != nil && account.IsOpenAIPassthroughEnabled() {
  762. upstreamModel := strings.TrimSpace(requestedModel)
  763. if upstreamModel == "" {
  764. return ""
  765. }
  766. if requireCompact {
  767. return resolveOpenAICompactForwardModel(account, upstreamModel)
  768. }
  769. return upstreamModel
  770. }
  771. // Compact mappings are keyed by the client-visible model. Prefer an exact
  772. // compact rule before ordinary account mapping; otherwise a normal alias can
  773. // hide the compact-specific rule and make scheduling disagree with Forward.
  774. if requireCompact && account != nil {
  775. if compactModel, matched := account.ResolveCompactMappedModel(strings.TrimSpace(requestedModel)); matched {
  776. if compactModel = strings.TrimSpace(compactModel); compactModel != "" {
  777. return compactModel
  778. }
  779. }
  780. }
  781. upstreamModel := resolveOpenAIForwardModel(account, requestedModel, "")
  782. if upstreamModel == "" {
  783. return ""
  784. }
  785. if requireCompact {
  786. compactModel := resolveOpenAICompactForwardModel(account, upstreamModel)
  787. if compactModel != upstreamModel {
  788. return compactModel
  789. }
  790. }
  791. return normalizeOpenAIModelForUpstream(account, upstreamModel)
  792. }
  793. // ResolveOpenAIAccountUpstreamModelForRequest exposes the scheduler's exact
  794. // account mapping chain to handler-side outcome reporting.
  795. func ResolveOpenAIAccountUpstreamModelForRequest(account *Account, requestedModel string, requireCompact bool) string {
  796. return resolveOpenAIAccountUpstreamModelForRequest(account, requestedModel, requireCompact)
  797. }
  798. // resolveOpenAIForwardMappedModels is the shared account mapping chain for
  799. // Forward callers. billingModel retains the ordinary mapping used for usage
  800. // accounting, while upstreamModel is the model the scheduler has admitted.
  801. func resolveOpenAIForwardMappedModels(account *Account, requestedModel string, requireCompact bool) (billingModel, upstreamModel string) {
  802. requestedModel = strings.TrimSpace(requestedModel)
  803. if account != nil && account.IsOpenAIPassthroughEnabled() {
  804. billingModel = requestedModel
  805. } else if account != nil {
  806. billingModel = strings.TrimSpace(account.GetMappedModel(requestedModel))
  807. }
  808. if billingModel == "" {
  809. billingModel = requestedModel
  810. }
  811. upstreamModel = resolveOpenAIAccountUpstreamModelForRequest(account, requestedModel, requireCompact)
  812. if strings.TrimSpace(upstreamModel) == "" {
  813. upstreamModel = billingModel
  814. }
  815. return billingModel, upstreamModel
  816. }
  817. func resolveOpenAIErrorSchedulingModel(billingModel, upstreamModel string) string {
  818. if upstreamModel = strings.TrimSpace(upstreamModel); upstreamModel != "" {
  819. return upstreamModel
  820. }
  821. return strings.TrimSpace(billingModel)
  822. }
  823. func (s *OpenAIGatewayService) selectAccountForModelWithExclusions(ctx context.Context, groupID *int64, platform string, sessionHash string, requestedModel string, excludedIDs map[int64]struct{}, requireCompact bool, stickyAccountID int64, requiredCapability OpenAIEndpointCapability, preferLowUpstreamRate bool) (*Account, error) {
  824. platform = NormalizeOpenAICompatiblePlatform(platform)
  825. if s.checkChannelPricingRestriction(ctx, groupID, requestedModel) {
  826. slog.Warn("channel pricing restriction blocked request",
  827. "group_id", derefGroupID(groupID),
  828. "model", requestedModel)
  829. return nil, fmt.Errorf("%w supporting model: %s (channel pricing restriction)", ErrNoAvailableAccounts, requestedModel)
  830. }
  831. // 1. 尝试粘性会话命中
  832. // Try sticky session hit
  833. if account := s.tryStickySessionHit(ctx, groupID, platform, sessionHash, requestedModel, excludedIDs, requireCompact, stickyAccountID, requiredCapability); account != nil {
  834. return account, nil
  835. }
  836. // 2. 获取可调度的 OpenAI 账号
  837. // Get schedulable OpenAI accounts
  838. accounts, err := s.listSchedulableAccounts(ctx, groupID, platform)
  839. if err != nil {
  840. return nil, fmt.Errorf("query accounts failed: %w", err)
  841. }
  842. // 3. 按优先级 + LRU 选择最佳账号
  843. // Select by priority + LRU
  844. selected, compactBlocked, filterStats := s.selectBestAccount(ctx, groupID, platform, accounts, requestedModel, excludedIDs, requireCompact, requiredCapability, preferLowUpstreamRate)
  845. if selected == nil {
  846. return nil, noAvailableOpenAISelectionError(requestedModel, compactBlocked, filterStats.summary(""))
  847. }
  848. // 4. 设置粘性会话绑定(利润门下推迟到 handler 终检通过后再绑定,
  849. // 终检否决的账号不得成为新的粘性目标;无门保持既有 eager 绑定与 TTL)
  850. // Set sticky session binding (deferred until terminal admission under a profit gate)
  851. if sessionHash != "" && !gatewayProfitControlGateActive(ctx) {
  852. _ = s.setStickySessionAccountID(ctx, groupID, sessionHash, selected.ID, openaiStickySessionTTL)
  853. }
  854. return selected, nil
  855. }
  856. // tryStickySessionHit 尝试从粘性会话获取账号。
  857. // 如果命中且账号可用则返回账号;如果账号不可用则清理会话并返回 nil。
  858. //
  859. // tryStickySessionHit attempts to get account from sticky session.
  860. // Returns account if hit and usable; clears session and returns nil if account is unavailable.
  861. func (s *OpenAIGatewayService) tryStickySessionHit(ctx context.Context, groupID *int64, platform string, sessionHash, requestedModel string, excludedIDs map[int64]struct{}, requireCompact bool, stickyAccountID int64, requiredCapability OpenAIEndpointCapability) *Account {
  862. if sessionHash == "" {
  863. return nil
  864. }
  865. platform = NormalizeOpenAICompatiblePlatform(platform)
  866. accountID := stickyAccountID
  867. if accountID <= 0 {
  868. var err error
  869. accountID, err = s.getStickySessionAccountID(ctx, groupID, sessionHash)
  870. if err != nil || accountID <= 0 {
  871. return nil
  872. }
  873. }
  874. if _, excluded := excludedIDs[accountID]; excluded {
  875. return nil
  876. }
  877. account, err := s.getSchedulableAccount(ctx, accountID)
  878. if err != nil {
  879. return nil
  880. }
  881. // 检查账号是否需要清理粘性会话
  882. // Check if sticky session should be cleared
  883. if shouldClearStickySession(account, requestedModel) {
  884. _ = s.deleteStickySessionAccountID(ctx, groupID, sessionHash)
  885. return nil
  886. }
  887. // 验证账号是否可用于当前请求
  888. // Verify account is usable for current request
  889. if !isOpenAICompatibleAccountEligibleForRequest(ctx, account, platform, requestedModel, false, requiredCapability) {
  890. return nil
  891. }
  892. if !parentHealthyForShadow(account, s.parentAccountLookup(ctx)) {
  893. _ = s.deleteStickySessionAccountID(ctx, groupID, sessionHash)
  894. return nil
  895. }
  896. if s.isOpenAIAccountRequestRuntimeBlocked(account, requestedModel) {
  897. _ = s.deleteStickySessionAccountID(ctx, groupID, sessionHash)
  898. return nil
  899. }
  900. account = s.recheckSelectedOpenAIAccountFromDB(ctx, account, groupID, platform, requestedModel, requireCompact, requiredCapability)
  901. if account == nil || !s.openAIAccountMatchesSchedulingGroup(account, groupID) {
  902. _ = s.deleteStickySessionAccountID(ctx, groupID, sessionHash)
  903. return nil
  904. }
  905. if groupID != nil && s.needsUpstreamChannelRestrictionCheck(ctx, groupID) &&
  906. s.isUpstreamModelRestrictedByChannel(ctx, *groupID, account, requestedModel, requireCompact) {
  907. _ = s.deleteStickySessionAccountID(ctx, groupID, sessionHash)
  908. return nil
  909. }
  910. // 刷新会话 TTL 并返回账号
  911. // Refresh session TTL and return account
  912. _ = s.refreshStickySessionTTL(ctx, groupID, sessionHash, openaiStickySessionTTL)
  913. return account
  914. }
  915. // selectBestAccount 从候选账号中选择最佳账号(优先级 + LRU)。
  916. // 返回 nil 表示无可用账号。
  917. //
  918. // selectBestAccount selects the best account from candidates (priority + LRU).
  919. // Returns nil if no available account. The second return reports whether at
  920. // least one candidate was filtered out solely because it lacks compact support
  921. // (only meaningful when the legacy /responses/compact requireCompact flag is
  922. // true); the third contains deterministic
  923. // exclusion diagnostics for the evaluated snapshot.
  924. func (s *OpenAIGatewayService) selectBestAccount(ctx context.Context, groupID *int64, platform string, accounts []Account, requestedModel string, excludedIDs map[int64]struct{}, requireCompact bool, requiredCapability OpenAIEndpointCapability, preferLowUpstreamRate bool) (*Account, bool, openAISelectionFilterStats) {
  925. platform = NormalizeOpenAICompatiblePlatform(platform)
  926. compactBlocked := false
  927. filterStats := openAISelectionFilterStats{pool: len(accounts)}
  928. needsUpstreamCheck := s.needsUpstreamChannelRestrictionCheck(ctx, groupID)
  929. eligible := make([]*Account, 0, len(accounts))
  930. compactTiers := make(map[int64]int, len(accounts))
  931. for i := range accounts {
  932. acc := &accounts[i]
  933. // 跳过被排除的账号
  934. // Skip excluded accounts
  935. if _, excluded := excludedIDs[acc.ID]; excluded {
  936. filterStats.exclude("excluded")
  937. continue
  938. }
  939. fresh := s.resolveFreshSchedulableOpenAIAccountBeforeProfit(ctx, acc, platform, requestedModel, false, requiredCapability)
  940. if fresh == nil {
  941. filterStats.exclude("ineligible")
  942. continue
  943. }
  944. fresh = s.recheckSelectedOpenAIAccountFromDBBeforeProfit(ctx, fresh, groupID, platform, requestedModel, false, requiredCapability)
  945. if fresh == nil {
  946. filterStats.exclude("ineligible")
  947. continue
  948. }
  949. if needsUpstreamCheck && s.isUpstreamModelRestrictedByChannel(ctx, *groupID, fresh, requestedModel, requireCompact) {
  950. filterStats.exclude("channel_restricted")
  951. continue
  952. }
  953. if vetoed, reason := openAIProfitControlVetoReason(ctx, fresh); vetoed {
  954. filterStats.exclude(reason)
  955. continue
  956. }
  957. compactTier := 0
  958. if requireCompact {
  959. compactTier = openAICompactSupportTier(fresh)
  960. if compactTier == 0 {
  961. compactBlocked = true
  962. filterStats.exclude("compact_unsupported")
  963. continue
  964. }
  965. }
  966. eligible = append(eligible, fresh)
  967. compactTiers[fresh.ID] = compactTier
  968. }
  969. if len(eligible) == 0 {
  970. return nil, compactBlocked, filterStats
  971. }
  972. rateOrder := openAILegacyUpstreamRateOrder{}
  973. if preferLowUpstreamRate {
  974. rateOrder = newOpenAILegacyUpstreamRateOrder(eligible, time.Now(), s.openAIOAuthSchedulingRateMultiplier(ctx))
  975. }
  976. sort.SliceStable(eligible, func(i, j int) bool {
  977. a, b := eligible[i], eligible[j]
  978. if requireCompact && compactTiers[a.ID] != compactTiers[b.ID] {
  979. return compactTiers[a.ID] > compactTiers[b.ID]
  980. }
  981. if rateCmp := rateOrder.compare(a, b); rateCmp != 0 {
  982. return rateCmp < 0
  983. }
  984. return s.isBetterAccount(a, b)
  985. })
  986. return eligible[0], compactBlocked, filterStats
  987. }
  988. // isBetterAccount 判断 candidate 是否比 current 更优。
  989. // 规则:优先级更高(数值更小)优先;同优先级时,未使用过的优先,其次是最久未使用的。
  990. //
  991. // isBetterAccount checks if candidate is better than current.
  992. // Rules: higher priority (lower value) wins; same priority: never used > least recently used.
  993. func (s *OpenAIGatewayService) isBetterAccount(candidate, current *Account) bool {
  994. // 优先级更高(数值更小)
  995. // Higher priority (lower value)
  996. if candidate.Priority < current.Priority {
  997. return true
  998. }
  999. if candidate.Priority > current.Priority {
  1000. return false
  1001. }
  1002. // 同优先级,比较最后使用时间
  1003. // Same priority, compare last used time
  1004. switch {
  1005. case candidate.LastUsedAt == nil && current.LastUsedAt != nil:
  1006. // candidate 从未使用,优先
  1007. return true
  1008. case candidate.LastUsedAt != nil && current.LastUsedAt == nil:
  1009. // current 从未使用,保持
  1010. return false
  1011. case candidate.LastUsedAt == nil && current.LastUsedAt == nil:
  1012. // 都未使用,保持
  1013. return false
  1014. default:
  1015. // 都使用过,选择最久未使用的
  1016. return candidate.LastUsedAt.Before(*current.LastUsedAt)
  1017. }
  1018. }
  1019. // SelectAccountWithLoadAwareness selects an account with load-awareness and wait plan.
  1020. func (s *OpenAIGatewayService) SelectAccountWithLoadAwareness(ctx context.Context, groupID *int64, sessionHash string, requestedModel string, excludedIDs map[int64]struct{}) (*AccountSelectionResult, error) {
  1021. ctx = s.withOpenAIQuotaAutoPauseContext(ctx)
  1022. ctx = s.withOpenAIGroupPrivacyRequirement(ctx, groupID)
  1023. // 分组利润控制:legacy 公共入口同样装门,保证不经
  1024. // selectAccountWithScheduler 的调用方也无法绕过利润准入。
  1025. ctx = s.withOpenAIProfitControlGate(ctx, groupID)
  1026. return s.selectAccountWithLoadAwareness(ctx, groupID, PlatformOpenAI, sessionHash, requestedModel, excludedIDs, false, "", true)
  1027. }
  1028. func (s *OpenAIGatewayService) selectAccountWithLoadAwareness(ctx context.Context, groupID *int64, platform string, sessionHash string, requestedModel string, excludedIDs map[int64]struct{}, requireCompact bool, requiredCapability OpenAIEndpointCapability, useUpstreamTokenCost bool) (*AccountSelectionResult, error) {
  1029. platform = NormalizeOpenAICompatiblePlatform(platform)
  1030. if s.checkChannelPricingRestriction(ctx, groupID, requestedModel) {
  1031. slog.Warn("channel pricing restriction blocked request",
  1032. "group_id", derefGroupID(groupID),
  1033. "model", requestedModel)
  1034. return nil, fmt.Errorf("%w supporting model: %s (channel pricing restriction)", ErrNoAvailableAccounts, requestedModel)
  1035. }
  1036. cfg := s.schedulingConfig()
  1037. preferLowUpstreamRate := useUpstreamTokenCost && s.isOpenAILowUpstreamRatePriorityEnabled(ctx)
  1038. needsUpstreamCheck := s.needsUpstreamChannelRestrictionCheck(ctx, groupID)
  1039. var stickyAccountID int64
  1040. if sessionHash != "" && s.cache != nil {
  1041. if accountID, err := s.getStickySessionAccountID(ctx, groupID, sessionHash); err == nil {
  1042. stickyAccountID = accountID
  1043. }
  1044. }
  1045. if s.concurrencyService == nil || !cfg.LoadBatchEnabled {
  1046. account, err := s.selectAccountForModelWithExclusions(ctx, groupID, platform, sessionHash, requestedModel, excludedIDs, requireCompact, stickyAccountID, requiredCapability, preferLowUpstreamRate)
  1047. if err != nil {
  1048. return nil, err
  1049. }
  1050. result, err := s.tryAcquireAccountSlot(ctx, account.ID, account.Mode1EffectiveConcurrency())
  1051. if err == nil && result != nil && result.Acquired {
  1052. return s.newAcquiredSelectionResult(ctx, account, result.ReleaseFunc)
  1053. }
  1054. if stickyAccountID > 0 && stickyAccountID == account.ID && s.concurrencyService != nil {
  1055. waitingCount, _ := s.concurrencyService.GetAccountWaitingCount(ctx, account.ID)
  1056. if waitingCount < cfg.StickySessionMaxWaiting {
  1057. return s.newSelectionResult(ctx, account, false, nil, &AccountWaitPlan{
  1058. AccountID: account.ID,
  1059. MaxConcurrency: account.Mode1EffectiveConcurrency(),
  1060. Timeout: cfg.StickySessionWaitTimeout,
  1061. MaxWaiting: cfg.StickySessionMaxWaiting,
  1062. })
  1063. }
  1064. }
  1065. return s.newSelectionResult(ctx, account, false, nil, &AccountWaitPlan{
  1066. AccountID: account.ID,
  1067. MaxConcurrency: account.Mode1EffectiveConcurrency(),
  1068. Timeout: cfg.FallbackWaitTimeout,
  1069. MaxWaiting: cfg.FallbackMaxWaiting,
  1070. })
  1071. }
  1072. accounts, err := s.listSchedulableAccounts(ctx, groupID, platform)
  1073. if err != nil {
  1074. return nil, err
  1075. }
  1076. if len(accounts) == 0 {
  1077. return nil, noAvailableOpenAISelectionError(requestedModel, false, openAISelectionFilterStats{}.summary(""))
  1078. }
  1079. isExcluded := func(accountID int64) bool {
  1080. if excludedIDs == nil {
  1081. return false
  1082. }
  1083. _, excluded := excludedIDs[accountID]
  1084. return excluded
  1085. }
  1086. // ============ Layer 1: Sticky session ============
  1087. // A healthy sticky account whose bounded wait queue is full may be used as a
  1088. // one-request capacity spillover in Layer 2. Keep that spillover temporary:
  1089. // rewriting the durable binding here would make a short burst migrate the
  1090. // whole conversation to a cache-cold account.
  1091. stickySpillover := false
  1092. if sessionHash != "" {
  1093. accountID := stickyAccountID
  1094. if accountID > 0 && !isExcluded(accountID) {
  1095. account, err := s.getSchedulableAccount(ctx, accountID)
  1096. if err == nil {
  1097. clearSticky := shouldClearStickySession(account, requestedModel)
  1098. if clearSticky {
  1099. _ = s.deleteStickySessionAccountID(ctx, groupID, sessionHash)
  1100. }
  1101. if !clearSticky && isOpenAICompatibleAccountEligibleForRequest(ctx, account, platform, requestedModel, false, requiredCapability) {
  1102. account = s.recheckSelectedOpenAIAccountFromDB(ctx, account, groupID, platform, requestedModel, requireCompact, requiredCapability)
  1103. if account == nil {
  1104. _ = s.deleteStickySessionAccountID(ctx, groupID, sessionHash)
  1105. } else if !s.openAIAccountMatchesSchedulingGroup(account, groupID) {
  1106. _ = s.deleteStickySessionAccountID(ctx, groupID, sessionHash)
  1107. } else if s.isOpenAIAccountRequestRuntimeBlocked(account, requestedModel) {
  1108. _ = s.deleteStickySessionAccountID(ctx, groupID, sessionHash)
  1109. } else if needsUpstreamCheck && s.isUpstreamModelRestrictedByChannel(ctx, *groupID, account, requestedModel, requireCompact) {
  1110. _ = s.deleteStickySessionAccountID(ctx, groupID, sessionHash)
  1111. } else if !parentHealthyForShadow(account, s.parentAccountLookup(ctx)) {
  1112. _ = s.deleteStickySessionAccountID(ctx, groupID, sessionHash)
  1113. } else {
  1114. result, err := s.tryAcquireAccountSlot(ctx, accountID, account.Mode1EffectiveConcurrency())
  1115. if err == nil && result != nil && result.Acquired {
  1116. selection, selectErr := s.newAcquiredSelectionResult(ctx, account, result.ReleaseFunc)
  1117. if selectErr != nil {
  1118. return nil, selectErr
  1119. }
  1120. _ = s.refreshStickySessionTTL(ctx, groupID, sessionHash, openaiStickySessionTTL)
  1121. return selection, nil
  1122. }
  1123. waitingCount, _ := s.concurrencyService.GetAccountWaitingCount(ctx, accountID)
  1124. if waitingCount < cfg.StickySessionMaxWaiting {
  1125. return s.newSelectionResult(ctx, account, false, nil, &AccountWaitPlan{
  1126. AccountID: accountID,
  1127. MaxConcurrency: account.Mode1EffectiveConcurrency(),
  1128. Timeout: cfg.StickySessionWaitTimeout,
  1129. MaxWaiting: cfg.StickySessionMaxWaiting,
  1130. })
  1131. }
  1132. stickySpillover = true
  1133. }
  1134. }
  1135. }
  1136. }
  1137. }
  1138. // ============ Layer 2: Load-aware selection ============
  1139. // Per-pass parent-health cache to avoid repeated DB calls when multiple shadow
  1140. // accounts share the same parent.
  1141. parentCacheL2 := make(map[int64]*Account)
  1142. parentLookupL2 := func(id int64) *Account {
  1143. if a, ok := parentCacheL2[id]; ok {
  1144. return a
  1145. }
  1146. if s.accountRepo == nil {
  1147. return nil
  1148. }
  1149. a, _ := s.accountRepo.GetByID(ctx, id)
  1150. parentCacheL2[id] = a
  1151. return a
  1152. }
  1153. baseCandidateCount := 0
  1154. filterStats := openAISelectionFilterStats{pool: len(accounts)}
  1155. candidates := make([]*Account, 0, len(accounts))
  1156. for i := range accounts {
  1157. acc := &accounts[i]
  1158. if isExcluded(acc.ID) {
  1159. filterStats.exclude("excluded")
  1160. continue
  1161. }
  1162. // Scheduler snapshots can be temporarily stale (bucket rebuild is throttled);
  1163. // re-check schedulability here so recently rate-limited/overloaded accounts
  1164. // are not selected again before the bucket is rebuilt.
  1165. if reason := openAICompatibleAccountEligibilityFailureReason(ctx, acc, platform, requestedModel, false, requiredCapability); reason != "" {
  1166. filterStats.exclude(reason)
  1167. continue
  1168. }
  1169. if !parentHealthyForShadow(acc, parentLookupL2) {
  1170. filterStats.exclude("shadow_parent_unhealthy")
  1171. continue
  1172. }
  1173. if s.isOpenAIAccountRequestRuntimeBlocked(acc, requestedModel) {
  1174. filterStats.exclude("runtime_blocked")
  1175. continue
  1176. }
  1177. if needsUpstreamCheck && s.isUpstreamModelRestrictedByChannel(ctx, *groupID, acc, requestedModel, requireCompact) {
  1178. filterStats.exclude("channel_upstream_restricted")
  1179. continue
  1180. }
  1181. baseCandidateCount++
  1182. candidates = append(candidates, acc)
  1183. }
  1184. if len(candidates) == 0 {
  1185. return nil, noAvailableOpenAISelectionError(requestedModel, false, filterStats.summary(""))
  1186. }
  1187. rateOrder := openAILegacyUpstreamRateOrder{}
  1188. if preferLowUpstreamRate {
  1189. rateOrder = newOpenAILegacyUpstreamRateOrder(candidates, time.Now(), s.openAIOAuthSchedulingRateMultiplier(ctx))
  1190. }
  1191. accountLoads := make([]AccountWithConcurrency, 0, len(candidates))
  1192. for _, acc := range candidates {
  1193. accountLoads = append(accountLoads, AccountWithConcurrency{
  1194. ID: acc.ID,
  1195. MaxConcurrency: acc.EffectiveLoadFactor(),
  1196. })
  1197. }
  1198. tryAcquireFromLoadMap := func(loadMap map[int64]*AccountLoadInfo) (*AccountSelectionResult, bool, error) {
  1199. var available []accountWithLoad
  1200. for _, acc := range candidates {
  1201. loadInfo := loadMap[acc.ID]
  1202. if loadInfo == nil {
  1203. loadInfo = &AccountLoadInfo{AccountID: acc.ID}
  1204. }
  1205. if loadInfo.LoadRate < 100 {
  1206. available = append(available, accountWithLoad{
  1207. account: acc,
  1208. loadInfo: loadInfo,
  1209. })
  1210. }
  1211. }
  1212. if len(available) == 0 {
  1213. return nil, false, nil
  1214. }
  1215. sort.SliceStable(available, func(i, j int) bool {
  1216. a, b := available[i], available[j]
  1217. if a.account.Priority != b.account.Priority {
  1218. return a.account.Priority < b.account.Priority
  1219. }
  1220. if a.loadInfo.LoadRate != b.loadInfo.LoadRate {
  1221. return a.loadInfo.LoadRate < b.loadInfo.LoadRate
  1222. }
  1223. switch {
  1224. case a.account.LastUsedAt == nil && b.account.LastUsedAt != nil:
  1225. return true
  1226. case a.account.LastUsedAt != nil && b.account.LastUsedAt == nil:
  1227. return false
  1228. case a.account.LastUsedAt == nil && b.account.LastUsedAt == nil:
  1229. return false
  1230. default:
  1231. return a.account.LastUsedAt.Before(*b.account.LastUsedAt)
  1232. }
  1233. })
  1234. shuffleWithinSortGroups(available)
  1235. if rateOrder.enabled {
  1236. sort.SliceStable(available, func(i, j int) bool {
  1237. return rateOrder.compare(available[i].account, available[j].account) < 0
  1238. })
  1239. }
  1240. selectionOrder := make([]accountWithLoad, 0, len(available))
  1241. if requireCompact {
  1242. appendTier := func(out []accountWithLoad, tier int) []accountWithLoad {
  1243. for _, item := range available {
  1244. if openAICompactSupportTier(item.account) == tier {
  1245. out = append(out, item)
  1246. }
  1247. }
  1248. return out
  1249. }
  1250. selectionOrder = appendTier(selectionOrder, 2)
  1251. selectionOrder = appendTier(selectionOrder, 1)
  1252. // tier 0 候选作为兜底追加:DB recheck 时若发现 cache tier 0 实际
  1253. // 已升级为 1/2(探测刚跑完,cache 尚未刷新),仍可正常命中。
  1254. selectionOrder = appendTier(selectionOrder, 0)
  1255. } else {
  1256. selectionOrder = append(selectionOrder, available...)
  1257. }
  1258. for _, item := range selectionOrder {
  1259. fresh := s.resolveFreshSchedulableOpenAIAccount(ctx, item.account, platform, requestedModel, false, requiredCapability)
  1260. if fresh == nil {
  1261. continue
  1262. }
  1263. fresh = s.recheckSelectedOpenAIAccountFromDB(ctx, fresh, groupID, platform, requestedModel, requireCompact, requiredCapability)
  1264. if fresh == nil {
  1265. continue
  1266. }
  1267. if needsUpstreamCheck && s.isUpstreamModelRestrictedByChannel(ctx, *groupID, fresh, requestedModel, requireCompact) {
  1268. continue
  1269. }
  1270. result, err := s.tryAcquireAccountSlot(ctx, fresh.ID, fresh.Mode1EffectiveConcurrency())
  1271. if err == nil && result != nil && result.Acquired {
  1272. selection, selectErr := s.newAcquiredSelectionResult(ctx, fresh, result.ReleaseFunc)
  1273. if selectErr != nil {
  1274. return nil, true, selectErr
  1275. }
  1276. if sessionHash != "" && !stickySpillover && !gatewayProfitControlGateActive(ctx) {
  1277. _ = s.setStickySessionAccountID(ctx, groupID, sessionHash, fresh.ID, openaiStickySessionTTL)
  1278. }
  1279. return selection, true, nil
  1280. }
  1281. }
  1282. return nil, true, nil
  1283. }
  1284. loadMap, err := s.concurrencyService.GetAccountsLoadBatch(ctx, accountLoads)
  1285. if err != nil {
  1286. ordered := append([]*Account(nil), candidates...)
  1287. sortAccountsByPriorityAndLastUsed(ordered, false)
  1288. if rateOrder.enabled {
  1289. sort.SliceStable(ordered, func(i, j int) bool {
  1290. return rateOrder.compare(ordered[i], ordered[j]) < 0
  1291. })
  1292. }
  1293. if requireCompact {
  1294. ordered = prioritizeOpenAICompactAccounts(ordered)
  1295. }
  1296. for _, acc := range ordered {
  1297. fresh := s.resolveFreshSchedulableOpenAIAccount(ctx, acc, platform, requestedModel, false, requiredCapability)
  1298. if fresh == nil {
  1299. continue
  1300. }
  1301. fresh = s.recheckSelectedOpenAIAccountFromDB(ctx, fresh, groupID, platform, requestedModel, requireCompact, requiredCapability)
  1302. if fresh == nil {
  1303. continue
  1304. }
  1305. if needsUpstreamCheck && s.isUpstreamModelRestrictedByChannel(ctx, *groupID, fresh, requestedModel, requireCompact) {
  1306. continue
  1307. }
  1308. result, err := s.tryAcquireAccountSlot(ctx, fresh.ID, fresh.Mode1EffectiveConcurrency())
  1309. if err == nil && result != nil && result.Acquired {
  1310. selection, selectErr := s.newAcquiredSelectionResult(ctx, fresh, result.ReleaseFunc)
  1311. if selectErr != nil {
  1312. return nil, selectErr
  1313. }
  1314. if sessionHash != "" && !stickySpillover && !gatewayProfitControlGateActive(ctx) {
  1315. _ = s.setStickySessionAccountID(ctx, groupID, sessionHash, fresh.ID, openaiStickySessionTTL)
  1316. }
  1317. return selection, nil
  1318. }
  1319. }
  1320. } else {
  1321. if selection, attempted, selectErr := tryAcquireFromLoadMap(loadMap); selectErr != nil {
  1322. return nil, selectErr
  1323. } else if selection != nil {
  1324. return selection, nil
  1325. } else if attempted {
  1326. if freshLoadMap, loadErr := s.concurrencyService.GetAccountsLoadBatchFresh(ctx, accountLoads); loadErr == nil {
  1327. if selection, _, selectErr := tryAcquireFromLoadMap(freshLoadMap); selectErr != nil {
  1328. return nil, selectErr
  1329. } else if selection != nil {
  1330. return selection, nil
  1331. }
  1332. }
  1333. }
  1334. }
  1335. // ============ Layer 3: Fallback wait ============
  1336. sortAccountsByPriorityAndLastUsed(candidates, false)
  1337. if rateOrder.enabled {
  1338. sort.SliceStable(candidates, func(i, j int) bool {
  1339. return rateOrder.compare(candidates[i], candidates[j]) < 0
  1340. })
  1341. }
  1342. if requireCompact {
  1343. candidates = prioritizeOpenAICompactAccounts(candidates)
  1344. }
  1345. for _, acc := range candidates {
  1346. fresh := s.resolveFreshSchedulableOpenAIAccount(ctx, acc, platform, requestedModel, false, requiredCapability)
  1347. if fresh == nil {
  1348. continue
  1349. }
  1350. fresh = s.recheckSelectedOpenAIAccountFromDB(ctx, fresh, groupID, platform, requestedModel, requireCompact, requiredCapability)
  1351. if fresh == nil {
  1352. continue
  1353. }
  1354. if needsUpstreamCheck && s.isUpstreamModelRestrictedByChannel(ctx, *groupID, fresh, requestedModel, requireCompact) {
  1355. continue
  1356. }
  1357. return s.newSelectionResult(ctx, fresh, false, nil, &AccountWaitPlan{
  1358. AccountID: fresh.ID,
  1359. MaxConcurrency: fresh.Mode1EffectiveConcurrency(),
  1360. Timeout: cfg.FallbackWaitTimeout,
  1361. MaxWaiting: cfg.FallbackMaxWaiting,
  1362. })
  1363. }
  1364. if requireCompact && baseCandidateCount > 0 {
  1365. return nil, ErrNoAvailableCompactAccounts
  1366. }
  1367. return nil, ErrNoAvailableAccounts
  1368. }
  1369. func (s *OpenAIGatewayService) listSchedulableAccounts(ctx context.Context, groupID *int64, platform string) ([]Account, error) {
  1370. platform = NormalizeOpenAICompatiblePlatform(platform)
  1371. if s.schedulerSnapshot != nil {
  1372. accounts, _, err := s.schedulerSnapshot.ListSchedulableAccounts(ctx, groupID, platform, false)
  1373. if err != nil {
  1374. return accounts, err
  1375. }
  1376. accounts = s.filterOpenAIAccountsBySchedulingThreshold(ctx, accounts)
  1377. if platform == PlatformGrok {
  1378. accounts = s.filterGrokFreeQuotaAccountsForOpenAI(ctx, accounts)
  1379. }
  1380. return FilterTierPoolAccounts(accounts, TierPoolFromContext(ctx)), nil
  1381. }
  1382. var accounts []Account
  1383. var err error
  1384. if s.cfg != nil && s.cfg.RunMode == config.RunModeSimple {
  1385. accounts, err = s.accountRepo.ListSchedulableByPlatform(ctx, platform)
  1386. } else if groupID != nil {
  1387. accounts, err = s.accountRepo.ListSchedulableByGroupIDAndPlatform(ctx, *groupID, platform)
  1388. } else {
  1389. accounts, err = s.accountRepo.ListSchedulableUngroupedByPlatform(ctx, platform)
  1390. }
  1391. if err != nil {
  1392. return nil, fmt.Errorf("query accounts failed: %w", err)
  1393. }
  1394. accounts = s.filterOpenAIAccountsBySchedulingThreshold(ctx, accounts)
  1395. if platform == PlatformGrok {
  1396. accounts = s.filterGrokFreeQuotaAccountsForOpenAI(ctx, accounts)
  1397. }
  1398. return FilterTierPoolAccounts(accounts, TierPoolFromContext(ctx)), nil
  1399. }
  1400. func (s *OpenAIGatewayService) tryAcquireAccountSlot(ctx context.Context, accountID int64, maxConcurrency int) (*AcquireResult, error) {
  1401. if s.concurrencyService == nil {
  1402. return &AcquireResult{Acquired: true, ReleaseFunc: func() {}}, nil
  1403. }
  1404. return s.concurrencyService.AcquireAccountSlot(ctx, accountID, maxConcurrency)
  1405. }
  1406. func (s *OpenAIGatewayService) resolveFreshSchedulableOpenAIAccount(ctx context.Context, account *Account, platform string, requestedModel string, requireCompact bool, requiredCapability OpenAIEndpointCapability) *Account {
  1407. fresh := s.resolveFreshSchedulableOpenAIAccountBeforeProfit(ctx, account, platform, requestedModel, requireCompact, requiredCapability)
  1408. if fresh == nil {
  1409. return nil
  1410. }
  1411. if vetoed, _ := openAIProfitControlVetoReason(ctx, fresh); vetoed {
  1412. return nil
  1413. }
  1414. return fresh
  1415. }
  1416. func (s *OpenAIGatewayService) resolveFreshSchedulableOpenAIAccountBeforeProfit(ctx context.Context, account *Account, platform string, requestedModel string, requireCompact bool, requiredCapability OpenAIEndpointCapability) *Account {
  1417. if account == nil {
  1418. return nil
  1419. }
  1420. platform = NormalizeOpenAICompatiblePlatform(platform)
  1421. fresh := account
  1422. if s.schedulerSnapshot != nil {
  1423. current, err := s.getSchedulableAccount(ctx, account.ID)
  1424. if err != nil || current == nil {
  1425. return nil
  1426. }
  1427. fresh = current
  1428. }
  1429. if !isOpenAICompatibleAccountEligibleForRequestBeforeProfit(ctx, fresh, platform, requestedModel, requireCompact, requiredCapability) {
  1430. return nil
  1431. }
  1432. if !parentHealthyForShadow(fresh, s.parentAccountLookup(ctx)) {
  1433. return nil
  1434. }
  1435. if s.isOpenAIAccountRequestRuntimeBlocked(fresh, requestedModel) {
  1436. return nil
  1437. }
  1438. if s.isOpenAIAccountBlockedBySchedulingThreshold(ctx, fresh) {
  1439. return nil
  1440. }
  1441. if s.isOpenAIProxyStreamQuarantined(ctx, fresh) {
  1442. return nil
  1443. }
  1444. return fresh
  1445. }
  1446. // parentAccountLookup 返回供 parentHealthyForShadow 使用的母账号解析闭包:经 accountRepo
  1447. // 按 ID 取当前 Account(repo 为空时 fail-closed 返回 nil)。统一调度/粘连各路径的母账号解析,
  1448. // 取代各调用点重复内联的同一闭包(历史上 recheck 等路径还漏写过 accountRepo==nil 守卫)。
  1449. // L2 候选循环改用带 per-pass 缓存的 parentLookupL2,不走此方法。
  1450. func (s *OpenAIGatewayService) parentAccountLookup(ctx context.Context) func(int64) *Account {
  1451. return func(id int64) *Account {
  1452. if s.accountRepo == nil {
  1453. return nil
  1454. }
  1455. a, _ := s.accountRepo.GetByID(ctx, id)
  1456. return a
  1457. }
  1458. }
  1459. func (s *OpenAIGatewayService) recheckSelectedOpenAIAccountFromDB(ctx context.Context, account *Account, groupID *int64, platform string, requestedModel string, requireCompact bool, requiredCapability OpenAIEndpointCapability) *Account {
  1460. latest := s.recheckSelectedOpenAIAccountFromDBBeforeProfit(ctx, account, groupID, platform, requestedModel, requireCompact, requiredCapability)
  1461. if latest == nil {
  1462. return nil
  1463. }
  1464. if vetoed, _ := openAIProfitControlVetoReason(ctx, latest); vetoed {
  1465. return nil
  1466. }
  1467. return latest
  1468. }
  1469. func (s *OpenAIGatewayService) recheckSelectedOpenAIAccountFromDBBeforeProfit(ctx context.Context, account *Account, groupID *int64, platform string, requestedModel string, requireCompact bool, requiredCapability OpenAIEndpointCapability) *Account {
  1470. if account == nil {
  1471. return nil
  1472. }
  1473. platform = NormalizeOpenAICompatiblePlatform(platform)
  1474. if s.schedulerSnapshot == nil || s.accountRepo == nil {
  1475. if s.openAIGroupRequiresPrivacySet(ctx, groupID) && !account.IsPrivacySet() {
  1476. return nil
  1477. }
  1478. if !isOpenAICompatibleAccountEligibleForRequestBeforeProfit(ctx, account, platform, requestedModel, requireCompact, requiredCapability) {
  1479. return nil
  1480. }
  1481. if s.isOpenAIAccountBlockedBySchedulingThreshold(ctx, account) {
  1482. return nil
  1483. }
  1484. if !parentHealthyForShadow(account, s.parentAccountLookup(ctx)) {
  1485. return nil
  1486. }
  1487. if s.isOpenAIProxyStreamQuarantined(ctx, account) {
  1488. return nil
  1489. }
  1490. return account
  1491. }
  1492. latest, err := s.accountRepo.GetByID(ctx, account.ID)
  1493. if err != nil || latest == nil {
  1494. return nil
  1495. }
  1496. if !s.openAIAccountMatchesSchedulingGroup(latest, groupID) {
  1497. return nil
  1498. }
  1499. if s.openAIGroupRequiresPrivacySet(ctx, groupID) && !latest.IsPrivacySet() {
  1500. return nil
  1501. }
  1502. if !isOpenAICompatibleAccountEligibleForRequestBeforeProfit(ctx, latest, platform, requestedModel, requireCompact, requiredCapability) {
  1503. return nil
  1504. }
  1505. if !parentHealthyForShadow(latest, s.parentAccountLookup(ctx)) {
  1506. return nil
  1507. }
  1508. if s.isOpenAIAccountRequestRuntimeBlocked(latest, requestedModel) {
  1509. return nil
  1510. }
  1511. if s.isOpenAIAccountBlockedBySchedulingThreshold(ctx, latest) {
  1512. return nil
  1513. }
  1514. if s.isOpenAIProxyStreamQuarantined(ctx, latest) {
  1515. return nil
  1516. }
  1517. return latest
  1518. }
  1519. func (s *OpenAIGatewayService) openAIAccountMatchesSchedulingGroup(account *Account, groupID *int64) bool {
  1520. if s != nil && s.cfg != nil && s.cfg.RunMode == config.RunModeSimple {
  1521. return account != nil
  1522. }
  1523. return openAIStickyAccountMatchesGroup(account, groupID)
  1524. }
  1525. func (s *OpenAIGatewayService) getSchedulableAccount(ctx context.Context, accountID int64) (*Account, error) {
  1526. var (
  1527. account *Account
  1528. err error
  1529. )
  1530. if s.schedulerSnapshot != nil {
  1531. account, err = s.schedulerSnapshot.GetAccount(ctx, accountID)
  1532. } else {
  1533. account, err = s.accountRepo.GetByID(ctx, accountID)
  1534. }
  1535. if err != nil || account == nil {
  1536. return account, err
  1537. }
  1538. if s.isOpenAIAccountBlockedBySchedulingThreshold(ctx, account) {
  1539. return nil, nil
  1540. }
  1541. // Legacy sticky (advanced scheduler off) must still free-gate Grok OAuth.
  1542. if account.IsGrok() {
  1543. if gated := s.filterGrokFreeQuotaAccountsForOpenAI(ctx, []Account{*account}); len(gated) == 0 {
  1544. return nil, nil
  1545. }
  1546. }
  1547. return account, nil
  1548. }
  1549. // filterGrokFreeQuotaAccountsForOpenAI applies the same local free soft-gate as
  1550. // GatewayService / advanced scheduler, for OpenAI-compatible legacy selection.
  1551. func (s *OpenAIGatewayService) filterGrokFreeQuotaAccountsForOpenAI(ctx context.Context, accounts []Account) []Account {
  1552. if s == nil {
  1553. return accounts
  1554. }
  1555. return filterGrokFreeQuotaAccountsCore(ctx, s.cfg, s.usageLogRepo, &openaiGrokFreeQuotaGateCache, accounts)
  1556. }
  1557. func (s *OpenAIGatewayService) filterOpenAIAccountsBySchedulingThreshold(ctx context.Context, accounts []Account) []Account {
  1558. if len(accounts) == 0 {
  1559. return accounts
  1560. }
  1561. filtered := make([]Account, 0, len(accounts))
  1562. for i := range accounts {
  1563. if s.isOpenAIAccountBlockedBySchedulingThreshold(ctx, &accounts[i]) {
  1564. continue
  1565. }
  1566. filtered = append(filtered, accounts[i])
  1567. }
  1568. return filtered
  1569. }
  1570. func (s *OpenAIGatewayService) isOpenAIAccountBlockedBySchedulingThreshold(ctx context.Context, account *Account) bool {
  1571. if s == nil || s.rateLimitService == nil || account == nil {
  1572. return false
  1573. }
  1574. return s.rateLimitService.ApplyAccountSchedulingThreshold(ctx, account)
  1575. }
  1576. func (s *OpenAIGatewayService) hydrateSelectedAccount(ctx context.Context, account *Account) (*Account, error) {
  1577. if account == nil {
  1578. return nil, nil
  1579. }
  1580. hydrated := account
  1581. if s.schedulerSnapshot != nil {
  1582. var err error
  1583. hydrated, err = s.schedulerSnapshot.GetAccount(ctx, account.ID)
  1584. if err != nil {
  1585. return nil, err
  1586. }
  1587. if hydrated == nil {
  1588. return nil, fmt.Errorf("selected openai account %d not found during hydration", account.ID)
  1589. }
  1590. }
  1591. if err := ResolveRandomProxyFromSource(ctx, hydrated, s.accountRepo); err != nil {
  1592. return nil, err
  1593. }
  1594. return hydrated, nil
  1595. }
  1596. func (s *OpenAIGatewayService) newSelectionResult(ctx context.Context, account *Account, acquired bool, release func(), waitPlan *AccountWaitPlan) (*AccountSelectionResult, error) {
  1597. hydrated, err := s.hydrateSelectedAccount(ctx, account)
  1598. if err != nil {
  1599. return nil, err
  1600. }
  1601. return attachSelectionProfitGate(ctx, &AccountSelectionResult{
  1602. Account: hydrated,
  1603. Acquired: acquired,
  1604. ReleaseFunc: release,
  1605. WaitPlan: waitPlan,
  1606. }), nil
  1607. }
  1608. func (s *OpenAIGatewayService) newAcquiredSelectionResult(ctx context.Context, account *Account, release func()) (*AccountSelectionResult, error) {
  1609. selection, err := s.newSelectionResult(ctx, account, true, release, nil)
  1610. if err != nil && release != nil {
  1611. release()
  1612. }
  1613. return selection, err
  1614. }
  1615. func (s *OpenAIGatewayService) schedulingConfig() config.GatewaySchedulingConfig {
  1616. if s.cfg != nil {
  1617. return s.cfg.Gateway.Scheduling
  1618. }
  1619. return config.GatewaySchedulingConfig{
  1620. StickySessionMaxWaiting: 3,
  1621. StickySessionWaitTimeout: 45 * time.Second,
  1622. FallbackWaitTimeout: 30 * time.Second,
  1623. FallbackMaxWaiting: 100,
  1624. LoadBatchEnabled: true,
  1625. SlotCleanupInterval: 30 * time.Second,
  1626. }
  1627. }