From fa4587041ccd7b6de8486c0d40da1028b66d5b28 Mon Sep 17 00:00:00 2001 From: Jarl Date: Sat, 22 Aug 2026 10:24:41 +0800 Subject: [PATCH] fix(openai): keep auto-review on parent account --- .../handler/openai_gateway_handler.go | 4 + .../service/openai_account_scheduler.go | 146 +++++- .../service/openai_gateway_scheduling.go | 7 + .../service/openai_guardian_affinity.go | 133 ++++++ .../service/openai_guardian_affinity_test.go | 450 ++++++++++++++++++ .../internal/service/openai_profit_control.go | 5 +- .../service/openai_ws_forwarder_support.go | 6 + 7 files changed, 734 insertions(+), 17 deletions(-) create mode 100644 backend/internal/service/openai_guardian_affinity.go create mode 100644 backend/internal/service/openai_guardian_affinity_test.go diff --git a/backend/internal/handler/openai_gateway_handler.go b/backend/internal/handler/openai_gateway_handler.go index a11113fe49..d9da0bef71 100644 --- a/backend/internal/handler/openai_gateway_handler.go +++ b/backend/internal/handler/openai_gateway_handler.go @@ -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 diff --git a/backend/internal/service/openai_account_scheduler.go b/backend/internal/service/openai_account_scheduler.go index 16ebde7b62..2f2fef563b 100644 --- a/backend/internal/service/openai_account_scheduler.go +++ b/backend/internal/service/openai_account_scheduler.go @@ -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, diff --git a/backend/internal/service/openai_gateway_scheduling.go b/backend/internal/service/openai_gateway_scheduling.go index 435ec6545a..42e1c3fac8 100644 --- a/backend/internal/service/openai_gateway_scheduling.go +++ b/backend/internal/service/openai_gateway_scheduling.go @@ -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 } diff --git a/backend/internal/service/openai_guardian_affinity.go b/backend/internal/service/openai_guardian_affinity.go new file mode 100644 index 0000000000..30ed388f0c --- /dev/null +++ b/backend/internal/service/openai_guardian_affinity.go @@ -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 +} diff --git a/backend/internal/service/openai_guardian_affinity_test.go b/backend/internal/service/openai_guardian_affinity_test.go new file mode 100644 index 0000000000..4901b36647 --- /dev/null +++ b/backend/internal/service/openai_guardian_affinity_test.go @@ -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() + } +} diff --git a/backend/internal/service/openai_profit_control.go b/backend/internal/service/openai_profit_control.go index 4153bc4094..a3418e6fcf 100644 --- a/backend/internal/service/openai_profit_control.go +++ b/backend/internal/service/openai_profit_control.go @@ -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) } diff --git a/backend/internal/service/openai_ws_forwarder_support.go b/backend/internal/service/openai_ws_forwarder_support.go index 11f54d1568..20ea5eb6b6 100644 --- a/backend/internal/service/openai_ws_forwarder_support.go +++ b/backend/internal/service/openai_ws_forwarder_support.go @@ -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