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
132 lines
3.6 KiB
Go
132 lines
3.6 KiB
Go
package config
|
|
|
|
import (
|
|
"os"
|
|
"path/filepath"
|
|
"strings"
|
|
"testing"
|
|
)
|
|
|
|
func TestQuotaAndResourcesLoad(t *testing.T) {
|
|
src := `
|
|
vertical = "t"
|
|
[loop]
|
|
[redmine]
|
|
url = "https://rm"
|
|
key_ref = "env:K"
|
|
scope_query = "p=1"
|
|
[litellm]
|
|
base_url = "http://l"
|
|
key_ref = "env:K"
|
|
[models]
|
|
a = "m"
|
|
default_tier = "a"
|
|
[models.classes]
|
|
x = "a"
|
|
[quota]
|
|
enabled = true
|
|
account = "zai-max-1"
|
|
plan_5h_credits = 28000
|
|
plan_weekly_credits = 140000
|
|
usage_url = "https://api.z.ai/api/coding/paas/v4/usage"
|
|
key_ref = "env:HARNESS_ZAI_KEY"
|
|
poll_interval_secs = 120
|
|
defer_at_pct = 80
|
|
block_at_pct = 92.5
|
|
peak_start = "22:30"
|
|
peak_end = "04:15"
|
|
timezone = "America/Chicago"
|
|
peak_weekdays_only = false
|
|
peak_classes = ["study", "read", "summarize"]
|
|
redis_url = "redis://192.168.3.78:6390/2"
|
|
[resources]
|
|
enabled = true
|
|
max_load_avg = 5.5
|
|
min_mem_available_mb = 1024
|
|
min_disk_free_mb = 2048
|
|
max_io_delay_pct = 75
|
|
`
|
|
path := filepath.Join(t.TempDir(), "harness.toml")
|
|
if err := os.WriteFile(path, []byte(src), 0o600); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
cfg, err := Load(path)
|
|
if err != nil {
|
|
t.Fatalf("Load: %v", err)
|
|
}
|
|
q := cfg.Quota
|
|
if !q.Enabled || q.Account != "zai-max-1" || q.Plan5hCredits != 28000 || q.PlanWeeklyCredits != 140000 {
|
|
t.Errorf("quota basics: %+v", q)
|
|
}
|
|
if q.UsageURL == "" || q.KeyRef != "env:HARNESS_ZAI_KEY" || q.PollIntervalSecs != 120 {
|
|
t.Errorf("poller config: %+v", q)
|
|
}
|
|
if q.DeferAtPct != 80 || q.BlockAtPct != 92.5 {
|
|
t.Errorf("thresholds: %+v", q)
|
|
}
|
|
if q.PeakStart != "22:30" || q.PeakEnd != "04:15" || q.Timezone != "America/Chicago" || q.PeakWeekdaysOnly {
|
|
t.Errorf("peak window: %+v", q)
|
|
}
|
|
if len(q.PeakClasses) != 3 || q.PeakClasses[0] != "study" {
|
|
t.Errorf("peak classes: %v", q.PeakClasses)
|
|
}
|
|
if q.RedisURL != "redis://192.168.3.78:6390/2" {
|
|
t.Errorf("redis url: %q", q.RedisURL)
|
|
}
|
|
r := cfg.Resources
|
|
if !r.Enabled || r.MaxLoadAvg != 5.5 || r.MinMemAvailableMB != 1024 || r.MinDiskFreeMB != 2048 || r.MaxIODelayPct != 75 {
|
|
t.Errorf("resources: %+v", r)
|
|
}
|
|
}
|
|
|
|
func TestQuotaValidateErrors(t *testing.T) {
|
|
base := `
|
|
vertical = "t"
|
|
[litellm]
|
|
base_url = "http://l"
|
|
key_ref = "env:K"
|
|
[models]
|
|
a = "m"
|
|
default_tier = "a"
|
|
`
|
|
cases := []struct {
|
|
name, extra, want string
|
|
}{
|
|
{"no account", "[quota]\nenabled = true\n", "account is required"},
|
|
{"bad bucket", "[quota]\nenabled = true\naccount = \"a\"\nplan_5h_credits = 0\n", "plan_5h_credits"},
|
|
{"defer > block", "[quota]\nenabled = true\naccount = \"a\"\ndefer_at_pct = 99\nblock_at_pct = 95\n", "defer_at_pct"},
|
|
{"bad window", "[quota]\nenabled = true\naccount = \"a\"\npeak_start = \"99:00\"\n", "peak_start"},
|
|
{"bad tz", "[quota]\nenabled = true\naccount = \"a\"\ntimezone = \"Mars/Olympus\"\n", "timezone"},
|
|
{"usage url without key", "[quota]\nenabled = true\naccount = \"a\"\nusage_url = \"https://x\"\n", "key_ref is required"},
|
|
{"bad threshold", "[resources]\nenabled = true\nmax_load_avg = -1\n", "thresholds must be > 0"},
|
|
}
|
|
for _, c := range cases {
|
|
t.Run(c.name, func(t *testing.T) {
|
|
path := filepath.Join(t.TempDir(), "harness.toml")
|
|
if err := os.WriteFile(path, []byte(base+c.extra), 0o600); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
_, err := Load(path)
|
|
if err == nil || !strings.Contains(err.Error(), c.want) {
|
|
t.Fatalf("err = %v, want containing %q", err, c.want)
|
|
}
|
|
})
|
|
}
|
|
}
|
|
|
|
func TestFloatTomLValues(t *testing.T) {
|
|
doc, err := ParseTOML("a = 6\nb = 6.25\nc = 1_000.5\n")
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if v, ok := doc.Float("a"); !ok || v != 6 {
|
|
t.Errorf("int as float = %v %v", v, ok)
|
|
}
|
|
if v, ok := doc.Float("b"); !ok || v != 6.25 {
|
|
t.Errorf("float = %v %v", v, ok)
|
|
}
|
|
if v, ok := doc.Float("c"); !ok || v != 1000.5 {
|
|
t.Errorf("underscore float = %v %v", v, ok)
|
|
}
|
|
}
|