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 } }