oadmin/sub2api
1
0
Fork 0
Code Issues Pull Requests 1 Actions Packages Projects Releases Wiki Activity

feat: Beta策略支持按模型区分处理（模型白名单）

Browse Source
This commit is contained in:
shaw
2026-04-07 20:28:14 +08:00
parent b2e379cf7a
commit 7c60ee3c85
8 changed files with 255 additions and 24 deletions
Download Patch File Download Diff File Expand all files Collapse all files

55
backend/internal/service/gateway_service.go
View File

@@ -3946,7 +3946,7 @@ func (s *GatewayService) Forward(ctx context.Context, c *gin.Context, account *A
// Beta policy: evaluate once; block check + cache filter set for buildUpstreamRequest.
// Always overwrite the cache to prevent stale values from a previous retry with a different account.
if account.Platform == PlatformAnthropic && c != nil {
policy := s.evaluateBetaPolicy(ctx, c.GetHeader("anthropic-beta"), account)
policy := s.evaluateBetaPolicy(ctx, c.GetHeader("anthropic-beta"), account, parsed.Model)
if policy.blockErr != nil {
return nil, policy.blockErr
}
@@ -5603,7 +5603,7 @@ func (s *GatewayService) buildUpstreamRequest(ctx context.Context, c *gin.Contex
}
// Build effective drop set: merge static defaults with dynamic beta policy filter rules
policyFilterSet := s.getBetaPolicyFilterSet(ctx, c, account)
policyFilterSet := s.getBetaPolicyFilterSet(ctx, c, account, modelID)
effectiveDropSet := mergeDropSets(policyFilterSet)
effectiveDropWithClaudeCodeSet := mergeDropSets(policyFilterSet, claude.BetaClaudeCode)
@@ -5843,7 +5843,7 @@ type betaPolicyResult struct {
}
// evaluateBetaPolicy loads settings once and evaluates all rules against the given request.
func (s *GatewayService) evaluateBetaPolicy(ctx context.Context, betaHeader string, account *Account) betaPolicyResult {
func (s *GatewayService) evaluateBetaPolicy(ctx context.Context, betaHeader string, account *Account, model string) betaPolicyResult {
if s.settingService == nil {
return betaPolicyResult{}
}
@@ -5858,10 +5858,11 @@ func (s *GatewayService) evaluateBetaPolicy(ctx context.Context, betaHeader stri
if !betaPolicyScopeMatches(rule.Scope, isOAuth, isBedrock) {
continue
}
switch rule.Action {
effectiveAction, effectiveErrMsg := resolveRuleAction(rule, model)
switch effectiveAction {
case BetaPolicyActionBlock:
if result.blockErr == nil && betaHeader != "" && containsBetaToken(betaHeader, rule.BetaToken) {
msg := rule.ErrorMessage
msg := effectiveErrMsg
if msg == "" {
msg = "beta feature " + rule.BetaToken + " is not allowed"
}
@@ -5903,7 +5904,7 @@ const betaPolicyFilterSetKey = "betaPolicyFilterSet"
// In the /v1/messages path, Forward() evaluates the policy first and caches the result;
// buildUpstreamRequest reuses it (zero extra DB calls). In the count_tokens path, this
// evaluates on demand (one DB call).
func (s *GatewayService) getBetaPolicyFilterSet(ctx context.Context, c *gin.Context, account *Account) map[string]struct{} {
func (s *GatewayService) getBetaPolicyFilterSet(ctx context.Context, c *gin.Context, account *Account, model string) map[string]struct{} {
if c != nil {
if v, ok := c.Get(betaPolicyFilterSetKey); ok {
if fs, ok := v.(map[string]struct{}); ok {
@@ -5911,7 +5912,7 @@ func (s *GatewayService) getBetaPolicyFilterSet(ctx context.Context, c *gin.Cont
}
}
}
return s.evaluateBetaPolicy(ctx, "", account).filterSet
return s.evaluateBetaPolicy(ctx, "", account, model).filterSet
}
// betaPolicyScopeMatches checks whether a rule's scope matches the current account type.
@@ -5930,6 +5931,33 @@ func betaPolicyScopeMatches(scope string, isOAuth bool, isBedrock bool) bool {
}
}
// matchModelWhitelist checks if a model matches any pattern in the whitelist.
// Reuses matchModelPattern from group.go which supports exact and wildcard prefix matching.
func matchModelWhitelist(model string, whitelist []string) bool {
for _, pattern := range whitelist {
if matchModelPattern(pattern, model) {
return true
}
}
return false
}
// resolveRuleAction determines the effective action and error message for a rule given the request model.
// When ModelWhitelist is empty, the rule's primary Action/ErrorMessage applies unconditionally.
// When non-empty, Action applies to matching models; FallbackAction/FallbackErrorMessage applies to others.
func resolveRuleAction(rule BetaPolicyRule, model string) (action, errorMessage string) {
if len(rule.ModelWhitelist) == 0 {
return rule.Action, rule.ErrorMessage
}
if matchModelWhitelist(model, rule.ModelWhitelist) {
return rule.Action, rule.ErrorMessage
}
if rule.FallbackAction != "" {
return rule.FallbackAction, rule.FallbackErrorMessage
}
return BetaPolicyActionPass, "" // default fallback: pass (fail-open)
}
// droppedBetaSet returns claude.DroppedBetas as a set, with optional extra tokens.
func droppedBetaSet(extra ...string) map[string]struct{} {
m := make(map[string]struct{}, len(defaultDroppedBetasSet)+len(extra))
@@ -5976,7 +6004,7 @@ func (s *GatewayService) resolveBedrockBetaTokensForRequest(
modelID string,
) ([]string, error) {
// 1. 对原始 header 中的 beta token 做 block 检查（快速失败）
policy := s.evaluateBetaPolicy(ctx, betaHeader, account)
policy := s.evaluateBetaPolicy(ctx, betaHeader, account, modelID)
if policy.blockErr != nil {
return nil, policy.blockErr
}
@@ -5988,7 +6016,7 @@ func (s *GatewayService) resolveBedrockBetaTokensForRequest(
// 例如：管理员 block 了 interleaved-thinking，客户端不在 header 中带该 token，
// 但请求体中包含 thinking 字段 → autoInjectBedrockBetaTokens 会自动补齐 →
// 如果不做此检查，block 规则会被绕过。
if blockErr := s.checkBetaPolicyBlockForTokens(ctx, betaTokens, account); blockErr != nil {
if blockErr := s.checkBetaPolicyBlockForTokens(ctx, betaTokens, account, modelID); blockErr != nil {
return nil, blockErr
}
@@ -5997,7 +6025,7 @@ func (s *GatewayService) resolveBedrockBetaTokensForRequest(
// checkBetaPolicyBlockForTokens 检查 token 列表中是否有被管理员 block 规则命中的 token。
// 用于补充 evaluateBetaPolicy 对 header 的检查，覆盖 body 自动注入的 token。
func (s *GatewayService) checkBetaPolicyBlockForTokens(ctx context.Context, tokens []string, account *Account) *BetaBlockedError {
func (s *GatewayService) checkBetaPolicyBlockForTokens(ctx context.Context, tokens []string, account *Account, model string) *BetaBlockedError {
if s.settingService == nil || len(tokens) == 0 {
return nil
}
@@ -6009,14 +6037,15 @@ func (s *GatewayService) checkBetaPolicyBlockForTokens(ctx context.Context, toke
isBedrock := account.IsBedrock()
tokenSet := buildBetaTokenSet(tokens)
for _, rule := range settings.Rules {
if rule.Action != BetaPolicyActionBlock {
effectiveAction, effectiveErrMsg := resolveRuleAction(rule, model)
if effectiveAction != BetaPolicyActionBlock {
continue
}
if !betaPolicyScopeMatches(rule.Scope, isOAuth, isBedrock) {
continue
}
if _, present := tokenSet[rule.BetaToken]; present {
msg := rule.ErrorMessage
msg := effectiveErrMsg
if msg == "" {
msg = "beta feature " + rule.BetaToken + " is not allowed"
}
@@ -8474,7 +8503,7 @@ func (s *GatewayService) buildCountTokensRequest(ctx context.Context, c *gin.Con
}
// Build effective drop set for count_tokens: merge static defaults with dynamic beta policy filter rules
ctEffectiveDropSet := mergeDropSets(s.getBetaPolicyFilterSet(ctx, c, account))
ctEffectiveDropSet := mergeDropSets(s.getBetaPolicyFilterSet(ctx, c, account, modelID))
// OAuth 账号：处理 anthropic-beta header
if tokenType == "oauth" {

12
backend/internal/service/setting_service.go
View File

@@ -1527,6 +1527,18 @@ func (s *SettingService) SetBetaPolicySettings(ctx context.Context, settings *Be
if !validScopes[rule.Scope] {
return fmt.Errorf("rule[%d]: invalid scope %q", i, rule.Scope)
}
// Validate model_whitelist patterns
for j, pattern := range rule.ModelWhitelist {
trimmed := strings.TrimSpace(pattern)
if trimmed == "" {
return fmt.Errorf("rule[%d]: model_whitelist[%d] cannot be empty", i, j)
}
settings.Rules[i].ModelWhitelist[j] = trimmed
}
// Validate fallback_action
if rule.FallbackAction != "" && !validActions[rule.FallbackAction] {
return fmt.Errorf("rule[%d]: invalid fallback_action %q", i, rule.FallbackAction)
}
}
data, err := json.Marshal(settings)

11
backend/internal/service/settings_view.go
View File

@@ -178,10 +178,13 @@ const (
// BetaPolicyRule 单条 Beta 策略规则
type BetaPolicyRule struct {
BetaToken string `json:"beta_token"` // beta token 值
Action string `json:"action"` // "pass" | "filter" | "block"
Scope string `json:"scope"` // "all" | "oauth" | "apikey" | "bedrock"
ErrorMessage string `json:"error_message,omitempty"` // 自定义错误消息 (action=block 时生效)
BetaToken string `json:"beta_token"` // beta token 值
Action string `json:"action"` // "pass" | "filter" | "block"
Scope string `json:"scope"` // "all" | "oauth" | "apikey" | "bedrock"
ErrorMessage string `json:"error_message,omitempty"` // 自定义错误消息 (action=block 时生效)
ModelWhitelist []string `json:"model_whitelist,omitempty"` // 模型匹配模式列表（为空=对所有模型生效）
FallbackAction string `json:"fallback_action,omitempty"` // 未匹配白名单的模型的处理方式
FallbackErrorMessage string `json:"fallback_error_message,omitempty"` // 未匹配白名单时的自定义错误消息 (fallback_action=block 时生效)
}
// BetaPolicySettings Beta 策略配置