From fc518c475ea13712ba4235d36499bb7b57153e7e Mon Sep 17 00:00:00 2001 From: reachableceo Date: Sat, 29 Aug 2026 05:37:15 -0500 Subject: [PATCH] quota: z.ai credit-bucket back-pressure + resource gate + usage accounting (Redmine 490+491) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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 --- README.md | 111 ++++++++- cmd/harness/main.go | 119 ++++++++- harness.toml.example | 39 +++ internal/config/config.go | 220 +++++++++++++++++ internal/config/quota_test.go | 131 ++++++++++ internal/config/toml.go | 19 +- internal/loop/daemon.go | 27 +- internal/loop/daemon_test.go | 18 +- internal/loop/gate.go | 165 +++++++++++++ internal/loop/gate_test.go | 293 ++++++++++++++++++++++ internal/loop/loop.go | 14 +- internal/loop/state.go | 62 ++++- internal/quota/client.go | 48 ++++ internal/quota/decision.go | 288 ++++++++++++++++++++++ internal/quota/estimate.go | 74 ++++++ internal/quota/resources.go | 143 +++++++++++ internal/quota/resources_test.go | 276 +++++++++++++++++++++ internal/quota/schedule.go | 77 ++++++ internal/quota/snapshot.go | 161 ++++++++++++ internal/quota/snapshot_test.go | 406 +++++++++++++++++++++++++++++++ internal/quota/state.go | 233 ++++++++++++++++++ 21 files changed, 2904 insertions(+), 20 deletions(-) create mode 100644 internal/config/quota_test.go create mode 100644 internal/loop/gate.go create mode 100644 internal/loop/gate_test.go create mode 100644 internal/quota/client.go create mode 100644 internal/quota/decision.go create mode 100644 internal/quota/estimate.go create mode 100644 internal/quota/resources.go create mode 100644 internal/quota/resources_test.go create mode 100644 internal/quota/schedule.go create mode 100644 internal/quota/snapshot.go create mode 100644 internal/quota/snapshot_test.go create mode 100644 internal/quota/state.go diff --git a/README.md b/README.md index c14ac97..6706b69 100644 --- a/README.md +++ b/README.md @@ -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 [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::{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) diff --git a/cmd/harness/main.go b/cmd/harness/main.go index 6f61714..6b9d707 100644 --- a/cmd/harness/main.go +++ b/cmd/harness/main.go @@ -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 [-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") diff --git a/harness.toml.example b/harness.toml.example index e10d232..7c9e287 100644 --- a/harness.toml.example +++ b/harness.toml.example @@ -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 diff --git a/internal/config/config.go b/internal/config/config.go index 78ed1e5..c234101 100644 --- a/internal/config/config.go +++ b/internal/config/config.go @@ -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 +} diff --git a/internal/config/quota_test.go b/internal/config/quota_test.go new file mode 100644 index 0000000..0e0ef3a --- /dev/null +++ b/internal/config/quota_test.go @@ -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) + } +} diff --git a/internal/config/toml.go b/internal/config/toml.go index ec29734..1c48593 100644 --- a/internal/config/toml.go +++ b/internal/config/toml.go @@ -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 } - 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) { diff --git a/internal/loop/daemon.go b/internal/loop/daemon.go index 84e0bac..0b8c180 100644 --- a/internal/loop/daemon.go +++ b/internal/loop/daemon.go @@ -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 { diff --git a/internal/loop/daemon_test.go b/internal/loop/daemon_test.go index ddab6fb..98902e2 100644 --- a/internal/loop/daemon_test.go +++ b/internal/loop/daemon_test.go @@ -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) { diff --git a/internal/loop/gate.go b/internal/loop/gate.go new file mode 100644 index 0000000..b1d209b --- /dev/null +++ b/internal/loop/gate.go @@ -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() +} diff --git a/internal/loop/gate_test.go b/internal/loop/gate_test.go new file mode 100644 index 0000000..864e609 --- /dev/null +++ b/internal/loop/gate_test.go @@ -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() +} diff --git a/internal/loop/loop.go b/internal/loop/loop.go index 72d95c5..49cd364 100644 --- a/internal/loop/loop.go +++ b/internal/loop/loop.go @@ -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 diff --git a/internal/loop/state.go b/internal/loop/state.go index ceaa409..f792f98 100644 --- a/internal/loop/state.go +++ b/internal/loop/state.go @@ -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,11 +33,18 @@ 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"` - Detail string `json:"detail,omitempty"` // human context; never secrets + // 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 } // loopState is the append-only loop log + dedup index: stateDir/loop.jsonl, @@ -47,7 +56,18 @@ type loopState struct { f *os.File path string seen map[string]string - readOnly bool // dry-run: dedup checks work, writes are refused + 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) { @@ -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() diff --git a/internal/quota/client.go b/internal/quota/client.go new file mode 100644 index 0000000..e8c1613 --- /dev/null +++ b/internal/quota/client.go @@ -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 +} diff --git a/internal/quota/decision.go b/internal/quota/decision.go new file mode 100644 index 0000000..85cb5f0 --- /dev/null +++ b/internal/quota/decision.go @@ -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 +} diff --git a/internal/quota/estimate.go b/internal/quota/estimate.go new file mode 100644 index 0000000..30158bb --- /dev/null +++ b/internal/quota/estimate.go @@ -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" +} diff --git a/internal/quota/resources.go b/internal/quota/resources.go new file mode 100644 index 0000000..1691752 --- /dev/null +++ b/internal/quota/resources.go @@ -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 +} diff --git a/internal/quota/resources_test.go b/internal/quota/resources_test.go new file mode 100644 index 0000000..0f34ba2 --- /dev/null +++ b/internal/quota/resources_test.go @@ -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") + } +} diff --git a/internal/quota/schedule.go b/internal/quota/schedule.go new file mode 100644 index 0000000..f0dd2f7 --- /dev/null +++ b/internal/quota/schedule.go @@ -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 +} diff --git a/internal/quota/snapshot.go b/internal/quota/snapshot.go new file mode 100644 index 0000000..709e2d0 --- /dev/null +++ b/internal/quota/snapshot.go @@ -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 +} diff --git a/internal/quota/snapshot_test.go b/internal/quota/snapshot_test.go new file mode 100644 index 0000000..2da4c7a --- /dev/null +++ b/internal/quota/snapshot_test.go @@ -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", `502 Bad Gateway`}, + {"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) diff --git a/internal/quota/state.go b/internal/quota/state.go new file mode 100644 index 0000000..ed239c5 --- /dev/null +++ b/internal/quota/state.go @@ -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::snapshot — latest polled QuotaSnapshot (JSON) +// mopac:quota::est:5h — estimated credits in the rolling 5h window +// mopac:quota::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" }