mrcharles
fc518c475e
The loop now consults quota and host state before every dispatch and
DEFERS gated work with a logged reason instead of letting turns die at
the provider (the 2026-08-28 19:00 quota-wall failure mode, replayed as
a test). Adds internal/quota: 5h/weekly credit buckets (provider poll
when z.ai ships an endpoint - fake-server tested - else locally
estimated from the documented credit formula), TZ-aware peak window
(default 01:00-05:00 America/Chicago weekdays, matching the documented
z.ai peak Mon-Fri 14:00-18:00 Singapore), block/defer thresholds, a
read-only load/mem/disk/IO-PSI monitor, an optional redis shared-state
hop (stdlib RESP2 mini-client) so all instances of an account
coordinate, per-class token+credit accounting in loop.jsonl, and
`harness quota status|probe|gate`. Config: [quota] + [resources]
sections; README runbook covers the redis container and deploy-time
cgroup enforcement.
💘 Generated with Crush
Assisted-by: Crush:glm-5.2
407 lines
14 KiB
Go
407 lines
14 KiB
Go
package quota
|
|
|
|
import (
|
|
"context"
|
|
"fmt"
|
|
"net/http"
|
|
"net/http/httptest"
|
|
"strings"
|
|
"testing"
|
|
"time"
|
|
|
|
"ukrrs.com/mopac/harness/internal/config"
|
|
)
|
|
|
|
func TestParseUsage(t *testing.T) {
|
|
ok := `{"usage":{"five_hour":{"used_credits":15400,"total_credits":28000,"reset_at":"2026-08-29T22:00:00Z"},
|
|
"weekly":{"used_credits":130200,"total_credits":140000,"reset_at":"2026-09-02T00:00:00Z"}}}`
|
|
alias := `{"usage":{"five_hour":{"credits_used":100,"limit_credits":200,"reset_time":"2026-08-29T22:00:00Z"},
|
|
"weekly":{"credits_used":"150","limit_credits":"200","reset_time":"2026-09-02T00:00:00Z"}}}`
|
|
bad := []struct {
|
|
name, body string
|
|
}{
|
|
{"html error page", `<html><body>502 Bad Gateway</body></html>`},
|
|
{"missing buckets", `{"usage":{}}`},
|
|
{"missing used", `{"usage":{"five_hour":{"total_credits":28000},"weekly":{"used_credits":1,"total_credits":2}}}`},
|
|
{"bad reset", `{"usage":{"five_hour":{"used_credits":1,"total_credits":2,"reset_at":"tomorrow"},"weekly":{"used_credits":1,"total_credits":2}}}`},
|
|
}
|
|
cases := []struct {
|
|
name string
|
|
body string
|
|
want5h float64
|
|
want5p float64
|
|
wantWkd float64
|
|
wantErr bool
|
|
}{
|
|
{"canonical shape", ok, 15400, 55.0, 93.0, false},
|
|
{"field aliases + string numbers", alias, 100, 50, 75, false},
|
|
}
|
|
for _, c := range cases {
|
|
t.Run(c.name, func(t *testing.T) {
|
|
snap, err := ParseUsage("zai-1", time.Now(), []byte(c.body))
|
|
if err != nil {
|
|
t.Fatalf("ParseUsage: %v", err)
|
|
}
|
|
b5, ok := snap.Bucket(Bucket5h)
|
|
if !ok || b5.Used != c.want5h || !nearly(b5.UsedPct(), c.want5p) {
|
|
t.Errorf("5h bucket = %+v, want used=%v pct=%v", b5, c.want5h, c.want5p)
|
|
}
|
|
bw, _ := snap.Bucket(BucketWeekly)
|
|
if !nearly(bw.UsedPct(), c.wantWkd) {
|
|
t.Errorf("weekly pct = %v, want %v", bw.UsedPct(), c.wantWkd)
|
|
}
|
|
if got := snap.MaxUsedPct(); !nearly(got, c.wantWkd) {
|
|
t.Errorf("MaxUsedPct = %v, want %v", got, c.wantWkd)
|
|
}
|
|
})
|
|
}
|
|
for _, c := range bad {
|
|
t.Run(c.name, func(t *testing.T) {
|
|
if _, err := ParseUsage("zai-1", time.Now(), []byte(c.body)); err == nil {
|
|
t.Fatal("expected error, got nil (a garbage body must never read as 'quota fine')")
|
|
}
|
|
})
|
|
}
|
|
}
|
|
|
|
func nearly(a, b float64) bool {
|
|
d := a - b
|
|
return d < 1e-9 && d > -1e-9
|
|
}
|
|
|
|
func TestParseUsageZeroLimitNeverBlocks(t *testing.T) {
|
|
snap, err := ParseUsage("a", time.Now(), []byte(`{"usage":{"five_hour":{"used_credits":999,"total_credits":0,"reset_at":""},"weekly":{"used_credits":1,"total_credits":2}}}`))
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if b, _ := snap.Bucket(Bucket5h); b.UsedPct() != 0 {
|
|
t.Errorf("zero-limit bucket must read 0%%, got %v", b.UsedPct())
|
|
}
|
|
}
|
|
|
|
func TestGateUsesBearerAndNeverLogsKey(t *testing.T) {
|
|
var gotAuth string
|
|
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
gotAuth = r.Header.Get("Authorization")
|
|
fmt.Fprintf(w, `{"usage":{"five_hour":{"used_credits":1,"total_credits":28000,"reset_at":"2026-08-29T22:00:00Z"},"weekly":{"used_credits":1,"total_credits":140000,"reset_at":"2026-09-02T00:00:00Z"}}}`)
|
|
}))
|
|
defer srv.Close()
|
|
|
|
cfg := gateTestCfg()
|
|
cfg.UsageURL = srv.URL
|
|
cfg.KeyRef = "literal:seekrit-zai-key"
|
|
gate := newGate(t, cfg)
|
|
snap := gate.Snapshot(context.Background())
|
|
if snap.Source != "usage_url" {
|
|
t.Fatalf("source = %s, want usage_url", snap.Source)
|
|
}
|
|
if gotAuth != "Bearer seekrit-zai-key" {
|
|
t.Errorf("auth header = %q", gotAuth)
|
|
}
|
|
// Poll failures (5xx) must fall back without panicking and never leak
|
|
// the key into the returned error strings.
|
|
srv.Config.Handler = http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
http.Error(w, "boom", http.StatusInternalServerError)
|
|
})
|
|
gate2 := newGate(t, cfg)
|
|
snap2 := gate2.Snapshot(context.Background())
|
|
if snap2 == nil {
|
|
t.Fatal("nil snapshot on poll failure; gate must degrade to estimates")
|
|
}
|
|
if strings.Contains(snap2.Source, "seekrit") {
|
|
t.Errorf("key leaked into snapshot source %q", snap2.Source)
|
|
}
|
|
}
|
|
|
|
func TestGateSharedSnapshotRoundTrip(t *testing.T) {
|
|
fake := newFakeRedis(t)
|
|
st, err := NewSharedState(fake.url(t))
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
defer st.Close()
|
|
|
|
snap := &QuotaSnapshot{Account: "zai-1", Source: "usage_url", FetchedAt: time.Now(),
|
|
Buckets: []Bucket{{ID: Bucket5h, Used: 1, Limit: 28000}}}
|
|
if err := st.PublishSnapshot(snap, time.Minute); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
got, err := st.LoadSnapshot("zai-1")
|
|
if err != nil || got == nil || got.Source != "usage_url" {
|
|
t.Fatalf("round trip: %v %+v", err, got)
|
|
}
|
|
if b, _ := got.Bucket(Bucket5h); b.Used != 1 {
|
|
t.Errorf("bucket lost: %+v", got.Buckets)
|
|
}
|
|
}
|
|
|
|
func TestEstimateCredits(t *testing.T) {
|
|
cases := []struct {
|
|
name string
|
|
model string
|
|
prompt, cached, comp int
|
|
peak bool
|
|
want float64
|
|
}{
|
|
// flagship: (10000*6.9 + 2000*1.7 + 10000*24)/10000 = 31.24
|
|
{"flagship peak", "glm-5.3", 10000, 2000, 10000, true, 31.24},
|
|
// off-peak halves: 15.62
|
|
{"flagship off-peak", "glm-5.3", 10000, 2000, 10000, false, 15.62},
|
|
// flash: (10000*2.3 + 2000*0.56 + 10000*8)/10000 = 10.412 -> /2
|
|
{"flash off-peak", "glm-4.7-flash", 10000, 2000, 10000, false, 5.206},
|
|
{"flash peak", "GLM-5.3-Flash", 10000, 0, 10000, true, 10.3},
|
|
}
|
|
for _, c := range cases {
|
|
t.Run(c.name, func(t *testing.T) {
|
|
got := EstimateCredits(c.model, c.prompt, c.cached, c.comp, c.peak)
|
|
if diff := got - c.want; diff > 1e-9 || diff < -1e-9 {
|
|
t.Errorf("credits = %v, want %v", got, c.want)
|
|
}
|
|
})
|
|
}
|
|
if MultipliersFor("glm-9.9-omega") != FlagshipMultipliers() {
|
|
t.Error("unknown model must read as flagship (conservative)")
|
|
}
|
|
}
|
|
|
|
func TestScheduleEdges(t *testing.T) {
|
|
cst, err := time.LoadLocation("America/Chicago")
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
at := func(day string, hm string) time.Time {
|
|
tm, err := time.ParseInLocation("2006-01-02 15:04", day+" "+hm, cst)
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
return tm
|
|
}
|
|
// Default window 01:00-05:00 CST, weekdays only.
|
|
sched, err := NewSchedule("01:00", "05:00", "America/Chicago", true)
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
cases := []struct {
|
|
name string
|
|
t time.Time
|
|
want bool
|
|
}{
|
|
{"just before start (Fri)", at("2026-08-28", "00:59"), false},
|
|
{"at start (Fri)", at("2026-08-28", "01:00"), true},
|
|
{"mid window (Fri)", at("2026-08-28", "03:00"), true},
|
|
{"last minute (Fri)", at("2026-08-28", "04:59"), true},
|
|
{"at end (Fri)", at("2026-08-28", "05:00"), false},
|
|
{"evening off-peak (Fri)", at("2026-08-28", "19:00"), false},
|
|
{"weekend inside window (Sat)", at("2026-08-29", "02:00"), false},
|
|
{"weekend inside window (Sun)", at("2026-08-30", "02:00"), false},
|
|
{"Monday early window", at("2026-08-31", "02:00"), true},
|
|
}
|
|
for _, c := range cases {
|
|
t.Run(c.name, func(t *testing.T) {
|
|
if got := sched.InPeak(c.t); got != c.want {
|
|
t.Errorf("InPeak(%s) = %v, want %v", c.t, got, c.want)
|
|
}
|
|
})
|
|
}
|
|
|
|
// Wrapped window 22:00-02:00, weekdays only: evening part carries the
|
|
// start day's weekday; morning part checks the PREVIOUS day.
|
|
wrap, err := NewSchedule("22:00", "02:00", "America/Chicago", true)
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
wrapCases := []struct {
|
|
name string
|
|
t time.Time
|
|
want bool
|
|
}{
|
|
{"Fri 23:00 (started Fri)", at("2026-08-28", "23:00"), true},
|
|
{"Sat 01:00 (window started Fri)", at("2026-08-29", "01:00"), true},
|
|
{"Sat 23:00 (started Sat)", at("2026-08-29", "23:00"), false},
|
|
{"Sun 01:00 (window started Sat)", at("2026-08-30", "01:00"), false},
|
|
}
|
|
for _, c := range wrapCases {
|
|
t.Run("wrap: "+c.name, func(t *testing.T) {
|
|
if got := wrap.InPeak(c.t); got != c.want {
|
|
t.Errorf("InPeak(%s) = %v, want %v", c.t, got, c.want)
|
|
}
|
|
})
|
|
}
|
|
|
|
// Every-day window ignores weekday flags.
|
|
daily, _ := NewSchedule("01:00", "05:00", "America/Chicago", false)
|
|
if !daily.InPeak(at("2026-08-29", "02:00")) { // Saturday
|
|
t.Error("weekdays_only=false must peak on Saturday too")
|
|
}
|
|
|
|
// UTC window evaluated on the INSTANT: 21:00 UTC is outside, but the
|
|
// same wall-clock digits in CST (21:00 CDT == 02:00 UTC) are inside.
|
|
utc, _ := NewSchedule("01:00", "05:00", "UTC", false)
|
|
utcT := time.Date(2026, 8, 28, 21, 0, 0, 0, time.UTC)
|
|
if utc.InPeak(utcT) {
|
|
t.Error("21:00 UTC must be outside the UTC window")
|
|
}
|
|
cstWall := time.Date(2026, 8, 28, 21, 0, 0, 0, cst)
|
|
if !utc.InPeak(cstWall) {
|
|
t.Error("21:00 CST (== 02:00 UTC) must be inside the UTC window")
|
|
}
|
|
|
|
if _, err := NewSchedule("25:00", "05:00", "UTC", false); err == nil {
|
|
t.Error("25:00 must fail to parse")
|
|
}
|
|
if _, err := NewSchedule("01:00", "05:00", "Mars/Olympus", false); err == nil {
|
|
t.Error("unknown timezone must fail")
|
|
}
|
|
}
|
|
|
|
func TestRecordTurnFeedsEstimateSnapshot(t *testing.T) {
|
|
fake := newFakeRedis(t)
|
|
cfg := gateTestCfg()
|
|
cfg.RedisURL = fake.url(t)
|
|
gate := newGate(t, cfg)
|
|
// Friday 03:00 CST: inside the default peak window.
|
|
fri := time.Date(2026, 8, 28, 3, 0, 0, 0, mustLoc(t, "America/Chicago"))
|
|
gate.SetClock(func() time.Time { return fri })
|
|
|
|
credits := gate.RecordTurn(EstimateTurnInput{Model: "glm-5.3", PromptTokens: 10000, CachedTokens: 2000, CompletionTokens: 10000, Peak: true})
|
|
if credits != 31.24 {
|
|
t.Fatalf("credits = %v, want 31.24", credits)
|
|
}
|
|
|
|
snap := gate.estimateSnapshot(fri)
|
|
b5, _ := snap.Bucket(Bucket5h)
|
|
bw, _ := snap.Bucket(BucketWeekly)
|
|
if b5.Used < 31.23 || b5.Used > 31.25 || bw.Used < 31.23 || bw.Used > 31.25 {
|
|
t.Errorf("estimate buckets after one turn: 5h=%v weekly=%v", b5.Used, bw.Used)
|
|
}
|
|
if b5.Limit != cfg.Plan5hCredits || bw.Limit != cfg.PlanWeeklyCredits {
|
|
t.Errorf("limits = %v/%v, want plan limits", b5.Limit, bw.Limit)
|
|
}
|
|
|
|
// Local-only mode (no redis) keeps its own accounting.
|
|
local := newGate(t, gateTestCfg())
|
|
local.SetClock(func() time.Time { return fri })
|
|
local.RecordTurn(EstimateTurnInput{Model: "glm-4.7-flash", PromptTokens: 10000, CompletionTokens: 10000, Peak: false})
|
|
ls := local.estimateSnapshot(fri)
|
|
lb, _ := ls.Bucket(Bucket5h)
|
|
if lb.Used < 5.14 || lb.Used > 5.16 {
|
|
t.Errorf("local estimate = %v, want ~5.15", lb.Used)
|
|
}
|
|
}
|
|
|
|
func mustLoc(t *testing.T, name string) *time.Location {
|
|
t.Helper()
|
|
l, err := time.LoadLocation(name)
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
return l
|
|
}
|
|
|
|
// TestDecideNineteenWall replays the 2026-08-28 19:00 quota wall: weekly
|
|
// bucket exhausted at 19:00 CST — the gate must DEFER every class with a
|
|
// surfaced reason instead of letting the turn die at the provider.
|
|
func TestDecideNineteenWall(t *testing.T) {
|
|
gate := newGate(t, gateTestCfg())
|
|
cst := mustLoc(t, "America/Chicago")
|
|
wall := time.Date(2026, 8, 28, 19, 0, 0, 0, cst) // Friday evening
|
|
gate.SetClock(func() time.Time { return wall })
|
|
gate.mu.Lock()
|
|
gate.lastSnap = &QuotaSnapshot{Account: "zai-1", Source: "usage_url", FetchedAt: wall,
|
|
Buckets: []Bucket{
|
|
{ID: Bucket5h, Used: 12000, Limit: 28000},
|
|
{ID: BucketWeekly, Used: 135800, Limit: 140000}, // 97%
|
|
}}
|
|
gate.mu.Unlock()
|
|
|
|
for _, class := range []string{"primary", "code", "study", "read", "review"} {
|
|
d := gate.Decide(context.Background(), class)
|
|
if d.Action != ActionDefer {
|
|
t.Fatalf("class %s at 97%% weekly: action = %s, want defer", class, d.Action)
|
|
}
|
|
if !strings.Contains(d.Reason, "weekly") || !strings.Contains(d.Reason, "97%") {
|
|
t.Errorf("defer reason must surface the bucket + ratio: %q", d.Reason)
|
|
}
|
|
}
|
|
}
|
|
|
|
func TestDecideLevels(t *testing.T) {
|
|
cfg := gateTestCfg()
|
|
gate := newGate(t, cfg)
|
|
// Wednesday noon CST: off-peak, mid-week.
|
|
noon := time.Date(2026, 9, 2, 12, 0, 0, 0, mustLoc(t, "America/Chicago"))
|
|
gate.SetClock(func() time.Time { return noon })
|
|
snap := func(pct5, pctW float64) {
|
|
gate.mu.Lock()
|
|
defer gate.mu.Unlock()
|
|
gate.lastSnap = &QuotaSnapshot{Account: cfg.Account, Source: "usage_url", FetchedAt: noon,
|
|
Buckets: []Bucket{
|
|
{ID: Bucket5h, Used: 28000 * pct5 / 100, Limit: 28000},
|
|
{ID: BucketWeekly, Used: 140000 * pctW / 100, Limit: 140000},
|
|
}}
|
|
}
|
|
|
|
cases := []struct {
|
|
name string
|
|
pct5, pctW float64
|
|
class string
|
|
want string
|
|
wantIn string
|
|
}{
|
|
{"healthy allows flagship", 10, 10, "primary", ActionAllow, ""},
|
|
{"healthy allows flash", 10, 10, "study", ActionAllow, ""},
|
|
{"defer level: heavy defers", 90, 50, "code", ActionDefer, ">= defer"},
|
|
{"defer level: flash continues", 90, 50, "study", ActionAllow, ""},
|
|
{"block level: flash defers too", 50, 96, "study", ActionDefer, ">= block"},
|
|
}
|
|
for _, c := range cases {
|
|
t.Run(c.name, func(t *testing.T) {
|
|
snap(c.pct5, c.pctW)
|
|
d := gate.Decide(context.Background(), c.class)
|
|
if d.Action != c.want {
|
|
t.Fatalf("action = %s (%s), want %s", d.Action, d.Reason, c.want)
|
|
}
|
|
if c.wantIn != "" && !strings.Contains(d.Reason, c.wantIn) {
|
|
t.Errorf("reason %q missing %q", d.Reason, c.wantIn)
|
|
}
|
|
})
|
|
}
|
|
|
|
// Peak window restriction (quota healthy): Friday 03:00 CST.
|
|
peak := time.Date(2026, 8, 28, 3, 0, 0, 0, mustLoc(t, "America/Chicago"))
|
|
gate.SetClock(func() time.Time { return peak })
|
|
snap(10, 10)
|
|
if d := gate.Decide(context.Background(), "primary"); d.Action != ActionDefer || !strings.Contains(d.Reason, "peak window") {
|
|
t.Errorf("peak + flagship class: %+v, want defer/peak reason", d)
|
|
}
|
|
if d := gate.Decide(context.Background(), "study"); d.Action != ActionAllow {
|
|
t.Errorf("peak + flash class must allow: %+v", d)
|
|
}
|
|
// Block beats peak: even flash defers when the wall is hit in-peak.
|
|
snap(10, 97)
|
|
if d := gate.Decide(context.Background(), "study"); d.Action != ActionDefer {
|
|
t.Errorf("peak + wall: %+v, want defer", d)
|
|
}
|
|
}
|
|
|
|
func gateTestCfg() config.QuotaConfig {
|
|
return config.QuotaConfig{
|
|
Enabled: true, Account: "zai-1",
|
|
Plan5hCredits: 28000, PlanWeeklyCredits: 140000,
|
|
PollIntervalSecs: 300, DeferAtPct: 85, BlockAtPct: 95,
|
|
PeakStart: "01:00", PeakEnd: "05:00", Timezone: "America/Chicago",
|
|
PeakWeekdaysOnly: true, PeakClasses: []string{"study", "read"},
|
|
}
|
|
}
|
|
|
|
func newGate(t *testing.T, cfg config.QuotaConfig) *Gate {
|
|
t.Helper()
|
|
gate, err := NewGate(cfg, nil, nil)
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
return gate
|
|
}
|
|
|
|
// (end of quota gate tests)
|