package quota import ( "context" "fmt" "net/http" "sort" "strings" "sync" "time" "ukrrs.com/mopac/harness/internal/config" ) // Decision is the back-pressure verdict for one candidate dispatch. type Decision struct { Action string // "allow" | "defer" Reason string // human/log line; empty on allow unless Peak is set Peak bool // decision was made inside the peak window (context) } // Action constants. const ( ActionAllow = "allow" ActionDefer = "defer" ) func allow(peak bool) Decision { return Decision{Action: ActionAllow, Peak: peak} } func defer_(reason string, peak bool) Decision { return Decision{Action: ActionDefer, Reason: reason, Peak: peak} } // Gate is the loop's back-pressure consultant: quota snapshot (polled or // estimated), peak schedule, thresholds. It never blocks startup on redis // or the provider being down; it degrades to local estimates and allows. type Gate struct { cfg config.QuotaConfig sched Schedule shared *SharedState keys *config.KeyResolver http *http.Client now func() time.Time mu sync.Mutex lastSnap *QuotaSnapshot // last successful poll (cache across failures) local5h float64 // process-local estimate fallbacks (redis down) localWeek float64 win5hID string winWeekID string } // NewGate builds the gate from the [quota] config. shared may be nil // (local-only estimates); keys may be nil when usage_url is unset. func NewGate(cfg config.QuotaConfig, keys *config.KeyResolver, shared *SharedState) (*Gate, error) { sched, err := NewSchedule(cfg.PeakStart, cfg.PeakEnd, cfg.Timezone, cfg.PeakWeekdaysOnly) if err != nil { return nil, fmt.Errorf("quota schedule: %w", err) } return &Gate{ cfg: cfg, sched: sched, shared: shared, keys: keys, http: &http.Client{Timeout: 15 * time.Second}, now: time.Now, }, nil } // SetClock overrides the gate's clock (tests inject a fake). func (g *Gate) SetClock(now func() time.Time) { g.now = now } // InPeak exposes the schedule verdict at the gate's clock. func (g *Gate) InPeak() bool { return g.sched.InPeak(g.now()) } // Schedule returns the parsed peak window (status output). func (g *Gate) Schedule() Schedule { return g.sched } // Snapshot returns the freshest quota state: a live poll when the provider // endpoint is configured and the cache is stale, else the redis-published // snapshot from any instance, else the local cache, else the estimate // synthesized from this instance's recorded consumption. The returned // snapshot is never nil when the gate is enabled. func (g *Gate) Snapshot(ctx context.Context) *QuotaSnapshot { now := g.now() if g.cfg.UsageURL != "" { if s := g.cachedPoll(); s != nil && now.Sub(s.FetchedAt) < time.Duration(g.cfg.PollIntervalSecs)*time.Second { return s } if s, err := g.poll(ctx, now); err == nil { return s } } if g.shared != nil { if s, err := g.shared.LoadSnapshot(g.cfg.Account); err == nil && s != nil && now.Sub(s.FetchedAt) < 2*time.Duration(g.cfg.PollIntervalSecs)*time.Second { return s } } if s := g.cachedPoll(); s != nil { return s } return g.estimateSnapshot(now) } func (g *Gate) cachedPoll() *QuotaSnapshot { g.mu.Lock() defer g.mu.Unlock() return g.lastSnap } func (g *Gate) poll(ctx context.Context, now time.Time) (*QuotaSnapshot, error) { key, err := g.keys.Resolve(ctx, g.cfg.KeyRef) if err != nil { return nil, fmt.Errorf("quota key: %w", err) } body, err := FetchUsage(ctx, g.http, g.cfg.UsageURL, key) if err != nil { return nil, err } snap, err := ParseUsage(g.cfg.Account, now, body) if err != nil { return nil, err } g.mu.Lock() g.lastSnap = snap g.mu.Unlock() _ = g.shared.PublishSnapshot(snap, time.Duration(g.cfg.PollIntervalSecs)*3*time.Second) return snap, nil } // window ids: the 5h bucket advances in fixed 5h steps from the epoch; the // weekly bucket is the Monday 00:00 of the schedule's TZ (z.ai weekly // credits reset 7 days after activation — an approximation until the usage // endpoint ships reset_at, noted in the REPORT). func (g *Gate) windowIDs(now time.Time) (w5h, week string) { five := 5 * time.Hour w5h = fmt.Sprintf("%d", now.Unix()/int64(five.Seconds())) mon := now.In(g.sched.Loc) for mon.Weekday() != time.Monday { mon = mon.AddDate(0, 0, -1) } y, m, d := mon.Date() week = fmt.Sprintf("%04d%02d%02d", y, int(m), d) return w5h, week } // RecordTurn records one completed turn's estimated credits into the shared // state (redis when configured; process-local maps always, so a down redis // never loses this instance's own accounting). func (g *Gate) RecordTurn(in EstimateTurnInput) float64 { credits := EstimateCredits(in.Model, in.PromptTokens, in.CachedTokens, in.CompletionTokens, in.Peak) now := g.now() w5h, week := g.windowIDs(now) g.mu.Lock() if g.win5hID != w5h { // 5h window rolled: reset the local counter g.local5h, g.win5hID = 0, w5h } if g.winWeekID != week { g.localWeek, g.winWeekID = 0, week } g.local5h += credits g.localWeek += credits g.mu.Unlock() if g.shared != nil { k5 := fmt.Sprintf("mopac:quota:%s:est:5h:%s", g.cfg.Account, w5h) kw := fmt.Sprintf("mopac:quota:%s:est:weekly:%s", g.cfg.Account, week) if _, err := g.shared.incrByFloat(k5, credits); err == nil { _, _ = g.shared.command("EXPIRE", k5, "21600") // 5h + 1h slack } _, _ = g.shared.incrByFloat(kw, credits) } return credits } // estimateSnapshot synthesizes quota state from recorded consumption against // the configured plan limits (the no-endpoint mode). func (g *Gate) estimateSnapshot(now time.Time) *QuotaSnapshot { w5h, week := g.windowIDs(now) var used5, usedW float64 if g.shared != nil { if v, err := g.shared.get(fmt.Sprintf("mopac:quota:%s:est:5h:%s", g.cfg.Account, w5h)); err == nil && v != "" { fmt.Sscanf(v, "%g", &used5) } if v, err := g.shared.get(fmt.Sprintf("mopac:quota:%s:est:weekly:%s", g.cfg.Account, week)); err == nil && v != "" { fmt.Sscanf(v, "%g", &usedW) } } if used5 == 0 || usedW == 0 { g.mu.Lock() if g.win5hID == w5h { used5 = max(used5, g.local5h) } if g.winWeekID == week { usedW = max(usedW, g.localWeek) } g.mu.Unlock() } src := "estimate" if g.shared != nil { src = "estimate+redis" } return &QuotaSnapshot{ Account: g.cfg.Account, Source: src, FetchedAt: now, Buckets: []Bucket{ {ID: Bucket5h, Used: used5, Limit: g.cfg.Plan5hCredits, WindowReset: now.Add(5 * time.Hour)}, {ID: BucketWeekly, Used: usedW, Limit: g.cfg.PlanWeeklyCredits, WindowReset: nextWeekStart(now, g.sched.Loc)}, }, } } func nextWeekStart(now time.Time, loc *time.Location) time.Time { t := now.In(loc) for t.Weekday() != time.Monday { t = t.AddDate(0, 0, 1) } y, m, d := t.Date() return time.Date(y, m, d, 0, 0, 0, 0, loc).UTC() } // Decide is the pre-dispatch consultation. Priority: hard quota wall first // (the 19:00 failure mode), then peak-window class restriction, then // soft-quota heavy-class deferral. Unknown quota state never defers. func (g *Gate) Decide(ctx context.Context, class string) Decision { peak := g.InPeak() snap := g.Snapshot(ctx) ratio := snap.MaxUsedPct() if ratio >= g.cfg.BlockAtPct { return defer_(fmt.Sprintf("quota: %s bucket at %.0f%% (>= block %.0f%%): all classes deferred until reset (source %s)", worstBucketID(snap), ratio, g.cfg.BlockAtPct, snap.Source), peak) } if peak && !g.peakClass(class) { return defer_(fmt.Sprintf("peak window %s-%s %s: class %q deferred to off-peak (flash-tier classes only: %s)", g.cfg.PeakStart, g.cfg.PeakEnd, g.cfg.Timezone, class, strings.Join(g.cfg.PeakClasses, ",")), peak) } if ratio >= g.cfg.DeferAtPct && !g.peakClass(class) { return defer_(fmt.Sprintf("quota: %s bucket at %.0f%% (>= defer %.0f%%): heavy class %q deferred; LLM-lite continues (source %s)", worstBucketID(snap), ratio, g.cfg.DeferAtPct, class, snap.Source), peak) } return allow(peak) } func (g *Gate) peakClass(class string) bool { for _, c := range g.cfg.PeakClasses { if c == class { return true } } return false } func worstBucketID(s *QuotaSnapshot) string { id, max := "", 0.0 for _, b := range s.Buckets { if p := b.UsedPct(); p > max { max, id = p, b.ID } } return id } // StatusLine renders the one-line quota status for loop startup / `harness // quota status`. func (g *Gate) StatusLine(snap *QuotaSnapshot) string { parts := make([]string, 0, len(snap.Buckets)+2) parts = append(parts, fmt.Sprintf("account=%s source=%s", snap.Account, snap.Source)) for _, b := range snap.Buckets { parts = append(parts, fmt.Sprintf("%s=%.0f/%.0f(%.0f%%)", b.ID, b.Used, b.Limit, b.UsedPct())) } if g.InPeak() { parts = append(parts, "PEAK") } return strings.Join(parts, " ") } // SortedClasses returns PeakClasses sorted (stable status output). func (g *Gate) SortedClasses() []string { out := append([]string{}, g.cfg.PeakClasses...) sort.Strings(out) return out }