diff --git a/README.md b/README.md index 22572369..fb65d500 100644 --- a/README.md +++ b/README.md @@ -645,6 +645,23 @@ subscription account, plan and key that answered a request through the gateway in the last 30 days, and `last: true` on the latest. It is kept in `served.json` beside `providers.json`, so a restart keeps it. +`magpie quota wait ` blocks until that subscription (any +of its accounts magpie has on) or that one account has allowance again — no +window that stops it used up — then exits 0, so a long task stopped by its +limit can go on unattended: + +```sh +until codex exec "…"; do magpie quota wait codex || break; done +``` + +It reads the vendors itself, whether or not the gateway runs, and again +shortly after the soonest reset it knows (every 1 to 10 minutes; an +allowance it can't read is asked again less often each time), saying on +stderr what it waits for and until when. Name an account by its email or +login, or as `/` when two subscriptions share it. +`--timeout 6h` exits 1 if it passes first, `--quiet` says nothing; an +unknown name exits 2 and Ctrl+C 130. + The *Gateway* tab in the app has this as copy buttons and ready-made snippets (shell, curl, Python, Node) for each API, the list of model ids, and the recent calls; `MAGPIE_DEBUG=1` logs every call to the terminal. @@ -910,6 +927,9 @@ magpie profiles magpie rm work magpie sync # refresh the models.dev catalog and every live model list + +magpie quota # what is left of every subscription, plan and key balance +magpie quota wait codex # block until a Codex account has allowance again ``` In the app, click any value to open a filtered list; type to search or to diff --git a/internal/provider/codex_switch.go b/internal/provider/codex_switch.go index 8952d9cf..ce904184 100644 --- a/internal/provider/codex_switch.go +++ b/internal/provider/codex_switch.go @@ -86,6 +86,35 @@ func usedUp(q SubscriptionQuota) bool { return usedPast(q, 100) } +// UsedUp is usedUp for the CLI's magpie quota wait, so what it waits out +// is what stops an account here. +func UsedUp(q SubscriptionQuota) bool { return usedUp(q) } + +// BackAt is when a used-up account has room again: when the last of the +// windows that stop it (usedUp's) starts again, a window's seconds to go +// counted from now. Zero when one of them doesn't say. +func BackAt(q SubscriptionQuota, now time.Time) time.Time { + var back time.Time + for _, w := range q.Windows { + if w.Aside || w.Model != "" || w.Used < 100 { + continue + } + at := time.Time{} + switch { + case w.ResetsAt != nil: + at = *w.ResetsAt + case w.ResetSecs > 0: + at = now.Add(time.Duration(w.ResetSecs) * time.Second) + default: + return time.Time{} + } + if at.After(back) { + back = at + } + } + return back +} + func usedPast(q SubscriptionQuota, share float64) bool { for _, w := range q.Windows { if !w.Aside && w.Model == "" && w.Used >= share { diff --git a/internal/provider/quotas.go b/internal/provider/quotas.go index 84f3dde4..cc5d4081 100644 --- a/internal/provider/quotas.go +++ b/internal/provider/quotas.go @@ -16,6 +16,14 @@ func Quotas(ctx context.Context) []SubscriptionQuota { return append(append(subs, plans...), balances...) } +// Allotted is Quotas without the keys' balances: the accounts and plans +// whose allowance runs out and starts again, which is what magpie quota +// wait waits on; no key's balance is asked for. +func Allotted(ctx context.Context) []SubscriptionQuota { + subs := SubscriptionUsage(ctx) + return append(subs, notShown(PlanQuotas(ctx), subs)...) +} + func quotas(ctx context.Context) (subs, plans, balances []SubscriptionQuota) { b := make(chan []SubscriptionQuota, 1) p := make(chan []SubscriptionQuota, 1) diff --git a/internal/provider/subscription_usage.go b/internal/provider/subscription_usage.go index 93054b57..1a6f5a9d 100644 --- a/internal/provider/subscription_usage.go +++ b/internal/provider/subscription_usage.go @@ -522,10 +522,7 @@ var claudeAsked atomic.Int64 // (claudeUsageWait); the next SubscriptionUsage waits for it. func AskClaudeUsage() { claudeAsked.Store(time.Now().UnixNano()) - c := &subscriptionUsageCache - c.Lock() - c.at, c.asked = time.Time{}, true - c.Unlock() + AskUsage() l := &loginUsageCache l.Lock() for k := range l.m { @@ -536,6 +533,17 @@ func AskClaudeUsage() { l.Unlock() } +// AskUsage has the next SubscriptionUsage wait for a new reading rather +// than hand back the last one, as AskClaudeUsage does, without having +// Claude Code's /usage run: magpie quota wait reads again each time it +// wakes, and only a wait for Claude needs Claude read. +func AskUsage() { + c := &subscriptionUsageCache + c.Lock() + c.at, c.asked = time.Time{}, true + c.Unlock() +} + // claudeWindows is the allowance of the Claude account user. Only the // account Claude Code is signed in to (active) is read, by Claude Code's // own /usage: when the user asked since it last was, or when the last diff --git a/main.go b/main.go index 95c9fdbc..f5087f51 100644 --- a/main.go +++ b/main.go @@ -3,6 +3,7 @@ package main import ( "context" + "errors" "fmt" "os" "strings" @@ -94,6 +95,8 @@ const usage = `magpie — one place to pick every agent's model magpie sessions --days N|today|all [--model ] [--folder ] [--json] what every session spent, day by day, with the top models and folders (7 days) magpie quota [] [--json] what is left of every subscription, plan and key balance + magpie quota wait [--timeout ] [--quiet] + block until that subscription (any of its accounts) or account has allowance again magpie sync refresh the model catalog and vendor model lists magpie agents list every supported agent magpie update [check] [--proxy ] [--mirror ] @@ -126,8 +129,16 @@ func main() { proc.EndProbes() // a CLI still being asked something isn't left to init sessions.Saved() // the session index kept, for the next run if err != nil { - fmt.Fprintln(os.Stderr, "magpie:", err) - os.Exit(1) + // a command with exit codes of its own (quota wait) says which + code := 1 + var e exitError + if errors.As(err, &e) { + code = e.code + } + if msg := err.Error(); msg != "" { + fmt.Fprintln(os.Stderr, "magpie:", msg) + } + os.Exit(code) } } diff --git a/quota_cli.go b/quota_cli.go index 3b6af215..2297328e 100644 --- a/quota_cli.go +++ b/quota_cli.go @@ -18,6 +18,7 @@ const quotaUsage = `usage: magpie quota […] [--json] magpie quota reset [] [--yes] magpie quota auto-reset [] [on|off] magpie quota alert [|off] [--balance |off] + magpie quota wait [--timeout ] [--quiet] what is left of every subscription, plan and key magpie has: each window's use and when it starts again, and each key's balance, asked of the vendors now (or less than a minute ago). --json is for scripts and agents, each entry with lastServedAt, when it @@ -34,7 +35,15 @@ const quotaUsage = `usage: magpie quota […] [--json] reaches 80% used, once each time the window runs (not windows set aside, such as on-demand spending); --balance 5 when a balance falls to 5 or under, in its own currency or credits, once until topped up past it. off turns either off, and alone - it says what is set. Both are off until set.` + it says what is set. Both are off until set. + quota wait codex returns once any Codex account magpie has on has allowance again + (a window that stops it no longer used up), and an account's email or login, or + /, waits for that one alone: it reads the vendors itself, gateway + running or not, again shortly after the soonest reset (every 1 to 10 minutes), and + says on stderr what it waits for. For a script to go on with work it stopped: + until codex exec …; do magpie quota wait codex || break; done. Exit 0 once there + is allowance, 1 when --timeout (30m, 6h) passes first, 2 for a name it doesn't + know, 130 on Ctrl+C.` // quotaCmd: magpie quota […] [--json] func quotaCmd(args []string) error { @@ -47,6 +56,9 @@ func quotaCmd(args []string) error { if len(args) > 1 && args[1] == "alert" { return quotaAlertCmd(args[2:]) } + if len(args) > 1 && args[1] == "wait" { + return quotaWaitCmd(args[2:]) + } asJSON := false var only []string for _, a := range args[1:] { diff --git a/quota_wait.go b/quota_wait.go new file mode 100644 index 00000000..9753e623 --- /dev/null +++ b/quota_wait.go @@ -0,0 +1,334 @@ +package main + +import ( + "cmp" + "context" + "errors" + "fmt" + "os" + "slices" + "strings" + "time" + + "github.com/yetone/magpie/internal/provider" +) + +// How often magpie quota wait reads the allowance again: shortly after the +// soonest reset it knows, but never sooner than a minute (the readings are +// kept that long anyway, and the vendors' usage endpoints are rate limited) +// nor later than ten, as a reset can come early (a Codex reset spent, a plan +// topped up) or not be told at all. +const ( + quotaWaitMin = time.Minute + quotaWaitMax = 10 * time.Minute + quotaWaitMargin = 30 * time.Second +) + +// exitError is an error that ends magpie with its own exit code rather +// than 1; with no err, nothing is printed. +type exitError struct { + code int + err error +} + +func (e exitError) Error() string { + if e.err == nil { + return "" + } + return e.err.Error() +} + +func (e exitError) Unwrap() error { return e.err } + +// errWaitTimeout is quota wait's --timeout running out, exit code 1. +var errWaitTimeout = errors.New("timed out") + +// quotaTarget is what magpie quota wait waits on: every account of a +// provider (user ""), or one account. +type quotaTarget struct { + provider string + user string + label string +} + +func (t quotaTarget) has(q provider.SubscriptionQuota) bool { + return strings.EqualFold(q.Provider, t.provider) && (t.user == "" || strings.EqualFold(q.User, t.user)) +} + +// resolveQuotaTarget finds what arg names among the readings, as magpie +// quota matches a provider (its id or name), or one account by its user, +// or / where two providers have the same user. find says +// whether arg is a provider at all, for an error that tells the two apart. +func resolveQuotaTarget(arg string, qs []provider.SubscriptionQuota, find func(string) bool) (quotaTarget, error) { + byProvider := func(name string) (string, bool) { + for _, q := range qs { + if strings.EqualFold(q.Provider, name) || strings.EqualFold(q.Name, name) { + return q.Provider, true + } + } + return "", false + } + if id, ok := byProvider(arg); ok { + return quotaTarget{provider: id, label: id}, nil + } + if p, u, ok := strings.Cut(arg, "/"); ok { + if id, ok := byProvider(p); ok { + t := quotaTarget{provider: id, user: u, label: id + " · " + u} + if slices.ContainsFunc(qs, t.has) { + return t, nil + } + return quotaTarget{}, fmt.Errorf("%s has no account %s · magpie quota %s lists its accounts", id, u, id) + } + } + var hits []quotaTarget + for _, q := range qs { + if q.User != "" && strings.EqualFold(q.User, arg) && !slices.ContainsFunc(hits, func(t quotaTarget) bool { return t.provider == q.Provider }) { + hits = append(hits, quotaTarget{provider: q.Provider, user: q.User, label: q.Provider + " · " + q.User}) + } + } + switch { + case len(hits) == 1: + return hits[0], nil + case len(hits) > 1: + var names []string + for _, t := range hits { + names = append(names, t.provider+"/"+t.user) + } + return quotaTarget{}, fmt.Errorf("%s is an account of more than one subscription: name one of %s", arg, strings.Join(names, ", ")) + } + if find != nil && find(arg) { + return quotaTarget{}, fmt.Errorf("%s has no subscription or plan whose allowance magpie reads: sign in to it first, or magpie quota lists what there is", arg) + } + return quotaTarget{}, fmt.Errorf("no subscription, plan or account %s · magpie quota lists them", arg) +} + +// quotaWait is magpie quota wait's loop, with what it reads, the clock and +// how it sleeps given, so tests run it on a fake clock. +type quotaWait struct { + // read is every account's allowance now; readable ones say Error "". + read func(ctx context.Context) []provider.SubscriptionQuota + // skip, when set, leaves out an account a provider's wait shouldn't + // count: one switched off or signed out in magpie, whose room the agent + // would never be moved to. + skip func(q provider.SubscriptionQuota) bool + now func() time.Time + sleep func(ctx context.Context, d time.Duration) error + say func(string) +} + +// readable is a reading made now that says the account's windows: not an +// error, nor the last one kept in its place (AsOf), which may be hours old. +func readable(q provider.SubscriptionQuota) bool { + return q.Error == "" && q.AsOf == nil && len(q.Windows) > 0 +} + +// run waits until an account of t is not used up, and returns it. The +// error is errWaitTimeout once timeout (when over 0) has passed, or ctx's +// once it is cancelled. +func (w quotaWait) run(ctx context.Context, t quotaTarget, timeout time.Duration) (provider.SubscriptionQuota, error) { + var deadline time.Time + if timeout > 0 { + deadline = w.now().Add(timeout) + } + backoff := quotaWaitMin + for { + qs := w.read(ctx) + if err := ctx.Err(); err != nil { + return provider.SubscriptionQuota{}, err + } + now := w.now() + known, unread, why := 0, 0, "" + var back time.Time + backWho := "" + for _, q := range qs { + if !t.has(q) || (t.user == "" && w.skip != nil && w.skip(q)) { + continue + } + if !readable(q) { + unread++ + why = cmp.Or(why, q.Error) + continue + } + if !provider.UsedUp(q) { + return q, nil + } + known++ + // the soonest any account is back; one that doesn't say when + // leaves it to the next look + if at := provider.BackAt(q, now); !at.IsZero() && (back.IsZero() || at.Before(back)) { + back, backWho = at, q.User + } + } + var d time.Duration + var line string + switch { + case known == 0: + // nothing read: try again, less often each time, as a vendor + // that is down or refusing isn't helped by asking every minute + d, backoff = backoff, min(backoff*2, quotaWaitMax) + line = fmt.Sprintf("%s: its allowance couldn't be read", t.label) + if why != "" { + line += " (" + why + ")" + } + line += fmt.Sprintf(" · trying again in %s", roundWait(d)) + case back.IsZero(): + backoff = quotaWaitMin + d = quotaWaitMax + line = fmt.Sprintf("%s is used up, and when it starts again isn't told · checking again in %s", t.label, roundWait(d)) + default: + backoff = quotaWaitMin + d = min(max(back.Sub(now)+quotaWaitMargin, quotaWaitMin), quotaWaitMax) + who := t.label + if t.user == "" && backWho != "" && known > 1 { + who += " (" + backWho + " first)" + } else if t.user == "" && known > 1 { + who += " (every account)" + } + line = fmt.Sprintf("%s is used up until %s · checking again at %s", who, provider.ResetClock(back, now), now.Add(d).Local().Format("15:04")) + } + if known > 0 && unread > 0 { + line += fmt.Sprintf(" · %d not read", unread) + } + if !deadline.IsZero() { + if !now.Before(deadline) { + return provider.SubscriptionQuota{}, errWaitTimeout + } + d = min(d, deadline.Sub(now)) + } + if w.say != nil { + w.say(line) + } + if err := w.sleep(ctx, d); err != nil { + return provider.SubscriptionQuota{}, err + } + } +} + +// roundWait is a wait in words, to the minute: "2m0s" reads as "2m". +func roundWait(d time.Duration) string { + s := d.Round(time.Minute).String() + s = strings.TrimSuffix(s, "0s") + return strings.TrimSuffix(s, "h0m") +} + +// sleepCtx sleeps d, or until ctx ends. +func sleepCtx(ctx context.Context, d time.Duration) error { + t := time.NewTimer(d) + defer t.Stop() + select { + case <-t.C: + return nil + case <-ctx.Done(): + return ctx.Err() + } +} + +// quotaWaitCmd: magpie quota wait [--timeout ] +// [--quiet] — blocks until the subscription (any of its accounts) or the +// account has allowance again, for a script to go on with the work it +// stopped (#720): until codex exec …; do magpie quota wait codex; done. +// Exits 0 then, 1 on --timeout, 2 on a name it doesn't know, 130 on Ctrl+C. +func quotaWaitCmd(args []string) error { + var arg string + var timeout time.Duration + quiet := false + for i := 0; i < len(args); i++ { + a := args[i] + switch { + case a == "help" || a == "-h" || a == "--help": + fmt.Println(quotaUsage) + return nil + case a == "--quiet" || a == "-q": + quiet = true + case a == "--timeout" || strings.HasPrefix(a, "--timeout="): + v, ok := strings.CutPrefix(a, "--timeout=") + if !ok { + if i+1 >= len(args) { + return exitError{2, fmt.Errorf("--timeout needs a duration, such as 30m or 6h")} + } + i++ + v = args[i] + } + d, err := time.ParseDuration(v) + if err != nil || d <= 0 { + return exitError{2, fmt.Errorf("--timeout is a duration such as 30m or 6h, not %q", v)} + } + timeout = d + case strings.HasPrefix(a, "-"): + return exitError{2, fmt.Errorf("unknown flag %s · %s", a, "magpie quota wait [--timeout ] [--quiet]")} + default: + if arg != "" { + return exitError{2, fmt.Errorf("one subscription or account at a time: %s or %s", arg, a)} + } + arg = a + } + } + if arg == "" { + return exitError{2, fmt.Errorf("wait for what? magpie quota wait , such as codex")} + } + ctx, stop := interruptContext() + defer stop() + claude := strings.HasPrefix(strings.ToLower(arg), "claude") + read := func(ctx context.Context) []provider.SubscriptionQuota { + // a new reading each time: the cached one is what the last look saw + if claude { + provider.AskClaudeUsage() + } else { + provider.AskUsage() + } + rctx, cancel := context.WithTimeout(ctx, 30*time.Second) + defer cancel() + return provider.Allotted(rctx) + } + qs := read(ctx) + if ctx.Err() != nil { + return exitError{130, nil} + } + t, err := resolveQuotaTarget(arg, qs, func(s string) bool { _, err := provider.Find(s); return err == nil }) + if err != nil { + return exitError{2, err} + } + claude = strings.EqualFold(t.provider, "claude") + skip := func(q provider.SubscriptionQuota) bool { + for _, l := range provider.Logins(q.Provider) { + if q.User != "" && strings.EqualFold(l.User, q.User) { + return !l.On || l.Lapsed != "" + } + } + return false + } + if t.user == "" && !slices.ContainsFunc(qs, func(q provider.SubscriptionQuota) bool { return t.has(q) && !skip(q) }) { + return exitError{2, fmt.Errorf("none of %s's accounts is on in magpie · turn one on, or name the account", t.provider)} + } + first := true + w := quotaWait{ + read: func(ctx context.Context) []provider.SubscriptionQuota { + if first { // what resolving read is new enough + first = false + return qs + } + return read(ctx) + }, + skip: skip, + now: time.Now, + sleep: sleepCtx, + } + if !quiet { + w.say = func(s string) { fmt.Fprintln(os.Stderr, muted.Render("…"), s) } + } + q, err := w.run(ctx, t, timeout) + switch { + case errors.Is(err, errWaitTimeout): + return fmt.Errorf("%s still has no allowance after %s", t.label, timeout) + case err != nil: + return exitError{130, nil} + } + if !quiet { + who := q.Provider + if q.User != "" { + who += " · " + q.User + } + fmt.Println(green.Render("✓"), who, "has allowance again") + } + return nil +} diff --git a/quota_wait_test.go b/quota_wait_test.go new file mode 100644 index 00000000..8b30a04e --- /dev/null +++ b/quota_wait_test.go @@ -0,0 +1,246 @@ +package main + +import ( + "context" + "errors" + "strings" + "testing" + "time" + + "github.com/yetone/magpie/internal/provider" +) + +// fakeWait runs quotaWait on a clock that moves only when it sleeps, and +// reads what at says for the time it is. +type fakeWait struct { + now time.Time + slept []time.Duration + reads int + said []string + at func(now time.Time) []provider.SubscriptionQuota + cancel func() // called on the nth sleep, when set + cancelAt int +} + +func (f *fakeWait) wait() quotaWait { + return quotaWait{ + read: func(context.Context) []provider.SubscriptionQuota { f.reads++; return f.at(f.now) }, + now: func() time.Time { return f.now }, + sleep: func(ctx context.Context, d time.Duration) error { + f.slept = append(f.slept, d) + if f.cancel != nil && len(f.slept) == f.cancelAt { + f.cancel() + return ctx.Err() + } + f.now = f.now.Add(d) + return nil + }, + say: func(s string) { f.said = append(f.said, s) }, + } +} + +func codexAt(user string, used float64, reset time.Time) provider.SubscriptionQuota { + r := reset + return provider.SubscriptionQuota{Provider: "codex", Name: "Codex", User: user, Windows: []provider.QuotaWindow{ + {Name: "5 hours", Used: 20, ResetsAt: &r}, + {Name: "weekly", Used: used, ResetsAt: &r}, + // on-demand spending used up stops nothing + {Name: "credits", Used: 100, Aside: true}, + }} +} + +var t0 = time.Date(2026, 10, 4, 12, 0, 0, 0, time.UTC) + +// An account with room: back at once, nothing slept. +func TestQuotaWaitAlreadyAvailable(t *testing.T) { + f := &fakeWait{now: t0, at: func(time.Time) []provider.SubscriptionQuota { + return []provider.SubscriptionQuota{codexAt("a@x", 100, t0.Add(time.Hour)), codexAt("b@x", 97, t0.Add(time.Hour))} + }} + q, err := f.wait().run(context.Background(), quotaTarget{provider: "codex", label: "codex"}, 0) + if err != nil || q.User != "b@x" || len(f.slept) != 0 { + t.Fatalf("got %q %v, slept %v", q.User, err, f.slept) + } +} + +// Every account used up: it sleeps to just past the soonest reset, in +// steps of at most ten minutes, and returns the account that came back. +func TestQuotaWaitUntilReset(t *testing.T) { + resetA, resetB := t0.Add(25*time.Minute), t0.Add(3*time.Hour) + f := &fakeWait{now: t0} + f.at = func(now time.Time) []provider.SubscriptionQuota { + a := codexAt("a@x", 100, resetA) + if !now.Before(resetA) { + a = codexAt("a@x", 0, resetA.Add(7*24*time.Hour)) + } + return []provider.SubscriptionQuota{codexAt("b@x", 100, resetB), a} + } + q, err := f.wait().run(context.Background(), quotaTarget{provider: "codex", label: "codex"}, 0) + if err != nil || q.User != "a@x" { + t.Fatalf("got %q %v", q.User, err) + } + want := []time.Duration{10 * time.Minute, 10 * time.Minute, 5*time.Minute + quotaWaitMargin} + if len(f.slept) != len(want) { + t.Fatalf("slept %v, want %v", f.slept, want) + } + for i := range want { + if f.slept[i] != want[i] { + t.Fatalf("slept %v, want %v", f.slept, want) + } + } + if len(f.said) != 3 || !strings.Contains(f.said[0], "codex (a@x first) is used up until") { + t.Errorf("said %q", f.said) + } +} + +// One account named: another's room doesn't end the wait, and a reset +// under a minute away is still a minute's sleep. +func TestQuotaWaitOneAccount(t *testing.T) { + reset := t0.Add(10 * time.Second) + f := &fakeWait{now: t0} + f.at = func(now time.Time) []provider.SubscriptionQuota { + a := codexAt("a@x", 100, reset) + if now.After(reset) { + a = codexAt("a@x", 5, reset.Add(time.Hour)) + } + return []provider.SubscriptionQuota{codexAt("b@x", 0, reset), a} + } + q, err := f.wait().run(context.Background(), quotaTarget{provider: "codex", user: "A@x", label: "codex · a@x"}, 0) + if err != nil || q.User != "a@x" || len(f.slept) != 1 || f.slept[0] != quotaWaitMin { + t.Fatalf("got %q %v, slept %v", q.User, err, f.slept) + } +} + +// --timeout: the last look is at the deadline, then errWaitTimeout. +func TestQuotaWaitTimeout(t *testing.T) { + f := &fakeWait{now: t0, at: func(time.Time) []provider.SubscriptionQuota { + return []provider.SubscriptionQuota{codexAt("a@x", 100, t0.Add(5*time.Hour))} + }} + _, err := f.wait().run(context.Background(), quotaTarget{provider: "codex", label: "codex"}, 25*time.Minute) + if !errors.Is(err, errWaitTimeout) { + t.Fatalf("err %v", err) + } + if f.now != t0.Add(25*time.Minute) || f.reads != 4 { + t.Fatalf("ended at %v after %d reads, slept %v", f.now.Sub(t0), f.reads, f.slept) + } +} + +// Usage that can't be read is asked again, less often each time up to +// ten minutes, and the wait goes on once it reads. +func TestQuotaWaitUnknownRetries(t *testing.T) { + f := &fakeWait{now: t0} + f.at = func(now time.Time) []provider.SubscriptionQuota { + switch { + case now.Sub(t0) < 20*time.Minute: + return []provider.SubscriptionQuota{{Provider: "codex", User: "a@x", Error: "HTTP 503"}} + case now.Sub(t0) < 40*time.Minute: + // the last reading kept in its place says nothing of now + at := t0 + q := codexAt("a@x", 0, t0) + q.AsOf = &at + return []provider.SubscriptionQuota{q} + } + return []provider.SubscriptionQuota{codexAt("a@x", 50, t0.Add(time.Hour))} + } + q, err := f.wait().run(context.Background(), quotaTarget{provider: "codex", label: "codex"}, 0) + if err != nil || q.User != "a@x" { + t.Fatalf("got %q %v", q.User, err) + } + want := []time.Duration{time.Minute, 2 * time.Minute, 4 * time.Minute, 8 * time.Minute, 10 * time.Minute, 10 * time.Minute, 10 * time.Minute} + if len(f.slept) != len(want) { + t.Fatalf("slept %v, want %v", f.slept, want) + } + for i := range want { + if f.slept[i] != want[i] { + t.Fatalf("slept %v, want %v", f.slept, want) + } + } + if !strings.Contains(f.said[0], "couldn't be read (HTTP 503)") { + t.Errorf("said %q", f.said[0]) + } +} + +// A reset that isn't told: looked at again every ten minutes. +func TestQuotaWaitResetUntold(t *testing.T) { + f := &fakeWait{now: t0} + f.at = func(now time.Time) []provider.SubscriptionQuota { + q := provider.SubscriptionQuota{Provider: "zcode", User: "u", Windows: []provider.QuotaWindow{{Name: "day", Used: 100}}} + if now.Sub(t0) >= 20*time.Minute { + q.Windows[0].Used = 0 + } + return []provider.SubscriptionQuota{q} + } + _, err := f.wait().run(context.Background(), quotaTarget{provider: "zcode", label: "zcode"}, 0) + if err != nil || len(f.slept) != 2 || f.slept[0] != quotaWaitMax { + t.Fatalf("%v, slept %v", err, f.slept) + } +} + +// An account switched off in magpie is passed over when waiting on its +// provider: its room would never be used. +func TestQuotaWaitSkipsOff(t *testing.T) { + f := &fakeWait{now: t0, at: func(now time.Time) []provider.SubscriptionQuota { + a := codexAt("a@x", 100, t0.Add(time.Minute)) + if now.After(t0) { + a = codexAt("a@x", 0, t0.Add(time.Hour)) + } + return []provider.SubscriptionQuota{codexAt("off@x", 0, t0), a} + }} + w := f.wait() + w.skip = func(q provider.SubscriptionQuota) bool { return q.User == "off@x" } + q, err := w.run(context.Background(), quotaTarget{provider: "codex", label: "codex"}, 0) + if err != nil || q.User != "a@x" || len(f.slept) != 1 { + t.Fatalf("got %q %v, slept %v", q.User, err, f.slept) + } +} + +// Ctrl+C mid-sleep ends the wait with the context's error. +func TestQuotaWaitCancelled(t *testing.T) { + ctx, cancel := context.WithCancel(context.Background()) + f := &fakeWait{now: t0, cancel: cancel, cancelAt: 2, at: func(time.Time) []provider.SubscriptionQuota { + return []provider.SubscriptionQuota{codexAt("a@x", 100, t0.Add(time.Hour))} + }} + _, err := f.wait().run(ctx, quotaTarget{provider: "codex", label: "codex"}, 0) + if !errors.Is(err, context.Canceled) { + t.Fatalf("err %v", err) + } +} + +func TestResolveQuotaTarget(t *testing.T) { + qs := []provider.SubscriptionQuota{ + codexAt("a@x", 0, t0), codexAt("b@x", 0, t0), + {Provider: "claude", Name: "Claude Code", User: "a@x"}, + {Provider: "zcode", Name: "ZCode", User: "solo@y"}, + } + known := func(s string) bool { return s == "kimi" } + for _, c := range []struct{ arg, provider, user, err string }{ + {arg: "codex", provider: "codex"}, + {arg: "Claude Code", provider: "claude"}, + {arg: "solo@y", provider: "zcode", user: "solo@y"}, + {arg: "codex/B@x", provider: "codex", user: "B@x"}, + {arg: "a@x", err: "codex/a@x, claude/a@x"}, + {arg: "codex/z@x", err: "codex has no account z@x"}, + {arg: "kimi", err: "no subscription or plan whose allowance"}, + {arg: "nope", err: "no subscription, plan or account nope"}, + } { + got, err := resolveQuotaTarget(c.arg, qs, known) + if c.err != "" { + if err == nil || !strings.Contains(err.Error(), c.err) { + t.Errorf("%s: %v, want %q", c.arg, err, c.err) + } + continue + } + if err != nil || got.provider != c.provider || got.user != c.user { + t.Errorf("%s: %+v %v", c.arg, got, err) + } + } +} + +// Unknown names and flags are exit code 2, before anything is read. +func TestQuotaWaitCmdUsageErrors(t *testing.T) { + for _, args := range [][]string{{}, {"--timeout", "soon", "codex"}, {"--frob", "codex"}, {"codex", "claude"}} { + var e exitError + if err := quotaWaitCmd(args); !errors.As(err, &e) || e.code != 2 { + t.Errorf("%q: %v", args, err) + } + } +}