mirror of
https://github.com/Wei-Shaw/sub2api.git
synced 2026-08-31 01:13:06 +08:00
fix(openai): keep auto-review on parent account
This commit is contained in:
@@ -514,6 +514,9 @@ func (h *OpenAIGatewayHandler) Responses(c *gin.Context) {
|
||||
if h.rejectIfCyberSessionBlocked(c, apiKey, sessionHashBody, reqModel, cyberBlockFormatResponses) {
|
||||
return
|
||||
}
|
||||
c.Request = c.Request.WithContext(service.WithOpenAIGuardianParentAffinity(
|
||||
c.Request.Context(), c, sessionHashBody, reqModel,
|
||||
))
|
||||
requireCompact := legacyCompact
|
||||
|
||||
maxAccountSwitches := h.maxAccountSwitches
|
||||
@@ -2004,6 +2007,7 @@ func (h *OpenAIGatewayHandler) ResponsesWebSocket(c *gin.Context) {
|
||||
firstMessage,
|
||||
openAIWSIngressFallbackSessionSeed(subject.UserID, apiKey.ID, apiKey.GroupID),
|
||||
)
|
||||
ctx = service.WithOpenAIGuardianParentAffinity(ctx, c, firstMessage, reqModel)
|
||||
maxAccountSwitches := h.maxAccountSwitches
|
||||
switchCount := 0
|
||||
profitVetoCount := 0
|
||||
|
||||
@@ -21,6 +21,7 @@ import (
|
||||
|
||||
const (
|
||||
openAIAccountScheduleLayerPreviousResponse = "previous_response_id"
|
||||
openAIAccountScheduleLayerGuardianParent = "guardian_parent"
|
||||
openAIAccountScheduleLayerSessionSticky = "session_hash"
|
||||
openAIAccountScheduleLayerLoadBalance = "load_balance"
|
||||
openAIAdvancedSchedulerSettingKey = "openai_advanced_scheduler_enabled"
|
||||
@@ -71,10 +72,12 @@ type OpenAIAccountScheduleRequest struct {
|
||||
Platform string
|
||||
SessionHash string
|
||||
StickyAccountID int64
|
||||
GuardianParentAccountID int64
|
||||
StickyPreviousAccountID int64
|
||||
StickyWeighted bool
|
||||
SubscriptionPriority bool
|
||||
PreserveStickyBinding bool
|
||||
RequirePrivacySet bool
|
||||
PreviousResponseID string
|
||||
PreviousResponseCanMove bool
|
||||
UseUpstreamTokenCost bool
|
||||
@@ -373,6 +376,9 @@ func (s *defaultOpenAIAccountScheduler) Select(
|
||||
ctx context.Context,
|
||||
req OpenAIAccountScheduleRequest,
|
||||
) (*AccountSelectionResult, OpenAIAccountScheduleDecision, error) {
|
||||
if s != nil && s.service != nil && s.service.openAIGroupRequiresPrivacySet(ctx, req.GroupID) {
|
||||
req.RequirePrivacySet = true
|
||||
}
|
||||
decision := OpenAIAccountScheduleDecision{}
|
||||
start := time.Now()
|
||||
defer func() {
|
||||
@@ -396,7 +402,14 @@ func (s *defaultOpenAIAccountScheduler) Select(
|
||||
return nil, decision, err
|
||||
}
|
||||
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 {
|
||||
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 {
|
||||
selection, escapedSticky, err := s.selectBySessionHash(ctx, req)
|
||||
if err != nil {
|
||||
@@ -465,6 +495,11 @@ func (s *defaultOpenAIAccountScheduler) selectBySessionHash(
|
||||
}
|
||||
|
||||
accountID := req.StickyAccountID
|
||||
clearBinding := func() {
|
||||
if !req.PreserveStickyBinding {
|
||||
_ = s.service.deleteStickySessionAccountID(ctx, req.GroupID, sessionHash)
|
||||
}
|
||||
}
|
||||
if accountID <= 0 {
|
||||
var err error
|
||||
accountID, err = s.service.getStickySessionAccountID(ctx, req.GroupID, sessionHash)
|
||||
@@ -483,40 +518,40 @@ func (s *defaultOpenAIAccountScheduler) selectBySessionHash(
|
||||
|
||||
account, err := s.service.getSchedulableAccount(ctx, accountID)
|
||||
if err != nil || account == nil {
|
||||
_ = s.service.deleteStickySessionAccountID(ctx, req.GroupID, sessionHash)
|
||||
clearBinding()
|
||||
return nil, false, nil
|
||||
}
|
||||
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
|
||||
}
|
||||
if !s.isAccountRequestCompatible(ctx, account, req) {
|
||||
return nil, false, nil
|
||||
}
|
||||
if !s.isAccountTransportCompatible(account, req.RequiredTransport) {
|
||||
_ = s.service.deleteStickySessionAccountID(ctx, req.GroupID, sessionHash)
|
||||
clearBinding()
|
||||
return nil, false, nil
|
||||
}
|
||||
account = s.service.recheckSelectedOpenAIAccountFromDB(ctx, account, req.GroupID, req.Platform, req.RequestedModel, req.RequireCompact, req.RequiredCapability)
|
||||
if account == nil || !s.service.openAIAccountMatchesSchedulingGroup(account, req.GroupID) || !s.isAccountTransportCompatible(account, req.RequiredTransport) {
|
||||
_ = s.service.deleteStickySessionAccountID(ctx, req.GroupID, sessionHash)
|
||||
if account == nil || !s.service.openAIAccountMatchesSchedulingGroup(account, req.GroupID) || !s.isAccountRequestCompatible(ctx, account, req) || !s.isAccountTransportCompatible(account, req.RequiredTransport) {
|
||||
clearBinding()
|
||||
return nil, false, nil
|
||||
}
|
||||
// Free-tier soft gate: sticky session must not pin an over-quota free OAuth account.
|
||||
// Admin QueryQuota / import probes do not use this path.
|
||||
if account != nil && len(s.filterGrokFreeQuotaAccounts(ctx, []Account{*account})) == 0 {
|
||||
_ = s.service.deleteStickySessionAccountID(ctx, req.GroupID, sessionHash)
|
||||
clearBinding()
|
||||
return nil, false, nil
|
||||
}
|
||||
// Team+model cool: sticky must not pin a sibling under the same team 429 window.
|
||||
now := time.Now()
|
||||
upstreamModel := canonicalOpenAIAccountSchedulingModel(account, req.RequestedModel)
|
||||
if account != nil && isGrokTeamModelRateLimited(account, upstreamModel, now) {
|
||||
_ = s.service.deleteStickySessionAccountID(ctx, req.GroupID, sessionHash)
|
||||
clearBinding()
|
||||
return nil, false, nil
|
||||
}
|
||||
if account != nil && isGrokModelQuotaBlocked(account.ID, upstreamModel, now) {
|
||||
_ = s.service.deleteStickySessionAccountID(ctx, req.GroupID, sessionHash)
|
||||
clearBinding()
|
||||
return nil, false, nil
|
||||
}
|
||||
escapeCfg := s.service.openAIStickyEscapeConfig()
|
||||
@@ -531,7 +566,9 @@ func (s *defaultOpenAIAccountScheduler) selectBySessionHash(
|
||||
}
|
||||
result, acquireErr := s.service.tryAcquireAccountSlot(ctx, accountID, account.Concurrency)
|
||||
if acquireErr == nil && result != nil && result.Acquired {
|
||||
_ = s.service.refreshStickySessionTTL(ctx, req.GroupID, sessionHash, s.service.openAIWSSessionStickyTTL())
|
||||
if !req.PreserveStickyBinding {
|
||||
_ = s.service.refreshStickySessionTTL(ctx, req.GroupID, sessionHash, s.service.openAIWSSessionStickyTTL())
|
||||
}
|
||||
return attachSelectionProfitGate(ctx, &AccountSelectionResult{
|
||||
Account: account,
|
||||
Acquired: true,
|
||||
@@ -1414,11 +1451,10 @@ func (s *defaultOpenAIAccountScheduler) selectByLoadBalance(
|
||||
filterStats.exclude("runtime_blocked")
|
||||
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() {
|
||||
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")
|
||||
continue
|
||||
}
|
||||
@@ -1730,6 +1766,9 @@ func (s *defaultOpenAIAccountScheduler) isAccountRequestCompatibleReason(ctx con
|
||||
if account == nil {
|
||||
return false, "account_nil"
|
||||
}
|
||||
if req.RequirePrivacySet && !account.IsPrivacySet() {
|
||||
return false, "privacy_not_set"
|
||||
}
|
||||
if s != nil && s.service != nil && s.service.isOpenAIAccountRequestRuntimeBlocked(account, req.RequestedModel) {
|
||||
return false, "runtime_blocked"
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
|
||||
type openAIGroupPrivacyRequirementContextKey struct{}
|
||||
|
||||
type openAIGroupPrivacyRequirement struct {
|
||||
groupID int64
|
||||
required bool
|
||||
}
|
||||
|
||||
func (s *OpenAIGatewayService) withOpenAIGroupPrivacyRequirement(ctx context.Context, groupID *int64) context.Context {
|
||||
return context.WithValue(ctx, openAIGroupPrivacyRequirementContextKey{}, openAIGroupPrivacyRequirement{
|
||||
groupID: derefGroupID(groupID),
|
||||
required: s.loadOpenAIGroupRequiresPrivacySet(ctx, groupID),
|
||||
})
|
||||
}
|
||||
|
||||
func (s *OpenAIGatewayService) openAIGroupRequiresPrivacySet(ctx context.Context, groupID *int64) bool {
|
||||
if cached, ok := ctx.Value(openAIGroupPrivacyRequirementContextKey{}).(openAIGroupPrivacyRequirement); ok && cached.groupID == derefGroupID(groupID) {
|
||||
return cached.required
|
||||
}
|
||||
return s.loadOpenAIGroupRequiresPrivacySet(ctx, groupID)
|
||||
}
|
||||
|
||||
func (s *OpenAIGatewayService) loadOpenAIGroupRequiresPrivacySet(ctx context.Context, groupID *int64) bool {
|
||||
if s == nil || groupID == nil || s.schedulerSnapshot == nil {
|
||||
return false
|
||||
}
|
||||
group, err := s.schedulerSnapshot.GetGroupByID(ctx, *groupID)
|
||||
if err != nil {
|
||||
return true
|
||||
}
|
||||
return group != nil && group.RequirePrivacySet
|
||||
}
|
||||
|
||||
func (s *OpenAIGatewayService) selectAccountWithSchedulerOnce(
|
||||
ctx context.Context,
|
||||
groupID *int64,
|
||||
@@ -2152,6 +2223,7 @@ func (s *OpenAIGatewayService) selectAccountWithSchedulerOnce(
|
||||
useUpstreamTokenCost bool,
|
||||
) (*AccountSelectionResult, OpenAIAccountScheduleDecision, error) {
|
||||
ctx = s.withOpenAIQuotaAutoPauseContext(ctx)
|
||||
ctx = s.withOpenAIGroupPrivacyRequirement(ctx, groupID)
|
||||
// 分组利润控制:唯一文本调度入口的防御性装门。handler 文本
|
||||
// 入口已在请求开始经 WithOpenAIRequestPricingContext 装门并固定 pricingAt,
|
||||
// 此处对同分组门直接复用(failover 重入阈值稳定),仅为不经 handler 装配的
|
||||
@@ -2163,13 +2235,52 @@ func (s *OpenAIGatewayService) selectAccountWithSchedulerOnce(
|
||||
}
|
||||
platform = NormalizeOpenAICompatiblePlatform(platform)
|
||||
decision := OpenAIAccountScheduleDecision{}
|
||||
preserveGuardianParentBinding := preserveOpenAIGuardianParentBinding(ctx, sessionHash)
|
||||
guardianParentAccountID := int64(0)
|
||||
if strings.TrimSpace(previousResponseID) == "" {
|
||||
guardianParentAccountID = s.resolveOpenAIGuardianParentAccountID(ctx, groupID)
|
||||
}
|
||||
scheduler := s.getOpenAIAccountScheduler(ctx)
|
||||
if scheduler == nil {
|
||||
decision.Layer = openAIAccountScheduleLayerLoadBalance
|
||||
if guardianParentAccountID > 0 {
|
||||
if s.checkChannelPricingRestriction(ctx, groupID, requestedModel) {
|
||||
return nil, decision, fmt.Errorf("%w supporting model: %s (channel pricing restriction)", ErrNoAvailableAccounts, requestedModel)
|
||||
}
|
||||
fallbackScheduler := &defaultOpenAIAccountScheduler{service: s, stats: newOpenAIAccountRuntimeStats()}
|
||||
selection, _, err := fallbackScheduler.selectBySessionHash(ctx, OpenAIAccountScheduleRequest{
|
||||
GroupID: groupID,
|
||||
Platform: platform,
|
||||
SessionHash: sessionHash,
|
||||
StickyAccountID: guardianParentAccountID,
|
||||
PreserveStickyBinding: true,
|
||||
RequestedModel: requestedModel,
|
||||
RequiredTransport: requiredTransport,
|
||||
RequiredCapability: requiredCapability,
|
||||
RequiredImageCapability: requiredImageCapability,
|
||||
RequireCompact: requireCompact,
|
||||
ExcludedIDs: excludedIDs,
|
||||
RequirePrivacySet: s.openAIGroupRequiresPrivacySet(ctx, groupID),
|
||||
})
|
||||
if err != nil {
|
||||
return nil, decision, err
|
||||
}
|
||||
if selection != nil && selection.Account != nil {
|
||||
decision.Layer = openAIAccountScheduleLayerGuardianParent
|
||||
decision.StickySessionHit = true
|
||||
decision.SelectedAccountID = selection.Account.ID
|
||||
decision.SelectedAccountType = selection.Account.Type
|
||||
return selection, decision, nil
|
||||
}
|
||||
}
|
||||
legacySessionHash := sessionHash
|
||||
if preserveGuardianParentBinding {
|
||||
legacySessionHash = ""
|
||||
}
|
||||
if requiredTransport == OpenAIUpstreamTransportAny || requiredTransport == OpenAIUpstreamTransportHTTPSSE {
|
||||
effectiveExcludedIDs := cloneExcludedAccountIDs(excludedIDs)
|
||||
for {
|
||||
selection, err := s.selectAccountWithLoadAwareness(ctx, groupID, platform, sessionHash, requestedModel, effectiveExcludedIDs, requireCompact, requiredCapability, useUpstreamTokenCost)
|
||||
selection, err := s.selectAccountWithLoadAwareness(ctx, groupID, platform, legacySessionHash, requestedModel, effectiveExcludedIDs, requireCompact, requiredCapability, useUpstreamTokenCost)
|
||||
if err != nil {
|
||||
return nil, decision, err
|
||||
}
|
||||
@@ -2194,7 +2305,7 @@ func (s *OpenAIGatewayService) selectAccountWithSchedulerOnce(
|
||||
|
||||
effectiveExcludedIDs := cloneExcludedAccountIDs(excludedIDs)
|
||||
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 {
|
||||
return nil, decision, err
|
||||
}
|
||||
@@ -2243,9 +2354,12 @@ func (s *OpenAIGatewayService) selectAccountWithSchedulerOnce(
|
||||
Platform: platform,
|
||||
SessionHash: sessionHash,
|
||||
StickyAccountID: stickyAccountID,
|
||||
GuardianParentAccountID: guardianParentAccountID,
|
||||
StickyPreviousAccountID: stickyPreviousAccountID,
|
||||
StickyWeighted: stickyWeighted,
|
||||
SubscriptionPriority: subscriptionPriority,
|
||||
PreserveStickyBinding: preserveGuardianParentBinding,
|
||||
RequirePrivacySet: s.openAIGroupRequiresPrivacySet(ctx, groupID),
|
||||
PreviousResponseID: previousResponseID,
|
||||
PreviousResponseCanMove: previousResponseCanMove,
|
||||
UseUpstreamTokenCost: useUpstreamTokenCost,
|
||||
|
||||
@@ -1043,6 +1043,7 @@ func (s *OpenAIGatewayService) isBetterAccount(candidate, current *Account) bool
|
||||
// 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)
|
||||
ctx = s.withOpenAIGroupPrivacyRequirement(ctx, groupID)
|
||||
// 分组利润控制:legacy 公共入口同样装门,保证不经
|
||||
// selectAccountWithScheduler 的调用方也无法绕过利润准入。
|
||||
ctx = s.withOpenAIProfitControlGate(ctx, groupID)
|
||||
@@ -1509,6 +1510,9 @@ func (s *OpenAIGatewayService) recheckSelectedOpenAIAccountFromDBBeforeProfit(ct
|
||||
}
|
||||
platform = NormalizeOpenAICompatiblePlatform(platform)
|
||||
if s.schedulerSnapshot == nil || s.accountRepo == nil {
|
||||
if s.openAIGroupRequiresPrivacySet(ctx, groupID) && !account.IsPrivacySet() {
|
||||
return nil
|
||||
}
|
||||
if !isOpenAICompatibleAccountEligibleForRequestBeforeProfit(ctx, account, platform, requestedModel, requireCompact, requiredCapability) {
|
||||
return nil
|
||||
}
|
||||
@@ -1531,6 +1535,9 @@ func (s *OpenAIGatewayService) recheckSelectedOpenAIAccountFromDBBeforeProfit(ct
|
||||
if !s.openAIAccountMatchesSchedulingGroup(latest, groupID) {
|
||||
return nil
|
||||
}
|
||||
if s.openAIGroupRequiresPrivacySet(ctx, groupID) && !latest.IsPrivacySet() {
|
||||
return nil
|
||||
}
|
||||
if !isOpenAICompatibleAccountEligibleForRequestBeforeProfit(ctx, latest, platform, requestedModel, requireCompact, requiredCapability) {
|
||||
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
|
||||
// refresh cannot become the new sticky target.
|
||||
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 s.BindStickySession(ctx, groupID, sessionHash, accountID)
|
||||
@@ -358,6 +358,9 @@ func (s *OpenAIGatewayService) BindStickySessionAfterProfitAdmission(ctx context
|
||||
if sessionHash == "" || accountID <= 0 {
|
||||
return nil
|
||||
}
|
||||
if preserveOpenAIGuardianParentBinding(ctx, sessionHash) {
|
||||
return nil
|
||||
}
|
||||
if !gatewayProfitControlGateActive(ctx) {
|
||||
return s.BindStickySession(ctx, groupID, sessionHash, accountID)
|
||||
}
|
||||
|
||||
@@ -596,6 +596,12 @@ func (s *OpenAIGatewayService) resolveAccountByPreviousResponseIDForCapability(
|
||||
_ = store.DeleteResponseAccount(ctx, derefGroupID(groupID), responseID)
|
||||
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)) {
|
||||
_ = store.DeleteResponseAccount(ctx, derefGroupID(groupID), responseID)
|
||||
return 0, nil, "", nil
|
||||
|
||||
Reference in New Issue
Block a user