Merge pull request #6068 from okbexx/fix/codex-guardian-parent-affinity

fix(openai): keep auto-review on parent account
This commit is contained in:
Wesley Liddick authored and GitHub committed 2026-08-22 13:41:42 +08:00
commit d45135d87d
7 files changed
+733 -16

No files matched your search

@@ -514,6 +514,9 @@ func (h *OpenAIGatewayHandler) Responses(c *gin.Context) {
if h.rejectIfCyberSessionBlocked(c, apiKey, sessionHashBody, reqModel, cyberBlockFormatResponses) { if h.rejectIfCyberSessionBlocked(c, apiKey, sessionHashBody, reqModel, cyberBlockFormatResponses) {
return return
} }
c.Request = c.Request.WithContext(service.WithOpenAIGuardianParentAffinity(
c.Request.Context(), c, sessionHashBody, reqModel,
))
requireCompact := legacyCompact requireCompact := legacyCompact
maxAccountSwitches := h.maxAccountSwitches maxAccountSwitches := h.maxAccountSwitches
@@ -2004,6 +2007,7 @@ func (h *OpenAIGatewayHandler) ResponsesWebSocket(c *gin.Context) {
firstMessage, firstMessage,
openAIWSIngressFallbackSessionSeed(subject.UserID, apiKey.ID, apiKey.GroupID), openAIWSIngressFallbackSessionSeed(subject.UserID, apiKey.ID, apiKey.GroupID),
) )
ctx = service.WithOpenAIGuardianParentAffinity(ctx, c, firstMessage, reqModel)
maxAccountSwitches := h.maxAccountSwitches maxAccountSwitches := h.maxAccountSwitches
switchCount := 0 switchCount := 0
profitVetoCount := 0 profitVetoCount := 0
@@ -21,6 +21,7 @@ import (
const ( const (
openAIAccountScheduleLayerPreviousResponse = "previous_response_id" openAIAccountScheduleLayerPreviousResponse = "previous_response_id"
openAIAccountScheduleLayerGuardianParent = "guardian_parent"
openAIAccountScheduleLayerSessionSticky = "session_hash" openAIAccountScheduleLayerSessionSticky = "session_hash"
openAIAccountScheduleLayerLoadBalance = "load_balance" openAIAccountScheduleLayerLoadBalance = "load_balance"
openAIAdvancedSchedulerSettingKey = "openai_advanced_scheduler_enabled" openAIAdvancedSchedulerSettingKey = "openai_advanced_scheduler_enabled"
@@ -71,10 +72,12 @@ type OpenAIAccountScheduleRequest struct {
Platform string Platform string
SessionHash string SessionHash string
StickyAccountID int64 StickyAccountID int64
GuardianParentAccountID int64
StickyPreviousAccountID int64 StickyPreviousAccountID int64
StickyWeighted bool StickyWeighted bool
SubscriptionPriority bool SubscriptionPriority bool
PreserveStickyBinding bool PreserveStickyBinding bool
RequirePrivacySet bool
PreviousResponseID string PreviousResponseID string
PreviousResponseCanMove bool PreviousResponseCanMove bool
UseUpstreamTokenCost bool UseUpstreamTokenCost bool
@@ -373,6 +376,9 @@ func (s *defaultOpenAIAccountScheduler) Select(
ctx context.Context, ctx context.Context,
req OpenAIAccountScheduleRequest, req OpenAIAccountScheduleRequest,
) (*AccountSelectionResult, OpenAIAccountScheduleDecision, error) { ) (*AccountSelectionResult, OpenAIAccountScheduleDecision, error) {
if s != nil && s.service != nil && s.service.openAIGroupRequiresPrivacySet(ctx, req.GroupID) {
req.RequirePrivacySet = true
}
decision := OpenAIAccountScheduleDecision{} decision := OpenAIAccountScheduleDecision{}
start := time.Now() start := time.Now()
defer func() { defer func() {
@@ -396,7 +402,14 @@ func (s *defaultOpenAIAccountScheduler) Select(
return nil, decision, err return nil, decision, err
} }
if selection != nil && selection.Account != nil { if selection != nil && selection.Account != nil {
if !s.isAccountTransportCompatible(selection.Account, req.RequiredTransport) { compatible, _ := s.isAccountRequestCompatibleReason(ctx, selection.Account, req)
hasGroupMetadata := len(selection.Account.GroupIDs) > 0 || len(selection.Account.AccountGroups) > 0
groupCompatible := !hasGroupMetadata || openAIStickyAccountMatchesGroup(selection.Account, req.GroupID)
if hasGroupMetadata && s.service != nil {
groupCompatible = s.service.openAIAccountMatchesSchedulingGroup(selection.Account, req.GroupID)
}
if !groupCompatible ||
!compatible || !s.isAccountTransportCompatible(selection.Account, req.RequiredTransport) {
if selection.ReleaseFunc != nil { if selection.ReleaseFunc != nil {
selection.ReleaseFunc() selection.ReleaseFunc()
} }
@@ -415,6 +428,23 @@ func (s *defaultOpenAIAccountScheduler) Select(
} }
} }
if req.GuardianParentAccountID > 0 {
parentReq := req
parentReq.StickyAccountID = req.GuardianParentAccountID
parentReq.PreserveStickyBinding = true
selection, _, err := s.selectBySessionHash(ctx, parentReq)
if err != nil {
return nil, decision, err
}
if selection != nil && selection.Account != nil {
decision.Layer = openAIAccountScheduleLayerGuardianParent
decision.StickySessionHit = true
decision.SelectedAccountID = selection.Account.ID
decision.SelectedAccountType = selection.Account.Type
return selection, decision, nil
}
}
if !req.StickyWeighted { if !req.StickyWeighted {
selection, escapedSticky, err := s.selectBySessionHash(ctx, req) selection, escapedSticky, err := s.selectBySessionHash(ctx, req)
if err != nil { if err != nil {
@@ -465,6 +495,11 @@ func (s *defaultOpenAIAccountScheduler) selectBySessionHash(
} }
accountID := req.StickyAccountID accountID := req.StickyAccountID
clearBinding := func() {
if !req.PreserveStickyBinding {
_ = s.service.deleteStickySessionAccountID(ctx, req.GroupID, sessionHash)
}
}
if accountID <= 0 { if accountID <= 0 {
var err error var err error
accountID, err = s.service.getStickySessionAccountID(ctx, req.GroupID, sessionHash) accountID, err = s.service.getStickySessionAccountID(ctx, req.GroupID, sessionHash)
@@ -483,40 +518,40 @@ func (s *defaultOpenAIAccountScheduler) selectBySessionHash(
account, err := s.service.getSchedulableAccount(ctx, accountID) account, err := s.service.getSchedulableAccount(ctx, accountID)
if err != nil || account == nil { if err != nil || account == nil {
_ = s.service.deleteStickySessionAccountID(ctx, req.GroupID, sessionHash) clearBinding()
return nil, false, nil return nil, false, nil
} }
if shouldClearStickySession(account, req.RequestedModel) || account.Platform != NormalizeOpenAICompatiblePlatform(req.Platform) || !account.IsOpenAICompatible() || !account.IsSchedulable() { if shouldClearStickySession(account, req.RequestedModel) || account.Platform != NormalizeOpenAICompatiblePlatform(req.Platform) || !account.IsOpenAICompatible() || !account.IsSchedulable() {
_ = s.service.deleteStickySessionAccountID(ctx, req.GroupID, sessionHash) clearBinding()
return nil, false, nil return nil, false, nil
} }
if !s.isAccountRequestCompatible(ctx, account, req) { if !s.isAccountRequestCompatible(ctx, account, req) {
return nil, false, nil return nil, false, nil
} }
if !s.isAccountTransportCompatible(account, req.RequiredTransport) { if !s.isAccountTransportCompatible(account, req.RequiredTransport) {
_ = s.service.deleteStickySessionAccountID(ctx, req.GroupID, sessionHash) clearBinding()
return nil, false, nil return nil, false, nil
} }
account = s.service.recheckSelectedOpenAIAccountFromDB(ctx, account, req.GroupID, req.Platform, req.RequestedModel, req.RequireCompact, req.RequiredCapability) account = s.service.recheckSelectedOpenAIAccountFromDB(ctx, account, req.GroupID, req.Platform, req.RequestedModel, req.RequireCompact, req.RequiredCapability)
if account == nil || !s.service.openAIAccountMatchesSchedulingGroup(account, req.GroupID) || !s.isAccountTransportCompatible(account, req.RequiredTransport) { if account == nil || !s.service.openAIAccountMatchesSchedulingGroup(account, req.GroupID) || !s.isAccountRequestCompatible(ctx, account, req) || !s.isAccountTransportCompatible(account, req.RequiredTransport) {
_ = s.service.deleteStickySessionAccountID(ctx, req.GroupID, sessionHash) clearBinding()
return nil, false, nil return nil, false, nil
} }
// Free-tier soft gate: sticky session must not pin an over-quota free OAuth account. // Free-tier soft gate: sticky session must not pin an over-quota free OAuth account.
// Admin QueryQuota / import probes do not use this path. // Admin QueryQuota / import probes do not use this path.
if account != nil && len(s.filterGrokFreeQuotaAccounts(ctx, []Account{*account})) == 0 { if account != nil && len(s.filterGrokFreeQuotaAccounts(ctx, []Account{*account})) == 0 {
_ = s.service.deleteStickySessionAccountID(ctx, req.GroupID, sessionHash) clearBinding()
return nil, false, nil return nil, false, nil
} }
// Team+model cool: sticky must not pin a sibling under the same team 429 window. // Team+model cool: sticky must not pin a sibling under the same team 429 window.
now := time.Now() now := time.Now()
upstreamModel := canonicalOpenAIAccountSchedulingModel(account, req.RequestedModel) upstreamModel := canonicalOpenAIAccountSchedulingModel(account, req.RequestedModel)
if account != nil && isGrokTeamModelRateLimited(account, upstreamModel, now) { if account != nil && isGrokTeamModelRateLimited(account, upstreamModel, now) {
_ = s.service.deleteStickySessionAccountID(ctx, req.GroupID, sessionHash) clearBinding()
return nil, false, nil return nil, false, nil
} }
if account != nil && isGrokModelQuotaBlocked(account.ID, upstreamModel, now) { if account != nil && isGrokModelQuotaBlocked(account.ID, upstreamModel, now) {
_ = s.service.deleteStickySessionAccountID(ctx, req.GroupID, sessionHash) clearBinding()
return nil, false, nil return nil, false, nil
} }
escapeCfg := s.service.openAIStickyEscapeConfig() escapeCfg := s.service.openAIStickyEscapeConfig()
@@ -531,7 +566,9 @@ func (s *defaultOpenAIAccountScheduler) selectBySessionHash(
} }
result, acquireErr := s.service.tryAcquireAccountSlot(ctx, accountID, account.Concurrency) result, acquireErr := s.service.tryAcquireAccountSlot(ctx, accountID, account.Concurrency)
if acquireErr == nil && result != nil && result.Acquired { if acquireErr == nil && result != nil && result.Acquired {
if !req.PreserveStickyBinding {
_ = s.service.refreshStickySessionTTL(ctx, req.GroupID, sessionHash, s.service.openAIWSSessionStickyTTL()) _ = s.service.refreshStickySessionTTL(ctx, req.GroupID, sessionHash, s.service.openAIWSSessionStickyTTL())
}
return attachSelectionProfitGate(ctx, &AccountSelectionResult{ return attachSelectionProfitGate(ctx, &AccountSelectionResult{
Account: account, Account: account,
Acquired: true, Acquired: true,
@@ -1414,11 +1451,10 @@ func (s *defaultOpenAIAccountScheduler) selectByLoadBalance(
filterStats.exclude("runtime_blocked") filterStats.exclude("runtime_blocked")
continue continue
} }
// require_privacy_set: 跳过 privacy 未设置的账号并标记异常 // require_privacy_set is a group-scoped eligibility gate. Do not mutate the
// shared account: another group may intentionally allow accounts whose
// upstream privacy setting has not been confirmed.
if schedGroup != nil && schedGroup.RequirePrivacySet && !account.IsPrivacySet() { if schedGroup != nil && schedGroup.RequirePrivacySet && !account.IsPrivacySet() {
s.service.BlockAccountScheduling(account, time.Time{}, "privacy_not_set")
_ = s.service.accountRepo.SetError(ctx, account.ID,
fmt.Sprintf("Privacy not set, required by group [%s]", schedGroup.Name))
filterStats.exclude("privacy_not_set") filterStats.exclude("privacy_not_set")
continue continue
} }
@@ -1730,6 +1766,9 @@ func (s *defaultOpenAIAccountScheduler) isAccountRequestCompatibleReason(ctx con
if account == nil { if account == nil {
return false, "account_nil" return false, "account_nil"
} }
if req.RequirePrivacySet && !account.IsPrivacySet() {
return false, "privacy_not_set"
}
if s != nil && s.service != nil && s.service.isOpenAIAccountRequestRuntimeBlocked(account, req.RequestedModel) { if s != nil && s.service != nil && s.service.isOpenAIAccountRequestRuntimeBlocked(account, req.RequestedModel) {
return false, "runtime_blocked" return false, "runtime_blocked"
} }
@@ -2136,6 +2175,38 @@ func (s *OpenAIGatewayService) selectAccountWithScheduler(
return s.selectAccountWithSchedulerOnce(withOpenAIProxyStreamQuarantineBypass(ctx), groupID, previousResponseID, sessionHash, requestedModel, excludedIDs, requiredTransport, requiredCapability, requiredImageCapability, requireCompact, platform, previousResponseCanMove, useUpstreamTokenCost) return s.selectAccountWithSchedulerOnce(withOpenAIProxyStreamQuarantineBypass(ctx), groupID, previousResponseID, sessionHash, requestedModel, excludedIDs, requiredTransport, requiredCapability, requiredImageCapability, requireCompact, platform, previousResponseCanMove, useUpstreamTokenCost)
} }
type openAIGroupPrivacyRequirementContextKey struct{}
type openAIGroupPrivacyRequirement struct {
groupID int64
required bool
}
func (s *OpenAIGatewayService) withOpenAIGroupPrivacyRequirement(ctx context.Context, groupID *int64) context.Context {
return context.WithValue(ctx, openAIGroupPrivacyRequirementContextKey{}, openAIGroupPrivacyRequirement{
groupID: derefGroupID(groupID),
required: s.loadOpenAIGroupRequiresPrivacySet(ctx, groupID),
})
}
func (s *OpenAIGatewayService) openAIGroupRequiresPrivacySet(ctx context.Context, groupID *int64) bool {
if cached, ok := ctx.Value(openAIGroupPrivacyRequirementContextKey{}).(openAIGroupPrivacyRequirement); ok && cached.groupID == derefGroupID(groupID) {
return cached.required
}
return s.loadOpenAIGroupRequiresPrivacySet(ctx, groupID)
}
func (s *OpenAIGatewayService) loadOpenAIGroupRequiresPrivacySet(ctx context.Context, groupID *int64) bool {
if s == nil || groupID == nil || s.schedulerSnapshot == nil {
return false
}
group, err := s.schedulerSnapshot.GetGroupByID(ctx, *groupID)
if err != nil {
return true
}
return group != nil && group.RequirePrivacySet
}
func (s *OpenAIGatewayService) selectAccountWithSchedulerOnce( func (s *OpenAIGatewayService) selectAccountWithSchedulerOnce(
ctx context.Context, ctx context.Context,
groupID *int64, groupID *int64,
@@ -2152,6 +2223,7 @@ func (s *OpenAIGatewayService) selectAccountWithSchedulerOnce(
useUpstreamTokenCost bool, useUpstreamTokenCost bool,
) (*AccountSelectionResult, OpenAIAccountScheduleDecision, error) { ) (*AccountSelectionResult, OpenAIAccountScheduleDecision, error) {
ctx = s.withOpenAIQuotaAutoPauseContext(ctx) ctx = s.withOpenAIQuotaAutoPauseContext(ctx)
ctx = s.withOpenAIGroupPrivacyRequirement(ctx, groupID)
// 分组利润控制:唯一文本调度入口的防御性装门。handler 文本 // 分组利润控制:唯一文本调度入口的防御性装门。handler 文本
// 入口已在请求开始经 WithOpenAIRequestPricingContext 装门并固定 pricingAt, // 入口已在请求开始经 WithOpenAIRequestPricingContext 装门并固定 pricingAt,
// 此处对同分组门直接复用(failover 重入阈值稳定),仅为不经 handler 装配的 // 此处对同分组门直接复用(failover 重入阈值稳定),仅为不经 handler 装配的
@@ -2163,13 +2235,52 @@ func (s *OpenAIGatewayService) selectAccountWithSchedulerOnce(
} }
platform = NormalizeOpenAICompatiblePlatform(platform) platform = NormalizeOpenAICompatiblePlatform(platform)
decision := OpenAIAccountScheduleDecision{} decision := OpenAIAccountScheduleDecision{}
preserveGuardianParentBinding := preserveOpenAIGuardianParentBinding(ctx, sessionHash)
guardianParentAccountID := int64(0)
if strings.TrimSpace(previousResponseID) == "" {
guardianParentAccountID = s.resolveOpenAIGuardianParentAccountID(ctx, groupID)
}
scheduler := s.getOpenAIAccountScheduler(ctx) scheduler := s.getOpenAIAccountScheduler(ctx)
if scheduler == nil { if scheduler == nil {
decision.Layer = openAIAccountScheduleLayerLoadBalance decision.Layer = openAIAccountScheduleLayerLoadBalance
if guardianParentAccountID > 0 {
if s.checkChannelPricingRestriction(ctx, groupID, requestedModel) {
return nil, decision, fmt.Errorf("%w supporting model: %s (channel pricing restriction)", ErrNoAvailableAccounts, requestedModel)
}
fallbackScheduler := &defaultOpenAIAccountScheduler{service: s, stats: newOpenAIAccountRuntimeStats()}
selection, _, err := fallbackScheduler.selectBySessionHash(ctx, OpenAIAccountScheduleRequest{
GroupID: groupID,
Platform: platform,
SessionHash: sessionHash,
StickyAccountID: guardianParentAccountID,
PreserveStickyBinding: true,
RequestedModel: requestedModel,
RequiredTransport: requiredTransport,
RequiredCapability: requiredCapability,
RequiredImageCapability: requiredImageCapability,
RequireCompact: requireCompact,
ExcludedIDs: excludedIDs,
RequirePrivacySet: s.openAIGroupRequiresPrivacySet(ctx, groupID),
})
if err != nil {
return nil, decision, err
}
if selection != nil && selection.Account != nil {
decision.Layer = openAIAccountScheduleLayerGuardianParent
decision.StickySessionHit = true
decision.SelectedAccountID = selection.Account.ID
decision.SelectedAccountType = selection.Account.Type
return selection, decision, nil
}
}
legacySessionHash := sessionHash
if preserveGuardianParentBinding {
legacySessionHash = ""
}
if requiredTransport == OpenAIUpstreamTransportAny || requiredTransport == OpenAIUpstreamTransportHTTPSSE { if requiredTransport == OpenAIUpstreamTransportAny || requiredTransport == OpenAIUpstreamTransportHTTPSSE {
effectiveExcludedIDs := cloneExcludedAccountIDs(excludedIDs) effectiveExcludedIDs := cloneExcludedAccountIDs(excludedIDs)
for { for {
selection, err := s.selectAccountWithLoadAwareness(ctx, groupID, platform, sessionHash, requestedModel, effectiveExcludedIDs, requireCompact, requiredCapability, useUpstreamTokenCost) selection, err := s.selectAccountWithLoadAwareness(ctx, groupID, platform, legacySessionHash, requestedModel, effectiveExcludedIDs, requireCompact, requiredCapability, useUpstreamTokenCost)
if err != nil { if err != nil {
return nil, decision, err return nil, decision, err
} }
@@ -2194,7 +2305,7 @@ func (s *OpenAIGatewayService) selectAccountWithSchedulerOnce(
effectiveExcludedIDs := cloneExcludedAccountIDs(excludedIDs) effectiveExcludedIDs := cloneExcludedAccountIDs(excludedIDs)
for { for {
selection, err := s.selectAccountWithLoadAwareness(ctx, groupID, platform, sessionHash, requestedModel, effectiveExcludedIDs, requireCompact, requiredCapability, useUpstreamTokenCost) selection, err := s.selectAccountWithLoadAwareness(ctx, groupID, platform, legacySessionHash, requestedModel, effectiveExcludedIDs, requireCompact, requiredCapability, useUpstreamTokenCost)
if err != nil { if err != nil {
return nil, decision, err return nil, decision, err
} }
@@ -2243,9 +2354,12 @@ func (s *OpenAIGatewayService) selectAccountWithSchedulerOnce(
Platform: platform, Platform: platform,
SessionHash: sessionHash, SessionHash: sessionHash,
StickyAccountID: stickyAccountID, StickyAccountID: stickyAccountID,
GuardianParentAccountID: guardianParentAccountID,
StickyPreviousAccountID: stickyPreviousAccountID, StickyPreviousAccountID: stickyPreviousAccountID,
StickyWeighted: stickyWeighted, StickyWeighted: stickyWeighted,
SubscriptionPriority: subscriptionPriority, SubscriptionPriority: subscriptionPriority,
PreserveStickyBinding: preserveGuardianParentBinding,
RequirePrivacySet: s.openAIGroupRequiresPrivacySet(ctx, groupID),
PreviousResponseID: previousResponseID, PreviousResponseID: previousResponseID,
PreviousResponseCanMove: previousResponseCanMove, PreviousResponseCanMove: previousResponseCanMove,
UseUpstreamTokenCost: useUpstreamTokenCost, UseUpstreamTokenCost: useUpstreamTokenCost,
@@ -1043,6 +1043,7 @@ func (s *OpenAIGatewayService) isBetterAccount(candidate, current *Account) bool
// SelectAccountWithLoadAwareness selects an account with load-awareness and wait plan. // 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) { func (s *OpenAIGatewayService) SelectAccountWithLoadAwareness(ctx context.Context, groupID *int64, sessionHash string, requestedModel string, excludedIDs map[int64]struct{}) (*AccountSelectionResult, error) {
ctx = s.withOpenAIQuotaAutoPauseContext(ctx) ctx = s.withOpenAIQuotaAutoPauseContext(ctx)
ctx = s.withOpenAIGroupPrivacyRequirement(ctx, groupID)
// 分组利润控制:legacy 公共入口同样装门,保证不经 // 分组利润控制:legacy 公共入口同样装门,保证不经
// selectAccountWithScheduler 的调用方也无法绕过利润准入。 // selectAccountWithScheduler 的调用方也无法绕过利润准入。
ctx = s.withOpenAIProfitControlGate(ctx, groupID) ctx = s.withOpenAIProfitControlGate(ctx, groupID)
@@ -1509,6 +1510,9 @@ func (s *OpenAIGatewayService) recheckSelectedOpenAIAccountFromDBBeforeProfit(ct
} }
platform = NormalizeOpenAICompatiblePlatform(platform) platform = NormalizeOpenAICompatiblePlatform(platform)
if s.schedulerSnapshot == nil || s.accountRepo == nil { if s.schedulerSnapshot == nil || s.accountRepo == nil {
if s.openAIGroupRequiresPrivacySet(ctx, groupID) && !account.IsPrivacySet() {
return nil
}
if !isOpenAICompatibleAccountEligibleForRequestBeforeProfit(ctx, account, platform, requestedModel, requireCompact, requiredCapability) { if !isOpenAICompatibleAccountEligibleForRequestBeforeProfit(ctx, account, platform, requestedModel, requireCompact, requiredCapability) {
return nil return nil
} }
@@ -1531,6 +1535,9 @@ func (s *OpenAIGatewayService) recheckSelectedOpenAIAccountFromDBBeforeProfit(ct
if !s.openAIAccountMatchesSchedulingGroup(latest, groupID) { if !s.openAIAccountMatchesSchedulingGroup(latest, groupID) {
return nil return nil
} }
if s.openAIGroupRequiresPrivacySet(ctx, groupID) && !latest.IsPrivacySet() {
return nil
}
if !isOpenAICompatibleAccountEligibleForRequestBeforeProfit(ctx, latest, platform, requestedModel, requireCompact, requiredCapability) { if !isOpenAICompatibleAccountEligibleForRequestBeforeProfit(ctx, latest, platform, requestedModel, requireCompact, requiredCapability) {
return nil return nil
} }
@@ -0,0 +1,133 @@
package service
import (
"context"
"strings"
"github.com/gin-gonic/gin"
"github.com/tidwall/gjson"
)
const (
codexAutoReviewModel = "codex-auto-review"
openAISubagentHeader = "x-openai-subagent"
codexParentThreadIDHeader = "x-codex-parent-thread-id"
codexTurnMetadataHeader = "x-codex-turn-metadata"
)
type openAIGuardianParentAffinityContextKey struct{}
type openAIGuardianParentAffinity struct {
currentSessionHash string
legacySessionHash string
}
// WithOpenAIGuardianParentAffinity records a Codex review request's parent
// thread as a routing hint. The hint is resolved against the current group's
// sticky-session namespace later; client headers never carry an account ID.
func WithOpenAIGuardianParentAffinity(ctx context.Context, c *gin.Context, body []byte, requestedModel string) context.Context {
if ctx == nil || c == nil || !strings.EqualFold(strings.TrimSpace(requestedModel), codexAutoReviewModel) {
return ctx
}
headerMetadata := c.GetHeader(codexTurnMetadataHeader)
bodyMetadata := openAIRequestPayloadView(body).Get("client_metadata.x-codex-turn-metadata").String()
if !hasUnambiguousOpenAICodexReviewSubagent(
c.GetHeader(openAISubagentHeader),
codexSubagentKindFromMetadata(headerMetadata),
codexSubagentKindFromMetadata(bodyMetadata),
) {
return ctx
}
parentID := ""
for _, candidate := range []string{
strings.TrimSpace(c.GetHeader(codexParentThreadIDHeader)),
codexParentThreadIDFromMetadata(headerMetadata),
codexParentThreadIDFromMetadata(bodyMetadata),
} {
if candidate == "" {
continue
}
if parentID != "" && parentID != candidate {
return ctx
}
parentID = candidate
}
if parentID == "" {
return ctx
}
currentHash, legacyHash := deriveOpenAISessionHashes(parentID)
if currentHash == "" {
return ctx
}
return context.WithValue(ctx, openAIGuardianParentAffinityContextKey{}, openAIGuardianParentAffinity{
currentSessionHash: currentHash,
legacySessionHash: legacyHash,
})
}
func codexParentThreadIDFromMetadata(raw string) string {
raw = strings.TrimSpace(raw)
if raw == "" || !gjson.Valid(raw) {
return ""
}
return strings.TrimSpace(gjson.Get(raw, "parent_thread_id").String())
}
func codexSubagentKindFromMetadata(raw string) string {
raw = strings.TrimSpace(raw)
if raw == "" || !gjson.Valid(raw) {
return ""
}
return strings.TrimSpace(gjson.Get(raw, "subagent_kind").String())
}
func hasUnambiguousOpenAICodexReviewSubagent(candidates ...string) bool {
subagent := ""
for _, candidate := range candidates {
candidate = strings.ToLower(strings.TrimSpace(candidate))
if candidate == "" {
continue
}
if subagent != "" && subagent != candidate {
return false
}
subagent = candidate
}
return subagent == "guardian" || subagent == "review"
}
func openAIGuardianParentAffinityFromContext(ctx context.Context) (openAIGuardianParentAffinity, bool) {
if ctx == nil {
return openAIGuardianParentAffinity{}, false
}
affinity, ok := ctx.Value(openAIGuardianParentAffinityContextKey{}).(openAIGuardianParentAffinity)
return affinity, ok && affinity.currentSessionHash != ""
}
func preserveOpenAIGuardianParentBinding(ctx context.Context, sessionHash string) bool {
affinity, ok := openAIGuardianParentAffinityFromContext(ctx)
if !ok {
return false
}
sessionHash = strings.TrimSpace(sessionHash)
return sessionHash != "" && (sessionHash == affinity.currentSessionHash || sessionHash == affinity.legacySessionHash)
}
func (s *OpenAIGatewayService) resolveOpenAIGuardianParentAccountID(ctx context.Context, groupID *int64) int64 {
if s == nil || s.cache == nil {
return 0
}
affinity, ok := openAIGuardianParentAffinityFromContext(ctx)
if !ok {
return 0
}
lookupCtx := withOpenAILegacySessionHash(ctx, affinity.legacySessionHash)
accountID, err := s.getStickySessionAccountID(lookupCtx, groupID, affinity.currentSessionHash)
if err != nil || accountID <= 0 {
return 0
}
return accountID
}
@@ -0,0 +1,450 @@
package service
import (
"context"
"errors"
"net/http"
"net/http/httptest"
"testing"
"time"
"github.com/Wei-Shaw/sub2api/internal/config"
"github.com/gin-gonic/gin"
"github.com/stretchr/testify/require"
)
type guardianAffinityGroupRepo struct {
GroupRepository
group *Group
err error
}
type guardianAffinityAccountRepo struct {
schedulerGroupAwareOpenAIAccountRepo
setErrorCalls int
}
func (r *guardianAffinityAccountRepo) SetError(context.Context, int64, string) error {
r.setErrorCalls++
return nil
}
func (r guardianAffinityGroupRepo) GetByID(context.Context, int64) (*Group, error) {
if r.err != nil {
return nil, r.err
}
return r.group, nil
}
func (r guardianAffinityGroupRepo) GetByIDLite(context.Context, int64) (*Group, error) {
if r.err != nil {
return nil, r.err
}
return r.group, nil
}
func guardianAffinityTestContext(t *testing.T, model, subagent, parentHeader, metadata string) context.Context {
t.Helper()
gin.SetMode(gin.TestMode)
rec := httptest.NewRecorder()
c, _ := gin.CreateTestContext(rec)
c.Request = httptest.NewRequest(http.MethodPost, "/openai/v1/responses", nil)
c.Request.Header.Set(openAISubagentHeader, subagent)
if parentHeader != "" {
c.Request.Header.Set(codexParentThreadIDHeader, parentHeader)
}
if metadata != "" {
c.Request.Header.Set(codexTurnMetadataHeader, metadata)
}
return WithOpenAIGuardianParentAffinity(context.Background(), c, nil, model)
}
func TestWithOpenAIGuardianParentAffinity_RequiresUnambiguousReviewLineage(t *testing.T) {
parentID := "11111111-1111-4111-8111-111111111111"
wantHash := DeriveSessionHashFromSeed(parentID)
for _, subagent := range []string{"guardian", "review", "GUARDIAN"} {
t.Run(subagent, func(t *testing.T) {
ctx := guardianAffinityTestContext(t, codexAutoReviewModel, subagent, parentID, `{"parent_thread_id":"`+parentID+`"}`)
affinity, ok := openAIGuardianParentAffinityFromContext(ctx)
require.True(t, ok)
require.Equal(t, wantHash, affinity.currentSessionHash)
})
}
t.Run("metadata only", func(t *testing.T) {
ctx := guardianAffinityTestContext(t, codexAutoReviewModel, "guardian", "", `{"parent_thread_id":"`+parentID+`"}`)
_, ok := openAIGuardianParentAffinityFromContext(ctx)
require.True(t, ok)
})
t.Run("websocket envelope metadata", func(t *testing.T) {
gin.SetMode(gin.TestMode)
rec := httptest.NewRecorder()
c, _ := gin.CreateTestContext(rec)
c.Request = httptest.NewRequest(http.MethodGet, "/openai/v1/responses", nil)
body := []byte(`{"type":"response.create","response":{"model":"codex-auto-review","client_metadata":{"x-codex-turn-metadata":"{\"parent_thread_id\":\"` + parentID + `\",\"subagent_kind\":\"guardian\"}"}}}`)
ctx := WithOpenAIGuardianParentAffinity(context.Background(), c, body, codexAutoReviewModel)
affinity, ok := openAIGuardianParentAffinityFromContext(ctx)
require.True(t, ok)
require.Equal(t, wantHash, affinity.currentSessionHash)
})
for name, ctx := range map[string]context.Context{
"ordinary model": guardianAffinityTestContext(t, "gpt-5.6-sol", "guardian", parentID, ""),
"ordinary subagent": guardianAffinityTestContext(t, codexAutoReviewModel, "collab_spawn", parentID, ""),
"missing parent": guardianAffinityTestContext(t, codexAutoReviewModel, "guardian", "", ""),
"conflicting lineage": guardianAffinityTestContext(t, codexAutoReviewModel, "guardian", parentID, `{"parent_thread_id":"different-parent"}`),
"conflicting subagent": guardianAffinityTestContext(t, codexAutoReviewModel, "guardian", parentID, `{"parent_thread_id":"`+parentID+`","subagent_kind":"collab_spawn"}`),
} {
t.Run(name, func(t *testing.T) {
_, ok := openAIGuardianParentAffinityFromContext(ctx)
require.False(t, ok)
})
}
}
func TestOpenAIGatewayService_GuardianParentAffinitySelectsParentAccountAcrossSchedulers(t *testing.T) {
parentID := "22222222-2222-4222-8222-222222222222"
parentHash := DeriveSessionHashFromSeed(parentID)
groupID := int64(102001)
for _, mode := range []struct {
name string
advanced string
stickyWeighted string
}{
{name: "legacy", advanced: "false"},
{name: "advanced", advanced: "true"},
{name: "advanced sticky weighted", advanced: "true", stickyWeighted: "true"},
} {
t.Run(mode.name, func(t *testing.T) {
accounts := []Account{
{
ID: 39001, Platform: PlatformOpenAI, Type: AccountTypeOAuth,
Status: StatusActive, Schedulable: true, Concurrency: 1, Priority: 10,
GroupIDs: []int64{groupID}, Credentials: map[string]any{"plan_type": "team"},
},
{
ID: 39002, Platform: PlatformOpenAI, Type: AccountTypeOAuth,
Status: StatusActive, Schedulable: true, Concurrency: 1, Priority: 0,
GroupIDs: []int64{groupID}, Credentials: map[string]any{"plan_type": "team"},
},
}
cfg := &config.Config{}
cfg.Gateway.OpenAIWS.LBTopK = 2
cache := &schedulerTestGatewayCache{sessionBindings: map[string]int64{"openai:" + parentHash: 39001}}
svc := &OpenAIGatewayService{
accountRepo: schedulerGroupAwareOpenAIAccountRepo{schedulerTestOpenAIAccountRepo{accounts: accounts}},
cache: cache,
cfg: cfg,
rateLimitService: newOpenAIAdvancedSchedulerRateLimitService(mode.advanced, mode.stickyWeighted),
concurrencyService: NewConcurrencyService(schedulerTestConcurrencyCache{acquireResults: map[int64]bool{39001: true, 39002: true}}),
}
ctx := guardianAffinityTestContext(t, codexAutoReviewModel, "guardian", parentID, "")
selection, decision, err := svc.SelectAccountWithScheduler(
ctx, &groupID, "", "guardian-child-session", codexAutoReviewModel,
nil, OpenAIUpstreamTransportAny, false,
)
require.NoError(t, err)
require.NotNil(t, selection)
require.Equal(t, int64(39001), selection.Account.ID)
require.Equal(t, openAIAccountScheduleLayerGuardianParent, decision.Layer)
require.Zero(t, cache.deletedSessions["openai:"+parentHash])
if selection.ReleaseFunc != nil {
selection.ReleaseFunc()
}
})
}
}
func TestOpenAIGatewayService_GuardianParentAffinityFallsBackWithoutCrossGroupOrFailoverBypass(t *testing.T) {
parentID := "33333333-3333-4333-8333-333333333333"
parentHash := DeriveSessionHashFromSeed(parentID)
groupID := int64(102011)
otherGroupID := int64(102012)
for name, excluded := range map[string]map[int64]struct{}{
"parent moved out of group": nil,
"parent excluded after upstream failure": {39011: {}},
} {
t.Run(name, func(t *testing.T) {
parentGroups := []int64{groupID}
if excluded == nil {
parentGroups = []int64{otherGroupID}
}
accounts := []Account{
{ID: 39011, Platform: PlatformOpenAI, Type: AccountTypeOAuth, Status: StatusActive, Schedulable: true, Concurrency: 1, Priority: 0, GroupIDs: parentGroups, Credentials: map[string]any{"plan_type": "team"}},
{ID: 39012, Platform: PlatformOpenAI, Type: AccountTypeOAuth, Status: StatusActive, Schedulable: true, Concurrency: 1, Priority: 5, GroupIDs: []int64{groupID}, Credentials: map[string]any{"plan_type": "team"}},
}
cache := &schedulerTestGatewayCache{sessionBindings: map[string]int64{"openai:" + parentHash: 39011}}
svc := &OpenAIGatewayService{
accountRepo: schedulerGroupAwareOpenAIAccountRepo{schedulerTestOpenAIAccountRepo{accounts: accounts}},
cache: cache,
cfg: &config.Config{},
rateLimitService: newOpenAIAdvancedSchedulerRateLimitService("true"),
concurrencyService: NewConcurrencyService(schedulerTestConcurrencyCache{acquireResults: map[int64]bool{39011: true, 39012: true}}),
}
ctx := guardianAffinityTestContext(t, codexAutoReviewModel, "guardian", parentID, "")
selection, _, err := svc.SelectAccountWithScheduler(
ctx, &groupID, "", "guardian-fallback-child", codexAutoReviewModel,
excluded, OpenAIUpstreamTransportAny, false,
)
require.NoError(t, err)
require.NotNil(t, selection)
require.Equal(t, int64(39012), selection.Account.ID)
require.Zero(t, cache.deletedSessions["openai:"+parentHash], "a child request must never delete its parent's binding")
if selection.ReleaseFunc != nil {
selection.ReleaseFunc()
}
})
}
}
func TestOpenAIGatewayService_GuardianParentHashCollisionPreservesParentBinding(t *testing.T) {
parentID := "44444444-4444-4444-8444-444444444444"
parentHash := DeriveSessionHashFromSeed(parentID)
groupID := int64(102021)
otherGroupID := int64(102022)
for _, advanced := range []string{"false", "true"} {
t.Run(map[string]string{"false": "legacy", "true": "advanced"}[advanced], func(t *testing.T) {
accounts := []Account{
{ID: 39021, Platform: PlatformOpenAI, Type: AccountTypeOAuth, Status: StatusActive, Schedulable: true, Concurrency: 1, Priority: 0, GroupIDs: []int64{otherGroupID}, Credentials: map[string]any{"plan_type": "team"}},
{ID: 39022, Platform: PlatformOpenAI, Type: AccountTypeOAuth, Status: StatusActive, Schedulable: true, Concurrency: 1, Priority: 5, GroupIDs: []int64{groupID}, Credentials: map[string]any{"plan_type": "team"}},
}
cache := &schedulerTestGatewayCache{sessionBindings: map[string]int64{"openai:" + parentHash: 39021}}
svc := &OpenAIGatewayService{
accountRepo: schedulerGroupAwareOpenAIAccountRepo{schedulerTestOpenAIAccountRepo{accounts: accounts}},
cache: cache,
cfg: &config.Config{},
rateLimitService: newOpenAIAdvancedSchedulerRateLimitService(advanced),
concurrencyService: NewConcurrencyService(schedulerTestConcurrencyCache{acquireResults: map[int64]bool{39021: true, 39022: true}}),
}
ctx := guardianAffinityTestContext(t, codexAutoReviewModel, "guardian", parentID, "")
selection, _, err := svc.SelectAccountWithScheduler(
ctx, &groupID, "", parentHash, codexAutoReviewModel,
nil, OpenAIUpstreamTransportAny, false,
)
require.NoError(t, err)
require.NotNil(t, selection)
require.Equal(t, int64(39022), selection.Account.ID)
require.NoError(t, svc.BindStickySessionAfterProfitAdmission(ctx, &groupID, parentHash, selection.Account.ID))
require.Equal(t, int64(39021), cache.sessionBindings["openai:"+parentHash])
require.Zero(t, cache.deletedSessions["openai:"+parentHash])
if selection.ReleaseFunc != nil {
selection.ReleaseFunc()
}
})
}
}
func TestOpenAIGatewayService_GuardianParentAffinityHonorsRequiredPrivacy(t *testing.T) {
parentID := "55555555-5555-4555-8555-555555555555"
parentHash := DeriveSessionHashFromSeed(parentID)
groupID := int64(102031)
for _, advanced := range []string{"false", "true"} {
t.Run(map[string]string{"false": "legacy", "true": "advanced"}[advanced], func(t *testing.T) {
accounts := []Account{
{ID: 39031, Platform: PlatformOpenAI, Type: AccountTypeOAuth, Status: StatusActive, Schedulable: true, Concurrency: 1, Priority: 0, GroupIDs: []int64{groupID}, Credentials: map[string]any{"plan_type": "team"}},
{ID: 39032, Platform: PlatformOpenAI, Type: AccountTypeOAuth, Status: StatusActive, Schedulable: true, Concurrency: 1, Priority: 5, GroupIDs: []int64{groupID}, Credentials: map[string]any{"plan_type": "team"}, Extra: map[string]any{"privacy_mode": PrivacyModeTrainingOff}},
}
repo := &guardianAffinityAccountRepo{schedulerGroupAwareOpenAIAccountRepo: schedulerGroupAwareOpenAIAccountRepo{schedulerTestOpenAIAccountRepo{accounts: accounts}}}
cache := &schedulerTestGatewayCache{sessionBindings: map[string]int64{"openai:" + parentHash: 39031}}
svc := &OpenAIGatewayService{
accountRepo: repo,
cache: cache,
cfg: &config.Config{},
rateLimitService: newOpenAIAdvancedSchedulerRateLimitService(advanced),
concurrencyService: NewConcurrencyService(schedulerTestConcurrencyCache{acquireResults: map[int64]bool{39031: true, 39032: true}}),
schedulerSnapshot: &SchedulerSnapshotService{
accountRepo: repo,
groupRepo: guardianAffinityGroupRepo{group: &Group{ID: groupID, Name: "privacy", RequirePrivacySet: true}},
},
}
ctx := guardianAffinityTestContext(t, codexAutoReviewModel, "guardian", parentID, "")
selection, _, err := svc.SelectAccountWithScheduler(
ctx, &groupID, "", "guardian-privacy-child", codexAutoReviewModel,
nil, OpenAIUpstreamTransportAny, false,
)
require.NoError(t, err)
require.NotNil(t, selection)
require.Equal(t, int64(39032), selection.Account.ID)
require.Zero(t, repo.setErrorCalls, "a group-scoped privacy gate must not globally error a shared account")
require.False(t, svc.isOpenAIAccountRequestRuntimeBlocked(&accounts[0], codexAutoReviewModel))
if selection.ReleaseFunc != nil {
selection.ReleaseFunc()
}
})
}
}
func TestOpenAIGatewayService_PreviousResponseHonorsGroupAndRequiredPrivacy(t *testing.T) {
groupID := int64(3904)
tests := []struct {
name string
boundAccount Account
groupErr error
}{
{
name: "privacy unset",
boundAccount: Account{
ID: 39041, Platform: PlatformOpenAI, Type: AccountTypeAPIKey,
Status: StatusActive, Schedulable: true, Concurrency: 1,
GroupIDs: []int64{groupID},
Extra: map[string]any{"openai_apikey_responses_websockets_v2_enabled": true},
},
},
{
name: "privacy policy lookup error fails closed",
boundAccount: Account{
ID: 39041, Platform: PlatformOpenAI, Type: AccountTypeAPIKey,
Status: StatusActive, Schedulable: true, Concurrency: 1,
GroupIDs: []int64{groupID},
Extra: map[string]any{"openai_apikey_responses_websockets_v2_enabled": true},
},
groupErr: errors.New("group repository unavailable"),
},
{
name: "different group",
boundAccount: Account{
ID: 39041, Platform: PlatformOpenAI, Type: AccountTypeAPIKey,
Status: StatusActive, Schedulable: true, Concurrency: 1,
GroupIDs: []int64{groupID + 1},
Extra: map[string]any{
"openai_apikey_responses_websockets_v2_enabled": true,
"privacy_mode": PrivacyModeTrainingOff,
},
},
},
}
for _, tc := range tests {
t.Run(tc.name, func(t *testing.T) {
fallback := Account{
ID: 39042, Platform: PlatformOpenAI, Type: AccountTypeAPIKey,
Status: StatusActive, Schedulable: true, Concurrency: 1, Priority: 5,
GroupIDs: []int64{groupID},
Extra: map[string]any{
"openai_apikey_responses_websockets_v2_enabled": true,
"privacy_mode": PrivacyModeTrainingOff,
},
}
accounts := []Account{tc.boundAccount, fallback}
repo := &guardianAffinityAccountRepo{schedulerGroupAwareOpenAIAccountRepo: schedulerGroupAwareOpenAIAccountRepo{schedulerTestOpenAIAccountRepo{accounts: accounts}}}
cache := &schedulerTestGatewayCache{}
store := NewOpenAIWSStateStore(cache)
groupRepo := guardianAffinityGroupRepo{
group: &Group{
ID: groupID, Name: "privacy-required", Platform: PlatformOpenAI,
Status: StatusActive, RequirePrivacySet: true,
},
err: tc.groupErr,
}
svc := &OpenAIGatewayService{
accountRepo: repo,
cache: cache,
cfg: &config.Config{},
rateLimitService: newOpenAIAdvancedSchedulerRateLimitService("true"),
concurrencyService: NewConcurrencyService(&schedulerTestConcurrencyCache{}),
openaiWSStateStore: store,
schedulerSnapshot: &SchedulerSnapshotService{
accountRepo: repo,
groupRepo: groupRepo,
},
}
responseID := "resp_privacy_guard"
require.NoError(t, store.BindResponseAccount(context.Background(), groupID, responseID, tc.boundAccount.ID, time.Hour))
directSelection, directErr := svc.SelectAccountByPreviousResponseID(
context.Background(), &groupID, responseID, codexAutoReviewModel, nil, false,
)
require.NoError(t, directErr)
require.Nil(t, directSelection, "the previous-response helper must enforce fresh group/privacy state")
selection, decision, err := svc.SelectAccountWithSchedulerForCapability(
context.Background(), &groupID, responseID, "", codexAutoReviewModel,
nil, OpenAIUpstreamTransportAny, OpenAIEndpointCapabilityResponses,
false, false, true,
)
require.NoError(t, err)
require.NotNil(t, selection)
require.Equal(t, fallback.ID, selection.Account.ID)
require.NotEqual(t, openAIAccountScheduleLayerPreviousResponse, decision.Layer)
require.Zero(t, repo.setErrorCalls)
require.False(t, svc.isOpenAIAccountRequestRuntimeBlocked(&accounts[0], codexAutoReviewModel))
boundAccountID, getErr := store.GetResponseAccount(context.Background(), groupID, responseID)
require.NoError(t, getErr)
require.Equal(t, tc.boundAccount.ID, boundAccountID, "transient policy misses must preserve the response binding")
if selection.ReleaseFunc != nil {
selection.ReleaseFunc()
}
})
}
}
func TestOpenAIGatewayService_PreviousResponseSimpleModeIgnoresGroupMembership(t *testing.T) {
groupID := int64(3905)
bound := Account{
ID: 39051, Platform: PlatformOpenAI, Type: AccountTypeAPIKey,
Status: StatusActive, Schedulable: true, Concurrency: 1,
GroupIDs: []int64{groupID + 1},
Extra: map[string]any{"openai_apikey_responses_websockets_v2_enabled": true},
}
fallback := Account{
ID: 39052, Platform: PlatformOpenAI, Type: AccountTypeAPIKey,
Status: StatusActive, Schedulable: true, Concurrency: 1, Priority: 10,
GroupIDs: []int64{groupID},
Extra: map[string]any{"openai_apikey_responses_websockets_v2_enabled": true},
}
accounts := []Account{bound, fallback}
repo := &guardianAffinityAccountRepo{schedulerGroupAwareOpenAIAccountRepo: schedulerGroupAwareOpenAIAccountRepo{schedulerTestOpenAIAccountRepo{accounts: accounts}}}
cache := &schedulerTestGatewayCache{}
store := NewOpenAIWSStateStore(cache)
cfg := &config.Config{RunMode: config.RunModeSimple}
svc := &OpenAIGatewayService{
accountRepo: repo,
cache: cache,
cfg: cfg,
rateLimitService: newOpenAIAdvancedSchedulerRateLimitService("true"),
concurrencyService: NewConcurrencyService(&schedulerTestConcurrencyCache{}),
openaiWSStateStore: store,
schedulerSnapshot: &SchedulerSnapshotService{
accountRepo: repo,
groupRepo: guardianAffinityGroupRepo{group: &Group{
ID: groupID, Name: "simple-mode", Platform: PlatformOpenAI, Status: StatusActive,
}},
},
}
responseID := "resp_simple_mode_cross_group"
require.NoError(t, store.BindResponseAccount(context.Background(), groupID, responseID, bound.ID, time.Hour))
directSelection, err := svc.SelectAccountByPreviousResponseID(
context.Background(), &groupID, responseID, codexAutoReviewModel, nil, false,
)
require.NoError(t, err)
require.NotNil(t, directSelection)
require.Equal(t, bound.ID, directSelection.Account.ID)
if directSelection.ReleaseFunc != nil {
directSelection.ReleaseFunc()
}
selection, decision, err := svc.SelectAccountWithSchedulerForCapability(
context.Background(), &groupID, responseID, "", codexAutoReviewModel,
nil, OpenAIUpstreamTransportAny, OpenAIEndpointCapabilityResponses,
false, false, true,
)
require.NoError(t, err)
require.NotNil(t, selection)
require.Equal(t, bound.ID, selection.Account.ID)
require.Equal(t, openAIAccountScheduleLayerPreviousResponse, decision.Layer)
if selection.ReleaseFunc != nil {
selection.ReleaseFunc()
}
}
@@ -342,7 +342,7 @@ func (s *OpenAIGatewayService) ProfitControlVetoLatest(ctx context.Context, sele
// only after the terminal post-slot check, so an account rejected after a rate // only after the terminal post-slot check, so an account rejected after a rate
// refresh cannot become the new sticky target. // refresh cannot become the new sticky target.
func (s *OpenAIGatewayService) bindOpenAIStickySessionDuringSelection(ctx context.Context, groupID *int64, sessionHash string, accountID int64) error { func (s *OpenAIGatewayService) bindOpenAIStickySessionDuringSelection(ctx context.Context, groupID *int64, sessionHash string, accountID int64) error {
if gatewayProfitControlGateActive(ctx) { if gatewayProfitControlGateActive(ctx) || preserveOpenAIGuardianParentBinding(ctx, sessionHash) {
return nil return nil
} }
return s.BindStickySession(ctx, groupID, sessionHash, accountID) return s.BindStickySession(ctx, groupID, sessionHash, accountID)
@@ -358,6 +358,9 @@ func (s *OpenAIGatewayService) BindStickySessionAfterProfitAdmission(ctx context
if sessionHash == "" || accountID <= 0 { if sessionHash == "" || accountID <= 0 {
return nil return nil
} }
if preserveOpenAIGuardianParentBinding(ctx, sessionHash) {
return nil
}
if !gatewayProfitControlGateActive(ctx) { if !gatewayProfitControlGateActive(ctx) {
return s.BindStickySession(ctx, groupID, sessionHash, accountID) return s.BindStickySession(ctx, groupID, sessionHash, accountID)
} }
@@ -596,6 +596,12 @@ func (s *OpenAIGatewayService) resolveAccountByPreviousResponseIDForCapability(
_ = store.DeleteResponseAccount(ctx, derefGroupID(groupID), responseID) _ = store.DeleteResponseAccount(ctx, derefGroupID(groupID), responseID)
return 0, nil, "", nil return 0, nil, "", nil
} }
if !s.openAIAccountMatchesSchedulingGroup(latest, groupID) {
return 0, nil, "", nil
}
if s.openAIGroupRequiresPrivacySet(ctx, groupID) && !latest.IsPrivacySet() {
return 0, nil, "", nil
}
if !parentHealthyForShadow(latest, s.parentAccountLookup(ctx)) { if !parentHealthyForShadow(latest, s.parentAccountLookup(ctx)) {
_ = store.DeleteResponseAccount(ctx, derefGroupID(groupID), responseID) _ = store.DeleteResponseAccount(ctx, derefGroupID(groupID), responseID)
return 0, nil, "", nil return 0, nil, "", nil