quota: z.ai credit-bucket back-pressure + resource gate + usage accounting (Redmine 490+491)

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
2026-08-29 05:37:15 -05:00
parent 9dbb20489c
commit fc518c475e
21 changed files with 2904 additions and 20 deletions
@@ -0,0 +1,293 @@
+package loop
+
+import (
+	"context"
+	"encoding/json"
+	"fmt"
+	"net/http"
+	"net/http/httptest"
+	"os"
+	"path/filepath"
+	"strings"
+	"sync/atomic"
+	"testing"
+	"time"
+
+	"ukrrs.com/mopac/harness/internal/config"
+)
+
+// fakeUsage is a mutable z.ai usage endpoint: the documented-shape JSON,
+// with the weekly bucket flippable between "wall" (97%) and healthy (10%).
+type fakeUsage struct {
+	wall atomic.Bool
+	srv  *httptest.Server
+}
+
+func newFakeUsage(t *testing.T) *fakeUsage {
+	t.Helper()
+	f := &fakeUsage{}
+	mux := http.NewServeMux()
+	mux.HandleFunc("/usage", func(w http.ResponseWriter, r *http.Request) {
+		if r.Header.Get("Authorization") == "" {
+			http.Error(w, "no bearer", http.StatusUnauthorized)
+			return
+		}
+		wk := 14000.0 // 10%
+		five := 2800.0
+		if f.wall.Load() {
+			wk = 135800.0 // 97%: the 2026-08-28 19:00 wall
+			five = 12000.0
+		}
+		fmt.Fprintf(w, `{"usage":{"five_hour":{"used_credits":%g,"total_credits":28000,"reset_at":"2026-08-29T22:00:00Z"},"weekly":{"used_credits":%g,"total_credits":140000,"reset_at":"2026-09-02T00:00:00Z"}}}`, five, wk)
+	})
+	f.srv = httptest.NewServer(mux)
+	t.Cleanup(f.srv.Close)
+	return f
+}
+
+func quotaTestConfig(t *testing.T, redmineURL, llmURL, usageURL string) *config.Config {
+	t.Helper()
+	cfg := loopTestConfig(t, redmineURL, llmURL)
+	cfg.Quota = quotaCfg(usageURL)
+	return cfg
+}
+
+func resCfg(procRoot string) config.ResourcesConfig {
+	return config.ResourcesConfig{
+		Enabled:           true,
+		MaxLoadAvg:        6.0,
+		MinMemAvailableMB: 2048,
+		MinDiskFreeMB:     5120,
+		MaxIODelayPct:     90,
+		ProcRoot:          procRoot,
+		SysRoot:           procRoot,
+	}
+}
+
+func quotaCfg(usageURL string) config.QuotaConfig {
+	c := config.QuotaConfig{
+		Enabled: true, Account: "zai-1",
+		Plan5hCredits: 28000, PlanWeeklyCredits: 140000,
+		UsageURL: usageURL, KeyRef: "literal:zai-key",
+		PollIntervalSecs: 300, DeferAtPct: 85, BlockAtPct: 95,
+		PeakStart: "01:00", PeakEnd: "05:00", Timezone: "America/Chicago",
+		PeakWeekdaysOnly: true, PeakClasses: []string{"study", "read"},
+	}
+	if usageURL == "" {
+		c.KeyRef = ""
+	}
+	return c
+}
+
+// TestLoopQuotaWallDefersAndRecovers replays 2026-08-28 ~19:00 CST: the
+// weekly bucket hit 97% and dispatched turns died at the provider. With the
+// gate: the scan DEFERS the turn with a surfaced reason, the loop exits
+// clean, the issue is NOT consumed (no note, not marked processed) — and
+// once quota recovers the very next scan dispatches it.
+func TestLoopQuotaWallDefersAndRecovers(t *testing.T) {
+	rm := newFakeRedmine(t)
+	rm.addIssue("490", "Big build task", "In Progress", "2026-08-28T21:30:00Z")
+	f := newFakeLLM(t, textMsg("the turn ran"))
+	usage := newFakeUsage(t)
+	usage.wall.Store(true) // the wall is up
+	cfg := quotaTestConfig(t, rm.srv.URL, f.srv.URL, usage.srv.URL+"/usage")
+
+	out := runOneScan(t, cfg, false)
+
+	if f.requestCount() != 0 {
+		t.Fatalf("LLM calls during the wall = %d, want 0 (turns must defer, not die)", f.requestCount())
+	}
+	for _, want := range []string{
+		"defer 490",
+		">= block 95%",
+		"weekly",
+	} {
+		if !strings.Contains(out, want) {
+			t.Errorf("output missing %q:\n%s", want, out)
+		}
+	}
+	if rm.noteCount("490") != 0 {
+		t.Errorf("deferred task must not be noted back")
+	}
+
+	// The defer is in the JSONL with the reason (json.Marshal HTML-escapes
+	// ">=", so match on the words); the task was NOT consumed.
+	statePath := filepath.Join(cfg.Loop.StateDir, "loop.jsonl")
+	data, err := os.ReadFile(statePath)
+	if err != nil {
+		t.Fatal(err)
+	}
+	if !strings.Contains(string(data), `"type":"defer"`) || !strings.Contains(string(data), "block 95") {
+		t.Errorf("JSONL defer event missing:\n%s", data)
+	}
+
+	// Quota recovers: next scan dispatches (task was never marked processed).
+	usage.wall.Store(false)
+	out2 := runOneScan(t, cfg, false)
+	if f.requestCount() != 1 {
+		t.Fatalf("LLM calls after recovery = %d, want 1 (defer must not consume the task)", f.requestCount())
+	}
+	if !strings.Contains(out2, "dispatch 490") || !strings.Contains(out2, "noted REPORT on #490") {
+		t.Errorf("recovery scan output wrong:\n%s", out2)
+	}
+}
+
+// TestLoopPeakWindowDefersHeavyClass: inside 01:00-05:00 CST the heavy
+// class defers with a peak reason while the flash-tier class runs.
+func TestLoopPeakWindowDefersHeavyClass(t *testing.T) {
+	rm := newFakeRedmine(t)
+	rm.addIssue("491", "Heavy flagship work", "In Progress", "2026-08-28T21:30:00Z")
+	f := newFakeLLM(t, textMsg("flash reply"))
+	cfg := quotaTestConfig(t, rm.srv.URL, f.srv.URL, "") // estimates only
+
+	// Friday 2026-08-28 03:00 CST = inside the peak window.
+	peakAt := time.Date(2026, 8, 28, 3, 0, 0, 0, cst(t))
+	var out strings.Builder
+	cond, err := New(cfg, &out)
+	if err != nil {
+		t.Fatal(err)
+	}
+	cond.gate.SetClock(func() time.Time { return peakAt })
+	if err := cond.RunLoop(context.Background(), LoopOpts{Once: true}); err != nil {
+		t.Fatal(err)
+	}
+	if f.requestCount() != 0 {
+		t.Fatalf("heavy class dispatched in peak window (%d calls)", f.requestCount())
+	}
+	if !strings.Contains(out.String(), "peak window") {
+		t.Errorf("defer reason must name the peak window:\n%s", out.String())
+	}
+
+	// A flash-tier class in the same peak window runs.
+	rm.addIssueClass("492", "Light study work", "In Progress", "2026-08-28T21:31:00Z", "study")
+	out2 := runScanAt(t, cfg, peakAt)
+	if f.requestCount() != 1 {
+		t.Fatalf("flash class must dispatch in peak: %d calls", f.requestCount())
+	}
+	if !strings.Contains(out2, "dispatch 492") {
+		t.Errorf("flash dispatch missing:\n%s", out2)
+	}
+
+	// Off-peak (Friday 12:00 CST): the heavy task now dispatches.
+	offPeak := time.Date(2026, 8, 28, 12, 0, 0, 0, cst(t))
+	_ = runScanAt(t, cfg, offPeak)
+	if f.requestCount() != 2 {
+		t.Fatalf("heavy class must dispatch off-peak: %d calls", f.requestCount())
+	}
+}
+
+// TestLoopResourceBusyDefers: a pinned-high load fixture defers dispatch;
+// the task stays unconsumed.
+func TestLoopResourceBusyDefers(t *testing.T) {
+	rm := newFakeRedmine(t)
+	rm.addIssue("492", "Task during a build storm", "In Progress", "2026-08-28T21:30:00Z")
+	f := newFakeLLM(t, textMsg("never"))
+	cfg := loopTestConfig(t, rm.srv.URL, f.srv.URL)
+
+	proc := filepath.Join(t.TempDir(), "proc")
+	os.MkdirAll(proc, 0o755)
+	os.WriteFile(filepath.Join(proc, "loadavg"), []byte("41.20 30.00 15.00 5/900 1\n"), 0o644)
+	cfg.Resources = resCfg(proc)
+
+	out := runOneScan(t, cfg, false)
+	if f.requestCount() != 0 {
+		t.Fatalf("busy host must defer dispatch (%d calls)", f.requestCount())
+	}
+	if !strings.Contains(out, "resources busy") || !strings.Contains(out, "load 41.20") {
+		t.Errorf("defer reason must surface the resource violation:\n%s", out)
+	}
+	if rm.noteCount("492") != 0 {
+		t.Errorf("deferred task must not be noted")
+	}
+}
+
+// TestLoopUsageAccounting: after a dispatched turn, the JSONL report event
+// carries class + tokens + estimated credits, and the usage table aggregates.
+func TestLoopUsageAccounting(t *testing.T) {
+	rm := newFakeRedmine(t)
+	rm.addIssue("493", "Accounted task", "In Progress", "2026-08-28T21:30:00Z")
+	f := newFakeLLM(t, textMsg("reply with 20/40 tokens"))
+	cfg := quotaTestConfig(t, rm.srv.URL, f.srv.URL, "")
+
+	runScanAt(t, cfg, time.Date(2026, 8, 28, 12, 0, 0, 0, cst(t))) // off-peak Friday
+
+	data, err := os.ReadFile(filepath.Join(cfg.Loop.StateDir, "loop.jsonl"))
+	if err != nil {
+		t.Fatal(err)
+	}
+	var reportEv map[string]any
+	for _, line := range strings.Split(strings.TrimSpace(string(data)), "\n") {
+		var ev map[string]any
+		if err := json.Unmarshal([]byte(line), &ev); err != nil {
+			t.Fatalf("bad JSONL line %q: %v", line, err)
+		}
+		if ev["type"] == "report" {
+			reportEv = ev
+		}
+	}
+	if reportEv == nil {
+		t.Fatal("no report event in JSONL")
+	}
+	if reportEv["class"] != "primary" {
+		t.Errorf("report event class = %v, want primary", reportEv["class"])
+	}
+	if reportEv["prompt_tokens"] != float64(20) || reportEv["completion_tokens"] != float64(40) {
+		t.Errorf("report tokens = %v", reportEv)
+	}
+	// This turn ran glm-5.3 off-peak: (20*6.9 + 40*24)/10000 / 2 = 0.0549
+	credits, _ := reportEv["credits"].(float64)
+	if credits < 0.0548 || credits > 0.0550 {
+		t.Errorf("credits = %v, want ~0.0549", credits)
+	}
+
+	report, err := UsageReport(cfg.Loop.StateDir)
+	if err != nil {
+		t.Fatal(err)
+	}
+	if !strings.Contains(report, "primary") || !strings.Contains(report, "TOTAL") || !strings.Contains(report, "1") {
+		t.Errorf("usage table wrong:\n%s", report)
+	}
+}
+
+// TestLoopGateOffByDefault: with [quota] disabled the loop dispatches
+// exactly as before (no behavior change for existing configs).
+func TestLoopGateOffByDefault(t *testing.T) {
+	rm := newFakeRedmine(t)
+	rm.addIssue("494", "Plain task", "In Progress", "2026-08-28T21:30:00Z")
+	f := newFakeLLM(t, textMsg("plain reply"))
+	cfg := loopTestConfig(t, rm.srv.URL, f.srv.URL)
+
+	out := runOneScan(t, cfg, false)
+	if f.requestCount() != 1 {
+		t.Fatalf("LLM calls = %d, want 1", f.requestCount())
+	}
+	if strings.Contains(out, "defer") || strings.Contains(out, "gate:") {
+		t.Errorf("disabled gate must be silent:\n%s", out)
+	}
+}
+
+func cst(t *testing.T) *time.Location {
+	t.Helper()
+	loc, err := time.LoadLocation("America/Chicago")
+	if err != nil {
+		t.Fatal(err)
+	}
+	return loc
+}
+
+// runScanAt runs a --once scan with the quota gate clock pinned.
+func runScanAt(t *testing.T, cfg *config.Config, at time.Time) string {
+	t.Helper()
+	var out strings.Builder
+	cond, err := New(cfg, &out)
+	if err != nil {
+		t.Fatal(err)
+	}
+	if cond.gate != nil {
+		cond.gate.SetClock(func() time.Time { return at })
+	}
+	if err := cond.RunLoop(context.Background(), LoopOpts{Once: true}); err != nil {
+		t.Fatalf("RunLoop: %v (output:\n%s)", err, out.String())
+	}
+	return out.String()
+}