easyai-ai-gateway/apps/api/internal/store/rate_limit_policy.go

package store

import (
	"errors"
	"fmt"
	"math"
	"strings"
)

var ErrInvalidRateLimitPolicy = errors.New("invalid rate limit policy")

func IsInvalidRateLimitPolicy(err error) bool {
	return errors.Is(err, ErrInvalidRateLimitPolicy)
}

const (
	RateLimitPolicyModeInherit  = "inherit"
	RateLimitPolicyModeOverride = "override"
	DefaultQueueMaxWaitSeconds  = 600
	MaxQueueSize                = 10000
	MaxQueueWaitSeconds         = 3600
)

type EffectiveRateLimitPolicyInput struct {
	BasePolicy            map[string]any
	PlatformPolicy        map[string]any
	RuntimePolicy         map[string]any
	RuntimePolicyExplicit bool
	RuntimePolicyOverride map[string]any
	ModelPolicy           map[string]any
	ModelPolicyMode       string
}

// EffectiveRateLimitPolicy resolves one authoritative policy. Platform
// policies are defaults for every bound model; explicit model/runtime settings
// replace the complete policy instead of merging individual metrics.
func EffectiveRateLimitPolicy(input EffectiveRateLimitPolicyInput) map[string]any {
	policy := input.BasePolicy
	if policySpecified(input.PlatformPolicy) {
		policy = input.PlatformPolicy
	}
	if input.RuntimePolicyExplicit {
		policy = input.RuntimePolicy
	}
	if raw, ok := input.RuntimePolicyOverride["rateLimitPolicy"]; ok {
		policy, _ = raw.(map[string]any)
	}
	mode := NormalizeRateLimitPolicyMode(input.ModelPolicyMode, input.ModelPolicy != nil)
	if mode == RateLimitPolicyModeOverride {
		policy = input.ModelPolicy
	}
	return NormalizeRateLimitPolicy(policy)
}

func NormalizeRateLimitPolicyMode(mode string, policyProvided bool) string {
	switch strings.ToLower(strings.TrimSpace(mode)) {
	case RateLimitPolicyModeOverride:
		return RateLimitPolicyModeOverride
	case RateLimitPolicyModeInherit:
		return RateLimitPolicyModeInherit
	default:
		if policyProvided {
			return RateLimitPolicyModeOverride
		}
		return RateLimitPolicyModeInherit
	}
}

func RateLimitPolicyMetric(policy map[string]any, metric string) (float64, bool) {
	rules, _ := NormalizeRateLimitPolicy(policy)["rules"].([]any)
	for _, rawRule := range rules {
		rule, _ := rawRule.(map[string]any)
		if strings.TrimSpace(stringValue(rule["metric"])) != metric {
			continue
		}
		value := floatValue(rule["limit"])
		return value, true
	}
	return 0, false
}

// NormalizeRateLimitPolicy keeps the canonical rules contract while accepting
// policies imported from server-main before that contract existed. Runtime
// enforcement and worker sizing must agree on these legacy limits; otherwise a
// configured max_concurrent_requests silently becomes unlimited.
func NormalizeRateLimitPolicy(policy map[string]any) map[string]any {
	if policy == nil {
		return nil
	}
	out := clonePolicy(policy)
	rules, _ := out["rules"].([]any)
	if len(rules) == 0 {
		legacyScopes := []map[string]any{policy}
		for _, key := range []string{"platformLimits", "modelLimits", "platform_limits", "model_limits"} {
			if nested, ok := policy[key].(map[string]any); ok {
				legacyScopes = append(legacyScopes, nested)
			}
		}
		if limit, ok := lowestPositiveLegacyLimit(legacyScopes,
			"max_concurrent_requests", "maxConcurrentRequests", "concurrent"); ok {
			rules = append(rules, map[string]any{
				"metric":          "concurrent",
				"limit":           limit,
				"leaseTtlSeconds": 120,
			})
		}
		if limit, ok := lowestPositiveLegacyLimit(legacyScopes,
			"max_request_per_minute", "maxRequestsPerMinute", "rpm"); ok {
			rules = append(rules, map[string]any{
				"metric":        "rpm",
				"limit":         limit,
				"windowSeconds": 60,
			})
		}
		if limit, ok := lowestPositiveLegacyLimit(legacyScopes,
			"max_tokens_per_minute", "maxTokensPerMinute", "tpm_total"); ok {
			rules = append(rules, map[string]any{
				"metric":        "tpm_total",
				"limit":         limit,
				"windowSeconds": 60,
			})
		}
	}
	normalizedRules := make([]any, 0, len(rules))
	for _, rawRule := range rules {
		rule, ok := rawRule.(map[string]any)
		if !ok {
			normalizedRules = append(normalizedRules, rawRule)
			continue
		}
		normalizedRule := clonePolicy(rule)
		if strings.TrimSpace(stringValue(normalizedRule["metric"])) == "queue_size" {
			if floatValue(normalizedRule["limit"]) <= 0 {
				continue
			}
			if floatValue(normalizedRule["maxWaitSeconds"]) <= 0 {
				normalizedRule["maxWaitSeconds"] = DefaultQueueMaxWaitSeconds
			}
		} else {
			delete(normalizedRule, "maxWaitSeconds")
		}
		normalizedRules = append(normalizedRules, normalizedRule)
	}
	out["rules"] = normalizedRules
	return out
}

type QueuePolicyRule struct {
	Limit          int
	MaxWaitSeconds int
	Policy         map[string]any
}

func QueueRuleFromPolicy(policy map[string]any) (QueuePolicyRule, bool) {
	normalized := NormalizeRateLimitPolicy(policy)
	rules, _ := normalized["rules"].([]any)
	for _, rawRule := range rules {
		rule, _ := rawRule.(map[string]any)
		if strings.TrimSpace(stringValue(rule["metric"])) != "queue_size" {
			continue
		}
		limit := int(math.Floor(floatValue(rule["limit"])))
		if limit <= 0 {
			return QueuePolicyRule{}, false
		}
		maxWaitSeconds := int(math.Floor(floatValue(rule["maxWaitSeconds"])))
		if maxWaitSeconds <= 0 {
			maxWaitSeconds = DefaultQueueMaxWaitSeconds
		}
		return QueuePolicyRule{
			Limit:          limit,
			MaxWaitSeconds: maxWaitSeconds,
			Policy:         normalized,
		}, true
	}
	return QueuePolicyRule{}, false
}

// NormalizeAndValidateRateLimitPolicy validates only the canonical queue
// extension. Existing rate-limit metrics intentionally keep their historical
// permissive parsing contract.
func NormalizeAndValidateRateLimitPolicy(policy map[string]any) (map[string]any, error) {
	if policy == nil {
		return nil, nil
	}
	rulesValue, rulesPresent := policy["rules"]
	if rulesPresent {
		if _, ok := rulesValue.([]any); !ok {
			return nil, fmt.Errorf("%w: rateLimitPolicy.rules must be an array", ErrInvalidRateLimitPolicy)
		}
	}
	rules, _ := rulesValue.([]any)
	for index, rawRule := range rules {
		rule, ok := rawRule.(map[string]any)
		if !ok {
			return nil, fmt.Errorf("%w: rateLimitPolicy.rules[%d] must be an object", ErrInvalidRateLimitPolicy, index)
		}
		metric := strings.TrimSpace(stringValue(rule["metric"]))
		maxWait, hasMaxWait := numericRuleValue(rule, "maxWaitSeconds")
		if metric != "queue_size" {
			if hasMaxWait {
				return nil, fmt.Errorf("%w: rateLimitPolicy.rules[%d].maxWaitSeconds is only valid for queue_size", ErrInvalidRateLimitPolicy, index)
			}
			continue
		}
		limit, hasLimit := numericRuleValue(rule, "limit")
		if !hasLimit {
			return nil, fmt.Errorf("%w: rateLimitPolicy.rules[%d].limit must be an integer", ErrInvalidRateLimitPolicy, index)
		}
		if limit != math.Trunc(limit) || limit < 0 || limit > MaxQueueSize {
			return nil, fmt.Errorf("%w: queue_size limit must be 0 or an integer between 1 and %d", ErrInvalidRateLimitPolicy, MaxQueueSize)
		}
		if limit == 0 {
			continue
		}
		if hasMaxWait && (maxWait != math.Trunc(maxWait) || maxWait < 1 || maxWait > MaxQueueWaitSeconds) {
			return nil, fmt.Errorf("%w: queue_size maxWaitSeconds must be an integer between 1 and %d", ErrInvalidRateLimitPolicy, MaxQueueWaitSeconds)
		}
	}
	return NormalizeRateLimitPolicy(policy), nil
}

func numericRuleValue(rule map[string]any, key string) (float64, bool) {
	raw, ok := rule[key]
	if !ok || raw == nil {
		return 0, false
	}
	switch raw.(type) {
	case float64, float32, int, int8, int16, int32, int64, uint, uint8, uint16, uint32, uint64:
		return floatValue(raw), true
	default:
		return 0, false
	}
}

func ConcurrentPolicyCapacity(policy map[string]any) (int, bool) {
	limit, ok := RateLimitPolicyMetric(policy, "concurrent")
	if !ok || limit <= 0 {
		return 0, false
	}
	if limit < 1 {
		return 1, true
	}
	if limit >= float64(math.MaxInt) {
		return math.MaxInt, true
	}
	return int(math.Floor(limit)), true
}

func policySpecified(policy map[string]any) bool {
	rules, _ := NormalizeRateLimitPolicy(policy)["rules"].([]any)
	return len(rules) > 0
}

func lowestPositiveLegacyLimit(scopes []map[string]any, keys ...string) (float64, bool) {
	limit := 0.0
	found := false
	for _, scope := range scopes {
		for _, key := range keys {
			value := floatValue(scope[key])
			if value <= 0 {
				continue
			}
			if !found || value < limit {
				limit = value
				found = true
			}
		}
	}
	return limit, found
}

func clonePolicy(policy map[string]any) map[string]any {
	if policy == nil {
		return nil
	}
	out := make(map[string]any, len(policy))
	for key, value := range policy {
		out[key] = value
	}
	return out
}