97 lines
2.6 KiB
Go
97 lines
2.6 KiB
Go
|
package usecase
|
||
|
|
|
||
|
|
import (
|
||
|
|
"sync"
|
||
|
|
"time"
|
||
|
|
|
||
|
|
"apps/backend/internal/module/usage/domain"
|
||
|
|
)
|
||
|
|
|
||
|
|
// 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) {}
|