feat: sync deployed v1.0.0-rc.38-custom.42

Sync production source including routing, billing, Astra detection and public uptime. Update custom build instructions and license packaging.

model/token_routing.go

@@ -0,0 +1,185 @@
+package model
+
+import (
+	"fmt"
+	"slices"
+	"strings"
+
+	"github.com/QuantumNous/new-api/common"
+)
+
+const TokenRoutingContextKey = "token_routing_policy"
+
+type ModelRoutingRule struct {
+	Model  string            `json:"model"`
+	Groups []TokenGroupEntry `json:"groups"`
+}
+
+type TokenRoutingPolicy struct {
+	Version    int                `json:"version"`
+	Mode       string             `json:"mode"`
+	Groups     []TokenGroupEntry  `json:"groups"`
+	ModelRules []ModelRoutingRule `json:"model_rules"`
+}
+
+func (p *TokenRoutingPolicy) Validate(maxGroups int, allowed func(string) bool) error {
+	if p.Version != 1 || (p.Mode != "custom" && p.Mode != "inherit" && p.Mode != "user") {
+		return fmt.Errorf("invalid routing policy version or mode")
+	}
+	if p.Mode == "custom" && len(p.Groups) == 0 {
+		return fmt.Errorf("a group pool must contain at least one group")
+	}
+	if p.Mode != "custom" && len(p.Groups) != 0 {
+		return fmt.Errorf("inherited pools cannot contain custom groups")
+	}
+	if len(p.Groups) > maxGroups || len(p.ModelRules) > 128 {
+		return fmt.Errorf("routing policy exceeds configured limits")
+	}
+	pool := map[string]bool{}
+	for _, entry := range p.Groups {
+		if entry.Group == "" || entry.Group == "auto" || entry.Group != strings.TrimSpace(entry.Group) || pool[entry.Group] || entry.Priority < 0 || entry.Priority > 10000 || !allowed(entry.Group) {
+			return fmt.Errorf("invalid or inaccessible pool group: %s", entry.Group)
+		}
+		pool[entry.Group] = true
+	}
+	models := map[string]bool{}
+	for _, rule := range p.ModelRules {
+		if rule.Model == "" || len(rule.Model) > 256 || strings.TrimSpace(rule.Model) != rule.Model || models[rule.Model] || len(rule.Groups) == 0 || len(rule.Groups) > maxGroups {
+			return fmt.Errorf("invalid or duplicate model routing rule: %s", rule.Model)
+		}
+		models[rule.Model] = true
+		seen := map[string]bool{}
+		for _, entry := range rule.Groups {
+			if !pool[entry.Group] || seen[entry.Group] || entry.Priority < 0 || entry.Priority > 10000 {
+				return fmt.Errorf("model rule references an invalid pool group: %s", entry.Group)
+			}
+			seen[entry.Group] = true
+		}
+	}
+	return nil
+}
+
+// Resolve is shared by routing and the management preview. Model preferences
+// precede remaining defaults; matching priorities stay in the same tier.
+func (p *TokenRoutingPolicy) Resolve(modelName string, inherited []string, allowed func(string) bool) []TokenGroupEntry {
+	defaults := slices.Clone(p.Groups)
+	if p.Mode != "custom" {
+		defaults = nil
+		for i, group := range inherited {
+			defaults = append(defaults, TokenGroupEntry{Group: group, Priority: i})
+		}
+	}
+	slices.SortStableFunc(defaults, func(a, b TokenGroupEntry) int { return a.Priority - b.Priority })
+	pool := map[string]bool{}
+	for _, entry := range defaults {
+		if allowed(entry.Group) {
+			pool[entry.Group] = true
+		}
+	}
+	var preferred []TokenGroupEntry
+	for _, rule := range p.ModelRules {
+		if rule.Model == modelName {
+			preferred = slices.Clone(rule.Groups)
+			break
+		}
+	}
+	slices.SortStableFunc(preferred, func(a, b TokenGroupEntry) int { return a.Priority - b.Priority })
+	result := []TokenGroupEntry{}
+	seen := map[string]bool{}
+	next := 0
+	for _, list := range [][]TokenGroupEntry{preferred, defaults} {
+		last, tier := 0, next
+		haveLast := false
+		for _, entry := range list {
+			if !pool[entry.Group] || seen[entry.Group] {
+				continue
+			}
+			if haveLast && last != entry.Priority {
+				tier++
+			}
+			result = append(result, TokenGroupEntry{Group: entry.Group, Priority: tier})
+			seen[entry.Group] = true
+			last = entry.Priority
+			haveLast = true
+		}
+		if haveLast {
+			next = tier + 1
+		}
+	}
+	return result
+}
+
+func (token *Token) GetRoutingPolicy() (*TokenRoutingPolicy, error) {
+	if token.RoutingPolicy == "" {
+		return nil, nil
+	}
+	var policy TokenRoutingPolicy
+	if err := common.UnmarshalJsonStr(token.RoutingPolicy, &policy); err != nil {
+		return nil, err
+	}
+	if policy.Version != 1 {
+		return nil, fmt.Errorf("unsupported routing policy version")
+	}
+	return &policy, nil
+}
+
+// LegacyRoutingPolicy preserves effective order, including an explicit old
+// pool which already contained the primary group. No model rule is invented.
+func (token *Token) LegacyRoutingPolicy() (*TokenRoutingPolicy, error) {
+	entries, err := token.GetGroupPool()
+	if err != nil {
+		return nil, err
+	}
+	policy := &TokenRoutingPolicy{Version: 1, Mode: "custom", Groups: []TokenGroupEntry{}, ModelRules: []ModelRoutingRule{}}
+	if token.Group == "auto" {
+		if len(entries) == 0 {
+			policy.Mode = "inherit"
+			return policy, nil
+		}
+		policy.Groups = entries
+		return policy, nil
+	}
+	if token.Group == "" {
+		policy.Mode = "user"
+		return policy, nil
+	}
+	if !token.CrossGroupRetry {
+		entries = nil
+	}
+	minPriority := 0
+	hasPrimary := false
+	for _, entry := range entries {
+		minPriority = min(minPriority, entry.Priority)
+		hasPrimary = hasPrimary || entry.Group == token.Group
+	}
+	if !hasPrimary {
+		policy.Groups = append(policy.Groups, TokenGroupEntry{Group: token.Group, Priority: minPriority - 1})
+	}
+	policy.Groups = append(policy.Groups, entries...)
+	policy.Groups = policy.Resolve("", nil, func(string) bool { return true })
+	return policy, nil
+}
+
+func MigrateTokenRoutingPolicies() (int, error) {
+	var tokens []Token
+	if err := DB.Where("routing_policy IS NULL OR routing_policy = ?", "").Find(&tokens).Error; err != nil {
+		return 0, err
+	}
+	count := 0
+	for _, token := range tokens {
+		policy, err := token.LegacyRoutingPolicy()
+		if err != nil {
+			return count, fmt.Errorf("token %d: %w", token.Id, err)
+		}
+		body, err := common.Marshal(policy)
+		if err != nil {
+			return count, err
+		}
+		r := DB.Model(&Token{}).Where("id = ? AND (routing_policy IS NULL OR routing_policy = ?)", token.Id, "").Update("routing_policy", string(body))
+		if r.Error != nil {
+			return count, r.Error
+		}
+		count += int(r.RowsAffected)
+	}
+	return count, nil
+}