diff --git a/README.md b/README.md index 297bd811..0047dc1a 100644 --- a/README.md +++ b/README.md @@ -588,6 +588,12 @@ so a UI can follow a turn with one request at a time. Only the session named is told; like `/v1/magpie/quotas`, it answers this machine, and another only with the key of a gateway shared on the local network. +To show the one in use without a session, `GET /v1/magpie/quotas` (and +`magpie quota --json`, the same list) has `lastServedAt` on each +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. + 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. diff --git a/internal/gateway/quotas_served_test.go b/internal/gateway/quotas_served_test.go new file mode 100644 index 00000000..c7d7edd2 --- /dev/null +++ b/internal/gateway/quotas_served_test.go @@ -0,0 +1,72 @@ +package gateway + +import ( + "encoding/json" + "io" + "net/http" + "net/http/httptest" + "strings" + "testing" + "time" + + "github.com/yetone/magpie/internal/provider" +) + +// GET /v1/magpie/quotas tells when each key last answered a request through +// the gateway, and marks the latest last (#570); one that never did has +// neither. +func TestQuotasLastServedAt(t *testing.T) { + fresh(t) + up := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.URL.Path == "/balance" { + io.WriteString(w, `{"balance":"7"}`) + return + } + servedBy(provider.Chat, "sol")(w, r) + })) + t.Cleanup(up.Close) + for _, id := range []string{"qa", "qb"} { + if err := provider.Save(provider.Provider{ID: id, Name: strings.ToUpper(id), Key: "k-" + id, Models: []string{"sol"}, + Chat: up.URL + "/v1", BalanceURL: up.URL + "/balance", BalancePath: "balance"}); err != nil { + t.Fatal(err) + } + } + provider.ForgetBalances() + s := New() + before := time.Now().Add(-time.Second) + rec := httptest.NewRecorder() + s.Handler().ServeHTTP(rec, httptest.NewRequest("POST", "/v1/chat/completions", strings.NewReader(`{"model":"qa/sol","messages":[{"role":"user","content":"hi"}]}`))) + if rec.Code != 200 { + t.Fatalf("%d %s", rec.Code, rec.Body) + } + rec = httptest.NewRecorder() + r := httptest.NewRequest("GET", "/v1/magpie/quotas", nil) + r.RemoteAddr = "127.0.0.1:5000" + s.Handler().ServeHTTP(rec, r) + var out struct { + Data []map[string]any `json:"data"` + } + if err := json.Unmarshal(rec.Body.Bytes(), &out); err != nil { + t.Fatal(err, rec.Body) + } + seen := 0 + for _, q := range out.Data { + switch q["provider"] { + case "qa": + seen++ + s, _ := q["lastServedAt"].(string) + at, err := time.Parse(time.RFC3339, s) + if err != nil || at.Before(before) || q["last"] != true { + t.Errorf("served key: %v (%v)", q, err) + } + case "qb": + seen++ + if _, ok := q["lastServedAt"]; ok || q["last"] != nil { + t.Errorf("a key that never answered: %v", q) + } + } + } + if seen != 2 { + t.Fatalf("quotas %s", rec.Body) + } +} diff --git a/internal/gateway/routing.go b/internal/gateway/routing.go index 80af7640..f7adc224 100644 --- a/internal/gateway/routing.go +++ b/internal/gateway/routing.go @@ -65,6 +65,7 @@ func served(rest, key string, tokens int) { // when it has one, the candidate's model-specific rest. func servedCandidate(c candidate, tokens int) { served(c.restKey(), c.restKey(), tokens) + provider.NoteServed(c.p, time.Now()) if id := c.restID(); id != c.restKey() { clearRest(id) } diff --git a/internal/provider/quotas.go b/internal/provider/quotas.go index 32639f9c..f11cfed8 100644 --- a/internal/provider/quotas.go +++ b/internal/provider/quotas.go @@ -79,6 +79,11 @@ type Quota struct { Renew string `json:"renew,omitempty"` // Resets are a Codex account's rate-limit resets, when it holds any. Resets *ResetCredits `json:"resets,omitempty"` + // LastServedAt is when the account, plan or key last answered a + // request through the gateway, nil when it hasn't in the last 30 days + // (served.go); Last is on the latest of them (#570). + LastServedAt *time.Time `json:"lastServedAt,omitempty"` + Last bool `json:"last,omitempty"` } // QuotaSpan is one window of an allowance: how much of it is used and @@ -114,7 +119,7 @@ func QuotaReport(ctx context.Context, now time.Time) []Quota { out = append(out, r) } } - return out + return withServed(out, LastServed()) } // ResetClock is when a window starts again, on the clock: "14:30" today, diff --git a/internal/provider/served.go b/internal/provider/served.go new file mode 100644 index 00000000..8ea0f6c5 --- /dev/null +++ b/internal/provider/served.go @@ -0,0 +1,167 @@ +package provider + +// When each subscription account, plan and key last answered a request +// through the gateway (#570): magpie quota --json and the gateway's GET +// /v1/magpie/quotas tell it as each Quota's lastServedAt, and mark the +// latest last, so a menu-bar widget can show the one in use — under +// routing Smart nobody is first for good. It is kept on disk (served.json, +// next to the providers), so the CLI, a process of its own, reads what the +// gateway noted, and a restart doesn't forget it. + +import ( + "encoding/json" + "os" + "path/filepath" + "strings" + "sync" + "time" +) + +// servedKeep is how long an answer is remembered: a key or account gone +// since is dropped from the file after it. +const servedKeep = 30 * 24 * time.Hour + +var servedAt = struct { + sync.Mutex + m map[string]time.Time // ServedID → when it last answered + path string // the file m goes with: another (a test's HOME) starts afresh +}{m: map[string]time.Time{}} + +// ours is servedAt.m for the file at path. Called with servedAt held. +func ours(path string) map[string]time.Time { + if servedAt.path != path { + servedAt.path, servedAt.m = path, map[string]time.Time{} + } + return servedAt.m +} + +func servedPath() string { return filepath.Join(filepath.Dir(Path()), "served.json") } + +// ServedID names who answers for p, as its Quota is found by: an account +// by its agent and user ("codex@a@b.c"), a key by its provider and the +// key's id ("deepseek#k1a2b3c"), else the provider. +func ServedID(p Provider) string { + if a := p.Account; a != nil { + return loginProvider(Login{Agent: a.Agent}) + "@" + strings.ToLower(a.User) + } + if p.Key != "" { + return p.ID + "#" + KeyID(p.Key) + } + return p.ID +} + +// NoteServed records that p just answered a request, at at. +func NoteServed(p Provider, at time.Time) { + id := ServedID(p) + path := servedPath() + servedAt.Lock() + defer servedAt.Unlock() + mine := ours(path) + if at.Before(mine[id]) { + return + } + mine[id] = at + m := readServed(path) + for k, t := range mine { + if t.After(m[k]) { + m[k] = t + } + } + for k, t := range m { + if at.Sub(t) > servedKeep { + delete(m, k) + } + } + if b, err := json.Marshal(m); err == nil && os.MkdirAll(filepath.Dir(path), 0o700) == nil { + writeFileAtomic(path, b) + } +} + +// LastServed is when each ServedID last answered: what this process noted +// and what the file holds, whichever gateway wrote it. +func LastServed() map[string]time.Time { + path := servedPath() + m := readServed(path) + servedAt.Lock() + for k, t := range ours(path) { + if t.After(m[k]) { + m[k] = t + } + } + servedAt.Unlock() + for k, t := range m { + if time.Since(t) > servedKeep { + delete(m, k) + } + } + return m +} + +func readServed(path string) map[string]time.Time { + m := map[string]time.Time{} + if b, err := os.ReadFile(path); err == nil { + json.Unmarshal(b, &m) + } + if m == nil { + m = map[string]time.Time{} + } + return m +} + +// withServed sets each quota's LastServedAt from served, and Last on the +// latest (on each of them, when a key's plan and balance tie). +func withServed(qs []Quota, served map[string]time.Time) []Quota { + if len(served) == 0 { + return qs + } + var keys map[string][]KeyInfo // provider → its keys, read once + var latest time.Time + for i := range qs { + q := &qs[i] + var at time.Time + latestOf := func(id string) { + if t, ok := served[id]; ok && t.After(at) { + at = t + } + } + user := strings.ToLower(q.User) + switch { + case q.Kind == "subscription" && user != "": + latestOf(q.Provider + "@" + user) + case q.User == "": + // the one account or key it has, or any of them: a balance + // one account has whichever key asks + for id, t := range served { + if (id == q.Provider || strings.HasPrefix(id, q.Provider+"@") || strings.HasPrefix(id, q.Provider+"#")) && t.After(at) { + at = t + } + } + default: + // a plan's or balance's key, by the name or mask it is told by + if keys == nil { + keys = map[string][]KeyInfo{} + for _, p := range All() { + keys[p.ID] = p.KeyList() + } + } + for _, k := range keys[q.Provider] { + if k.Name == q.User || k.Name == "" && k.Masked == q.User { + latestOf(q.Provider + "#" + k.ID) + } + } + } + if !at.IsZero() { + t := at.UTC() + q.LastServedAt = &t + if at.After(latest) { + latest = at + } + } + } + for i := range qs { + if qs[i].LastServedAt != nil && qs[i].LastServedAt.Equal(latest) { + qs[i].Last = true + } + } + return qs +} diff --git a/internal/provider/served_test.go b/internal/provider/served_test.go new file mode 100644 index 00000000..44a3bdba --- /dev/null +++ b/internal/provider/served_test.go @@ -0,0 +1,61 @@ +package provider + +import ( + "encoding/json" + "os" + "strings" + "testing" + "time" +) + +// Each quota is told when its account, plan or key last answered (#570): +// an account by its user, a key by the name or mask its card has, a card +// with no user by any of its provider's; the latest is marked last. +func TestQuotasLastServed(t *testing.T) { + home := t.TempDir() + t.Setenv("HOME", home) + t.Setenv("USERPROFILE", home) + t.Setenv("XDG_CONFIG_HOME", t.TempDir()) + if err := Save(Provider{ID: "dk", Name: "DK", Key: "sk-first-key-0001", KeyName: "work", Chat: "http://127.0.0.1:1/v1", + Keys: []KeyAccount{{Key: "sk-other-key-0002"}}}); err != nil { + t.Fatal(err) + } + now := time.Now().Truncate(time.Second) + NoteServed(Provider{ID: "codex", Account: &Account{Agent: "codex", User: "A@x.com"}}, now.Add(-time.Hour)) + NoteServed(Provider{ID: "dk", Key: "sk-other-key-0002"}, now.Add(-time.Minute)) + NoteServed(Provider{ID: "solo", Key: "sk-solo"}, now.Add(-2*time.Hour)) + // the CLI is a process of its own: what it is told is read from the file + servedAt.Lock() + servedAt.m = map[string]time.Time{} + servedAt.Unlock() + qs := withServed([]Quota{ + {Provider: "codex", Kind: "subscription", User: "a@x.com"}, + {Provider: "codex", Kind: "subscription", User: "b@x.com"}, + {Provider: "dk", Kind: "plan", User: "work"}, + {Provider: "dk", Kind: "plan", User: Mask("sk-other-key-0002")}, + {Provider: "dk", Kind: "balance"}, + {Provider: "solo", Kind: "balance"}, + {Provider: "none", Kind: "balance"}, + }, LastServed()) + want := []*time.Time{ptr(now.Add(-time.Hour)), nil, nil, ptr(now.Add(-time.Minute)), ptr(now.Add(-time.Minute)), ptr(now.Add(-2 * time.Hour)), nil} + for i, q := range qs { + if (q.LastServedAt == nil) != (want[i] == nil) || q.LastServedAt != nil && !q.LastServedAt.Equal(*want[i]) { + t.Errorf("%d %s %s %q: lastServedAt %v, want %v", i, q.Provider, q.Kind, q.User, q.LastServedAt, want[i]) + } + if last := i == 3 || i == 4; q.Last != last { + t.Errorf("%d %s %q: last %v", i, q.Provider, q.User, q.Last) + } + } + b, _ := json.Marshal(qs[0]) + if !strings.Contains(string(b), `"lastServedAt":"`+now.Add(-time.Hour).UTC().Format(time.RFC3339)+`"`) || strings.Contains(string(b), `"last"`) { + t.Errorf("json %s", b) + } + if b, _ := json.Marshal(qs[1]); strings.Contains(string(b), "lastServedAt") || strings.Contains(string(b), `"last"`) { + t.Errorf("never served: %s", b) + } + if _, err := os.Stat(servedPath()); err != nil { + t.Fatal("not kept on disk:", err) + } +} + +func ptr(t time.Time) *time.Time { return &t } diff --git a/quota_cli.go b/quota_cli.go index 05e45bf9..3b6af215 100644 --- a/quota_cli.go +++ b/quota_cli.go @@ -20,7 +20,8 @@ const quotaUsage = `usage: magpie quota […] [--json] magpie quota alert [|off] [--balance |off] 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; the gateway answers the same at + a minute ago). --json is for scripts and agents, each entry with lastServedAt, when it + last answered through the gateway, and last: true on the latest; the gateway answers the same at GET http://127.0.0.1:3425/v1/magpie/quotas, and to another machine only while magpie is shared on the local network, with its key as the API key (Authorization: Bearer or x-api-key) A Codex account that holds rate-limit resets says how many; quota reset spends one, diff --git a/quota_served_cli_test.go b/quota_served_cli_test.go new file mode 100644 index 00000000..b026831e --- /dev/null +++ b/quota_served_cli_test.go @@ -0,0 +1,84 @@ +package main + +import ( + "encoding/json" + "io" + "net/http" + "net/http/httptest" + "os" + "path/filepath" + "testing" + "time" + + "github.com/yetone/magpie/internal/agentenv" + "github.com/yetone/magpie/internal/provider" +) + +// magpie quota --json tells when each key last answered through the gateway +// — a process of its own, which wrote it down — and marks the latest last, +// as GET /v1/magpie/quotas does (#570). +func TestQuotaJSONLastServedAt(t *testing.T) { + groupsHome(t) + for _, k := range agentenv.Vars { + t.Setenv(k, "") + os.Unsetenv(k) + } + // no agent's CLI nor the keychain is asked: what answers is inert + bin := t.TempDir() + for _, name := range []string{"security", "secret-tool", "claude", "codex", "cursor-agent", "devin", "grok", "kiro-cli"} { + if err := os.WriteFile(filepath.Join(bin, name), []byte("#!/bin/sh\nexit 1\n"), 0o755); err != nil { + t.Fatal(err) + } + } + t.Setenv("PATH", bin+string(os.PathListSeparator)+os.Getenv("PATH")) + up := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { io.WriteString(w, `{"balance":"7"}`) })) + t.Cleanup(up.Close) + for _, id := range []string{"a", "b"} { + p, err := provider.Find(id) + if err != nil { + t.Fatal(err) + } + p.BalanceURL, p.BalancePath = up.URL, "balance" + if err := provider.Save(*p); err != nil { + t.Fatal(err) + } + } + provider.ForgetBalances() + at := time.Now().Add(-time.Minute).Truncate(time.Second) + a, _ := provider.Find("a") + provider.NoteServed(*a, at) + r, w, err := os.Pipe() + if err != nil { + t.Fatal(err) + } + old := os.Stdout + os.Stdout = w + done := make(chan []byte) + go func() { b, _ := io.ReadAll(r); done <- b }() + err = quotaCmd([]string{"quota", "a", "b", "--json"}) + w.Close() + os.Stdout = old + out := <-done + if err != nil { + t.Fatal(err) + } + var qs []map[string]any + if err := json.Unmarshal(out, &qs); err != nil { + t.Fatal(err, string(out)) + } + if len(qs) != 2 { + t.Fatalf("%s", out) + } + for _, q := range qs { + switch q["provider"] { + case "a": + if q["lastServedAt"] != at.UTC().Format(time.RFC3339) || q["last"] != true { + t.Errorf("served: %v", q) + } + case "b": + if _, ok := q["lastServedAt"]; ok || q["last"] != nil { + t.Errorf("never served: %v", q) + } + } + } +}