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
This commit is contained in:
@@ -9,12 +9,15 @@ other agents interact with a stack only through Redmine (SoR), Discourse
|
||||
(docs) and Gitea (code), never by attaching to the loop.
|
||||
|
||||
Status: 2026-08-29 — skeleton + event receiver + self-host loop + **OWUI
|
||||
front door live**: `harness serve` answers OpenAI-compatible
|
||||
`/v1/models` + `/v1/chat/completions` on LAN port 8090 (bearer vkey,
|
||||
stateless bounded turns, live-proven through LiteLLM to glm-5.3 and
|
||||
glm-4.7-flash). Also live: `harness loop` (Redmine SoR self-hosting, fake-
|
||||
Redmine e2e test-asserted), the MVP demo path, and `harness events` on port
|
||||
4100.
|
||||
front door live** + **quota/resource gates** (Redmine 490+491): `harness
|
||||
loop` consults the z.ai credit buckets (5h + weekly, polled or locally
|
||||
estimated), a TZ-aware peak window (default 01:00-05:00 CST weekdays) and
|
||||
host load/mem/disk/IO before every dispatch — gated work DEFERS with a
|
||||
logged reason and is reconsidered next scan, never hard-failed; per-class
|
||||
token+credit accounting lands in the loop JSONL (`harness quota status`).
|
||||
`harness serve` (OWUI front door, LAN 8090), `harness loop` (fake-Redmine
|
||||
e2e test-asserted), the MVP demo path and `harness events` (port 4100)
|
||||
all live.
|
||||
|
||||
## Quickstart
|
||||
|
||||
@@ -226,6 +229,57 @@ Live on the LAN 2026-08-29: catalog of 9 models, 401/400/200 paths, and real
|
||||
turns through LiteLLM (`mopac-primary` -> glm-5.3, `mopac-study` ->
|
||||
glm-4.7-flash) with usage accounting, driven by python urllib.
|
||||
|
||||
### Quota + resource gates (`[quota]` / `[resources]`, Redmine 490+491)
|
||||
|
||||
The 2026-08-28 ~19:00 quota wall killed dispatched turns mid-flight; the
|
||||
gates turn that failure mode into a logged throttle. Before EVERY dispatch
|
||||
the loop consults, in order:
|
||||
|
||||
1. **Host resources** (`[resources]`, read-only `/proc` + statfs): loadavg,
|
||||
mem available, work-root disk free, IO pressure (`/proc/pressure/io`,
|
||||
skipped when PSI is absent). Any violation defers with all reasons
|
||||
surfaced.
|
||||
2. **Quota buckets** (`[quota]`): the z.ai coding plan's 5-hour and weekly
|
||||
credit windows. `>= block_at_pct` (default 95%) defers EVERYTHING with
|
||||
the bucket/ratio in the reason — the wall, caught early.
|
||||
3. **Peak window**: z.ai peak hours (documented Mon-Fri 14:00-18:00
|
||||
Singapore == 01:00-05:00 CST in winter) charge full rate; inside the
|
||||
window only `peak_classes` (the flash/LLM-lite tier) dispatch, heavy
|
||||
classes defer to off-peak (50% credit cost).
|
||||
4. **Soft quota**: `>= defer_at_pct` (default 85%) defers heavy classes
|
||||
while LLM-lite continues.
|
||||
|
||||
A deferred task is NOT consumed: no turn, no note, no dedup marker — the
|
||||
next scan reconsiders it (the 19:00-wall scenario is replayed as a test:
|
||||
wall up -> defer + loop clean -> quota recovers -> dispatch). Defer events
|
||||
land in `loop.jsonl` (`"type":"defer"` + reason), deduped per task+reason.
|
||||
|
||||
**Quota state** comes from two sources: the provider endpoint (`usage_url`,
|
||||
bearer `key_ref`, parsed into buckets with reset times) and, when that is
|
||||
unconfigured/unreachable, locally ESTIMATED consumption — per-turn credits
|
||||
computed from the documented z.ai formula (input x 6.9 + cached x 1.7 +
|
||||
output x 24, per 10k tokens; flash 2.3/0.56/8; off-peak 50% off) against
|
||||
the configured plan limits. z.ai documents the buckets but publishes no
|
||||
usage REST route today (probed 2026-08-29 — see the REPORT); the parser
|
||||
targets the documented shape and is fake-server-tested, so flipping
|
||||
`usage_url` on when z.ai ships it is a config edit. LIVE VERIFICATION open.
|
||||
|
||||
**Shared state**: with `redis_url` set, the latest snapshot and the credit
|
||||
estimates live in one redis container so all harness instances of an
|
||||
account (9 accounts across 2 hosts) coordinate — see the runbook below.
|
||||
Redis down = this instance's local estimate; the loop never stops for it.
|
||||
|
||||
**Usage accounting**: every dispatched turn appends class + tokens +
|
||||
estimated credits to its `report` event in `loop.jsonl`; `harness quota
|
||||
status` renders the per-class table (the feed for the per-instance
|
||||
Discourse usage reports).
|
||||
|
||||
```sh
|
||||
./bin/harness quota status # snapshot + peak window + resources + usage table
|
||||
./bin/harness quota gate # the allow/defer verdict per class, right now
|
||||
./bin/harness quota probe # one usage_url poll; parsed buckets or the error
|
||||
```
|
||||
|
||||
### Help
|
||||
|
||||
```sh
|
||||
@@ -317,9 +371,10 @@ Subcommands (from `harness help`):
|
||||
|---|---|
|
||||
| `harness help` | print usage (also `-h`, `--help`) |
|
||||
| `harness once` | run ONE conductor iteration, then exit |
|
||||
| `harness loop` | run the self-host daemon until SIGINT (poll -> turn -> note/status writeback) |
|
||||
| `harness loop` | run the self-host daemon until SIGINT (poll -> gate -> turn -> note/status writeback) |
|
||||
| `harness events` | run the webhook receiver until SIGINT/SIGTERM |
|
||||
| `harness serve` | run the OpenAI-compatible front door until SIGINT/SIGTERM (the OWUI connection) |
|
||||
| `harness quota` | gate surface: `status` (snapshot + usage accounting), `gate` (per-class verdicts), `probe` (one usage poll) |
|
||||
|
||||
Flags for `once`:
|
||||
|
||||
@@ -353,6 +408,12 @@ Flags for `serve`:
|
||||
| `-config PATH` | config file (default `$HARNESS_CONFIG`, then `./harness.toml`) |
|
||||
| `-listen ADDR` | bind address (overrides `[serve]` listen) |
|
||||
|
||||
Flags for `quota` (`harness quota <status|probe|gate> [flags]`):
|
||||
|
||||
| Flag | Meaning |
|
||||
|---|---|
|
||||
| `-config PATH` | config file (default `$HARNESS_CONFIG`, then `./harness.toml`) |
|
||||
|
||||
Exit codes:
|
||||
|
||||
| Code | Meaning |
|
||||
@@ -458,9 +519,12 @@ Sourced from [REPORT.md](REPORT.md) — keep both in sync.
|
||||
| Status transitions | Works | `[redmine.status_map]` names → ids via `/issue_statuses.json`; refresh advances the dedup marker past its own writes |
|
||||
| `mpk:` key refs | Works | `[keyproxy]` hop (POST `/v1/resolve`, bearer, cached); local refs unaffected |
|
||||
| Gitea REPORT commit | Works (off) | `[gitea] commit_reports`: contents-API create-or-update right after the REPORT lands |
|
||||
| Quota gate (490) | Works (off) | `[quota]`: 5h/weekly credit buckets (poll or estimate), block/defer/peak back-pressure, defer-not-fail; usage accounting per class in loop.jsonl; z.ai usage endpoint LIVE VERIFICATION open |
|
||||
| Shared quota state | Works (off) | `redis_url`: one redis container, all instances of an account share snapshot + estimates; stdlib RESP2 mini-client, fail-soft |
|
||||
| Resource gate (491) | Works (off) | `[resources]`: loadavg/mem/disk/IO-PSI thresholds, read-only, defer-not-fail |
|
||||
| Event → turn dispatch | Stubbed | conductor `DispatchEvent` prints what it would do; wiring is phase 3 |
|
||||
| Tests | Works | table-driven, stdlib only; build/vet/test clean on go1.26 |
|
||||
| Budget/semaphore gate | Stubbed | tokens in REPORT, no cost/spend enforcement yet |
|
||||
| Budget/semaphore gate | Partial | credit-bucket back-pressure is live (rows above); LiteLLM-$-spend budget keys remain open |
|
||||
| `bw:` key refs | Stubbed | error until the bitwarden wrapper (phase 3) |
|
||||
| Loop concurrency | Stubbed | v0 = one turn at a time; in-process concurrency knob later |
|
||||
| Streaming + turn resume | Stubbed | retry is request-level today; serve door is non-streaming by design v0 |
|
||||
@@ -470,6 +534,37 @@ Sourced from [REPORT.md](REPORT.md) — keep both in sync.
|
||||
| Local inbox intake | Stubbed | not started |
|
||||
| Next (phase 3) | Next | bitwarden wrapper, full permission layer, event → turn wiring, budget gate via LiteLLM spend APIs, serve streaming + serve tools, loop concurrency knob |
|
||||
|
||||
## Runbook: shared quota state + resource control (deploy-time)
|
||||
|
||||
**Redis container** (the shared-state hop for `[quota] redis_url`; one per
|
||||
host-pair, container only, no host packages):
|
||||
|
||||
```sh
|
||||
docker run -d --name mopac-quota-redis --restart unless-stopped \
|
||||
-p 192.168.3.78:6390:6379 \
|
||||
-v /srv/mopac-quota-redis:/data \
|
||||
redis:7-alpine --appendonly yes
|
||||
```
|
||||
|
||||
Every harness instance of the same z.ai account then sets the same
|
||||
`account` + `redis_url` in `[quota]`; keys are namespaced
|
||||
`mopac:quota:<account>:{snapshot,est:*}`. Redis unreachable = local
|
||||
estimates only (fail-soft, logged). The harness speaks RESP2 directly —
|
||||
no client library, no host redis-cli needed.
|
||||
|
||||
**cgroup enforcement** (ticket 491, deploy-time): the in-harness gate is
|
||||
read-only and advisory — it defers dispatch when the HOST is busy. To keep
|
||||
builds/turns from making the host busy in the first place, run each loop
|
||||
container under cgroup limits at deploy:
|
||||
|
||||
```sh
|
||||
docker run ... --memory 4g --cpus 2 --pids-limit 512 \
|
||||
--io-max bandwidth=/data:100mb ... # device-specific; see docker run(1)
|
||||
```
|
||||
|
||||
or a systemd slice for non-container deploys (`CPUQuota=200%`,
|
||||
`MemoryMax=4G`, `IOWeight`). The gate catches what the limits don't.
|
||||
|
||||
## Docs and links
|
||||
|
||||
- [DESIGN.md](DESIGN.md) — design spec (hard rules, build order, org model)
|
||||
|
||||
+118
-1
@@ -10,6 +10,7 @@ import (
|
||||
"net/http"
|
||||
"os"
|
||||
"os/signal"
|
||||
"sort"
|
||||
"strings"
|
||||
"syscall"
|
||||
"time"
|
||||
@@ -17,6 +18,7 @@ import (
|
||||
"ukrrs.com/mopac/harness/internal/config"
|
||||
"ukrrs.com/mopac/harness/internal/events"
|
||||
"ukrrs.com/mopac/harness/internal/loop"
|
||||
"ukrrs.com/mopac/harness/internal/quota"
|
||||
"ukrrs.com/mopac/harness/internal/serve"
|
||||
)
|
||||
|
||||
@@ -27,6 +29,7 @@ Usage:
|
||||
harness loop [-config PATH] [-interval DUR] [-once] [-dry-run]
|
||||
harness events [-config PATH] [-listen ADDR]
|
||||
harness serve [-config PATH] [-listen ADDR]
|
||||
harness quota <status|probe|gate> [-config PATH]
|
||||
|
||||
once runs ONE conductor iteration and exits (chain by re-invoking; no daemon):
|
||||
intake (Redmine scope, or the [demo] issue) -> plan/model routing ->
|
||||
@@ -52,6 +55,15 @@ serve runs the OpenAI-compatible front door until SIGINT/SIGTERM (the
|
||||
non-streaming v0; tools off. Own port - coexists with events.
|
||||
Routes: POST /v1/chat/completions, GET /v1/models, GET /healthz.
|
||||
|
||||
quota is the Redmine 490/491 gate surface:
|
||||
status one-shot: quota snapshot (polled or estimated), peak window,
|
||||
host resources, and the per-class token+credit usage table from
|
||||
the loop state (the Discourse usage-report feed)
|
||||
probe poll [quota] usage_url once and print the parsed buckets (or the
|
||||
raw error; bearer key never printed)
|
||||
gate evaluate the back-pressure gates NOW: the decision for every
|
||||
[models.classes] class + resource thresholds
|
||||
|
||||
Flags:
|
||||
-config PATH config file (default $HARNESS_CONFIG or ./harness.toml)
|
||||
-dry-run once/loop: intake + plan only; no LLM call, no REPORT,
|
||||
@@ -61,7 +73,6 @@ Flags:
|
||||
-interval DUR loop: poll interval (overrides [loop] poll_interval_secs)
|
||||
-once loop: single scan then exit (cron-able)
|
||||
-listen ADDR events/serve: bind address (overrides [events]/[serve] listen)
|
||||
|
||||
Exit codes:
|
||||
0 ok (including "no tasks in scope"; loop: clean SIGINT stop)
|
||||
1 usage / config / routing / writeback error
|
||||
@@ -90,6 +101,8 @@ func run(args []string) int {
|
||||
return runEvents(args[1:])
|
||||
case "serve":
|
||||
return runServe(args[1:])
|
||||
case "quota":
|
||||
return runQuota(args[1:])
|
||||
default:
|
||||
fmt.Fprintf(os.Stderr, "harness: unknown command %q\n\n%s", args[0], usage)
|
||||
return 1
|
||||
@@ -199,6 +212,110 @@ func runLoop(args []string) int {
|
||||
return 0
|
||||
}
|
||||
|
||||
func runQuota(args []string) int {
|
||||
sub := "status"
|
||||
var rest []string
|
||||
if len(args) > 0 && !strings.HasPrefix(args[0], "-") {
|
||||
sub = args[0]
|
||||
rest = args[1:]
|
||||
}
|
||||
switch sub {
|
||||
case "status", "probe", "gate":
|
||||
default:
|
||||
fmt.Fprintf(os.Stderr, "harness: quota: unknown subcommand %q (want status|probe|gate)\n", sub)
|
||||
return 1
|
||||
}
|
||||
fs := flag.NewFlagSet("quota "+sub, flag.ContinueOnError)
|
||||
cfgPath := fs.String("config", "", "config file path")
|
||||
if err := fs.Parse(rest); err != nil {
|
||||
return 1
|
||||
}
|
||||
if fs.NArg() > 0 {
|
||||
fmt.Fprintf(os.Stderr, "harness: unexpected argument %q\n", fs.Arg(0))
|
||||
return 1
|
||||
}
|
||||
if *cfgPath == "" {
|
||||
*cfgPath = os.Getenv("HARNESS_CONFIG")
|
||||
}
|
||||
if *cfgPath == "" {
|
||||
*cfgPath = "harness.toml"
|
||||
}
|
||||
cfg, err := config.Load(*cfgPath)
|
||||
if err != nil {
|
||||
fmt.Fprintf(os.Stderr, "harness: %v\n", err)
|
||||
return 1
|
||||
}
|
||||
conductor, err := loop.New(cfg, os.Stdout)
|
||||
if err != nil {
|
||||
fmt.Fprintf(os.Stderr, "harness: %v\n", err)
|
||||
return 1
|
||||
}
|
||||
gate := conductor.Gate()
|
||||
|
||||
switch sub {
|
||||
case "probe":
|
||||
if gate == nil || cfg.Quota.UsageURL == "" {
|
||||
fmt.Println("harness: quota: no [quota] usage_url configured (running on estimates)")
|
||||
return 0
|
||||
}
|
||||
snap := gate.Snapshot(context.Background())
|
||||
fmt.Printf("harness: quota: account=%s source=%s fetched=%s\n", snap.Account, snap.Source, snap.FetchedAt.Format(time.RFC3339))
|
||||
for _, b := range snap.Buckets {
|
||||
reset := "-"
|
||||
if !b.WindowReset.IsZero() {
|
||||
reset = b.WindowReset.Format(time.RFC3339)
|
||||
}
|
||||
fmt.Printf(" %-7s used=%.0f limit=%.0f (%.1f%%) resets=%s\n", b.ID, b.Used, b.Limit, b.UsedPct(), reset)
|
||||
}
|
||||
return 0
|
||||
|
||||
case "gate":
|
||||
if gate == nil {
|
||||
fmt.Println("harness: quota gate off ([quota] enabled = false)")
|
||||
return 0
|
||||
}
|
||||
fmt.Printf("harness: quota gate: %s\n", conductor.GateStatusLine(context.Background()))
|
||||
for _, class := range sortedClasses(cfg) {
|
||||
d := gate.Decide(context.Background(), class)
|
||||
tag := "ALLOW"
|
||||
if d.Action == quota.ActionDefer {
|
||||
tag = "DEFER"
|
||||
}
|
||||
fmt.Printf(" %-6s class=%-12s %s\n", tag, class, d.Reason)
|
||||
}
|
||||
return 0
|
||||
}
|
||||
|
||||
// status
|
||||
if gate == nil && !cfg.Resources.Enabled {
|
||||
fmt.Println("harness: quota gate off ([quota]/[resources] enabled = false)")
|
||||
} else if gate != nil {
|
||||
fmt.Printf("harness: quota: %s\n", conductor.GateStatusLine(context.Background()))
|
||||
fmt.Printf(" peak window: %s-%s %s weekdays_only=%v currently=%v\n",
|
||||
cfg.Quota.PeakStart, cfg.Quota.PeakEnd, cfg.Quota.Timezone, cfg.Quota.PeakWeekdaysOnly, gate.InPeak())
|
||||
}
|
||||
if st, ok := conductor.ResourceSample(); ok {
|
||||
fmt.Printf("harness: resources: load=%.2f memAvail=%.0fMB diskFree=%.0fMB ioDelay=%.1f%%(psi=%v)\n",
|
||||
st.LoadAvg1m, st.MemAvailable, st.DiskFree, st.IODelayPct, st.HasPSI)
|
||||
}
|
||||
report, err := loop.UsageReport(cfg.Loop.StateDir)
|
||||
if err != nil {
|
||||
fmt.Fprintf(os.Stderr, "harness: usage report: %v\n", err)
|
||||
return 0
|
||||
}
|
||||
fmt.Printf("harness: usage by class (%s):\n%s", cfg.Loop.StateDir, report)
|
||||
return 0
|
||||
}
|
||||
|
||||
func sortedClasses(cfg *config.Config) []string {
|
||||
out := make([]string, 0, len(cfg.Models.Classes))
|
||||
for class := range cfg.Models.Classes {
|
||||
out = append(out, class)
|
||||
}
|
||||
sort.Strings(out)
|
||||
return out
|
||||
}
|
||||
|
||||
func runEvents(args []string) int {
|
||||
fs := flag.NewFlagSet("events", flag.ContinueOnError)
|
||||
cfgPath := fs.String("config", "", "config file path")
|
||||
|
||||
@@ -165,3 +165,42 @@ secret_ref = "env:HARNESS_GITEA_WEBHOOK_SECRET"
|
||||
listen = ":8090" # publish on the LAN via docker -p
|
||||
vkey_ref = "env:HARNESS_SERVE_VKEY"
|
||||
# enabled_models = ["mopac-primary", "mopac-study"] # optional subset
|
||||
|
||||
# QUOTA GATE (Redmine 490, off by default until verified live): z.ai coding
|
||||
# plan credit buckets (5h + weekly), back-pressure thresholds, TZ-aware peak
|
||||
# window, and the redis shared-state hop so all harness instances of one
|
||||
# account coordinate. The loop DEFERS gated work with a logged reason; it
|
||||
# never hard-fails. With no usage_url it runs on locally estimated
|
||||
# consumption against the configured plan limits (see README: z.ai has no
|
||||
# public usage endpoint yet — LIVE VERIFICATION open).
|
||||
# [quota]
|
||||
# enabled = true
|
||||
# account = "zai-max-1" # plan label shared across instances
|
||||
# plan_5h_credits = 28000 # Max plan; Lite 2000/10000, Pro 12000/60000
|
||||
# plan_weekly_credits = 140000
|
||||
# usage_url = "" # set when z.ai ships the endpoint
|
||||
# key_ref = "env:HARNESS_ZAI_KEY" # bearer for usage_url; never logged
|
||||
# poll_interval_secs = 300
|
||||
# defer_at_pct = 85 # heavy classes defer, flash tier continues
|
||||
# block_at_pct = 95 # everything defers until reset (the wall)
|
||||
# # Peak = z.ai Mon-Fri 14:00-18:00 Singapore == 01:00-05:00 America/Chicago
|
||||
# # in winter (00:00-04:00 during US DST — adjust in March/November).
|
||||
# peak_start = "01:00"
|
||||
# peak_end = "05:00"
|
||||
# timezone = "America/Chicago"
|
||||
# peak_weekdays_only = true
|
||||
# peak_classes = ["study", "read"] # flash/LLM-lite classes allowed in peak
|
||||
# # Shared state for the 9 instances across 2 hosts (redis docker container
|
||||
# # on 192.168.3.78; see README runbook). Empty = local-only estimates.
|
||||
# redis_url = "redis://192.168.3.78:6390/0"
|
||||
|
||||
# RESOURCE GATE (Redmine 491, off by default): read-only host monitor —
|
||||
# loadavg, mem available, work_root disk free, IO pressure (/proc/pressure
|
||||
# io, skipped when PSI is absent). The loop defers dispatch while busy.
|
||||
# cgroup enforcement is deploy-time (README runbook).
|
||||
# [resources]
|
||||
# enabled = true
|
||||
# max_load_avg = 6.0
|
||||
# min_mem_available_mb = 2048
|
||||
# min_disk_free_mb = 5120
|
||||
# max_io_delay_pct = 90.0
|
||||
|
||||
@@ -2,7 +2,9 @@ package config
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"strconv"
|
||||
"strings"
|
||||
"time"
|
||||
)
|
||||
|
||||
// Config is the full harness.toml surface. Defaults live in Default();
|
||||
@@ -21,6 +23,8 @@ type Config struct {
|
||||
Serve ServeConfig
|
||||
KeyProxy KeyProxyConfig
|
||||
Gitea GiteaConfig
|
||||
Quota QuotaConfig
|
||||
Resources ResourcesConfig
|
||||
}
|
||||
|
||||
// LoopConfig is both the per-turn bound and the `harness loop` daemon
|
||||
@@ -105,6 +109,65 @@ type GiteaConfig struct {
|
||||
CommitReports bool
|
||||
}
|
||||
|
||||
// QuotaConfig is the z.ai coding-plan quota gate: bucket limits, polling,
|
||||
// back-pressure thresholds, the peak-hour schedule, and the optional
|
||||
// shared-state hop (redis container) for multi-instance coordination.
|
||||
// The quota gate NEVER hard-fails the loop; when it cannot see quota state
|
||||
// it degrades to permissive (allow) with a logged caveat.
|
||||
type QuotaConfig struct {
|
||||
Enabled bool
|
||||
|
||||
// Account labels this instance's z.ai plan in the shared state so the
|
||||
// nine harness instances across two hosts do not cross-contaminate.
|
||||
Account string
|
||||
|
||||
// Plan credit buckets (docs: Lite 2k/10k, Pro 12k/60k, Max 28k/140k).
|
||||
Plan5hCredits float64
|
||||
PlanWeeklyCredits float64
|
||||
|
||||
// UsageURL + KeyRef: the z.ai usage/limits endpoint (bearer auth; the
|
||||
// key is a ref, resolved at call time, never logged). Empty = no
|
||||
// upstream polling; the gate runs on locally estimated consumption.
|
||||
UsageURL string
|
||||
KeyRef string
|
||||
PollIntervalSecs int
|
||||
|
||||
// Back-pressure thresholds, percent of bucket used. defer_at: heavy
|
||||
// classes defer, flash-tier continue. block_at: everything defers.
|
||||
DeferAtPct float64
|
||||
BlockAtPct float64
|
||||
|
||||
// Peak window (z.ai charges 3x-ish during peak; off-peak is 50% off).
|
||||
// Defaults verified against the z.ai docs: peak = Mon-Fri 14:00-18:00
|
||||
// Singapore (UTC+8) == 01:00-05:00 America/Chicago in winter.
|
||||
PeakStart string // "HH:MM" in Timezone
|
||||
PeakEnd string
|
||||
Timezone string // IANA name; default America/Chicago (CST)
|
||||
PeakWeekdaysOnly bool
|
||||
// PeakClasses: task classes allowed to run inside the peak window
|
||||
// (the LLM-lite / flash tier); every other class defers with a reason.
|
||||
PeakClasses []string
|
||||
|
||||
// RedisURL: shared quota state for all instances of this account
|
||||
// (docker container; no host packages). Empty = local-only estimate.
|
||||
RedisURL string
|
||||
}
|
||||
|
||||
// ResourcesConfig is the read-only system resource gate (Redmine 491):
|
||||
// load average, memory available, work-root disk free, IO pressure. The
|
||||
// loop defers dispatch while the host is busy; enforcement via cgroups is
|
||||
// a deploy-time concern (see the runbook in README).
|
||||
type ResourcesConfig struct {
|
||||
Enabled bool
|
||||
MaxLoadAvg float64 // 1m load average
|
||||
MinMemAvailableMB float64
|
||||
MinDiskFreeMB float64 // work_root filesystem
|
||||
MaxIODelayPct float64 // /proc/pressure/io "some avg60" percent
|
||||
// ProcRoot/SysRoot are seams for tests; default /proc and /sys.
|
||||
ProcRoot string
|
||||
SysRoot string
|
||||
}
|
||||
|
||||
// EventsConfig is the `harness events` webhook receiver surface.
|
||||
type EventsConfig struct {
|
||||
Listen string // bind address (publish via docker -p)
|
||||
@@ -157,6 +220,28 @@ func Default() *Config {
|
||||
},
|
||||
Serve: ServeConfig{Listen: ":8090"},
|
||||
KeyProxy: KeyProxyConfig{CacheTTLSecs: 60},
|
||||
Quota: QuotaConfig{
|
||||
Enabled: false, // opt-in per instance until verified live
|
||||
Plan5hCredits: 28000, // Max plan
|
||||
PlanWeeklyCredits: 140000,
|
||||
PollIntervalSecs: 300,
|
||||
DeferAtPct: 85,
|
||||
BlockAtPct: 95,
|
||||
PeakStart: "01:00",
|
||||
PeakEnd: "05:00",
|
||||
Timezone: "America/Chicago",
|
||||
PeakWeekdaysOnly: true,
|
||||
PeakClasses: []string{"study", "read"},
|
||||
},
|
||||
Resources: ResourcesConfig{
|
||||
Enabled: false,
|
||||
MaxLoadAvg: 6.0,
|
||||
MinMemAvailableMB: 2048,
|
||||
MinDiskFreeMB: 5120,
|
||||
MaxIODelayPct: 90.0,
|
||||
ProcRoot: "/proc",
|
||||
SysRoot: "/sys",
|
||||
},
|
||||
Demo: DemoConfig{
|
||||
ID: "demo-1",
|
||||
Subject: "MVP demo: GLM self-description",
|
||||
@@ -343,6 +428,76 @@ func (c *Config) apply(doc TOMLDoc) error {
|
||||
if v, ok := sv.StringList("enabled_models"); ok {
|
||||
c.Serve.EnabledModels = v
|
||||
}
|
||||
|
||||
qt := doc.Table("quota")
|
||||
if v, ok := qt.Bool("enabled"); ok {
|
||||
c.Quota.Enabled = v
|
||||
}
|
||||
if v, ok := qt.String("account"); ok {
|
||||
c.Quota.Account = v
|
||||
}
|
||||
if v, ok := qt.Float("plan_5h_credits"); ok {
|
||||
c.Quota.Plan5hCredits = v
|
||||
}
|
||||
if v, ok := qt.Float("plan_weekly_credits"); ok {
|
||||
c.Quota.PlanWeeklyCredits = v
|
||||
}
|
||||
if v, ok := qt.String("usage_url"); ok {
|
||||
c.Quota.UsageURL = v
|
||||
}
|
||||
if v, ok := qt.String("key_ref"); ok {
|
||||
c.Quota.KeyRef = v
|
||||
}
|
||||
if v, ok := qt.Int("poll_interval_secs"); ok {
|
||||
c.Quota.PollIntervalSecs = int(v)
|
||||
}
|
||||
if v, ok := qt.Float("defer_at_pct"); ok {
|
||||
c.Quota.DeferAtPct = v
|
||||
}
|
||||
if v, ok := qt.Float("block_at_pct"); ok {
|
||||
c.Quota.BlockAtPct = v
|
||||
}
|
||||
if v, ok := qt.String("peak_start"); ok {
|
||||
c.Quota.PeakStart = v
|
||||
}
|
||||
if v, ok := qt.String("peak_end"); ok {
|
||||
c.Quota.PeakEnd = v
|
||||
}
|
||||
if v, ok := qt.String("timezone"); ok {
|
||||
c.Quota.Timezone = v
|
||||
}
|
||||
if v, ok := qt.Bool("peak_weekdays_only"); ok {
|
||||
c.Quota.PeakWeekdaysOnly = v
|
||||
}
|
||||
if v, ok := qt.StringList("peak_classes"); ok {
|
||||
c.Quota.PeakClasses = v
|
||||
}
|
||||
if v, ok := qt.String("redis_url"); ok {
|
||||
c.Quota.RedisURL = v
|
||||
}
|
||||
|
||||
rs := doc.Table("resources")
|
||||
if v, ok := rs.Bool("enabled"); ok {
|
||||
c.Resources.Enabled = v
|
||||
}
|
||||
if v, ok := rs.Float("max_load_avg"); ok {
|
||||
c.Resources.MaxLoadAvg = v
|
||||
}
|
||||
if v, ok := rs.Float("min_mem_available_mb"); ok {
|
||||
c.Resources.MinMemAvailableMB = v
|
||||
}
|
||||
if v, ok := rs.Float("min_disk_free_mb"); ok {
|
||||
c.Resources.MinDiskFreeMB = v
|
||||
}
|
||||
if v, ok := rs.Float("max_io_delay_pct"); ok {
|
||||
c.Resources.MaxIODelayPct = v
|
||||
}
|
||||
if v, ok := rs.String("proc_root"); ok {
|
||||
c.Resources.ProcRoot = v
|
||||
}
|
||||
if v, ok := rs.String("sys_root"); ok {
|
||||
c.Resources.SysRoot = v
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
@@ -484,5 +639,70 @@ func (c *Config) Validate() error {
|
||||
return fmt.Errorf("[gitea]: %w", err)
|
||||
}
|
||||
}
|
||||
|
||||
// The quota gate is optional; when on, buckets, thresholds and the peak
|
||||
// window must be coherent so the gate never divides by zero or wraps
|
||||
// nonsense windows. usage_url without key_ref is an error (bearer-only).
|
||||
if c.Quota.Enabled {
|
||||
if c.Quota.Plan5hCredits <= 0 || c.Quota.PlanWeeklyCredits <= 0 {
|
||||
return fmt.Errorf("[quota]: plan_5h_credits and plan_weekly_credits must be > 0")
|
||||
}
|
||||
if c.Quota.PollIntervalSecs < 1 {
|
||||
return fmt.Errorf("[quota]: poll_interval_secs must be >= 1")
|
||||
}
|
||||
if !(c.Quota.DeferAtPct > 0 && c.Quota.DeferAtPct < 100) || !(c.Quota.BlockAtPct > 0 && c.Quota.BlockAtPct <= 100) {
|
||||
return fmt.Errorf("[quota]: defer_at_pct and block_at_pct must be in (0,100]")
|
||||
}
|
||||
if c.Quota.DeferAtPct > c.Quota.BlockAtPct {
|
||||
return fmt.Errorf("[quota]: defer_at_pct (%.0f) must be <= block_at_pct (%.0f)", c.Quota.DeferAtPct, c.Quota.BlockAtPct)
|
||||
}
|
||||
if _, err := ParseHHMM(c.Quota.PeakStart); err != nil {
|
||||
return fmt.Errorf("[quota]: peak_start: %w", err)
|
||||
}
|
||||
if _, err := ParseHHMM(c.Quota.PeakEnd); err != nil {
|
||||
return fmt.Errorf("[quota]: peak_end: %w", err)
|
||||
}
|
||||
if _, err := time.LoadLocation(c.Quota.Timezone); err != nil {
|
||||
return fmt.Errorf("[quota]: unknown timezone %q", c.Quota.Timezone)
|
||||
}
|
||||
if c.Quota.Account == "" {
|
||||
return fmt.Errorf("[quota]: account is required when the quota gate is enabled (shared-state bucket label)")
|
||||
}
|
||||
if c.Quota.UsageURL != "" && c.Quota.KeyRef == "" {
|
||||
return fmt.Errorf("[quota]: key_ref is required when usage_url is set (bearer auth)")
|
||||
}
|
||||
if c.Quota.KeyRef != "" {
|
||||
if err := CheckKeyRef(c.Quota.KeyRef); err != nil {
|
||||
return fmt.Errorf("[quota]: %w", err)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// The resource gate is optional; when on, thresholds must be sane.
|
||||
if c.Resources.Enabled {
|
||||
if c.Resources.MaxLoadAvg <= 0 || c.Resources.MinMemAvailableMB <= 0 || c.Resources.MinDiskFreeMB <= 0 {
|
||||
return fmt.Errorf("[resources]: thresholds must be > 0")
|
||||
}
|
||||
if !(c.Resources.MaxIODelayPct > 0 && c.Resources.MaxIODelayPct <= 100) {
|
||||
return fmt.Errorf("[resources]: max_io_delay_pct must be in (0,100]")
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// ParseHHMM parses a wall-clock time of day, "HH:MM" (24h).
|
||||
func ParseHHMM(s string) (time.Duration, error) {
|
||||
parts := strings.Split(s, ":")
|
||||
if len(parts) != 2 || len(parts[0]) != 2 || len(parts[1]) != 2 {
|
||||
return 0, fmt.Errorf("want HH:MM, got %q", s)
|
||||
}
|
||||
h, err := strconv.Atoi(parts[0])
|
||||
if err != nil || h < 0 || h > 23 {
|
||||
return 0, fmt.Errorf("bad hour in %q", s)
|
||||
}
|
||||
m, err := strconv.Atoi(parts[1])
|
||||
if err != nil || m < 0 || m > 59 {
|
||||
return 0, fmt.Errorf("bad minute in %q", s)
|
||||
}
|
||||
return time.Duration(h)*time.Hour + time.Duration(m)*time.Minute, nil
|
||||
}
|
||||
|
||||
@@ -0,0 +1,131 @@
|
||||
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)
|
||||
}
|
||||
}
|
||||
+16
-3
@@ -10,6 +10,7 @@ import (
|
||||
"fmt"
|
||||
"os"
|
||||
"sort"
|
||||
"strconv"
|
||||
"strings"
|
||||
)
|
||||
|
||||
@@ -48,6 +49,15 @@ func (d TOMLDoc) Bool(key string) (bool, bool) {
|
||||
return v, ok
|
||||
}
|
||||
|
||||
func (d TOMLDoc) Float(key string) (float64, bool) {
|
||||
if v, ok := d[key].(float64); ok {
|
||||
return v, true
|
||||
}
|
||||
// Bare integers read as floats too (max_load_avg = 6 parses as int64).
|
||||
v, ok := d[key].(int64)
|
||||
return float64(v), ok
|
||||
}
|
||||
|
||||
func (d TOMLDoc) StringList(key string) ([]string, bool) {
|
||||
raw, ok := d[key].([]any)
|
||||
if !ok {
|
||||
@@ -280,10 +290,13 @@ func parseScalar(v string, no int) (any, error) {
|
||||
return false, nil
|
||||
}
|
||||
n, err := parseTOMLInt(v)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("harness.toml:%d: unsupported value %q (want string, int, bool, or array)", no, v)
|
||||
}
|
||||
if err == nil {
|
||||
return n, nil
|
||||
}
|
||||
if f, ferr := strconv.ParseFloat(strings.ReplaceAll(v, "_", ""), 64); ferr == nil {
|
||||
return f, nil
|
||||
}
|
||||
return nil, fmt.Errorf("harness.toml:%d: unsupported value %q (want string, int, float, bool, or array)", no, v)
|
||||
}
|
||||
|
||||
func parseTOMLInt(v string) (int64, error) {
|
||||
|
||||
+26
-1
@@ -15,6 +15,7 @@ import (
|
||||
"time"
|
||||
|
||||
"ukrrs.com/mopac/harness/internal/models"
|
||||
"ukrrs.com/mopac/harness/internal/quota"
|
||||
"ukrrs.com/mopac/harness/internal/task"
|
||||
"ukrrs.com/mopac/harness/internal/writeback"
|
||||
)
|
||||
@@ -69,6 +70,9 @@ func (c *Conductor) RunLoop(ctx context.Context, opts LoopOpts) error {
|
||||
|
||||
fmt.Fprintf(c.out, "harness: loop: vertical=%s poll=%s state=%s redmine=%s (SIGINT to stop)\n",
|
||||
c.cfg.Vertical, interval, c.cfg.Loop.StateDir, c.redmineHost())
|
||||
if status := c.gateStatusLine(ctx); status != "off" {
|
||||
fmt.Fprintf(c.out, "harness: loop: gate: %s\n", status)
|
||||
}
|
||||
if gitea != nil {
|
||||
fmt.Fprintf(c.out, "harness: loop: gitea report commit on (%s/%s)\n", c.cfg.Gitea.Owner, c.cfg.Gitea.Repo)
|
||||
}
|
||||
@@ -121,6 +125,13 @@ func (c *Conductor) scanOnce(ctx context.Context, state *loopState, writer *writ
|
||||
fmt.Fprintf(c.out, "harness: loop: would dispatch %s (updated %s): %q\n", t.ID, t.UpdatedOn, t.Subject)
|
||||
continue
|
||||
}
|
||||
// Back-pressure: quota buckets + peak schedule + host resources. A
|
||||
// defer skips the turn WITHOUT consuming the task — the next scan
|
||||
// reconsiders it (the 19:00-wall class of failure becomes a logged
|
||||
// throttle instead of dead turns).
|
||||
if !c.gateConsult(ctx, state, t) {
|
||||
continue
|
||||
}
|
||||
c.dispatchTask(ctx, state, writer, gitea, t)
|
||||
}
|
||||
return nil
|
||||
@@ -136,9 +147,23 @@ func (c *Conductor) dispatchTask(ctx context.Context, state *loopState, writer *
|
||||
|
||||
run, err := c.runTask(ctx, t)
|
||||
if run != nil && run.turn != nil && run.reportPath != "" {
|
||||
// Usage accounting: tokens + estimated z.ai credits, per class,
|
||||
// into the JSONL (feeds the Discourse usage reports) and the
|
||||
// shared quota state (all instances' estimates).
|
||||
credits := 0.0
|
||||
if c.gate != nil {
|
||||
credits = c.gate.RecordTurn(quota.EstimateTurnInput{
|
||||
Model: run.decision.Model,
|
||||
PromptTokens: run.turn.PromptTokens,
|
||||
CompletionTokens: run.turn.CompletionTokens,
|
||||
Peak: c.gate.InPeak(),
|
||||
})
|
||||
}
|
||||
_ = state.log(loopEvent{
|
||||
Type: evReport, TaskID: t.ID, Model: run.decision.Model,
|
||||
Type: evReport, TaskID: t.ID, Model: run.decision.Model, Class: t.Class,
|
||||
ReportPath: run.reportPath, StopReason: run.turn.StopReason,
|
||||
PromptTokens: run.turn.PromptTokens, CompletionTokens: run.turn.CompletionTokens,
|
||||
TotalTokens: run.turn.TotalTokens, Credits: credits,
|
||||
})
|
||||
}
|
||||
if err != nil {
|
||||
|
||||
@@ -27,6 +27,7 @@ type fakeRedmine struct {
|
||||
status string
|
||||
updatedOn string
|
||||
notes []string
|
||||
class string
|
||||
}
|
||||
statusIDs map[string]int
|
||||
noted []string // issue ids that got a note, in order
|
||||
@@ -41,6 +42,7 @@ func newFakeRedmine(t *testing.T) *fakeRedmine {
|
||||
status string
|
||||
updatedOn string
|
||||
notes []string
|
||||
class string
|
||||
}{},
|
||||
statusIDs: map[string]int{"New": 1, "In Progress": 2, "Done": 3},
|
||||
}
|
||||
@@ -56,8 +58,12 @@ func newFakeRedmine(t *testing.T) *fakeRedmine {
|
||||
b.WriteString(",")
|
||||
}
|
||||
first = false
|
||||
fmt.Fprintf(&b, `{"id":%s,"subject":%q,"description":"do the thing","updated_on":%q,"status":{"name":%q},"custom_fields":[{"name":"Class","value":"primary"}]}`,
|
||||
id, is.subject, is.updatedOn, is.status)
|
||||
class := is.class
|
||||
if class == "" {
|
||||
class = "primary"
|
||||
}
|
||||
fmt.Fprintf(&b, `{"id":%s,"subject":%q,"description":"do the thing","updated_on":%q,"status":{"name":%q},"custom_fields":[{"name":"Class","value":%q}]}`,
|
||||
id, is.subject, is.updatedOn, is.status, class)
|
||||
}
|
||||
b.WriteString(`]}`)
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
@@ -131,6 +137,11 @@ func newFakeRedmine(t *testing.T) *fakeRedmine {
|
||||
}
|
||||
|
||||
func (f *fakeRedmine) addIssue(id, subject, status, updatedOn string) {
|
||||
f.addIssueClass(id, subject, status, updatedOn, "primary")
|
||||
}
|
||||
|
||||
// addIssueClass registers an issue with an explicit Class custom field.
|
||||
func (f *fakeRedmine) addIssueClass(id, subject, status, updatedOn, class string) {
|
||||
f.mu.Lock()
|
||||
defer f.mu.Unlock()
|
||||
f.issues[id] = struct {
|
||||
@@ -138,7 +149,8 @@ func (f *fakeRedmine) addIssue(id, subject, status, updatedOn string) {
|
||||
status string
|
||||
updatedOn string
|
||||
notes []string
|
||||
}{subject: subject, status: status, updatedOn: updatedOn}
|
||||
class string
|
||||
}{subject: subject, status: status, updatedOn: updatedOn, class: class}
|
||||
}
|
||||
|
||||
func (f *fakeRedmine) setUpdatedOn(id, updatedOn string) {
|
||||
|
||||
@@ -0,0 +1,165 @@
|
||||
package loop
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"sort"
|
||||
"strings"
|
||||
|
||||
"ukrrs.com/mopac/harness/internal/config"
|
||||
"ukrrs.com/mopac/harness/internal/quota"
|
||||
"ukrrs.com/mopac/harness/internal/task"
|
||||
)
|
||||
|
||||
// buildGates constructs the quota gate and resource config for the
|
||||
// conductor. Quota/resources are opt-in ([quota] enabled / [resources]
|
||||
// enabled); failures to build (bad schedule, unreachable redis) are fatal
|
||||
// at startup because a silently missing gate is a silently burning plan.
|
||||
func buildGates(cfg *config.Config, keys *config.KeyResolver) (*quota.Gate, error) {
|
||||
if !cfg.Quota.Enabled && !cfg.Resources.Enabled {
|
||||
return nil, nil
|
||||
}
|
||||
var shared *quota.SharedState
|
||||
if cfg.Quota.RedisURL != "" {
|
||||
var err error
|
||||
shared, err = quota.NewSharedState(cfg.Quota.RedisURL)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("quota redis: %w", err)
|
||||
}
|
||||
}
|
||||
if !cfg.Quota.Enabled {
|
||||
return nil, nil
|
||||
}
|
||||
gate, err := quota.NewGate(cfg.Quota, keys, shared)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return gate, nil
|
||||
}
|
||||
|
||||
// gateConsult is the pre-dispatch back-pressure check: quota snapshot +
|
||||
// schedule + system resources. A defer never marks the task processed, so
|
||||
// the next scan reconsiders it (defer is throttle, not rejection); the
|
||||
// (task, updated_on, reason) tuple is logged once per process to keep the
|
||||
// JSONL free of per-scan spam.
|
||||
func (c *Conductor) gateConsult(ctx context.Context, state *loopState, t task.Task) bool {
|
||||
if c.gate == nil && !c.resEnabled() {
|
||||
return true
|
||||
}
|
||||
if reasons := c.resourceReasons(); len(reasons) > 0 {
|
||||
c.logDefer(state, t, "resources busy: "+strings.Join(reasons, "; "))
|
||||
return false
|
||||
}
|
||||
if c.gate == nil {
|
||||
return true
|
||||
}
|
||||
d := c.gate.Decide(ctx, t.Class)
|
||||
if d.Action == quota.ActionDefer {
|
||||
c.logDefer(state, t, d.Reason)
|
||||
return false
|
||||
}
|
||||
return true
|
||||
}
|
||||
|
||||
// resourceReasons samples the host and returns threshold violations ([] =
|
||||
// not busy, or the monitor is off/unreadable — read failures never defer).
|
||||
func (c *Conductor) resourceReasons() []string {
|
||||
rc := c.cfg.Resources
|
||||
if !rc.Enabled {
|
||||
return nil
|
||||
}
|
||||
st := quota.ReadResourceStats(rc.ProcRoot, rc.SysRoot, c.cfg.WorkRoot)
|
||||
return quota.BusyCheck{
|
||||
MaxLoadAvg: rc.MaxLoadAvg,
|
||||
MinMemAvailableMB: rc.MinMemAvailableMB,
|
||||
MinDiskFreeMB: rc.MinDiskFreeMB,
|
||||
MaxIODelayPct: rc.MaxIODelayPct,
|
||||
}.Evaluate(st)
|
||||
}
|
||||
|
||||
func (c *Conductor) resEnabled() bool { return c.cfg.Resources.Enabled }
|
||||
|
||||
// logDefer records one defer event (deduped per task+reason in-process).
|
||||
func (c *Conductor) logDefer(state *loopState, t task.Task, reason string) {
|
||||
key := t.ID + "|" + t.UpdatedOn + "|" + reason
|
||||
if c.deferLogged == nil {
|
||||
c.deferLogged = map[string]bool{}
|
||||
}
|
||||
if c.deferLogged[key] {
|
||||
return
|
||||
}
|
||||
c.deferLogged[key] = true
|
||||
_ = state.log(loopEvent{Type: evDefer, TaskID: t.ID, UpdatedOn: t.UpdatedOn, Subject: t.Subject, Detail: reason})
|
||||
fmt.Fprintf(c.out, "harness: loop: defer %s: %s (reconsidered next scan; never marked processed)\n", t.ID, reason)
|
||||
}
|
||||
|
||||
// gateStatusLine is the startup/`quota status` one-liner.
|
||||
func (c *Conductor) gateStatusLine(ctx context.Context) string {
|
||||
if c.gate == nil && !c.resEnabled() {
|
||||
return "off"
|
||||
}
|
||||
parts := []string{}
|
||||
if c.gate != nil {
|
||||
snap := c.gate.Snapshot(ctx)
|
||||
parts = append(parts, c.gate.StatusLine(snap))
|
||||
}
|
||||
if c.resEnabled() {
|
||||
st := quota.ReadResourceStats(c.cfg.Resources.ProcRoot, c.cfg.Resources.SysRoot, c.cfg.WorkRoot)
|
||||
parts = append(parts, fmt.Sprintf("load=%.2f memAvail=%.0fMB diskFree=%.0fMB", st.LoadAvg1m, st.MemAvailable, st.DiskFree))
|
||||
}
|
||||
return strings.Join(parts, " | ")
|
||||
}
|
||||
|
||||
// usageRow is one line of the per-class usage accounting (feeds the
|
||||
// per-instance Discourse usage reports).
|
||||
type usageRow struct {
|
||||
Class string `json:"class"`
|
||||
Turns int `json:"turns"`
|
||||
PromptTokens int `json:"prompt_tokens"`
|
||||
CompletionTokens int `json:"completion_tokens"`
|
||||
TotalTokens int `json:"total_tokens"`
|
||||
Credits float64 `json:"credits"`
|
||||
}
|
||||
|
||||
// UsageReport renders the per-class token+credit accounting table from a
|
||||
// loop state dir (read-only; `harness quota status`).
|
||||
func UsageReport(stateDir string) (string, error) {
|
||||
s, err := openLoopStateReadOnly(stateDir)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
defer s.Close()
|
||||
return usageReport(s.usageRows()), nil
|
||||
}
|
||||
|
||||
// GateStatusLine is the exported one-liner for `harness quota` output.
|
||||
func (c *Conductor) GateStatusLine(ctx context.Context) string {
|
||||
return c.gateStatusLine(ctx)
|
||||
}
|
||||
|
||||
// ResourceSample reads one host sample for the CLI (nil-safe when off).
|
||||
func (c *Conductor) ResourceSample() (quota.ResourceStats, bool) {
|
||||
if !c.resEnabled() {
|
||||
return quota.ResourceStats{}, false
|
||||
}
|
||||
rc := c.cfg.Resources
|
||||
return quota.ReadResourceStats(rc.ProcRoot, rc.SysRoot, c.cfg.WorkRoot), true
|
||||
}
|
||||
|
||||
// usageReport renders the accounting table from the loop state.
|
||||
func usageReport(rows []usageRow) string {
|
||||
sort.Slice(rows, func(i, j int) bool { return rows[i].Class < rows[j].Class })
|
||||
var b strings.Builder
|
||||
fmt.Fprintf(&b, "%-14s %6s %12s %12s %12s %10s\n", "CLASS", "TURNS", "PROMPT", "COMPLETION", "TOTAL", "CREDITS")
|
||||
var tot usageRow
|
||||
for _, r := range rows {
|
||||
fmt.Fprintf(&b, "%-14s %6d %12d %12d %12d %10.2f\n", r.Class, r.Turns, r.PromptTokens, r.CompletionTokens, r.TotalTokens, r.Credits)
|
||||
tot.Turns += r.Turns
|
||||
tot.PromptTokens += r.PromptTokens
|
||||
tot.CompletionTokens += r.CompletionTokens
|
||||
tot.TotalTokens += r.TotalTokens
|
||||
tot.Credits += r.Credits
|
||||
}
|
||||
fmt.Fprintf(&b, "%-14s %6d %12d %12d %12d %10.2f\n", "TOTAL", tot.Turns, tot.PromptTokens, tot.CompletionTokens, tot.TotalTokens, tot.Credits)
|
||||
return b.String()
|
||||
}
|
||||
@@ -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()
|
||||
}
|
||||
+13
-1
@@ -17,6 +17,7 @@ import (
|
||||
"ukrrs.com/mopac/harness/internal/intake"
|
||||
"ukrrs.com/mopac/harness/internal/llm"
|
||||
"ukrrs.com/mopac/harness/internal/models"
|
||||
"ukrrs.com/mopac/harness/internal/quota"
|
||||
"ukrrs.com/mopac/harness/internal/task"
|
||||
"ukrrs.com/mopac/harness/internal/tools"
|
||||
"ukrrs.com/mopac/harness/internal/writeback"
|
||||
@@ -35,6 +36,9 @@ type Conductor struct {
|
||||
bash *tools.BashTool
|
||||
keys *config.KeyResolver
|
||||
out io.Writer
|
||||
gate *quota.Gate
|
||||
// deferLogged dedups defer log lines per (task, updated_on, reason).
|
||||
deferLogged map[string]bool
|
||||
}
|
||||
|
||||
// New validates routing config and prepares the tool palette.
|
||||
@@ -57,9 +61,17 @@ func New(cfg *config.Config, out io.Writer) (*Conductor, error) {
|
||||
return nil, err
|
||||
}
|
||||
}
|
||||
return &Conductor{cfg: cfg, router: router, bash: bash, keys: config.NewKeyResolver(cfg), out: out}, nil
|
||||
keys := config.NewKeyResolver(cfg)
|
||||
gate, err := buildGates(cfg, keys)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return &Conductor{cfg: cfg, router: router, bash: bash, keys: keys, out: out, gate: gate, deferLogged: map[string]bool{}}, nil
|
||||
}
|
||||
|
||||
// Gate exposes the quota gate (nil when off) for `harness quota status`.
|
||||
func (c *Conductor) Gate() *quota.Gate { return c.gate }
|
||||
|
||||
// OnceOpts controls a single iteration.
|
||||
type OnceOpts struct {
|
||||
DryRun bool // intake + plan only, no LLM call
|
||||
|
||||
+57
-1
@@ -6,6 +6,7 @@ import (
|
||||
"fmt"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"sort"
|
||||
"sync"
|
||||
"time"
|
||||
)
|
||||
@@ -19,6 +20,7 @@ const (
|
||||
evCommit = "commit" // REPORT committed to gitea (optional step)
|
||||
evRefresh = "refresh" // dedup marker advanced to post-writeback updated_on
|
||||
evError = "error" // a step failed; loop continues
|
||||
evDefer = "defer" // dispatch deferred by the quota/resource gate (throttle, not rejection)
|
||||
)
|
||||
|
||||
// loopEvent is one line of the append-only loop log. It doubles as the
|
||||
@@ -31,10 +33,17 @@ type loopEvent struct {
|
||||
UpdatedOn string `json:"updated_on,omitempty"`
|
||||
Subject string `json:"subject,omitempty"`
|
||||
Model string `json:"model,omitempty"`
|
||||
Class string `json:"class,omitempty"` // usage accounting key
|
||||
ReportPath string `json:"report_path,omitempty"`
|
||||
StopReason string `json:"stop_reason,omitempty"`
|
||||
StatusFrom string `json:"status_from,omitempty"`
|
||||
StatusTo string `json:"status_to,omitempty"`
|
||||
// Per-turn usage accounting (report events): tokens + estimated z.ai
|
||||
// credits. Feeds the per-instance usage tables / Discourse reports.
|
||||
PromptTokens int `json:"prompt_tokens,omitempty"`
|
||||
CompletionTokens int `json:"completion_tokens,omitempty"`
|
||||
TotalTokens int `json:"total_tokens,omitempty"`
|
||||
Credits float64 `json:"credits,omitempty"`
|
||||
Detail string `json:"detail,omitempty"` // human context; never secrets
|
||||
}
|
||||
|
||||
@@ -47,9 +56,20 @@ type loopState struct {
|
||||
f *os.File
|
||||
path string
|
||||
seen map[string]string
|
||||
usage map[string]usageAcc // class -> per-class totals (report events)
|
||||
readOnly bool // dry-run: dedup checks work, writes are refused
|
||||
}
|
||||
|
||||
// usageAcc accumulates per-class turn accounting; it is rebuilt from the
|
||||
// JSONL at open so usage reporting survives restarts.
|
||||
type usageAcc struct {
|
||||
Turns int `json:"turns"`
|
||||
PromptTokens int `json:"prompt_tokens"`
|
||||
CompletionTokens int `json:"completion_tokens"`
|
||||
TotalTokens int `json:"total_tokens"`
|
||||
Credits float64 `json:"credits"`
|
||||
}
|
||||
|
||||
func openLoopState(stateDir string) (*loopState, error) {
|
||||
if err := os.MkdirAll(stateDir, 0o700); err != nil {
|
||||
return nil, fmt.Errorf("loop state dir: %w", err)
|
||||
@@ -59,7 +79,7 @@ func openLoopState(stateDir string) (*loopState, error) {
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("open loop log: %w", err)
|
||||
}
|
||||
s := &loopState{path: path, f: f, seen: make(map[string]string)}
|
||||
s := &loopState{path: path, f: f, seen: make(map[string]string), usage: make(map[string]usageAcc)}
|
||||
if err := s.loadIndex(); err != nil {
|
||||
f.Close()
|
||||
return nil, err
|
||||
@@ -73,6 +93,7 @@ func openLoopStateReadOnly(stateDir string) (*loopState, error) {
|
||||
s := &loopState{
|
||||
path: filepath.Join(stateDir, "loop.jsonl"),
|
||||
seen: make(map[string]string),
|
||||
usage: make(map[string]usageAcc),
|
||||
readOnly: true,
|
||||
}
|
||||
if err := s.loadIndex(); err != nil {
|
||||
@@ -100,6 +121,9 @@ func (s *loopState) loadIndex() error {
|
||||
if ev.Type == evDispatch || ev.Type == evRefresh {
|
||||
s.seen[ev.TaskID] = ev.UpdatedOn
|
||||
}
|
||||
if ev.Type == evReport {
|
||||
s.addUsage(ev.Class, ev)
|
||||
}
|
||||
}
|
||||
return sc.Err()
|
||||
}
|
||||
@@ -131,9 +155,41 @@ func (s *loopState) log(ev loopEvent) error {
|
||||
if ev.Type == evDispatch || ev.Type == evRefresh {
|
||||
s.seen[ev.TaskID] = ev.UpdatedOn
|
||||
}
|
||||
if ev.Type == evReport {
|
||||
s.addUsage(ev.Class, ev)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (s *loopState) addUsage(class string, ev loopEvent) {
|
||||
if class == "" {
|
||||
class = "(unknown)"
|
||||
}
|
||||
acc := s.usage[class]
|
||||
acc.Turns++
|
||||
acc.PromptTokens += ev.PromptTokens
|
||||
acc.CompletionTokens += ev.CompletionTokens
|
||||
acc.TotalTokens += ev.TotalTokens
|
||||
acc.Credits += ev.Credits
|
||||
s.usage[class] = acc
|
||||
}
|
||||
|
||||
// usageRows snapshots the per-class accounting (sorted by class).
|
||||
func (s *loopState) usageRows() []usageRow {
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
rows := make([]usageRow, 0, len(s.usage))
|
||||
for class, acc := range s.usage {
|
||||
rows = append(rows, usageRow{
|
||||
Class: class, Turns: acc.Turns,
|
||||
PromptTokens: acc.PromptTokens, CompletionTokens: acc.CompletionTokens,
|
||||
TotalTokens: acc.TotalTokens, Credits: acc.Credits,
|
||||
})
|
||||
}
|
||||
sort.Slice(rows, func(i, j int) bool { return rows[i].Class < rows[j].Class })
|
||||
return rows
|
||||
}
|
||||
|
||||
func (s *loopState) Close() error {
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
|
||||
@@ -0,0 +1,48 @@
|
||||
package quota
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"io"
|
||||
"net/http"
|
||||
"strings"
|
||||
)
|
||||
|
||||
// FetchUsage GETs the provider usage endpoint with bearer auth. The key is
|
||||
// resolved per call from the [quota] key_ref and never logged; error paths
|
||||
// carry only status codes and redacted/truncated bodies (same discipline as
|
||||
// the keyproxy hop).
|
||||
func FetchUsage(ctx context.Context, client *http.Client, usageURL, key string) ([]byte, error) {
|
||||
req, err := http.NewRequestWithContext(ctx, http.MethodGet, usageURL, nil)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("usage request: %w", err)
|
||||
}
|
||||
req.Header.Set("Authorization", "Bearer "+key)
|
||||
req.Header.Set("Accept", "application/json")
|
||||
resp, err := client.Do(req)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("usage poll: %w", redactQuery(err))
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
body, _ := io.ReadAll(io.LimitReader(resp.Body, 1<<20))
|
||||
if resp.StatusCode < 200 || resp.StatusCode >= 300 {
|
||||
return nil, fmt.Errorf("usage poll: HTTP %d: %s", resp.StatusCode, truncateBody(string(body)))
|
||||
}
|
||||
return body, nil
|
||||
}
|
||||
|
||||
func redactQuery(err error) error {
|
||||
msg := err.Error()
|
||||
if i := strings.Index(msg, "?"); i >= 0 {
|
||||
msg = msg[:i] + "?..."
|
||||
}
|
||||
return fmt.Errorf("%s", msg)
|
||||
}
|
||||
|
||||
func truncateBody(s string) string {
|
||||
s = strings.TrimSpace(s)
|
||||
if len(s) > 200 {
|
||||
s = s[:200] + "..."
|
||||
}
|
||||
return s
|
||||
}
|
||||
@@ -0,0 +1,288 @@
|
||||
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
|
||||
}
|
||||
@@ -0,0 +1,74 @@
|
||||
package quota
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"math"
|
||||
"strings"
|
||||
)
|
||||
|
||||
// Multipliers are the z.ai credit formula coefficients (per 10k tokens,
|
||||
// docs.z.ai/devpack/overview). Family selects by concrete model name; the
|
||||
// flash tier (GLM-5.3-Flash, routed from glm-4.7*) is the cheap one.
|
||||
type Multipliers struct {
|
||||
Input float64
|
||||
CachedInput float64
|
||||
Output float64
|
||||
}
|
||||
|
||||
// Flagship multipliers: GLM-5.3 (and everything auto-routed to it).
|
||||
func FlagshipMultipliers() Multipliers { return Multipliers{Input: 6.9, CachedInput: 1.7, Output: 24} }
|
||||
|
||||
// Flash multipliers: GLM-5.3-Flash (and glm-4.7* routing).
|
||||
func FlashMultipliers() Multipliers { return Multipliers{Input: 2.3, CachedInput: 0.56, Output: 8} }
|
||||
|
||||
// MultipliersFor maps a concrete proxy model name to its credit multiplier
|
||||
// family. Unknown models read as flagship (conservative: overestimate cost
|
||||
// rather than silently burn quota).
|
||||
func MultipliersFor(model string) Multipliers {
|
||||
m := strings.ToLower(model)
|
||||
switch {
|
||||
case strings.Contains(m, "flash"):
|
||||
return FlashMultipliers()
|
||||
default:
|
||||
return FlagshipMultipliers()
|
||||
}
|
||||
}
|
||||
|
||||
// EstimateCredits computes the z.ai credits one turn consumed from its token
|
||||
// usage: (input*in + cached*cache + output*out) / 10000, halved when the
|
||||
// turn ran off-peak (z.ai charges 50% outside peak hours).
|
||||
func EstimateCredits(model string, promptTokens, cachedTokens, completionTokens int, peak bool) float64 {
|
||||
mult := MultipliersFor(model)
|
||||
credits := (float64(promptTokens)*mult.Input +
|
||||
float64(cachedTokens)*mult.CachedInput +
|
||||
float64(completionTokens)*mult.Output) / 10000
|
||||
if !peak {
|
||||
credits /= 2
|
||||
}
|
||||
// Round to 6 decimals: keeps redis INCRBYFLOAT values readable and the
|
||||
// JSONL compact; sub-microcredit noise is meaningless.
|
||||
return math.Round(credits*1e6) / 1e6
|
||||
}
|
||||
|
||||
// EstimateTurnInput is the per-turn usage record the loop hands the gate.
|
||||
type EstimateTurnInput struct {
|
||||
Model string
|
||||
PromptTokens int
|
||||
CachedTokens int
|
||||
CompletionTokens int
|
||||
Peak bool // turn ran inside the peak window
|
||||
}
|
||||
|
||||
// Describe renders a human summary of one turn's credit cost (logs, REPORTs).
|
||||
func (e EstimateTurnInput) Describe() string {
|
||||
return fmt.Sprintf("%s: %d/%d/%d tokens (in/cached/out) %s = %.4f credits",
|
||||
e.Model, e.PromptTokens, e.CachedTokens, e.CompletionTokens,
|
||||
peakTag(e.Peak), EstimateCredits(e.Model, e.PromptTokens, e.CachedTokens, e.CompletionTokens, e.Peak))
|
||||
}
|
||||
|
||||
func peakTag(peak bool) string {
|
||||
if peak {
|
||||
return "peak"
|
||||
}
|
||||
return "off-peak"
|
||||
}
|
||||
@@ -0,0 +1,143 @@
|
||||
package quota
|
||||
|
||||
import (
|
||||
"bufio"
|
||||
"fmt"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strconv"
|
||||
"strings"
|
||||
"syscall"
|
||||
)
|
||||
|
||||
// ResourceStats is one read-only sample of host load: the numbers the
|
||||
// resource gate compares against its thresholds. Zero fields mean "could
|
||||
// not read" and never trigger a busy verdict on their own.
|
||||
type ResourceStats struct {
|
||||
LoadAvg1m float64
|
||||
MemAvailable float64 // MB
|
||||
DiskFree float64 // MB, at path
|
||||
IODelayPct float64 // /proc/pressure/io "some avg60" percent; 0 if PSI absent
|
||||
HasPSI bool
|
||||
}
|
||||
|
||||
// ReadResourceStats samples /proc/loadavg, /proc/meminfo, /proc/pressure/io
|
||||
// and statfs(path) — all read-only, no host packages. procRoot/sysRoot are
|
||||
// seams (tests point them at fixture trees; production passes /proc, /sys).
|
||||
func ReadResourceStats(procRoot, sysRoot, path string) ResourceStats {
|
||||
var st ResourceStats
|
||||
if la, ok := readLoadAvg(filepath.Join(procRoot, "loadavg")); ok {
|
||||
st.LoadAvg1m = la
|
||||
}
|
||||
if mb, ok := readMemAvailable(filepath.Join(procRoot, "meminfo")); ok {
|
||||
st.MemAvailable = mb
|
||||
}
|
||||
// PSI lives under /proc/pressure (io); older kernels expose none — the
|
||||
// gate then skips the IO check rather than failing.
|
||||
if p, ok := readIODelay(filepath.Join(procRoot, "pressure", "io")); ok {
|
||||
st.IODelayPct, st.HasPSI = p, true
|
||||
}
|
||||
_ = sysRoot // reserved: /sys/class/... sources when PSI is absent
|
||||
if df, ok := readDiskFree(path); ok {
|
||||
st.DiskFree = df
|
||||
}
|
||||
return st
|
||||
}
|
||||
|
||||
func readLoadAvg(path string) (float64, bool) {
|
||||
data, err := os.ReadFile(path)
|
||||
if err != nil {
|
||||
return 0, false
|
||||
}
|
||||
fields := strings.Fields(string(data))
|
||||
if len(fields) < 1 {
|
||||
return 0, false
|
||||
}
|
||||
f, err := strconv.ParseFloat(fields[0], 64)
|
||||
if err != nil {
|
||||
return 0, false
|
||||
}
|
||||
return f, true
|
||||
}
|
||||
|
||||
func readMemAvailable(path string) (float64, bool) {
|
||||
f, err := os.Open(path)
|
||||
if err != nil {
|
||||
return 0, false
|
||||
}
|
||||
defer f.Close()
|
||||
sc := bufio.NewScanner(f)
|
||||
for sc.Scan() {
|
||||
line := sc.Text()
|
||||
if strings.HasPrefix(line, "MemAvailable:") {
|
||||
fields := strings.Fields(line)
|
||||
if len(fields) < 2 {
|
||||
return 0, false
|
||||
}
|
||||
kb, err := strconv.ParseFloat(fields[1], 64)
|
||||
if err != nil {
|
||||
return 0, false
|
||||
}
|
||||
return kb / 1024, true
|
||||
}
|
||||
}
|
||||
return 0, false
|
||||
}
|
||||
|
||||
// readIODelay parses /proc/pressure/io, e.g.
|
||||
// "some avg10=0.00 avg60=0.12 avg300=0.05 total=123456789" — avg60 is the
|
||||
// steady-state signal (a build spiking IO shows here within a minute).
|
||||
func readIODelay(path string) (float64, bool) {
|
||||
data, err := os.ReadFile(path)
|
||||
if err != nil {
|
||||
return 0, false
|
||||
}
|
||||
for _, line := range strings.Split(string(data), "\n") {
|
||||
if !strings.HasPrefix(line, "some ") {
|
||||
continue
|
||||
}
|
||||
for _, field := range strings.Fields(line)[1:] {
|
||||
if strings.HasPrefix(field, "avg60=") {
|
||||
if v, err := strconv.ParseFloat(strings.TrimPrefix(field, "avg60="), 64); err == nil {
|
||||
return v, true
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
return 0, false
|
||||
}
|
||||
|
||||
func readDiskFree(path string) (float64, bool) {
|
||||
var fs syscall.Statfs_t
|
||||
if err := syscall.Statfs(path, &fs); err != nil {
|
||||
return 0, false
|
||||
}
|
||||
return float64(fs.Bavail) * float64(fs.Bsize) / (1024 * 1024), true
|
||||
}
|
||||
|
||||
// BusyCheck compares a sample against the resource thresholds. Violations
|
||||
// are collected (all reported, not just the first) — the loop defers while
|
||||
// any threshold trips, with every reason surfaced in the defer log line.
|
||||
type BusyCheck struct {
|
||||
MaxLoadAvg float64
|
||||
MinMemAvailableMB float64
|
||||
MinDiskFreeMB float64
|
||||
MaxIODelayPct float64
|
||||
}
|
||||
|
||||
func (c BusyCheck) Evaluate(st ResourceStats) []string {
|
||||
var reasons []string
|
||||
if c.MaxLoadAvg > 0 && st.LoadAvg1m > c.MaxLoadAvg {
|
||||
reasons = append(reasons, fmt.Sprintf("load %.2f > %.2f", st.LoadAvg1m, c.MaxLoadAvg))
|
||||
}
|
||||
if c.MinMemAvailableMB > 0 && st.MemAvailable > 0 && st.MemAvailable < c.MinMemAvailableMB {
|
||||
reasons = append(reasons, fmt.Sprintf("mem available %.0fMB < %.0fMB", st.MemAvailable, c.MinMemAvailableMB))
|
||||
}
|
||||
if c.MinDiskFreeMB > 0 && st.DiskFree > 0 && st.DiskFree < c.MinDiskFreeMB {
|
||||
reasons = append(reasons, fmt.Sprintf("disk free %.0fMB < %.0fMB", st.DiskFree, c.MinDiskFreeMB))
|
||||
}
|
||||
if c.MaxIODelayPct > 0 && st.HasPSI && st.IODelayPct > c.MaxIODelayPct {
|
||||
reasons = append(reasons, fmt.Sprintf("io delay %.1f%% > %.1f%%", st.IODelayPct, c.MaxIODelayPct))
|
||||
}
|
||||
return reasons
|
||||
}
|
||||
@@ -0,0 +1,276 @@
|
||||
package quota
|
||||
|
||||
import (
|
||||
"bufio"
|
||||
"fmt"
|
||||
"net"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strconv"
|
||||
"strings"
|
||||
"sync"
|
||||
"testing"
|
||||
"time"
|
||||
)
|
||||
|
||||
func TestResourceStatsFromFixture(t *testing.T) {
|
||||
root := t.TempDir()
|
||||
proc := filepath.Join(root, "proc")
|
||||
os.MkdirAll(filepath.Join(proc, "pressure"), 0o755)
|
||||
os.WriteFile(filepath.Join(proc, "loadavg"), []byte("7.32 5.10 2.20 3/900 12345\n"), 0o644)
|
||||
os.WriteFile(filepath.Join(proc, "meminfo"), []byte("MemTotal: 16000000 kB\nMemAvailable: 3145728 kB\n"), 0o644)
|
||||
os.WriteFile(filepath.Join(proc, "pressure", "io"), []byte("some avg10=1.00 avg60=2.50 avg300=0.10 total=123456789\nfull avg10=0.00 avg60=0.00 avg300=0.00 total=0\n"), 0o644)
|
||||
work := filepath.Join(root, "work")
|
||||
os.MkdirAll(work, 0o755)
|
||||
|
||||
st := ReadResourceStats(proc, filepath.Join(root, "sys"), work)
|
||||
if st.LoadAvg1m != 7.32 {
|
||||
t.Errorf("load = %v, want 7.32", st.LoadAvg1m)
|
||||
}
|
||||
if st.MemAvailable != 3072 {
|
||||
t.Errorf("mem = %v MB, want 3072", st.MemAvailable)
|
||||
}
|
||||
if !st.HasPSI || st.IODelayPct != 2.5 {
|
||||
t.Errorf("io = %v (psi=%v), want 2.5", st.IODelayPct, st.HasPSI)
|
||||
}
|
||||
if st.DiskFree <= 0 {
|
||||
t.Errorf("disk free = %v, want > 0 (statfs on tempdir)", st.DiskFree)
|
||||
}
|
||||
|
||||
check := BusyCheck{MaxLoadAvg: 6, MinMemAvailableMB: 4096, MinDiskFreeMB: 1, MaxIODelayPct: 90}
|
||||
reasons := check.Evaluate(st)
|
||||
joined := strings.Join(reasons, ";")
|
||||
for _, want := range []string{"load 7.32 > 6", "mem available 3072MB < 4096MB"} {
|
||||
if !strings.Contains(joined, want) {
|
||||
t.Errorf("reasons %q missing %q", joined, want)
|
||||
}
|
||||
}
|
||||
// IO under threshold: must NOT appear.
|
||||
if strings.Contains(joined, "io delay") {
|
||||
t.Errorf("io below threshold must not trip: %q", joined)
|
||||
}
|
||||
|
||||
// Missing PSI file: the IO check is skipped, not an error.
|
||||
os.Remove(filepath.Join(proc, "pressure", "io"))
|
||||
st2 := ReadResourceStats(proc, filepath.Join(root, "sys"), work)
|
||||
if st2.HasPSI {
|
||||
t.Error("missing PSI must read HasPSI=false")
|
||||
}
|
||||
if got := (BusyCheck{MaxIODelayPct: 1}).Evaluate(st2); len(got) != 0 {
|
||||
t.Errorf("no PSI must never trip IO: %v", got)
|
||||
}
|
||||
}
|
||||
|
||||
// TestResourceStatsNeverBusyOnUnreadable mirrors the deploy invariant: read
|
||||
// failures defer nothing (the gate is advisory, /proc layout changes must
|
||||
// not stop work).
|
||||
func TestResourceStatsNeverBusyOnUnreadable(t *testing.T) {
|
||||
root := t.TempDir() // empty: no loadavg, no meminfo
|
||||
st := ReadResourceStats(filepath.Join(root, "proc"), filepath.Join(root, "sys"), root)
|
||||
if got := (BusyCheck{MaxLoadAvg: 1, MinMemAvailableMB: 1, MinDiskFreeMB: 1, MaxIODelayPct: 1}).Evaluate(st); len(got) > 1 {
|
||||
t.Errorf("unreadable sources must not trip load/mem (disk may read): %v", got)
|
||||
}
|
||||
}
|
||||
|
||||
// fakeRedis is a minimal RESP2 server: enough command surface for the
|
||||
// SharedState client (SET/GET/INCRBYFLOAT/EXPIRE/SELECT), backed by a map.
|
||||
type fakeRedis struct {
|
||||
mu sync.Mutex
|
||||
data map[string]string
|
||||
ttl map[string]int64
|
||||
srv net.Listener
|
||||
}
|
||||
|
||||
func newFakeRedis(t *testing.T) *fakeRedis {
|
||||
t.Helper()
|
||||
f := &fakeRedis{data: map[string]string{}, ttl: map[string]int64{}}
|
||||
srv, err := net.Listen("tcp", "127.0.0.1:0")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
f.srv = srv
|
||||
go f.serve()
|
||||
t.Cleanup(func() { srv.Close() })
|
||||
return f
|
||||
}
|
||||
|
||||
func (f *fakeRedis) url(t *testing.T) string {
|
||||
t.Helper()
|
||||
return "redis://" + f.srv.Addr().String() + "/0"
|
||||
}
|
||||
|
||||
func (f *fakeRedis) serve() {
|
||||
for {
|
||||
conn, err := f.srv.Accept()
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
go f.handle(conn)
|
||||
}
|
||||
}
|
||||
|
||||
func (f *fakeRedis) handle(conn net.Conn) {
|
||||
defer conn.Close()
|
||||
r := bufio.NewReader(conn)
|
||||
w := bufio.NewWriter(conn)
|
||||
for {
|
||||
args, err := readCommand(r)
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
if len(args) == 0 {
|
||||
continue
|
||||
}
|
||||
f.mu.Lock()
|
||||
reply := f.exec(args)
|
||||
f.mu.Unlock()
|
||||
w.WriteString(reply + "\r\n")
|
||||
if w.Flush() != nil {
|
||||
return
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func readCommand(r *bufio.Reader) ([]string, error) {
|
||||
line, err := r.ReadString('\n')
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
line = strings.TrimRight(line, "\r\n")
|
||||
if !strings.HasPrefix(line, "*") {
|
||||
return strings.Fields(line), nil
|
||||
}
|
||||
n, _ := strconv.Atoi(line[1:])
|
||||
args := make([]string, 0, n)
|
||||
for i := 0; i < n; i++ {
|
||||
hl, err := r.ReadString('\n')
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
hl = strings.TrimRight(hl, "\r\n")
|
||||
ln, _ := strconv.Atoi(strings.TrimPrefix(hl, "$"))
|
||||
buf := make([]byte, ln+2)
|
||||
if _, err := readFull(r, buf); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
args = append(args, string(buf[:ln]))
|
||||
}
|
||||
return args, nil
|
||||
}
|
||||
|
||||
func readFull(r *bufio.Reader, buf []byte) (int, error) {
|
||||
total := 0
|
||||
for total < len(buf) {
|
||||
n, err := r.Read(buf[total:])
|
||||
total += n
|
||||
if err != nil {
|
||||
return total, err
|
||||
}
|
||||
}
|
||||
return total, nil
|
||||
}
|
||||
|
||||
func (f *fakeRedis) exec(args []string) string {
|
||||
cmd := strings.ToUpper(args[0])
|
||||
switch cmd {
|
||||
case "PING":
|
||||
return "+PONG"
|
||||
case "SELECT":
|
||||
return "+OK"
|
||||
case "SET":
|
||||
f.data[args[1]] = args[2]
|
||||
if len(args) >= 5 && strings.ToUpper(args[3]) == "EX" {
|
||||
f.ttl[args[1]], _ = strconv.ParseInt(args[4], 10, 64)
|
||||
}
|
||||
return "+OK"
|
||||
case "GET":
|
||||
if v, ok := f.data[args[1]]; ok {
|
||||
return fmt.Sprintf("$%d\r\n%s", len(v), v)
|
||||
}
|
||||
return "$-1"
|
||||
case "EXPIRE":
|
||||
if _, ok := f.data[args[1]]; ok {
|
||||
f.ttl[args[1]], _ = strconv.ParseInt(args[2], 10, 64)
|
||||
return ":1"
|
||||
}
|
||||
return ":0"
|
||||
case "INCRBYFLOAT":
|
||||
cur := 0.0
|
||||
if v, ok := f.data[args[1]]; ok {
|
||||
cur, _ = strconv.ParseFloat(v, 64)
|
||||
}
|
||||
delta, err := strconv.ParseFloat(args[2], 64)
|
||||
if err != nil {
|
||||
return "-ERR bad delta"
|
||||
}
|
||||
cur += delta
|
||||
f.data[args[1]] = strconv.FormatFloat(cur, 'f', 6, 64)
|
||||
return "$" + strconv.Itoa(len(f.data[args[1]])) + "\r\n" + f.data[args[1]]
|
||||
default:
|
||||
return "-ERR unknown command '" + cmd + "'"
|
||||
}
|
||||
}
|
||||
|
||||
func TestSharedStateRoundTrip(t *testing.T) {
|
||||
fake := newFakeRedis(t)
|
||||
st, err := NewSharedState(fake.url(t))
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer st.Close()
|
||||
|
||||
if err := st.setEx("k", "v1", time.Second); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if v, err := st.get("k"); err != nil || v != "v1" {
|
||||
t.Fatalf("get = %q err=%v", v, err)
|
||||
}
|
||||
if v, err := st.incrByFloat("ctr", 1.5); err != nil || v != 1.5 {
|
||||
t.Fatalf("incr = %v err=%v", v, err)
|
||||
}
|
||||
if v, err := st.incrByFloat("ctr", 2.25); err != nil || v != 3.75 {
|
||||
t.Fatalf("incr2 = %v err=%v (want 3.75)", v, err)
|
||||
}
|
||||
if v, err := st.get("missing"); err != nil || v != "" {
|
||||
t.Fatalf("missing key = %q err=%v (want empty, no error)", v, err)
|
||||
}
|
||||
|
||||
// snapshot publish/load through the typed helpers
|
||||
snap := &QuotaSnapshot{Account: "zai-9", Source: "usage_url", FetchedAt: time.Now(),
|
||||
Buckets: []Bucket{{ID: Bucket5h, Used: 3.5, Limit: 28000}}}
|
||||
if err := st.PublishSnapshot(snap, time.Minute); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
got, err := st.LoadSnapshot("zai-9")
|
||||
if err != nil || got == nil {
|
||||
t.Fatalf("load = %v err=%v", got, err)
|
||||
}
|
||||
if b, _ := got.Bucket(Bucket5h); b.Used != 3.5 {
|
||||
t.Errorf("bucket round trip: %+v", got.Buckets)
|
||||
}
|
||||
}
|
||||
|
||||
func TestSharedStateDownNeverFatal(t *testing.T) {
|
||||
// Nothing listens on this port: every op must error softly (nil state
|
||||
// ops are no-ops; a dead server returns errors, never panics).
|
||||
dead, err := NewSharedState("redis://127.0.0.1:1/0")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := dead.setEx("k", "v", time.Second); err == nil {
|
||||
t.Error("dead redis SET must error")
|
||||
}
|
||||
if _, err := dead.incrByFloat("k", 1); err == nil {
|
||||
t.Error("dead redis INCR must error")
|
||||
}
|
||||
var nilState *SharedState
|
||||
if err := nilState.setEx("k", "v", time.Second); err != nil {
|
||||
t.Errorf("nil state must no-op: %v", err)
|
||||
}
|
||||
if st, err := NewSharedState(""); err != nil || st != nil {
|
||||
t.Errorf("empty url must return nil state, got %v %v", st, err)
|
||||
}
|
||||
if _, err := NewSharedState("redis://"); err == nil {
|
||||
t.Error("hostless url must error")
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,77 @@
|
||||
package quota
|
||||
|
||||
import (
|
||||
_ "time/tzdata" // embedded zone database: TZ-aware windows work in scratch containers
|
||||
|
||||
"time"
|
||||
|
||||
"ukrrs.com/mopac/harness/internal/config"
|
||||
)
|
||||
|
||||
// Schedule is the TZ-aware peak window. z.ai peak hours are documented as
|
||||
// Monday-Friday 14:00-18:00 Singapore (UTC+8), which is 01:00-05:00
|
||||
// America/Chicago in winter (CST) / 00:00-04:00 during US DST — the default
|
||||
// matches Charles's "0100 to 0500 CST" and is configurable to the minute.
|
||||
// The window may wrap midnight (start > end): then it covers evenings of
|
||||
// the start day plus early mornings of the following day.
|
||||
type Schedule struct {
|
||||
Start, End time.Duration // minutes-since-midnight in Loc
|
||||
Loc *time.Location
|
||||
WeekdaysOnly bool
|
||||
}
|
||||
|
||||
// NewSchedule parses the window out of the [quota] config fields.
|
||||
func NewSchedule(peakStart, peakEnd, tzName string, weekdaysOnly bool) (Schedule, error) {
|
||||
start, err := config.ParseHHMM(peakStart)
|
||||
if err != nil {
|
||||
return Schedule{}, err
|
||||
}
|
||||
end, err := config.ParseHHMM(peakEnd)
|
||||
if err != nil {
|
||||
return Schedule{}, err
|
||||
}
|
||||
loc, err := time.LoadLocation(tzName)
|
||||
if err != nil {
|
||||
return Schedule{}, err
|
||||
}
|
||||
return Schedule{Start: start, End: end, Loc: loc, WeekdaysOnly: weekdaysOnly}, nil
|
||||
}
|
||||
|
||||
// SameDay reports whether t is inside the window without crossing midnight
|
||||
// (start <= t < end, single-day window).
|
||||
func (s Schedule) inWindow(mins time.Duration) bool {
|
||||
if s.Start <= s.End {
|
||||
return mins >= s.Start && mins < s.End
|
||||
}
|
||||
// Wrapped window: [start, 24h) of the start day, [0, end) of the next.
|
||||
return mins >= s.Start || mins < s.End
|
||||
}
|
||||
|
||||
// InPeak reports whether t falls inside the peak window, evaluated in the
|
||||
// schedule's timezone. For wrapped windows the weekday check applies to the
|
||||
// day the window STARTED on (a Sunday 22:00-02:00 window is off all week
|
||||
// when weekdays-only, because it starts on Sunday).
|
||||
func (s Schedule) InPeak(t time.Time) bool {
|
||||
local := t.In(s.Loc)
|
||||
mins := time.Duration(local.Hour())*time.Hour + time.Duration(local.Minute())*time.Minute
|
||||
if s.Start <= s.End {
|
||||
if !s.inWindow(mins) {
|
||||
return false
|
||||
}
|
||||
return !s.WeekdaysOnly || isWeekday(local)
|
||||
}
|
||||
// Evening half: today carries the window.
|
||||
if mins >= s.Start {
|
||||
return !s.WeekdaysOnly || isWeekday(local)
|
||||
}
|
||||
// Morning half: the window started yesterday.
|
||||
if mins < s.End {
|
||||
return !s.WeekdaysOnly || isWeekday(local.AddDate(0, 0, -1))
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
func isWeekday(t time.Time) bool {
|
||||
wd := t.Weekday()
|
||||
return wd >= time.Monday && wd <= time.Friday
|
||||
}
|
||||
@@ -0,0 +1,161 @@
|
||||
// Package quota is the z.ai coding-plan quota gate (Redmine 490) plus the
|
||||
// read-only system resource monitor (Redmine 491): it tracks plan credit
|
||||
// buckets (5-hour + weekly), computes per-turn credit consumption from the
|
||||
// documented z.ai formula, applies TZ-aware peak-hour scheduling and
|
||||
// back-pressure thresholds, and shares quota state across harness instances
|
||||
// through an optional redis container. The gate NEVER hard-fails the loop:
|
||||
// unknown state is permissive, exhaustion defers work with a logged reason.
|
||||
//
|
||||
// z.ai credit model (docs.z.ai/devpack/overview, 2026-08-29):
|
||||
//
|
||||
// credits = (input*in_mult + cached_input*cache_mult + output*out_mult) / 10000
|
||||
//
|
||||
// GLM-5.3: 6.9/1.7/24; GLM-5.3-Flash: 2.3/0.56/8. Off-peak hours charge 50%.
|
||||
// Plans: Lite 2,000/10,000, Pro 12,000/60,000, Max 28,000/140,000 credits
|
||||
// (5-hour / weekly). LIVE VERIFICATION of the usage endpoint is OPEN: z.ai
|
||||
// documents the buckets but no public usage REST route exists today
|
||||
// (probed 2026-08-29, see REPORT-20260829-0500-quota); the parser targets
|
||||
// the documented shape and runs against a fake server until z.ai ships it.
|
||||
package quota
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"time"
|
||||
)
|
||||
|
||||
// Bucket IDs for the z.ai coding-plan windows.
|
||||
const (
|
||||
Bucket5h = "5h"
|
||||
BucketWeekly = "weekly"
|
||||
)
|
||||
|
||||
// Bucket is one usage window: credits used against the plan limit and when
|
||||
// the window resets. Used/Limit are in z.ai credits, not tokens or dollars.
|
||||
type Bucket struct {
|
||||
ID string `json:"id"`
|
||||
Used float64 `json:"used"`
|
||||
Limit float64 `json:"limit"`
|
||||
WindowReset time.Time `json:"window_reset"`
|
||||
}
|
||||
|
||||
// UsedPct is the bucket's consumption ratio in percent of the limit.
|
||||
// A zero/negative limit reads as 0 (never blocks).
|
||||
func (b Bucket) UsedPct() float64 {
|
||||
if b.Limit <= 0 {
|
||||
return 0
|
||||
}
|
||||
return b.Used / b.Limit * 100
|
||||
}
|
||||
|
||||
// QuotaSnapshot is one account's quota state at a point in time, either
|
||||
// polled from the provider (Source "usage_url") or synthesized from locally
|
||||
// estimated consumption against the configured plan limits (Source
|
||||
// "estimate").
|
||||
type QuotaSnapshot struct {
|
||||
Account string `json:"account"`
|
||||
Source string `json:"source"`
|
||||
FetchedAt time.Time `json:"fetched_at"`
|
||||
Buckets []Bucket `json:"buckets"`
|
||||
}
|
||||
|
||||
// Bucket returns the bucket with the given id, if present.
|
||||
func (s *QuotaSnapshot) Bucket(id string) (Bucket, bool) {
|
||||
for _, b := range s.Buckets {
|
||||
if b.ID == id {
|
||||
return b, true
|
||||
}
|
||||
}
|
||||
return Bucket{}, false
|
||||
}
|
||||
|
||||
// MaxUsedPct is the highest consumption ratio across all buckets — the
|
||||
// number back-pressure thresholds compare against.
|
||||
func (s *QuotaSnapshot) MaxUsedPct() float64 {
|
||||
if s == nil {
|
||||
return 0
|
||||
}
|
||||
max := 0.0
|
||||
for _, b := range s.Buckets {
|
||||
if p := b.UsedPct(); p > max {
|
||||
max = p
|
||||
}
|
||||
}
|
||||
return max
|
||||
}
|
||||
|
||||
// providerUsage is the z.ai usage response shape this parser targets. z.ai
|
||||
// has not published the route yet; field aliases (used_credits|credits_used,
|
||||
// total_credits|limit_credits) keep the parser tolerant of either naming.
|
||||
type providerUsage struct {
|
||||
Usage struct {
|
||||
FiveHour providerBucket `json:"five_hour"`
|
||||
Weekly providerBucket `json:"weekly"`
|
||||
} `json:"usage"`
|
||||
}
|
||||
|
||||
type providerBucket struct {
|
||||
Used any `json:"used_credits"`
|
||||
UsedAlias any `json:"credits_used"`
|
||||
Limit any `json:"total_credits"`
|
||||
LimitAlias any `json:"limit_credits"`
|
||||
ResetAt string `json:"reset_at"`
|
||||
ResetAlias string `json:"reset_time"`
|
||||
}
|
||||
|
||||
// ParseUsage decodes a provider usage response into a QuotaSnapshot for the
|
||||
// account. It is strict about structure (unknown shapes error so a silent
|
||||
// HTML error page never reads as "quota fine") and lenient about field
|
||||
// naming within the documented buckets.
|
||||
func ParseUsage(account string, fetchedAt time.Time, body []byte) (*QuotaSnapshot, error) {
|
||||
var p providerUsage
|
||||
if err := json.Unmarshal(body, &p); err != nil {
|
||||
return nil, fmt.Errorf("usage decode: %w", err)
|
||||
}
|
||||
num := func(v any) (float64, bool) {
|
||||
switch n := v.(type) {
|
||||
case float64:
|
||||
return n, true
|
||||
case string:
|
||||
var f float64
|
||||
if _, err := fmt.Sscanf(n, "%g", &f); err == nil {
|
||||
return f, true
|
||||
}
|
||||
}
|
||||
return 0, false
|
||||
}
|
||||
snap := &QuotaSnapshot{Account: account, Source: "usage_url", FetchedAt: fetchedAt}
|
||||
for id, pb := range map[string]providerBucket{Bucket5h: p.Usage.FiveHour, BucketWeekly: p.Usage.Weekly} {
|
||||
used, ok := num(pb.Used)
|
||||
if !ok {
|
||||
used, ok = num(pb.UsedAlias)
|
||||
}
|
||||
if !ok {
|
||||
return nil, fmt.Errorf("usage decode: bucket %q has no used_credits", id)
|
||||
}
|
||||
limit, ok := num(pb.Limit)
|
||||
if !ok {
|
||||
limit, ok = num(pb.LimitAlias)
|
||||
}
|
||||
if !ok {
|
||||
return nil, fmt.Errorf("usage decode: bucket %q has no total_credits", id)
|
||||
}
|
||||
reset := pb.ResetAt
|
||||
if reset == "" {
|
||||
reset = pb.ResetAlias
|
||||
}
|
||||
rt := time.Time{}
|
||||
if reset != "" {
|
||||
var err error
|
||||
rt, err = time.Parse(time.RFC3339, reset)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("usage decode: bucket %q reset_at: %w", id, err)
|
||||
}
|
||||
}
|
||||
snap.Buckets = append(snap.Buckets, Bucket{ID: id, Used: used, Limit: limit, WindowReset: rt})
|
||||
}
|
||||
if len(snap.Buckets) != 2 {
|
||||
return nil, fmt.Errorf("usage decode: want five_hour and weekly buckets, got %d", len(snap.Buckets))
|
||||
}
|
||||
return snap, nil
|
||||
}
|
||||
@@ -0,0 +1,406 @@
|
||||
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)
|
||||
@@ -0,0 +1,233 @@
|
||||
package quota
|
||||
|
||||
import (
|
||||
"bufio"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"net"
|
||||
"net/url"
|
||||
"strconv"
|
||||
"strings"
|
||||
"sync"
|
||||
"time"
|
||||
)
|
||||
|
||||
// SharedState is the cross-instance quota state bus: one redis container on
|
||||
// the LAN (DECISION 2026-08-29, see REPORT-20260829-0500-quota) so the nine
|
||||
// harness instances across two hosts read/write the same quota snapshot and
|
||||
// credit estimates. It is a deliberately tiny stdlib-only RESP2 client
|
||||
// (SET/GET/INCRBYFLOAT/EXPIRE — the whole surface this package needs), no
|
||||
// host packages and no external Go deps. Every op fails soft: with redis
|
||||
// down the gate degrades to this instance's local estimate and the loop
|
||||
// keeps running.
|
||||
type SharedState struct {
|
||||
addr string
|
||||
db int
|
||||
mu sync.Mutex
|
||||
conn net.Conn
|
||||
rw *bufio.ReadWriter
|
||||
dialTO time.Duration
|
||||
}
|
||||
|
||||
// NewSharedState parses redis://host:port[/db]; empty URL returns nil (the
|
||||
// nil state is valid and fully functional as a no-op).
|
||||
func NewSharedState(redisURL string) (*SharedState, error) {
|
||||
if redisURL == "" {
|
||||
return nil, nil
|
||||
}
|
||||
u, err := url.Parse(redisURL)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("redis url: %w", err)
|
||||
}
|
||||
host := u.Host
|
||||
if host == "" {
|
||||
return nil, fmt.Errorf("redis url: missing host")
|
||||
}
|
||||
if !strings.Contains(host, ":") {
|
||||
host += ":6379"
|
||||
}
|
||||
db := 0
|
||||
if s := strings.TrimPrefix(u.Path, "/"); s != "" {
|
||||
if db, err = strconv.Atoi(s); err != nil {
|
||||
return nil, fmt.Errorf("redis url db: %w", err)
|
||||
}
|
||||
}
|
||||
return &SharedState{addr: host, db: db, dialTO: 3 * time.Second}, nil
|
||||
}
|
||||
|
||||
// Close releases the connection, if any.
|
||||
func (s *SharedState) Close() error {
|
||||
if s == nil {
|
||||
return nil
|
||||
}
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
if s.conn != nil {
|
||||
err := s.conn.Close()
|
||||
s.conn, s.rw = nil, nil
|
||||
return err
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// command runs one RESP command and returns the reply. Transport failures
|
||||
// reset the connection so the next op redials once.
|
||||
func (s *SharedState) command(args ...string) (string, error) {
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
if s.conn == nil {
|
||||
conn, err := net.DialTimeout("tcp", s.addr, s.dialTO)
|
||||
if err != nil {
|
||||
return "", fmt.Errorf("redis dial %s: %w", s.addr, err)
|
||||
}
|
||||
s.conn = conn
|
||||
s.rw = bufio.NewReadWriter(bufio.NewReader(conn), bufio.NewWriter(conn))
|
||||
if s.db != 0 {
|
||||
if _, err := s.execLocked("SELECT", strconv.Itoa(s.db)); err != nil {
|
||||
s.resetLocked()
|
||||
return "", err
|
||||
}
|
||||
}
|
||||
}
|
||||
reply, err := s.execLocked(args...)
|
||||
if err != nil {
|
||||
s.resetLocked()
|
||||
}
|
||||
return reply, err
|
||||
}
|
||||
|
||||
// execLocked writes one command and reads one reply (mutex held).
|
||||
func (s *SharedState) execLocked(args ...string) (string, error) {
|
||||
var b strings.Builder
|
||||
fmt.Fprintf(&b, "*%d\r\n", len(args))
|
||||
for _, a := range args {
|
||||
fmt.Fprintf(&b, "$%d\r\n%s\r\n", len(a), a)
|
||||
}
|
||||
if _, err := s.rw.WriteString(b.String()); err != nil {
|
||||
return "", fmt.Errorf("redis write: %w", err)
|
||||
}
|
||||
if err := s.rw.Flush(); err != nil {
|
||||
return "", fmt.Errorf("redis flush: %w", err)
|
||||
}
|
||||
return readReply(s.rw.Reader)
|
||||
}
|
||||
|
||||
func (s *SharedState) resetLocked() {
|
||||
if s.conn != nil {
|
||||
_ = s.conn.Close()
|
||||
}
|
||||
s.conn, s.rw = nil, nil
|
||||
}
|
||||
|
||||
// readReply parses one RESP2 reply: +simple / -error / :integer / $bulk.
|
||||
func readReply(r *bufio.Reader) (string, error) {
|
||||
line, err := r.ReadString('\n')
|
||||
if err != nil {
|
||||
return "", fmt.Errorf("redis read: %w", err)
|
||||
}
|
||||
line = strings.TrimRight(line, "\r\n")
|
||||
if line == "" {
|
||||
return "", errors.New("redis: empty reply")
|
||||
}
|
||||
switch line[0] {
|
||||
case '+', ':':
|
||||
return line[1:], nil
|
||||
case '-':
|
||||
return "", fmt.Errorf("redis: %s", line[1:])
|
||||
case '$':
|
||||
n, err := strconv.Atoi(line[1:])
|
||||
if err != nil {
|
||||
return "", fmt.Errorf("redis bulk len: %w", err)
|
||||
}
|
||||
if n < 0 {
|
||||
return "", nil // nil bulk: missing key
|
||||
}
|
||||
buf := make([]byte, n+2)
|
||||
if _, err := ioReadFull(r, buf); err != nil {
|
||||
return "", fmt.Errorf("redis bulk read: %w", err)
|
||||
}
|
||||
return string(buf[:n]), nil
|
||||
default:
|
||||
return "", fmt.Errorf("redis: unexpected reply %q", line)
|
||||
}
|
||||
}
|
||||
|
||||
func ioReadFull(r *bufio.Reader, buf []byte) (int, error) {
|
||||
total := 0
|
||||
for total < len(buf) {
|
||||
n, err := r.Read(buf[total:])
|
||||
total += n
|
||||
if err != nil {
|
||||
return total, err
|
||||
}
|
||||
}
|
||||
return total, nil
|
||||
}
|
||||
|
||||
// Key layout (one account's plan is one shared universe):
|
||||
//
|
||||
// mopac:quota:<account>:snapshot — latest polled QuotaSnapshot (JSON)
|
||||
// mopac:quota:<account>:est:5h — estimated credits in the rolling 5h window
|
||||
// mopac:quota:<account>:est:weekly — estimated credits in the plan week
|
||||
//
|
||||
// Estimate keys carry the window id in the VALUE-side bookkeeping done by
|
||||
// the caller (Gate): 5h keys expire after 5h+slack; weekly keys are keyed
|
||||
// by ISO week via the caller and need no expiry.
|
||||
|
||||
func (s *SharedState) get(key string) (string, error) {
|
||||
if s == nil {
|
||||
return "", nil
|
||||
}
|
||||
return s.command("GET", key)
|
||||
}
|
||||
|
||||
func (s *SharedState) setEx(key, val string, ttl time.Duration) error {
|
||||
if s == nil {
|
||||
return nil
|
||||
}
|
||||
_, err := s.command("SET", key, val, "EX", strconv.Itoa(int(ttl.Seconds())))
|
||||
return err
|
||||
}
|
||||
|
||||
// incrByFloat adds delta to key (creating at delta) and returns the new value.
|
||||
func (s *SharedState) incrByFloat(key string, delta float64) (float64, error) {
|
||||
if s == nil {
|
||||
return 0, nil
|
||||
}
|
||||
reply, err := s.command("INCRBYFLOAT", key, strconv.FormatFloat(delta, 'f', 6, 64))
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
return strconv.ParseFloat(strings.TrimSpace(reply), 64)
|
||||
}
|
||||
|
||||
// PublishSnapshot stores the polled snapshot for all instances to read.
|
||||
func (s *SharedState) PublishSnapshot(snap *QuotaSnapshot, ttl time.Duration) error {
|
||||
if s == nil || snap == nil {
|
||||
return nil
|
||||
}
|
||||
body, err := json.Marshal(snap)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
return s.setEx(snapshotKey(snap.Account), string(body), ttl)
|
||||
}
|
||||
|
||||
// LoadSnapshot returns the last published snapshot, or nil when absent.
|
||||
func (s *SharedState) LoadSnapshot(account string) (*QuotaSnapshot, error) {
|
||||
body, err := s.get(snapshotKey(account))
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if body == "" {
|
||||
return nil, nil
|
||||
}
|
||||
var snap QuotaSnapshot
|
||||
if err := json.Unmarshal([]byte(body), &snap); err != nil {
|
||||
return nil, fmt.Errorf("shared snapshot decode: %w", err)
|
||||
}
|
||||
return &snap, nil
|
||||
}
|
||||
|
||||
func snapshotKey(account string) string { return "mopac:quota:" + account + ":snapshot" }
|
||||
Reference in New Issue
Block a user