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

package store

import (
	"math"
	"strings"
)

const (
	RateLimitPolicyModeInherit  = "inherit"
	RateLimitPolicyModeOverride = "override"
)

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 {
		return out
	}

	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,
		})
	}
	if len(rules) > 0 {
		out["rules"] = rules
	}
	return out
}

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
}