fix(openai): keep auto-review on parent account

This commit is contained in:
Jarl
2026-08-22 10:24:41 +08:00
parent 67380eafd5
commit fa4587041c
7 changed files with 734 additions and 17 deletions
@@ -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