thread-master/apps/backend/internal/module/usage/usecase/platform_gate.go

package usecase

import (
	"fmt"
	"sync"
	"time"

	"apps/backend/internal/module/usage/domain"

	"github.com/zeromicro/go-zero/core/stores/redis"
)

// PlatformGate decides whether platform-paid keys still have capacity.
// When capacity is exhausted, only BYOK users can proceed (Resolve already prefers BYOK).
type PlatformGate interface {
	// AllowPlatform reports whether a new platform call for meter may start.
	// BYOK path must never call this.
	AllowPlatform(meter string) bool
	// MarkLimited marks platform as limited until until (e.g. after upstream 429).
	// until zero = clear limit.
	MarkLimited(until time.Time)
}

// MemoryPlatformGate — process-local RPM + optional global cool-down.
// Multi-instance deployments should later share via Redis; this covers single-node
// and "ops flip" cool-down after provider rate limits.
type MemoryPlatformGate struct {
	mu sync.Mutex

	// Per-minute caps; 0 = unlimited for that class.
	AIPerMin     int
	SearchPerMin int

	// Global cool-down: when set and now < limitedUntil, all platform denied.
	limitedUntil time.Time

	aiWinStart     time.Time
	aiCount        int
	searchWinStart time.Time
	searchCount    int
}

// NewMemoryPlatformGate builds a gate. Zero limits mean unlimited RPM (cool-down still applies).
func NewMemoryPlatformGate(aiPerMin, searchPerMin int) *MemoryPlatformGate {
	return &MemoryPlatformGate{AIPerMin: aiPerMin, SearchPerMin: searchPerMin}
}

func (g *MemoryPlatformGate) MarkLimited(until time.Time) {
	if g == nil {
		return
	}
	g.mu.Lock()
	defer g.mu.Unlock()
	g.limitedUntil = until
}

func (g *MemoryPlatformGate) AllowPlatform(meter string) bool {
	if g == nil {
		return true
	}
	g.mu.Lock()
	defer g.mu.Unlock()
	now := time.Now()
	if !g.limitedUntil.IsZero() && now.Before(g.limitedUntil) {
		return false
	}
	// clear expired cool-down
	if !g.limitedUntil.IsZero() && !now.Before(g.limitedUntil) {
		g.limitedUntil = time.Time{}
	}
	switch meter {
	case domain.MeterWebSearch:
		return g.takeWindow(now, &g.searchWinStart, &g.searchCount, g.SearchPerMin)
	default:
		// ai_copy / ai_research / ai_image share AI pool
		return g.takeWindow(now, &g.aiWinStart, &g.aiCount, g.AIPerMin)
	}
}

func (g *MemoryPlatformGate) takeWindow(now time.Time, start *time.Time, count *int, limit int) bool {
	if limit <= 0 {
		return true
	}
	if start.IsZero() || now.Sub(*start) >= time.Minute {
		*start = now
		*count = 0
	}
	if *count >= limit {
		return false
	}
	*count++
	return true
}

// AlwaysAllowPlatform is a no-op gate (tests / unlimited).
type AlwaysAllowPlatform struct{}

func (AlwaysAllowPlatform) AllowPlatform(string) bool { return true }
func (AlwaysAllowPlatform) MarkLimited(time.Time)     {}

// RedisPlatformGate coordinates rate limits and provider cooldowns across all
// gateway and worker instances. Redis failures deny platform-paid calls rather
// than bypassing a shared safety limit.
type RedisPlatformGate struct {
	client       *redis.Redis
	aiPerMin     int
	searchPerMin int
}

func NewRedisPlatformGate(conf redis.RedisConf, aiPerMin, searchPerMin int) *RedisPlatformGate {
	return &RedisPlatformGate{
		client: redis.MustNewRedis(conf), aiPerMin: aiPerMin, searchPerMin: searchPerMin,
	}
}

func (g *RedisPlatformGate) MarkLimited(until time.Time) {
	if g == nil || g.client == nil {
		return
	}
	key := "haixun:usage:platform:limited"
	if until.IsZero() || !until.After(time.Now()) {
		_, _ = g.client.Del(key)
		return
	}
	seconds := int(time.Until(until).Seconds())
	if seconds < 1 {
		seconds = 1
	}
	_ = g.client.Setex(key, "1", seconds)
}

func (g *RedisPlatformGate) AllowPlatform(meter string) bool {
	if g == nil || g.client == nil {
		return false
	}
	limit, bucket := g.aiPerMin, "ai"
	if meter == domain.MeterWebSearch {
		limit, bucket = g.searchPerMin, "search"
	}
	if limit <= 0 {
		limited, err := g.client.Exists("haixun:usage:platform:limited")
		return err == nil && !limited
	}
	window := time.Now().UTC().Format("200601021504")
	result, err := g.client.Eval(`
if redis.call('EXISTS', KEYS[1]) == 1 then return 0 end
local n = redis.call('INCR', KEYS[2])
if n == 1 then redis.call('EXPIRE', KEYS[2], 120) end
if n > tonumber(ARGV[1]) then return 0 end
return 1
`, []string{"haixun:usage:platform:limited", fmt.Sprintf("haixun:usage:platform:%s:%s", bucket, window)}, limit)
	if err != nil {
		return false
	}
	switch v := result.(type) {
	case int64:
		return v == 1
	case string:
		return v == "1"
	default:
		return false
	}
}