李建琦
6d655c9903
Release / update-version (push) Has been cancelled
Release / build-frontend (push) Has been cancelled
Release / release (push) Has been cancelled
Release / sync-version-file (push) Has been cancelled
CI / shell (push) Canceled after 0s
CI / test (push) Canceled after 0s
CI / frontend (push) Canceled after 0s
CI / golangci-lint (push) Canceled after 0s
Security Scan / backend-security (push) Canceled after 0s
Security Scan / frontend-security (push) Canceled after 0s
1578 lines
58 KiB
Go
1578 lines
58 KiB
Go
package service
|
|
|
|
// 本文件由 openai_gateway_service.go 纯移动拆分而来：粘性会话哈希、账号选择与
|
|
// 负载感知调度、配额自动暂停判定、并发槽位获取。仅做代码搬迁，无任何行为变更。
|
|
|
|
import (
|
|
"context"
|
|
"encoding/json"
|
|
"fmt"
|
|
"log/slog"
|
|
"sort"
|
|
"strconv"
|
|
"strings"
|
|
"time"
|
|
|
|
"github.com/Wei-Shaw/sub2api/internal/config"
|
|
"github.com/Wei-Shaw/sub2api/internal/pkg/openai"
|
|
"github.com/Wei-Shaw/sub2api/internal/pkg/xai"
|
|
"github.com/gin-gonic/gin"
|
|
"github.com/tidwall/gjson"
|
|
)
|
|
|
|
const (
|
|
openCodeSessionAffinityHeader = "X-Session-Affinity"
|
|
openCodeSessionIDHeader = "X-Session-Id"
|
|
openCodeNativeSessionHeader = "X-OpenCode-Session"
|
|
codeBuddyConversationHeader = "X-Conversation-ID"
|
|
)
|
|
|
|
var explicitOpenAIHeaderSessionNames = []string{
|
|
"session_id",
|
|
"conversation_id",
|
|
openCodeSessionAffinityHeader,
|
|
openCodeSessionIDHeader,
|
|
openCodeNativeSessionHeader,
|
|
codeBuddyConversationHeader,
|
|
}
|
|
|
|
// explicitOpenAIHeaderSessionID resolves stable conversation identifiers sent
|
|
// by OpenAI-compatible clients. Keep this list limited to session-scoped
|
|
// fields: request/message IDs rotate every turn and would defeat sticky routing
|
|
// and upstream prompt caching.
|
|
func explicitOpenAIHeaderSessionID(c *gin.Context) string {
|
|
if c == nil {
|
|
return ""
|
|
}
|
|
|
|
for _, header := range explicitOpenAIHeaderSessionNames {
|
|
if sessionID := strings.TrimSpace(c.GetHeader(header)); sessionID != "" {
|
|
return sessionID
|
|
}
|
|
}
|
|
return ""
|
|
}
|
|
|
|
// ExtractSessionID extracts the raw session ID from headers or body without hashing.
|
|
// Used by ForwardAsAnthropic to pass as prompt_cache_key for upstream cache.
|
|
func (s *OpenAIGatewayService) ExtractSessionID(c *gin.Context, body []byte) string {
|
|
return explicitOpenAIRequestSessionID(c, body)
|
|
}
|
|
|
|
func explicitOpenAISessionID(c *gin.Context, body []byte) string {
|
|
if c == nil {
|
|
return ""
|
|
}
|
|
|
|
sessionID := explicitOpenAIHeaderSessionID(c)
|
|
if sessionID == "" && len(body) > 0 {
|
|
sessionID = strings.TrimSpace(gjson.GetBytes(body, "prompt_cache_key").String())
|
|
}
|
|
return sessionID
|
|
}
|
|
|
|
// explicitOpenAIRequestSessionID extends the common OpenAI session signals
|
|
// with Grok's native conversation header only for requests authenticated to a
|
|
// Grok group. This keeps an unrelated x-grok-conv-id header from changing
|
|
// scheduling or upstream session behavior for non-Grok groups.
|
|
//
|
|
// For Grok groups only, previous_response_id is a last-resort sticky seed so
|
|
// multi-turn Responses chains stay on the same OAuth account when no explicit
|
|
// session/conversation/prompt_cache_key is present. Non-Grok groups omit this
|
|
// so HTTP OpenAI paths that delete previous_response_id before upstream are
|
|
// unchanged.
|
|
func explicitOpenAIRequestSessionID(c *gin.Context, body []byte) string {
|
|
if c == nil {
|
|
return ""
|
|
}
|
|
|
|
sessionID := explicitOpenAIHeaderSessionID(c)
|
|
if sessionID == "" && isGrokRequestContext(c) {
|
|
sessionID = strings.TrimSpace(c.GetHeader(grokConversationIDHeader))
|
|
}
|
|
if sessionID == "" && len(body) > 0 {
|
|
sessionID = strings.TrimSpace(gjson.GetBytes(body, "prompt_cache_key").String())
|
|
}
|
|
if sessionID == "" && isGrokRequestContext(c) && len(body) > 0 {
|
|
sessionID = grokPreviousResponseSessionSeed(body)
|
|
}
|
|
return sessionID
|
|
}
|
|
|
|
// grokPreviousResponseSessionSeed returns a stable sticky seed from a Responses
|
|
// previous_response_id. Only resp_* response ids are accepted; message ids and
|
|
// unknown shapes must not pin sticky routing or prompt-cache identity.
|
|
func grokPreviousResponseSessionSeed(body []byte) string {
|
|
id := strings.TrimSpace(gjson.GetBytes(body, "previous_response_id").String())
|
|
if id == "" {
|
|
return ""
|
|
}
|
|
if ClassifyOpenAIPreviousResponseIDKind(id) != OpenAIPreviousResponseIDKindResponseID {
|
|
return ""
|
|
}
|
|
// Namespace so content-derived seeds never collide with response ids.
|
|
return "grok-prev-resp:" + id
|
|
}
|
|
|
|
// GenerateExplicitSessionHash generates a sticky-session hash only from explicit
|
|
// client session signals. It intentionally skips content-derived fallback and is
|
|
// used by stateless endpoints such as /v1/images.
|
|
func (s *OpenAIGatewayService) GenerateExplicitSessionHash(c *gin.Context, body []byte) string {
|
|
sessionID := explicitOpenAIRequestSessionID(c, body)
|
|
if sessionID == "" {
|
|
return ""
|
|
}
|
|
|
|
currentHash, legacyHash := deriveOpenAISessionHashes(sessionID)
|
|
attachOpenAILegacySessionHashToGin(c, legacyHash)
|
|
return currentHash
|
|
}
|
|
|
|
// GenerateSessionHash generates a sticky-session hash for OpenAI requests.
|
|
//
|
|
// Priority:
|
|
// 1. Header: session_id
|
|
// 2. Header: conversation_id
|
|
// 3. Header: x-session-affinity / x-session-id / x-opencode-session (OpenCode)
|
|
// 4. Header: x-conversation-id (CodeBuddy)
|
|
// 5. Header: x-grok-conv-id (Grok groups only)
|
|
// 6. Body: prompt_cache_key
|
|
// 7. Body: content-based fallback (model + system + tools + first user message)
|
|
//
|
|
// Grok sticky affinity is intentionally separate from the upstream
|
|
// prompt_cache_key identity (resolveGrokCacheIdentity): sticky pins an OAuth
|
|
// account for multi-turn routing, while the cache identity is tenant+model
|
|
// isolated for xAI server-side prompt cache. For Grok groups we scope the
|
|
// sticky seed with the client-requested model so switching models does not
|
|
// inherit a stale account binding (grok2api affinityKey pattern).
|
|
func (s *OpenAIGatewayService) GenerateSessionHash(c *gin.Context, body []byte) string {
|
|
if c == nil {
|
|
return ""
|
|
}
|
|
|
|
sessionID := explicitOpenAIRequestSessionID(c, body)
|
|
if sessionID == "" && len(body) > 0 {
|
|
sessionID = deriveOpenAIContentSessionSeed(body)
|
|
}
|
|
if sessionID == "" {
|
|
return ""
|
|
}
|
|
|
|
if isGrokRequestContext(c) {
|
|
sessionID = grokStickyAffinitySeed(sessionID, body)
|
|
}
|
|
|
|
currentHash, legacyHash := deriveOpenAISessionHashes(sessionID)
|
|
attachOpenAILegacySessionHashToGin(c, legacyHash)
|
|
return currentHash
|
|
}
|
|
|
|
// grokStickyAffinitySeed scopes sticky routing by model without changing the
|
|
// upstream prompt_cache_key written by applyGrokResponsesCacheIdentity.
|
|
func grokStickyAffinitySeed(sessionID string, body []byte) string {
|
|
sessionID = strings.TrimSpace(sessionID)
|
|
if sessionID == "" {
|
|
return ""
|
|
}
|
|
model := ""
|
|
if len(body) > 0 {
|
|
model = strings.ToLower(strings.TrimSpace(gjson.GetBytes(body, "model").String()))
|
|
}
|
|
if model == "" {
|
|
return "grok-affinity:v1:" + sessionID
|
|
}
|
|
return "grok-affinity:v1:" + model + ":" + sessionID
|
|
}
|
|
|
|
// GenerateSessionHashWithFallback 先按常规信号生成会话哈希；
|
|
// 当未携带 session_id/conversation_id/prompt_cache_key 时，使用 fallbackSeed 生成稳定哈希。
|
|
// 该方法用于 WS ingress，避免会话信号缺失时发生跨账号漂移。
|
|
func (s *OpenAIGatewayService) GenerateSessionHashWithFallback(c *gin.Context, body []byte, fallbackSeed string) string {
|
|
sessionHash := s.GenerateSessionHash(c, body)
|
|
if sessionHash != "" {
|
|
return sessionHash
|
|
}
|
|
|
|
seed := strings.TrimSpace(fallbackSeed)
|
|
if seed == "" {
|
|
return ""
|
|
}
|
|
|
|
currentHash, legacyHash := deriveOpenAISessionHashes(seed)
|
|
attachOpenAILegacySessionHashToGin(c, legacyHash)
|
|
return currentHash
|
|
}
|
|
|
|
func resolveOpenAIUpstreamOriginator(c *gin.Context, isOfficialClient bool) string {
|
|
if c != nil {
|
|
if originator := strings.TrimSpace(c.GetHeader("originator")); originator != "" {
|
|
return originator
|
|
}
|
|
}
|
|
if isOfficialClient {
|
|
return openai.CodexDefaultOriginator
|
|
}
|
|
return "opencode"
|
|
}
|
|
|
|
// BindStickySession sets session -> account binding with standard TTL.
|
|
func (s *OpenAIGatewayService) BindStickySession(ctx context.Context, groupID *int64, sessionHash string, accountID int64) error {
|
|
if sessionHash == "" || accountID <= 0 {
|
|
return nil
|
|
}
|
|
ttl := openaiStickySessionTTL
|
|
if s != nil && s.cfg != nil && s.cfg.Gateway.OpenAIWS.StickySessionTTLSeconds > 0 {
|
|
ttl = time.Duration(s.cfg.Gateway.OpenAIWS.StickySessionTTLSeconds) * time.Second
|
|
}
|
|
return s.setStickySessionAccountID(ctx, groupID, sessionHash, accountID, ttl)
|
|
}
|
|
|
|
// SelectAccount selects an OpenAI account with sticky session support
|
|
func (s *OpenAIGatewayService) SelectAccount(ctx context.Context, groupID *int64, sessionHash string) (*Account, error) {
|
|
return s.SelectAccountForModel(ctx, groupID, sessionHash, "")
|
|
}
|
|
|
|
// SelectAccountForModel selects an account supporting the requested model
|
|
func (s *OpenAIGatewayService) SelectAccountForModel(ctx context.Context, groupID *int64, sessionHash string, requestedModel string) (*Account, error) {
|
|
return s.SelectAccountForModelWithExclusions(ctx, groupID, sessionHash, requestedModel, nil)
|
|
}
|
|
|
|
// SelectAccountForModelWithExclusions selects an account supporting the requested model while excluding specified accounts.
|
|
// SelectAccountForModelWithExclusions 选择支持指定模型的账号，同时排除指定的账号。
|
|
func (s *OpenAIGatewayService) SelectAccountForModelWithExclusions(ctx context.Context, groupID *int64, sessionHash string, requestedModel string, excludedIDs map[int64]struct{}) (*Account, error) {
|
|
return s.selectAccountForModelWithExclusions(s.withOpenAIQuotaAutoPauseContext(ctx), groupID, PlatformOpenAI, sessionHash, requestedModel, excludedIDs, false, 0, "", false)
|
|
}
|
|
|
|
// NormalizeOpenAICompatiblePlatform 保留 grok 与国产 OpenAI 兼容供应商（kimi/zhipu/
|
|
// deepseek）的原值，其他值一律归一为 openai。调度器据此对账号与请求做精确平台匹配：
|
|
// kimi 分组请求只命中 kimi 账号，语义与 openai/grok 一致。
|
|
// （upstream 曾将本函数改为未导出 normalizeOpenAICompatiblePlatform，本分支的
|
|
// handler 调度入口仍需导出，保持导出名。）
|
|
func NormalizeOpenAICompatiblePlatform(platform string) string {
|
|
switch platform {
|
|
case PlatformGrok, PlatformKimi, PlatformZhipu, PlatformDeepseek:
|
|
return platform
|
|
default:
|
|
return PlatformOpenAI
|
|
}
|
|
}
|
|
|
|
// noAvailableOpenAISelectionError builds the standard "no account available" error
|
|
// while preserving the legacy /responses/compact error when applicable.
|
|
// details carries an optional machine-parseable exclusion summary (e.g.
|
|
// "pool=2, filtered: quota_auto_pause_7d=1 runtime_blocked=1") appended in
|
|
// parentheses. It is for server-side logs / ops diagnostics only: handlers
|
|
// never forward this error text to OpenAI-platform clients (they respond with
|
|
// the generic classification message). Callers that must preserve the legacy
|
|
// message pass "".
|
|
func noAvailableOpenAISelectionError(requestedModel string, compactBlocked bool, details string) error {
|
|
if compactBlocked {
|
|
return ErrNoAvailableCompactAccounts
|
|
}
|
|
message := "no available OpenAI accounts"
|
|
if requestedModel != "" {
|
|
message = fmt.Sprintf("no available OpenAI accounts supporting model: %s", requestedModel)
|
|
}
|
|
if details != "" {
|
|
message += " (" + details + ")"
|
|
}
|
|
return openAINoAvailableSelectionError{message: message}
|
|
}
|
|
|
|
type openAINoAvailableSelectionError struct {
|
|
message string
|
|
}
|
|
|
|
func (e openAINoAvailableSelectionError) Error() string {
|
|
return e.message
|
|
}
|
|
|
|
func (e openAINoAvailableSelectionError) Unwrap() error {
|
|
return ErrNoAvailableAccounts
|
|
}
|
|
|
|
// openAICompactSupportTier classifies an OpenAI-compatible account by compact capability.
|
|
// 0 = explicitly unsupported, 1 = unknown / not yet probed, 2 = explicitly supported.
|
|
func openAICompactSupportTier(account *Account) int {
|
|
if account == nil {
|
|
return 0
|
|
}
|
|
if account.IsGrok() {
|
|
return 2
|
|
}
|
|
if !account.IsOpenAI() {
|
|
return 0
|
|
}
|
|
supported, known := account.OpenAICompactSupportKnown()
|
|
if !known {
|
|
return 1
|
|
}
|
|
if supported {
|
|
return 2
|
|
}
|
|
return 0
|
|
}
|
|
|
|
// isOpenAICompatibleAccountEligibleForRequest 判断 OpenAI 兼容账号是否满足本次请求的调度条件。
|
|
// 检查内容包括：平台匹配、账号可用性、quota 自动暂停、spark 路由限制、模型支持及端点能力。
|
|
//
|
|
// 注意：对 spark 影子账号，调用方还须额外调用 parentHealthyForShadow(account, lookup)
|
|
// 检查母账号凭据可用性；该检查未内置于本函数，以避免注入 DB 依赖。
|
|
func isOpenAICompatibleAccountEligibleForRequest(ctx context.Context, account *Account, platform string, requestedModel string, requireCompact bool, requiredCapability OpenAIEndpointCapability) bool {
|
|
if !isOpenAICompatibleAccountEligibleForRequestBeforeProfit(ctx, account, platform, requestedModel, requireCompact, requiredCapability) {
|
|
return false
|
|
}
|
|
// 分组利润控制：legacy 引擎的粘性/候选循环与 DB recheck 共用
|
|
// 本判定，任何 fallback 都不能把利润不合格账号重新放回候选。
|
|
if vetoed, _ := openAIProfitControlVetoReason(ctx, account); vetoed {
|
|
return false
|
|
}
|
|
return true
|
|
}
|
|
|
|
// isOpenAICompatibleAccountEligibleForRequestBeforeProfit applies every
|
|
// ordinary scheduling gate. Legacy selection uses it before classifying the
|
|
// profit veto so earlier failures retain their actual reason.
|
|
func isOpenAICompatibleAccountEligibleForRequestBeforeProfit(ctx context.Context, account *Account, platform string, requestedModel string, requireCompact bool, requiredCapability OpenAIEndpointCapability) bool {
|
|
platform = NormalizeOpenAICompatiblePlatform(platform)
|
|
if account == nil || account.Platform != platform || !account.IsOpenAICompatible() || !account.IsSchedulableForModelWithContext(ctx, requestedModel) {
|
|
return false
|
|
}
|
|
if account.IsOpenAI() {
|
|
if paused, reason := shouldAutoPauseOpenAIAccountByQuota(ctx, account); paused {
|
|
// Debug level: this fires per-candidate on the scheduling hot path, so Info
|
|
// would amplify into log spam once several accounts cross the threshold.
|
|
slog.Debug("account_auto_paused_by_quota",
|
|
"account_id", account.ID,
|
|
"window", reason.window,
|
|
"threshold", reason.threshold,
|
|
"utilization", reason.utilization,
|
|
)
|
|
return false
|
|
}
|
|
}
|
|
if account.IsGrok() {
|
|
if paused, reason := shouldAutoPauseGrokAccountByQuota(account); paused {
|
|
slog.Debug("grok_account_auto_paused_by_quota",
|
|
"account_id", account.ID,
|
|
"window", reason.window,
|
|
"threshold", reason.threshold,
|
|
"utilization", reason.utilization,
|
|
)
|
|
return false
|
|
}
|
|
}
|
|
if requestedModel != "" && !account.IsModelSupported(requestedModel) {
|
|
return false
|
|
}
|
|
if !account.SupportsOpenAIEndpointCapability(requiredCapability) {
|
|
if account.IsGrok() && requiredCapability == OpenAIEndpointCapabilityGrokMediaGeneration {
|
|
_, reason := account.GrokMediaGenerationEligibility()
|
|
slog.Debug("grok_media_account_ineligible", "account_id", account.ID, "reason", reason)
|
|
}
|
|
return false
|
|
}
|
|
if requireCompact && openAICompactSupportTier(account) == 0 {
|
|
return false
|
|
}
|
|
return true
|
|
}
|
|
|
|
type openAIQuotaAutoPauseDecision struct {
|
|
window string
|
|
threshold float64
|
|
utilization float64
|
|
}
|
|
|
|
func shouldAutoPauseGrokAccountByQuota(account *Account) (bool, openAIQuotaAutoPauseDecision) {
|
|
if account == nil || !account.IsGrok() || account.Type != AccountTypeOAuth {
|
|
return false, openAIQuotaAutoPauseDecision{}
|
|
}
|
|
snapshot, err := grokQuotaSnapshotFromExtra(account.Extra)
|
|
if err != nil || snapshot == nil {
|
|
return false, openAIQuotaAutoPauseDecision{}
|
|
}
|
|
now := time.Now()
|
|
if grokQuotaSnapshotStaleForPause(snapshot, now) {
|
|
return false, openAIQuotaAutoPauseDecision{}
|
|
}
|
|
if grokQuotaRetryAfterActive(snapshot, now) {
|
|
return true, openAIQuotaAutoPauseDecision{window: "retry_after", threshold: 1, utilization: 1}
|
|
}
|
|
if paused, decision := shouldAutoPauseGrokQuotaWindow("requests", snapshot.Requests, now); paused {
|
|
return true, decision
|
|
}
|
|
if paused, decision := shouldAutoPauseGrokQuotaWindow("tokens", snapshot.Tokens, now); paused {
|
|
return true, decision
|
|
}
|
|
return false, openAIQuotaAutoPauseDecision{}
|
|
}
|
|
|
|
func grokQuotaRetryAfterActive(snapshot *xai.QuotaSnapshot, now time.Time) bool {
|
|
if snapshot == nil || snapshot.RetryAfterSeconds == nil || *snapshot.RetryAfterSeconds <= 0 {
|
|
return false
|
|
}
|
|
if strings.TrimSpace(snapshot.UpdatedAt) == "" {
|
|
return true
|
|
}
|
|
updatedAt, err := parseTime(snapshot.UpdatedAt)
|
|
if err != nil {
|
|
return true
|
|
}
|
|
retryAfterUntil := updatedAt.Add(time.Duration(*snapshot.RetryAfterSeconds) * time.Second)
|
|
return now.Before(retryAfterUntil)
|
|
}
|
|
|
|
func shouldAutoPauseGrokQuotaWindow(name string, window *xai.QuotaWindow, now time.Time) (bool, openAIQuotaAutoPauseDecision) {
|
|
if window == nil || window.Limit == nil || window.Remaining == nil || *window.Limit <= 0 {
|
|
return false, openAIQuotaAutoPauseDecision{}
|
|
}
|
|
if window.ResetUnix != nil && *window.ResetUnix > 0 && !now.Before(time.Unix(*window.ResetUnix, 0)) {
|
|
return false, openAIQuotaAutoPauseDecision{}
|
|
}
|
|
utilization := float64(*window.Limit-*window.Remaining) / float64(*window.Limit)
|
|
if *window.Remaining <= 0 || utilization >= 1 {
|
|
return true, openAIQuotaAutoPauseDecision{window: name, threshold: 1, utilization: utilization}
|
|
}
|
|
return false, openAIQuotaAutoPauseDecision{}
|
|
}
|
|
|
|
func grokQuotaSnapshotStaleForPause(snapshot *xai.QuotaSnapshot, now time.Time) bool {
|
|
if snapshot == nil || strings.TrimSpace(snapshot.UpdatedAt) == "" {
|
|
return false
|
|
}
|
|
updatedAt, err := parseTime(snapshot.UpdatedAt)
|
|
if err != nil {
|
|
return false
|
|
}
|
|
return now.Sub(updatedAt) >= openAICodexAutoPauseStaleAfter
|
|
}
|
|
|
|
func shouldAutoPauseOpenAIAccountByQuota(ctx context.Context, account *Account) (bool, openAIQuotaAutoPauseDecision) {
|
|
if account == nil || !account.IsOpenAI() {
|
|
return false, openAIQuotaAutoPauseDecision{}
|
|
}
|
|
// Per-account explicit-disable flags must take precedence over the global default.
|
|
// Without these, leaving the account threshold blank means "use global default",
|
|
// so an admin has no way to exempt a single account from auto-pause once a global
|
|
// default exists. The disable flag is per-window so an account can opt out of
|
|
// only 5h or only 7d auto-pause.
|
|
disabled5h := resolveAccountExtraBool(account.Extra, "auto_pause_5h_disabled")
|
|
disabled7d := resolveAccountExtraBool(account.Extra, "auto_pause_7d_disabled")
|
|
threshold5h, threshold7d := resolveOpenAIQuotaAutoPauseThresholds(ctx, account)
|
|
now := time.Now()
|
|
if !disabled5h && threshold5h > 0 {
|
|
if utilization, ok := resolveOpenAIQuotaUtilization(account.Extra, "5h", now); ok && utilization >= threshold5h {
|
|
return true, openAIQuotaAutoPauseDecision{window: "5h", threshold: threshold5h, utilization: utilization}
|
|
}
|
|
}
|
|
if !disabled7d && threshold7d > 0 {
|
|
if utilization, ok := resolveOpenAIQuotaUtilization(account.Extra, "7d", now); ok && utilization >= threshold7d {
|
|
return true, openAIQuotaAutoPauseDecision{window: "7d", threshold: threshold7d, utilization: utilization}
|
|
}
|
|
}
|
|
return false, openAIQuotaAutoPauseDecision{}
|
|
}
|
|
|
|
// resolveAccountExtraBool reads a bool-like value from account extra, tolerating
|
|
// the few shapes JSON unmarshalling may produce (real bool, "true"/"false"
|
|
// strings, 0/1 numbers).
|
|
func resolveAccountExtraBool(extra map[string]any, key string) bool {
|
|
if len(extra) == 0 {
|
|
return false
|
|
}
|
|
value, ok := extra[key]
|
|
if !ok || value == nil {
|
|
return false
|
|
}
|
|
switch v := value.(type) {
|
|
case bool:
|
|
return v
|
|
case string:
|
|
parsed, err := strconv.ParseBool(strings.TrimSpace(v))
|
|
return err == nil && parsed
|
|
case float64:
|
|
return v != 0
|
|
case float32:
|
|
return v != 0
|
|
case int:
|
|
return v != 0
|
|
case int64:
|
|
return v != 0
|
|
case json.Number:
|
|
if i, err := v.Int64(); err == nil {
|
|
return i != 0
|
|
}
|
|
}
|
|
return false
|
|
}
|
|
|
|
func resolveOpenAIQuotaAutoPauseThresholds(ctx context.Context, account *Account) (float64, float64) {
|
|
threshold5h, _ := resolveAccountExtraNumber(account.Extra, "auto_pause_5h_threshold")
|
|
threshold7d, _ := resolveAccountExtraNumber(account.Extra, "auto_pause_7d_threshold")
|
|
threshold5h = clamp01(threshold5h)
|
|
threshold7d = clamp01(threshold7d)
|
|
if threshold5h > 0 && threshold7d > 0 {
|
|
return threshold5h, threshold7d
|
|
}
|
|
settings := openAIQuotaAutoPauseSettingsFromContext(ctx)
|
|
if threshold5h <= 0 {
|
|
threshold5h = clamp01(settings.DefaultThreshold5h)
|
|
}
|
|
if threshold7d <= 0 {
|
|
threshold7d = clamp01(settings.DefaultThreshold7d)
|
|
}
|
|
return threshold5h, threshold7d
|
|
}
|
|
|
|
func resolveAccountExtraNumber(extra map[string]any, keys ...string) (float64, bool) {
|
|
if len(extra) == 0 {
|
|
return 0, false
|
|
}
|
|
for _, key := range keys {
|
|
value, ok := extra[key]
|
|
if !ok || value == nil {
|
|
continue
|
|
}
|
|
switch v := value.(type) {
|
|
case float64:
|
|
return v, true
|
|
case float32:
|
|
return float64(v), true
|
|
case int:
|
|
return float64(v), true
|
|
case int64:
|
|
return float64(v), true
|
|
case json.Number:
|
|
parsed, err := v.Float64()
|
|
if err == nil {
|
|
return parsed, true
|
|
}
|
|
case string:
|
|
parsed, err := strconv.ParseFloat(strings.TrimSpace(v), 64)
|
|
if err == nil {
|
|
return parsed, true
|
|
}
|
|
}
|
|
}
|
|
return 0, false
|
|
}
|
|
|
|
// resolveOpenAIQuotaUtilization returns the current utilization ratio (0..1) for the
|
|
// given Codex usage window. ok=false means there is no usable signal to pause on:
|
|
// either no snapshot exists, or the window has already rolled over so the cached
|
|
// percentage is stale. The stale guard matters because a paused account stops
|
|
// receiving requests, so its snapshot is never refreshed from upstream headers —
|
|
// without this check an old used_percent would keep the account paused forever even
|
|
// after the real window reset.
|
|
func resolveOpenAIQuotaUtilization(extra map[string]any, window string, now time.Time) (float64, bool) {
|
|
usedPercent := readOpenAIQuotaUsedPercent(extra, window)
|
|
if usedPercent <= 0 {
|
|
return 0, false
|
|
}
|
|
if openAIQuotaWindowReset(extra, window, now) {
|
|
return 0, false
|
|
}
|
|
// 快照过于陈旧（账号长期未收到流量刷新）时，不再据此暂停。放行后下一次响应头
|
|
// 会刷新快照实现自愈，避免账号在错误/过期的 used% 上被永久跳过（issue #2994）。
|
|
if openAICodexSnapshotStaleForPause(extra, now) {
|
|
return 0, false
|
|
}
|
|
return usedPercent / 100, true
|
|
}
|
|
|
|
// openAICodexSnapshotStaleForPause reports whether the Codex usage snapshot is stale
|
|
// enough that it should no longer keep an account auto-paused. It anchors on
|
|
// codex_usage_updated_at (always written by buildCodexUsageExtraUpdates). A missing or
|
|
// unparseable timestamp returns false (treated as fresh, so the account stays paused) —
|
|
// this is deliberate: it prevents any snapshot without a write time from silently escaping
|
|
// auto-pause, and a genuinely-exhausted account that is actively served refreshes the
|
|
// timestamp on every response so it never crosses the staleness bound.
|
|
func openAICodexSnapshotStaleForPause(extra map[string]any, now time.Time) bool {
|
|
if len(extra) == 0 {
|
|
return false
|
|
}
|
|
updatedRaw, ok := extra["codex_usage_updated_at"]
|
|
if !ok {
|
|
return false
|
|
}
|
|
updatedAt, err := parseTime(fmt.Sprint(updatedRaw))
|
|
if err != nil {
|
|
return false
|
|
}
|
|
return now.Sub(updatedAt) >= openAICodexAutoPauseStaleAfter
|
|
}
|
|
|
|
// openAIQuotaWindowReset reports whether the Codex usage window's reset time has
|
|
// already passed relative to now. It prefers the absolute codex_<window>_reset_at
|
|
// timestamp and falls back to codex_<window>_reset_after_seconds anchored at
|
|
// codex_usage_updated_at, mirroring AccountUsageService's window-progress logic.
|
|
func openAIQuotaWindowReset(extra map[string]any, window string, now time.Time) bool {
|
|
if len(extra) == 0 {
|
|
return false
|
|
}
|
|
if resetAtRaw, ok := extra["codex_"+window+"_reset_at"]; ok {
|
|
if resetAt, err := parseTime(fmt.Sprint(resetAtRaw)); err == nil {
|
|
return !now.Before(resetAt)
|
|
}
|
|
}
|
|
resetAfter := parseExtraInt(extra["codex_"+window+"_reset_after_seconds"])
|
|
if resetAfter <= 0 {
|
|
return false
|
|
}
|
|
base := now
|
|
if updatedRaw, ok := extra["codex_usage_updated_at"]; ok {
|
|
if updatedAt, err := parseTime(fmt.Sprint(updatedRaw)); err == nil {
|
|
base = updatedAt
|
|
}
|
|
}
|
|
resetAt := base.Add(time.Duration(resetAfter) * time.Second)
|
|
return !now.Before(resetAt)
|
|
}
|
|
|
|
func readOpenAIQuotaUsedPercent(extra map[string]any, window string) float64 {
|
|
if len(extra) == 0 {
|
|
return 0
|
|
}
|
|
if value, ok := resolveAccountExtraNumber(extra, "codex_"+window+"_used_percent"); ok {
|
|
return value
|
|
}
|
|
return 0
|
|
}
|
|
|
|
type openAIQuotaAutoPauseCtxKey struct{}
|
|
|
|
func withOpenAIQuotaAutoPauseSettings(ctx context.Context, settings OpsOpenAIAccountQuotaAutoPauseSettings) context.Context {
|
|
if ctx == nil {
|
|
ctx = context.Background()
|
|
}
|
|
return context.WithValue(ctx, openAIQuotaAutoPauseCtxKey{}, settings)
|
|
}
|
|
|
|
func openAIQuotaAutoPauseSettingsFromContext(ctx context.Context) OpsOpenAIAccountQuotaAutoPauseSettings {
|
|
if ctx == nil {
|
|
return OpsOpenAIAccountQuotaAutoPauseSettings{}
|
|
}
|
|
settings, _ := ctx.Value(openAIQuotaAutoPauseCtxKey{}).(OpsOpenAIAccountQuotaAutoPauseSettings)
|
|
return settings
|
|
}
|
|
|
|
func (s *OpenAIGatewayService) withOpenAIQuotaAutoPauseContext(ctx context.Context) context.Context {
|
|
if s == nil || s.settingService == nil {
|
|
return ctx
|
|
}
|
|
return withOpenAIQuotaAutoPauseSettings(ctx, s.settingService.GetOpenAIQuotaAutoPauseSettings(ctx))
|
|
}
|
|
|
|
// prioritizeOpenAICompactAccounts re-orders a slice so that accounts with known
|
|
// compact support are tried first, followed by unknown, then explicitly unsupported.
|
|
// The relative order within each tier is preserved.
|
|
func prioritizeOpenAICompactAccounts(accounts []*Account) []*Account {
|
|
if len(accounts) == 0 {
|
|
return nil
|
|
}
|
|
supported := make([]*Account, 0, len(accounts))
|
|
unknown := make([]*Account, 0, len(accounts))
|
|
unsupported := make([]*Account, 0, len(accounts))
|
|
for _, account := range accounts {
|
|
switch openAICompactSupportTier(account) {
|
|
case 2:
|
|
supported = append(supported, account)
|
|
case 1:
|
|
unknown = append(unknown, account)
|
|
default:
|
|
unsupported = append(unsupported, account)
|
|
}
|
|
}
|
|
out := make([]*Account, 0, len(accounts))
|
|
out = append(out, supported...)
|
|
out = append(out, unknown...)
|
|
out = append(out, unsupported...)
|
|
return out
|
|
}
|
|
|
|
// resolveOpenAIAccountUpstreamModelForRequest resolves the upstream model that
|
|
// would be sent for a given request, honoring the legacy compact-only mapping
|
|
// when the caller is on the /responses/compact path.
|
|
func resolveOpenAIAccountUpstreamModelForRequest(account *Account, requestedModel string, requireCompact bool) string {
|
|
// Forward checks the raw Chat Completions fallback before passthrough.
|
|
// These API-key accounts therefore apply normal account model_mapping and
|
|
// upstream normalization, but never compact_model_mapping.
|
|
if shouldForwardOpenAIResponsesViaRawChatCompletions(account) {
|
|
upstreamModel := resolveOpenAIForwardModel(account, requestedModel, "")
|
|
return normalizeOpenAIModelForUpstream(account, upstreamModel)
|
|
}
|
|
|
|
// Passthrough accounts only replace authentication. Their Forward path
|
|
// keeps the channel-mapped model in the request body and does not apply the
|
|
// account's normal model_mapping. Legacy /responses/compact is the one
|
|
// exception: forwardOpenAIPassthrough applies compact_model_mapping
|
|
// directly to that channel-mapped model.
|
|
if account != nil && account.IsOpenAIPassthroughEnabled() {
|
|
upstreamModel := strings.TrimSpace(requestedModel)
|
|
if upstreamModel == "" {
|
|
return ""
|
|
}
|
|
if requireCompact {
|
|
return resolveOpenAICompactForwardModel(account, upstreamModel)
|
|
}
|
|
return upstreamModel
|
|
}
|
|
|
|
upstreamModel := resolveOpenAIForwardModel(account, requestedModel, "")
|
|
if upstreamModel == "" {
|
|
return ""
|
|
}
|
|
if requireCompact {
|
|
compactModel := resolveOpenAICompactForwardModel(account, upstreamModel)
|
|
if compactModel != upstreamModel {
|
|
return compactModel
|
|
}
|
|
}
|
|
return normalizeOpenAIModelForUpstream(account, upstreamModel)
|
|
}
|
|
|
|
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) {
|
|
platform = NormalizeOpenAICompatiblePlatform(platform)
|
|
if s.checkChannelPricingRestriction(ctx, groupID, requestedModel) {
|
|
slog.Warn("channel pricing restriction blocked request",
|
|
"group_id", derefGroupID(groupID),
|
|
"model", requestedModel)
|
|
return nil, fmt.Errorf("%w supporting model: %s (channel pricing restriction)", ErrNoAvailableAccounts, requestedModel)
|
|
}
|
|
|
|
// 1. 尝试粘性会话命中
|
|
// Try sticky session hit
|
|
if account := s.tryStickySessionHit(ctx, groupID, platform, sessionHash, requestedModel, excludedIDs, requireCompact, stickyAccountID, requiredCapability); account != nil {
|
|
return account, nil
|
|
}
|
|
|
|
// 2. 获取可调度的 OpenAI 账号
|
|
// Get schedulable OpenAI accounts
|
|
accounts, err := s.listSchedulableAccounts(ctx, groupID, platform)
|
|
if err != nil {
|
|
return nil, fmt.Errorf("query accounts failed: %w", err)
|
|
}
|
|
|
|
// 3. 按优先级 + LRU 选择最佳账号
|
|
// Select by priority + LRU
|
|
selected, compactBlocked, filterStats := s.selectBestAccount(ctx, groupID, platform, accounts, requestedModel, excludedIDs, requireCompact, requiredCapability, preferLowUpstreamRate)
|
|
|
|
if selected == nil {
|
|
return nil, noAvailableOpenAISelectionError(requestedModel, compactBlocked, filterStats.summary(""))
|
|
}
|
|
|
|
hydrated, err := s.hydrateSelectedAccount(ctx, selected)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
// 4. 设置粘性会话绑定（利润门下推迟到 handler 终检通过后再绑定，
|
|
// 终检否决的账号不得成为新的粘性目标；无门保持既有 eager 绑定与 TTL）
|
|
// Set sticky session binding (deferred until terminal admission under a profit gate)
|
|
if sessionHash != "" && !gatewayProfitControlGateActive(ctx) {
|
|
_ = s.setStickySessionAccountID(ctx, groupID, sessionHash, selected.ID, openaiStickySessionTTL)
|
|
}
|
|
|
|
return hydrated, nil
|
|
}
|
|
|
|
// tryStickySessionHit 尝试从粘性会话获取账号。
|
|
// 如果命中且账号可用则返回账号；如果账号不可用则清理会话并返回 nil。
|
|
//
|
|
// tryStickySessionHit attempts to get account from sticky session.
|
|
// Returns account if hit and usable; clears session and returns nil if account is unavailable.
|
|
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 {
|
|
if sessionHash == "" {
|
|
return nil
|
|
}
|
|
platform = NormalizeOpenAICompatiblePlatform(platform)
|
|
|
|
accountID := stickyAccountID
|
|
if accountID <= 0 {
|
|
var err error
|
|
accountID, err = s.getStickySessionAccountID(ctx, groupID, sessionHash)
|
|
if err != nil || accountID <= 0 {
|
|
return nil
|
|
}
|
|
}
|
|
|
|
if _, excluded := excludedIDs[accountID]; excluded {
|
|
return nil
|
|
}
|
|
|
|
account, err := s.getSchedulableAccount(ctx, accountID)
|
|
if err != nil {
|
|
return nil
|
|
}
|
|
|
|
// 检查账号是否需要清理粘性会话
|
|
// Check if sticky session should be cleared
|
|
if shouldClearStickySession(account, requestedModel) {
|
|
_ = s.deleteStickySessionAccountID(ctx, groupID, sessionHash)
|
|
return nil
|
|
}
|
|
|
|
// 验证账号是否可用于当前请求
|
|
// Verify account is usable for current request
|
|
if !isOpenAICompatibleAccountEligibleForRequest(ctx, account, platform, requestedModel, false, requiredCapability) {
|
|
return nil
|
|
}
|
|
if !parentHealthyForShadow(account, s.parentAccountLookup(ctx)) {
|
|
_ = s.deleteStickySessionAccountID(ctx, groupID, sessionHash)
|
|
return nil
|
|
}
|
|
if s.isOpenAIAccountRequestRuntimeBlocked(account, requestedModel) {
|
|
_ = s.deleteStickySessionAccountID(ctx, groupID, sessionHash)
|
|
return nil
|
|
}
|
|
account = s.recheckSelectedOpenAIAccountFromDB(ctx, account, groupID, platform, requestedModel, requireCompact, requiredCapability)
|
|
if account == nil || !s.openAIAccountMatchesSchedulingGroup(account, groupID) {
|
|
_ = s.deleteStickySessionAccountID(ctx, groupID, sessionHash)
|
|
return nil
|
|
}
|
|
if groupID != nil && s.needsUpstreamChannelRestrictionCheck(ctx, groupID) &&
|
|
s.isUpstreamModelRestrictedByChannel(ctx, *groupID, account, requestedModel, requireCompact) {
|
|
_ = s.deleteStickySessionAccountID(ctx, groupID, sessionHash)
|
|
return nil
|
|
}
|
|
|
|
// 刷新会话 TTL 并返回账号
|
|
// Refresh session TTL and return account
|
|
_ = s.refreshStickySessionTTL(ctx, groupID, sessionHash, openaiStickySessionTTL)
|
|
return account
|
|
}
|
|
|
|
// selectBestAccount 从候选账号中选择最佳账号（优先级 + LRU）。
|
|
// 返回 nil 表示无可用账号。
|
|
//
|
|
// selectBestAccount selects the best account from candidates (priority + LRU).
|
|
// Returns nil if no available account. The second return reports whether at
|
|
// least one candidate was filtered out solely because it lacks compact support
|
|
// (only meaningful when the legacy /responses/compact requireCompact flag is
|
|
// true); the third contains deterministic
|
|
// exclusion diagnostics for the evaluated snapshot.
|
|
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) {
|
|
platform = NormalizeOpenAICompatiblePlatform(platform)
|
|
compactBlocked := false
|
|
filterStats := openAISelectionFilterStats{pool: len(accounts)}
|
|
needsUpstreamCheck := s.needsUpstreamChannelRestrictionCheck(ctx, groupID)
|
|
eligible := make([]*Account, 0, len(accounts))
|
|
compactTiers := make(map[int64]int, len(accounts))
|
|
|
|
for i := range accounts {
|
|
acc := &accounts[i]
|
|
|
|
// 跳过被排除的账号
|
|
// Skip excluded accounts
|
|
if _, excluded := excludedIDs[acc.ID]; excluded {
|
|
filterStats.exclude("excluded")
|
|
continue
|
|
}
|
|
|
|
fresh := s.resolveFreshSchedulableOpenAIAccountBeforeProfit(ctx, acc, platform, requestedModel, false, requiredCapability)
|
|
if fresh == nil {
|
|
filterStats.exclude("ineligible")
|
|
continue
|
|
}
|
|
fresh = s.recheckSelectedOpenAIAccountFromDBBeforeProfit(ctx, fresh, groupID, platform, requestedModel, false, requiredCapability)
|
|
if fresh == nil {
|
|
filterStats.exclude("ineligible")
|
|
continue
|
|
}
|
|
if needsUpstreamCheck && s.isUpstreamModelRestrictedByChannel(ctx, *groupID, fresh, requestedModel, requireCompact) {
|
|
filterStats.exclude("channel_restricted")
|
|
continue
|
|
}
|
|
if vetoed, reason := openAIProfitControlVetoReason(ctx, fresh); vetoed {
|
|
filterStats.exclude(reason)
|
|
continue
|
|
}
|
|
compactTier := 0
|
|
if requireCompact {
|
|
compactTier = openAICompactSupportTier(fresh)
|
|
if compactTier == 0 {
|
|
compactBlocked = true
|
|
filterStats.exclude("compact_unsupported")
|
|
continue
|
|
}
|
|
}
|
|
|
|
eligible = append(eligible, fresh)
|
|
compactTiers[fresh.ID] = compactTier
|
|
}
|
|
|
|
if len(eligible) == 0 {
|
|
return nil, compactBlocked, filterStats
|
|
}
|
|
rateOrder := openAILegacyUpstreamRateOrder{}
|
|
if preferLowUpstreamRate {
|
|
rateOrder = newOpenAILegacyUpstreamRateOrder(eligible, time.Now(), s.openAIOAuthSchedulingRateMultiplier(ctx))
|
|
}
|
|
sort.SliceStable(eligible, func(i, j int) bool {
|
|
a, b := eligible[i], eligible[j]
|
|
if requireCompact && compactTiers[a.ID] != compactTiers[b.ID] {
|
|
return compactTiers[a.ID] > compactTiers[b.ID]
|
|
}
|
|
if rateCmp := rateOrder.compare(a, b); rateCmp != 0 {
|
|
return rateCmp < 0
|
|
}
|
|
return s.isBetterAccount(a, b)
|
|
})
|
|
return eligible[0], compactBlocked, filterStats
|
|
}
|
|
|
|
// isBetterAccount 判断 candidate 是否比 current 更优。
|
|
// 规则：优先级更高（数值更小）优先；同优先级时，未使用过的优先，其次是最久未使用的。
|
|
//
|
|
// isBetterAccount checks if candidate is better than current.
|
|
// Rules: higher priority (lower value) wins; same priority: never used > least recently used.
|
|
func (s *OpenAIGatewayService) isBetterAccount(candidate, current *Account) bool {
|
|
// 优先级更高（数值更小）
|
|
// Higher priority (lower value)
|
|
if candidate.Priority < current.Priority {
|
|
return true
|
|
}
|
|
if candidate.Priority > current.Priority {
|
|
return false
|
|
}
|
|
|
|
// 同优先级，比较最后使用时间
|
|
// Same priority, compare last used time
|
|
switch {
|
|
case candidate.LastUsedAt == nil && current.LastUsedAt != nil:
|
|
// candidate 从未使用，优先
|
|
return true
|
|
case candidate.LastUsedAt != nil && current.LastUsedAt == nil:
|
|
// current 从未使用，保持
|
|
return false
|
|
case candidate.LastUsedAt == nil && current.LastUsedAt == nil:
|
|
// 都未使用，保持
|
|
return false
|
|
default:
|
|
// 都使用过，选择最久未使用的
|
|
return candidate.LastUsedAt.Before(*current.LastUsedAt)
|
|
}
|
|
}
|
|
|
|
// SelectAccountWithLoadAwareness selects an account with load-awareness and wait plan.
|
|
func (s *OpenAIGatewayService) SelectAccountWithLoadAwareness(ctx context.Context, groupID *int64, sessionHash string, requestedModel string, excludedIDs map[int64]struct{}) (*AccountSelectionResult, error) {
|
|
ctx = s.withOpenAIQuotaAutoPauseContext(ctx)
|
|
// 分组利润控制：legacy 公共入口同样装门，保证不经
|
|
// selectAccountWithScheduler 的调用方也无法绕过利润准入。
|
|
ctx = s.withOpenAIProfitControlGate(ctx, groupID)
|
|
return s.selectAccountWithLoadAwareness(ctx, groupID, PlatformOpenAI, sessionHash, requestedModel, excludedIDs, false, "", true)
|
|
}
|
|
|
|
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) {
|
|
platform = NormalizeOpenAICompatiblePlatform(platform)
|
|
if s.checkChannelPricingRestriction(ctx, groupID, requestedModel) {
|
|
slog.Warn("channel pricing restriction blocked request",
|
|
"group_id", derefGroupID(groupID),
|
|
"model", requestedModel)
|
|
return nil, fmt.Errorf("%w supporting model: %s (channel pricing restriction)", ErrNoAvailableAccounts, requestedModel)
|
|
}
|
|
|
|
cfg := s.schedulingConfig()
|
|
preferLowUpstreamRate := useUpstreamTokenCost && s.isOpenAILowUpstreamRatePriorityEnabled(ctx)
|
|
needsUpstreamCheck := s.needsUpstreamChannelRestrictionCheck(ctx, groupID)
|
|
var stickyAccountID int64
|
|
if sessionHash != "" && s.cache != nil {
|
|
if accountID, err := s.getStickySessionAccountID(ctx, groupID, sessionHash); err == nil {
|
|
stickyAccountID = accountID
|
|
}
|
|
}
|
|
if s.concurrencyService == nil || !cfg.LoadBatchEnabled {
|
|
account, err := s.selectAccountForModelWithExclusions(ctx, groupID, platform, sessionHash, requestedModel, excludedIDs, requireCompact, stickyAccountID, requiredCapability, preferLowUpstreamRate)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
result, err := s.tryAcquireAccountSlot(ctx, account.ID, account.Concurrency)
|
|
if err == nil && result != nil && result.Acquired {
|
|
return s.newAcquiredSelectionResult(ctx, account, result.ReleaseFunc)
|
|
}
|
|
if stickyAccountID > 0 && stickyAccountID == account.ID && s.concurrencyService != nil {
|
|
waitingCount, _ := s.concurrencyService.GetAccountWaitingCount(ctx, account.ID)
|
|
if waitingCount < cfg.StickySessionMaxWaiting {
|
|
return s.newSelectionResult(ctx, account, false, nil, &AccountWaitPlan{
|
|
AccountID: account.ID,
|
|
MaxConcurrency: account.Concurrency,
|
|
Timeout: cfg.StickySessionWaitTimeout,
|
|
MaxWaiting: cfg.StickySessionMaxWaiting,
|
|
})
|
|
}
|
|
}
|
|
return s.newSelectionResult(ctx, account, false, nil, &AccountWaitPlan{
|
|
AccountID: account.ID,
|
|
MaxConcurrency: account.Concurrency,
|
|
Timeout: cfg.FallbackWaitTimeout,
|
|
MaxWaiting: cfg.FallbackMaxWaiting,
|
|
})
|
|
}
|
|
|
|
accounts, err := s.listSchedulableAccounts(ctx, groupID, platform)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
if len(accounts) == 0 {
|
|
return nil, ErrNoAvailableAccounts
|
|
}
|
|
|
|
isExcluded := func(accountID int64) bool {
|
|
if excludedIDs == nil {
|
|
return false
|
|
}
|
|
_, excluded := excludedIDs[accountID]
|
|
return excluded
|
|
}
|
|
|
|
// ============ Layer 1: Sticky session ============
|
|
if sessionHash != "" {
|
|
accountID := stickyAccountID
|
|
if accountID > 0 && !isExcluded(accountID) {
|
|
account, err := s.getSchedulableAccount(ctx, accountID)
|
|
if err == nil {
|
|
clearSticky := shouldClearStickySession(account, requestedModel)
|
|
if clearSticky {
|
|
_ = s.deleteStickySessionAccountID(ctx, groupID, sessionHash)
|
|
}
|
|
if !clearSticky && isOpenAICompatibleAccountEligibleForRequest(ctx, account, platform, requestedModel, false, requiredCapability) {
|
|
account = s.recheckSelectedOpenAIAccountFromDB(ctx, account, groupID, platform, requestedModel, requireCompact, requiredCapability)
|
|
if account == nil {
|
|
_ = s.deleteStickySessionAccountID(ctx, groupID, sessionHash)
|
|
} else if !s.openAIAccountMatchesSchedulingGroup(account, groupID) {
|
|
_ = s.deleteStickySessionAccountID(ctx, groupID, sessionHash)
|
|
} else if s.isOpenAIAccountRequestRuntimeBlocked(account, requestedModel) {
|
|
_ = s.deleteStickySessionAccountID(ctx, groupID, sessionHash)
|
|
} else if needsUpstreamCheck && s.isUpstreamModelRestrictedByChannel(ctx, *groupID, account, requestedModel, requireCompact) {
|
|
_ = s.deleteStickySessionAccountID(ctx, groupID, sessionHash)
|
|
} else if !parentHealthyForShadow(account, s.parentAccountLookup(ctx)) {
|
|
_ = s.deleteStickySessionAccountID(ctx, groupID, sessionHash)
|
|
} else {
|
|
result, err := s.tryAcquireAccountSlot(ctx, accountID, account.Concurrency)
|
|
if err == nil && result != nil && result.Acquired {
|
|
selection, selectErr := s.newAcquiredSelectionResult(ctx, account, result.ReleaseFunc)
|
|
if selectErr != nil {
|
|
return nil, selectErr
|
|
}
|
|
_ = s.refreshStickySessionTTL(ctx, groupID, sessionHash, openaiStickySessionTTL)
|
|
return selection, nil
|
|
}
|
|
|
|
waitingCount, _ := s.concurrencyService.GetAccountWaitingCount(ctx, accountID)
|
|
if waitingCount < cfg.StickySessionMaxWaiting {
|
|
return s.newSelectionResult(ctx, account, false, nil, &AccountWaitPlan{
|
|
AccountID: accountID,
|
|
MaxConcurrency: account.Concurrency,
|
|
Timeout: cfg.StickySessionWaitTimeout,
|
|
MaxWaiting: cfg.StickySessionMaxWaiting,
|
|
})
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
// ============ Layer 2: Load-aware selection ============
|
|
// Per-pass parent-health cache to avoid repeated DB calls when multiple shadow
|
|
// accounts share the same parent.
|
|
parentCacheL2 := make(map[int64]*Account)
|
|
parentLookupL2 := func(id int64) *Account {
|
|
if a, ok := parentCacheL2[id]; ok {
|
|
return a
|
|
}
|
|
if s.accountRepo == nil {
|
|
return nil
|
|
}
|
|
a, _ := s.accountRepo.GetByID(ctx, id)
|
|
parentCacheL2[id] = a
|
|
return a
|
|
}
|
|
baseCandidateCount := 0
|
|
candidates := make([]*Account, 0, len(accounts))
|
|
for i := range accounts {
|
|
acc := &accounts[i]
|
|
if isExcluded(acc.ID) {
|
|
continue
|
|
}
|
|
// Scheduler snapshots can be temporarily stale (bucket rebuild is throttled);
|
|
// re-check schedulability here so recently rate-limited/overloaded accounts
|
|
// are not selected again before the bucket is rebuilt.
|
|
if !isOpenAICompatibleAccountEligibleForRequest(ctx, acc, platform, requestedModel, false, requiredCapability) {
|
|
continue
|
|
}
|
|
if !parentHealthyForShadow(acc, parentLookupL2) {
|
|
continue
|
|
}
|
|
if s.isOpenAIAccountRequestRuntimeBlocked(acc, requestedModel) {
|
|
continue
|
|
}
|
|
if needsUpstreamCheck && s.isUpstreamModelRestrictedByChannel(ctx, *groupID, acc, requestedModel, requireCompact) {
|
|
continue
|
|
}
|
|
baseCandidateCount++
|
|
candidates = append(candidates, acc)
|
|
}
|
|
|
|
if len(candidates) == 0 {
|
|
return nil, ErrNoAvailableAccounts
|
|
}
|
|
rateOrder := openAILegacyUpstreamRateOrder{}
|
|
if preferLowUpstreamRate {
|
|
rateOrder = newOpenAILegacyUpstreamRateOrder(candidates, time.Now(), s.openAIOAuthSchedulingRateMultiplier(ctx))
|
|
}
|
|
|
|
accountLoads := make([]AccountWithConcurrency, 0, len(candidates))
|
|
for _, acc := range candidates {
|
|
accountLoads = append(accountLoads, AccountWithConcurrency{
|
|
ID: acc.ID,
|
|
MaxConcurrency: acc.EffectiveLoadFactor(),
|
|
})
|
|
}
|
|
|
|
tryAcquireFromLoadMap := func(loadMap map[int64]*AccountLoadInfo) (*AccountSelectionResult, bool, error) {
|
|
var available []accountWithLoad
|
|
for _, acc := range candidates {
|
|
loadInfo := loadMap[acc.ID]
|
|
if loadInfo == nil {
|
|
loadInfo = &AccountLoadInfo{AccountID: acc.ID}
|
|
}
|
|
if loadInfo.LoadRate < 100 {
|
|
available = append(available, accountWithLoad{
|
|
account: acc,
|
|
loadInfo: loadInfo,
|
|
})
|
|
}
|
|
}
|
|
|
|
if len(available) == 0 {
|
|
return nil, false, nil
|
|
}
|
|
|
|
sort.SliceStable(available, func(i, j int) bool {
|
|
a, b := available[i], available[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
|
|
}
|
|
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)
|
|
}
|
|
})
|
|
shuffleWithinSortGroups(available)
|
|
if rateOrder.enabled {
|
|
sort.SliceStable(available, func(i, j int) bool {
|
|
return rateOrder.compare(available[i].account, available[j].account) < 0
|
|
})
|
|
}
|
|
|
|
selectionOrder := make([]accountWithLoad, 0, len(available))
|
|
if requireCompact {
|
|
appendTier := func(out []accountWithLoad, tier int) []accountWithLoad {
|
|
for _, item := range available {
|
|
if openAICompactSupportTier(item.account) == tier {
|
|
out = append(out, item)
|
|
}
|
|
}
|
|
return out
|
|
}
|
|
selectionOrder = appendTier(selectionOrder, 2)
|
|
selectionOrder = appendTier(selectionOrder, 1)
|
|
// tier 0 候选作为兜底追加：DB recheck 时若发现 cache tier 0 实际
|
|
// 已升级为 1/2（探测刚跑完，cache 尚未刷新），仍可正常命中。
|
|
selectionOrder = appendTier(selectionOrder, 0)
|
|
} else {
|
|
selectionOrder = append(selectionOrder, available...)
|
|
}
|
|
|
|
for _, item := range selectionOrder {
|
|
fresh := s.resolveFreshSchedulableOpenAIAccount(ctx, item.account, platform, requestedModel, false, requiredCapability)
|
|
if fresh == nil {
|
|
continue
|
|
}
|
|
fresh = s.recheckSelectedOpenAIAccountFromDB(ctx, fresh, groupID, platform, requestedModel, requireCompact, requiredCapability)
|
|
if fresh == nil {
|
|
continue
|
|
}
|
|
if needsUpstreamCheck && s.isUpstreamModelRestrictedByChannel(ctx, *groupID, fresh, requestedModel, requireCompact) {
|
|
continue
|
|
}
|
|
result, err := s.tryAcquireAccountSlot(ctx, fresh.ID, fresh.Concurrency)
|
|
if err == nil && result != nil && result.Acquired {
|
|
selection, selectErr := s.newAcquiredSelectionResult(ctx, fresh, result.ReleaseFunc)
|
|
if selectErr != nil {
|
|
return nil, true, selectErr
|
|
}
|
|
if sessionHash != "" && !gatewayProfitControlGateActive(ctx) {
|
|
_ = s.setStickySessionAccountID(ctx, groupID, sessionHash, fresh.ID, openaiStickySessionTTL)
|
|
}
|
|
return selection, true, nil
|
|
}
|
|
}
|
|
return nil, true, nil
|
|
}
|
|
|
|
loadMap, err := s.concurrencyService.GetAccountsLoadBatch(ctx, accountLoads)
|
|
if err != nil {
|
|
ordered := append([]*Account(nil), candidates...)
|
|
sortAccountsByPriorityAndLastUsed(ordered, false)
|
|
if rateOrder.enabled {
|
|
sort.SliceStable(ordered, func(i, j int) bool {
|
|
return rateOrder.compare(ordered[i], ordered[j]) < 0
|
|
})
|
|
}
|
|
if requireCompact {
|
|
ordered = prioritizeOpenAICompactAccounts(ordered)
|
|
}
|
|
for _, acc := range ordered {
|
|
fresh := s.resolveFreshSchedulableOpenAIAccount(ctx, acc, platform, requestedModel, false, requiredCapability)
|
|
if fresh == nil {
|
|
continue
|
|
}
|
|
fresh = s.recheckSelectedOpenAIAccountFromDB(ctx, fresh, groupID, platform, requestedModel, requireCompact, requiredCapability)
|
|
if fresh == nil {
|
|
continue
|
|
}
|
|
if needsUpstreamCheck && s.isUpstreamModelRestrictedByChannel(ctx, *groupID, fresh, requestedModel, requireCompact) {
|
|
continue
|
|
}
|
|
result, err := s.tryAcquireAccountSlot(ctx, fresh.ID, fresh.Concurrency)
|
|
if err == nil && result != nil && result.Acquired {
|
|
selection, selectErr := s.newAcquiredSelectionResult(ctx, fresh, result.ReleaseFunc)
|
|
if selectErr != nil {
|
|
return nil, selectErr
|
|
}
|
|
if sessionHash != "" && !gatewayProfitControlGateActive(ctx) {
|
|
_ = s.setStickySessionAccountID(ctx, groupID, sessionHash, fresh.ID, openaiStickySessionTTL)
|
|
}
|
|
return selection, nil
|
|
}
|
|
}
|
|
} else {
|
|
if selection, attempted, selectErr := tryAcquireFromLoadMap(loadMap); selectErr != nil {
|
|
return nil, selectErr
|
|
} else if selection != nil {
|
|
return selection, nil
|
|
} else if attempted {
|
|
if freshLoadMap, loadErr := s.concurrencyService.GetAccountsLoadBatchFresh(ctx, accountLoads); loadErr == nil {
|
|
if selection, _, selectErr := tryAcquireFromLoadMap(freshLoadMap); selectErr != nil {
|
|
return nil, selectErr
|
|
} else if selection != nil {
|
|
return selection, nil
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
// ============ Layer 3: Fallback wait ============
|
|
sortAccountsByPriorityAndLastUsed(candidates, false)
|
|
if rateOrder.enabled {
|
|
sort.SliceStable(candidates, func(i, j int) bool {
|
|
return rateOrder.compare(candidates[i], candidates[j]) < 0
|
|
})
|
|
}
|
|
if requireCompact {
|
|
candidates = prioritizeOpenAICompactAccounts(candidates)
|
|
}
|
|
for _, acc := range candidates {
|
|
fresh := s.resolveFreshSchedulableOpenAIAccount(ctx, acc, platform, requestedModel, false, requiredCapability)
|
|
if fresh == nil {
|
|
continue
|
|
}
|
|
fresh = s.recheckSelectedOpenAIAccountFromDB(ctx, fresh, groupID, platform, requestedModel, requireCompact, requiredCapability)
|
|
if fresh == nil {
|
|
continue
|
|
}
|
|
if needsUpstreamCheck && s.isUpstreamModelRestrictedByChannel(ctx, *groupID, fresh, requestedModel, requireCompact) {
|
|
continue
|
|
}
|
|
return s.newSelectionResult(ctx, fresh, false, nil, &AccountWaitPlan{
|
|
AccountID: fresh.ID,
|
|
MaxConcurrency: fresh.Concurrency,
|
|
Timeout: cfg.FallbackWaitTimeout,
|
|
MaxWaiting: cfg.FallbackMaxWaiting,
|
|
})
|
|
}
|
|
|
|
if requireCompact && baseCandidateCount > 0 {
|
|
return nil, ErrNoAvailableCompactAccounts
|
|
}
|
|
return nil, ErrNoAvailableAccounts
|
|
}
|
|
|
|
func (s *OpenAIGatewayService) listSchedulableAccounts(ctx context.Context, groupID *int64, platform string) ([]Account, error) {
|
|
platform = NormalizeOpenAICompatiblePlatform(platform)
|
|
if s.schedulerSnapshot != nil {
|
|
accounts, _, err := s.schedulerSnapshot.ListSchedulableAccounts(ctx, groupID, platform, false)
|
|
if err != nil {
|
|
return accounts, err
|
|
}
|
|
accounts = s.filterOpenAIAccountsBySchedulingThreshold(ctx, accounts)
|
|
if platform == PlatformGrok {
|
|
accounts = s.filterGrokFreeQuotaAccountsForOpenAI(ctx, accounts)
|
|
}
|
|
return accounts, nil
|
|
}
|
|
var accounts []Account
|
|
var err error
|
|
if s.cfg != nil && s.cfg.RunMode == config.RunModeSimple {
|
|
accounts, err = s.accountRepo.ListSchedulableByPlatform(ctx, platform)
|
|
} else if groupID != nil {
|
|
accounts, err = s.accountRepo.ListSchedulableByGroupIDAndPlatform(ctx, *groupID, platform)
|
|
} else {
|
|
accounts, err = s.accountRepo.ListSchedulableUngroupedByPlatform(ctx, platform)
|
|
}
|
|
if err != nil {
|
|
return nil, fmt.Errorf("query accounts failed: %w", err)
|
|
}
|
|
accounts = s.filterOpenAIAccountsBySchedulingThreshold(ctx, accounts)
|
|
if platform == PlatformGrok {
|
|
accounts = s.filterGrokFreeQuotaAccountsForOpenAI(ctx, accounts)
|
|
}
|
|
return accounts, nil
|
|
}
|
|
|
|
func (s *OpenAIGatewayService) tryAcquireAccountSlot(ctx context.Context, accountID int64, maxConcurrency int) (*AcquireResult, error) {
|
|
if s.concurrencyService == nil {
|
|
return &AcquireResult{Acquired: true, ReleaseFunc: func() {}}, nil
|
|
}
|
|
return s.concurrencyService.AcquireAccountSlot(ctx, accountID, maxConcurrency)
|
|
}
|
|
|
|
func (s *OpenAIGatewayService) resolveFreshSchedulableOpenAIAccount(ctx context.Context, account *Account, platform string, requestedModel string, requireCompact bool, requiredCapability OpenAIEndpointCapability) *Account {
|
|
fresh := s.resolveFreshSchedulableOpenAIAccountBeforeProfit(ctx, account, platform, requestedModel, requireCompact, requiredCapability)
|
|
if fresh == nil {
|
|
return nil
|
|
}
|
|
if vetoed, _ := openAIProfitControlVetoReason(ctx, fresh); vetoed {
|
|
return nil
|
|
}
|
|
return fresh
|
|
}
|
|
|
|
func (s *OpenAIGatewayService) resolveFreshSchedulableOpenAIAccountBeforeProfit(ctx context.Context, account *Account, platform string, requestedModel string, requireCompact bool, requiredCapability OpenAIEndpointCapability) *Account {
|
|
if account == nil {
|
|
return nil
|
|
}
|
|
platform = NormalizeOpenAICompatiblePlatform(platform)
|
|
|
|
fresh := account
|
|
if s.schedulerSnapshot != nil {
|
|
current, err := s.getSchedulableAccount(ctx, account.ID)
|
|
if err != nil || current == nil {
|
|
return nil
|
|
}
|
|
fresh = current
|
|
}
|
|
|
|
if !isOpenAICompatibleAccountEligibleForRequestBeforeProfit(ctx, fresh, platform, requestedModel, requireCompact, requiredCapability) {
|
|
return nil
|
|
}
|
|
if !parentHealthyForShadow(fresh, s.parentAccountLookup(ctx)) {
|
|
return nil
|
|
}
|
|
if s.isOpenAIAccountRequestRuntimeBlocked(fresh, requestedModel) {
|
|
return nil
|
|
}
|
|
if s.isOpenAIAccountBlockedBySchedulingThreshold(ctx, fresh) {
|
|
return nil
|
|
}
|
|
if s.isOpenAIProxyStreamQuarantined(ctx, fresh) {
|
|
return nil
|
|
}
|
|
return fresh
|
|
}
|
|
|
|
// parentAccountLookup 返回供 parentHealthyForShadow 使用的母账号解析闭包:经 accountRepo
|
|
// 按 ID 取当前 Account(repo 为空时 fail-closed 返回 nil)。统一调度/粘连各路径的母账号解析,
|
|
// 取代各调用点重复内联的同一闭包(历史上 recheck 等路径还漏写过 accountRepo==nil 守卫)。
|
|
// L2 候选循环改用带 per-pass 缓存的 parentLookupL2,不走此方法。
|
|
func (s *OpenAIGatewayService) parentAccountLookup(ctx context.Context) func(int64) *Account {
|
|
return func(id int64) *Account {
|
|
if s.accountRepo == nil {
|
|
return nil
|
|
}
|
|
a, _ := s.accountRepo.GetByID(ctx, id)
|
|
return a
|
|
}
|
|
}
|
|
|
|
func (s *OpenAIGatewayService) recheckSelectedOpenAIAccountFromDB(ctx context.Context, account *Account, groupID *int64, platform string, requestedModel string, requireCompact bool, requiredCapability OpenAIEndpointCapability) *Account {
|
|
latest := s.recheckSelectedOpenAIAccountFromDBBeforeProfit(ctx, account, groupID, platform, requestedModel, requireCompact, requiredCapability)
|
|
if latest == nil {
|
|
return nil
|
|
}
|
|
if vetoed, _ := openAIProfitControlVetoReason(ctx, latest); vetoed {
|
|
return nil
|
|
}
|
|
return latest
|
|
}
|
|
|
|
func (s *OpenAIGatewayService) recheckSelectedOpenAIAccountFromDBBeforeProfit(ctx context.Context, account *Account, groupID *int64, platform string, requestedModel string, requireCompact bool, requiredCapability OpenAIEndpointCapability) *Account {
|
|
if account == nil {
|
|
return nil
|
|
}
|
|
platform = NormalizeOpenAICompatiblePlatform(platform)
|
|
if s.schedulerSnapshot == nil || s.accountRepo == nil {
|
|
if !isOpenAICompatibleAccountEligibleForRequestBeforeProfit(ctx, account, platform, requestedModel, requireCompact, requiredCapability) {
|
|
return nil
|
|
}
|
|
if s.isOpenAIAccountBlockedBySchedulingThreshold(ctx, account) {
|
|
return nil
|
|
}
|
|
if !parentHealthyForShadow(account, s.parentAccountLookup(ctx)) {
|
|
return nil
|
|
}
|
|
if s.isOpenAIProxyStreamQuarantined(ctx, account) {
|
|
return nil
|
|
}
|
|
return account
|
|
}
|
|
|
|
latest, err := s.accountRepo.GetByID(ctx, account.ID)
|
|
if err != nil || latest == nil {
|
|
return nil
|
|
}
|
|
if !s.openAIAccountMatchesSchedulingGroup(latest, groupID) {
|
|
return nil
|
|
}
|
|
if !isOpenAICompatibleAccountEligibleForRequestBeforeProfit(ctx, latest, platform, requestedModel, requireCompact, requiredCapability) {
|
|
return nil
|
|
}
|
|
if !parentHealthyForShadow(latest, s.parentAccountLookup(ctx)) {
|
|
return nil
|
|
}
|
|
if s.isOpenAIAccountRequestRuntimeBlocked(latest, requestedModel) {
|
|
return nil
|
|
}
|
|
if s.isOpenAIAccountBlockedBySchedulingThreshold(ctx, latest) {
|
|
return nil
|
|
}
|
|
if s.isOpenAIProxyStreamQuarantined(ctx, latest) {
|
|
return nil
|
|
}
|
|
return latest
|
|
}
|
|
|
|
func (s *OpenAIGatewayService) openAIAccountMatchesSchedulingGroup(account *Account, groupID *int64) bool {
|
|
if s != nil && s.cfg != nil && s.cfg.RunMode == config.RunModeSimple {
|
|
return account != nil
|
|
}
|
|
return openAIStickyAccountMatchesGroup(account, groupID)
|
|
}
|
|
|
|
func (s *OpenAIGatewayService) getSchedulableAccount(ctx context.Context, accountID int64) (*Account, error) {
|
|
var (
|
|
account *Account
|
|
err error
|
|
)
|
|
if s.schedulerSnapshot != nil {
|
|
account, err = s.schedulerSnapshot.GetAccount(ctx, accountID)
|
|
} else {
|
|
account, err = s.accountRepo.GetByID(ctx, accountID)
|
|
}
|
|
if err != nil || account == nil {
|
|
return account, err
|
|
}
|
|
if s.isOpenAIAccountBlockedBySchedulingThreshold(ctx, account) {
|
|
return nil, nil
|
|
}
|
|
// Legacy sticky (advanced scheduler off) must still free-gate Grok OAuth.
|
|
if account.IsGrok() {
|
|
if gated := s.filterGrokFreeQuotaAccountsForOpenAI(ctx, []Account{*account}); len(gated) == 0 {
|
|
return nil, nil
|
|
}
|
|
}
|
|
return account, nil
|
|
}
|
|
|
|
// filterGrokFreeQuotaAccountsForOpenAI applies the same local free soft-gate as
|
|
// GatewayService / advanced scheduler, for OpenAI-compatible legacy selection.
|
|
func (s *OpenAIGatewayService) filterGrokFreeQuotaAccountsForOpenAI(ctx context.Context, accounts []Account) []Account {
|
|
if s == nil {
|
|
return accounts
|
|
}
|
|
return filterGrokFreeQuotaAccountsCore(ctx, s.cfg, s.usageLogRepo, &openaiGrokFreeQuotaGateCache, accounts)
|
|
}
|
|
|
|
func (s *OpenAIGatewayService) filterOpenAIAccountsBySchedulingThreshold(ctx context.Context, accounts []Account) []Account {
|
|
if len(accounts) == 0 {
|
|
return accounts
|
|
}
|
|
|
|
filtered := make([]Account, 0, len(accounts))
|
|
for i := range accounts {
|
|
if s.isOpenAIAccountBlockedBySchedulingThreshold(ctx, &accounts[i]) {
|
|
continue
|
|
}
|
|
filtered = append(filtered, accounts[i])
|
|
}
|
|
return filtered
|
|
}
|
|
|
|
func (s *OpenAIGatewayService) isOpenAIAccountBlockedBySchedulingThreshold(ctx context.Context, account *Account) bool {
|
|
if s == nil || s.rateLimitService == nil || account == nil {
|
|
return false
|
|
}
|
|
return s.rateLimitService.ApplyAccountSchedulingThreshold(ctx, account)
|
|
}
|
|
|
|
func (s *OpenAIGatewayService) hydrateSelectedAccount(ctx context.Context, account *Account) (*Account, error) {
|
|
if account == nil || s.schedulerSnapshot == nil {
|
|
return account, nil
|
|
}
|
|
hydrated, err := s.schedulerSnapshot.GetAccount(ctx, account.ID)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
if hydrated == nil {
|
|
return nil, fmt.Errorf("selected openai account %d not found during hydration", account.ID)
|
|
}
|
|
return hydrated, nil
|
|
}
|
|
|
|
func (s *OpenAIGatewayService) newSelectionResult(ctx context.Context, account *Account, acquired bool, release func(), waitPlan *AccountWaitPlan) (*AccountSelectionResult, error) {
|
|
hydrated, err := s.hydrateSelectedAccount(ctx, account)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
return attachSelectionProfitGate(ctx, &AccountSelectionResult{
|
|
Account: hydrated,
|
|
Acquired: acquired,
|
|
ReleaseFunc: release,
|
|
WaitPlan: waitPlan,
|
|
}), nil
|
|
}
|
|
|
|
func (s *OpenAIGatewayService) newAcquiredSelectionResult(ctx context.Context, account *Account, release func()) (*AccountSelectionResult, error) {
|
|
selection, err := s.newSelectionResult(ctx, account, true, release, nil)
|
|
if err != nil && release != nil {
|
|
release()
|
|
}
|
|
return selection, err
|
|
}
|
|
|
|
func (s *OpenAIGatewayService) schedulingConfig() config.GatewaySchedulingConfig {
|
|
if s.cfg != nil {
|
|
return s.cfg.Gateway.Scheduling
|
|
}
|
|
return config.GatewaySchedulingConfig{
|
|
StickySessionMaxWaiting: 3,
|
|
StickySessionWaitTimeout: 45 * time.Second,
|
|
FallbackWaitTimeout: 30 * time.Second,
|
|
FallbackMaxWaiting: 100,
|
|
LoadBatchEnabled: true,
|
|
SlotCleanupInterval: 30 * time.Second,
|
|
}
|
|
}
|