gateway: each gateway key can have its own limit, a token total and/or an estimated cost per day, week or month, and a spent key is refused before any provider is asked (margbug01 on #585, v0.1.671: 每个网关密钥都能设置自己的用量上限) — a gateway key named who called and the Usage page grouped its calls, but nothing held a key to an amount, so one client shared on the network could spend a subscription's whole window. A key now carries limit {period day|week|month, tokens, cost, cacheReads} in caller-keys.json, set with the new "limit-key" change. Windows are calendar ones in local time (from midnight, from Monday, from the 1st) and a call counts in the window it started in. Tokens are a call's uncached input, output and cache writes, plus cache reads when the limit says so; cost is an estimate at the Usage page's effective prices (usage.NewCoster), a call with no known price adding none and counted as unpriced. What a key used is summed from usage.jsonl (usage.Visit), so it survives a restart; the new internal/budget keeps the sums in memory per window and appendUsage writes a key's record through budget.Append, which appends and counts under one lock so a window read meanwhile never counts a call twice. Before a keyed POST reaches the mux, keyLimited reserves for it: its body's size in tokens (4 bytes a token) plus the key's mean output per call this window, at the window's cost per token or the named model's input price; it is let in while used + reserved is under the cap, so requests sent at once overshoot by about one call, and the reservation is released when the request ends, a streamed reply being settled then with the usage its vendor reported. A spent key gets a 429 in its API's error shape (Anthropic, OpenAI chat/responses, Gemini) naming the key, the cap and the reset time, with Retry-After, X-Should-Retry: false and X-Magpie-Limit-Reset, recorded as a rejected call (gateway_key_limit); other keys go on. Requests from this computer without a gateway key are not limited; a key used from this computer is. count_tokens and bridge callbacks are never held. GET /v1/magpie/limit tells a key its own status. Gateway keys rows show a Limit badge (what is used of it, per what, Spent in red) and a line under a limited key (used of limit, left, resets / refused until, in flight; unpriced calls in its title); Limit opens an editor under the row (window, tokens with k/m/万/亿, cost in US$, Count cache reads too, what is counted) staged until Save, marked unsaved, Cancel/Escape dropping it, a bad amount disabling Save, No limit removing it; en and zh; no left-border accent; no click scrolls. The CLI gains magpie gateway-key limit <id> [off | day|week|month --tokens N --cost USD --cache-reads] and a LIMIT column in list. TestKeyLimitRefusesBeforeUpstream (two calls within 700 tokens, the third a 429 the fake upstream never sees, message/type/headers, the rejected record, Anthropic's shape on /v1/messages, another key and a keyless loopback request let in, still refused after budget.Forget as after a restart, a raised and a removed limit let in, /v1/magpie/limit), TestKeyLimitHoldsParallelRequests (20 at once against 2000 tokens: some refused, at most one call over, no reservation left; without the reservation all 20 got in and 7350 were used), TestKeyLimitWindowResets, TestLimitWindow, TestLimitValid and TestGatewayKeyLimitCLI are new; gateway-key-limit.test.cjs (WebKit, en and zh) is new and passes, as do gateway-caller-keys, gateway-connect-keys, api-key-theme and click-scroll (whose request-pick case flaked once under a parallel run and passes alone, as on main). go vet ./internal/..., go build -tags nogui (darwin, GOOS=linux), GOOS=windows go build, go test -tags nogui -count=1 ./... under a temp HOME pass (TestFXInState's known TempDir-cleanup flake aside).
This commit is contained in:
yetone committed
1 parent
b2ac45b0bc
commit
0879894af3
21 files changed
+1342
-15
No files matched your search
@@ -187,8 +187,30 @@ magpie gateway-key add "Remote laptop" # prints the new credential once
|
||||
magpie gateway-key list # ids, names, enabled state and masked keys
|
||||
magpie gateway-key rotate <id> # prints the replacement; identity stays the same
|
||||
magpie gateway-key remove <id> # revokes remote access
|
||||
magpie gateway-key limit <id> week --tokens 2m --cost 5 # its own limit
|
||||
magpie gateway-key limit <id> # limit, used, left and reset
|
||||
magpie gateway-key limit <id> off # no limit
|
||||
```
|
||||
|
||||
Each gateway key can have its own **limit**: a token total, an estimated
|
||||
cost in US$, or both, per day, week or month (calendar windows in local
|
||||
time: from midnight, from Monday, from the 1st). Set it with **Limit** on
|
||||
the key's row (saved with Save) or `magpie gateway-key limit`; the row shows
|
||||
what the key has used, what is left and when it resets. Tokens counted are a
|
||||
call's uncached input, output and cache writes, plus cache reads when **Count
|
||||
cache reads too** is on. Cost is an estimate at the Usage page's prices; a
|
||||
call with no known price adds none. A call counts in the window it started
|
||||
in. Once a key is spent, its requests are refused before any provider is
|
||||
asked, with a 429 in the API's own error shape that names the key, the
|
||||
limit and the reset time, plus `Retry-After`; other keys are unaffected. The
|
||||
counts are read from the usage log, so they survive a restart. A request in
|
||||
flight holds a reservation (its body's size in tokens plus the key's mean
|
||||
output per call), so requests sent at once overshoot by about one call; a
|
||||
streamed reply is settled when it ends with the usage its vendor reported.
|
||||
A key can read its own status with `GET /v1/magpie/limit`. Requests from
|
||||
this computer that send no gateway key are not limited; a gateway key used
|
||||
from this computer is.
|
||||
|
||||
While LAN sharing is enabled, remote requests require an enabled gateway key
|
||||
sent as Bearer, `x-api-key`, `x-goog-api-key` or `?key=`. Loopback remains
|
||||
permissive: any token works, including a stale or disabled gateway key.
|
||||
|
||||
+143
-4
@@ -4,10 +4,14 @@ import (
|
||||
"fmt"
|
||||
"io"
|
||||
"os"
|
||||
"slices"
|
||||
"strconv"
|
||||
"strings"
|
||||
"text/tabwriter"
|
||||
"time"
|
||||
|
||||
"github.com/yetone/magpie/internal/access"
|
||||
"github.com/yetone/magpie/internal/budget"
|
||||
)
|
||||
|
||||
func gatewayKeys(args []string) error { return gatewayKeysTo(os.Stdout, args) }
|
||||
@@ -18,8 +22,12 @@ func gatewayKeysTo(out io.Writer, args []string) error {
|
||||
args = []string{"list"}
|
||||
}
|
||||
action := args[0]
|
||||
if action == "limit" {
|
||||
access.MigrateLegacyLANKeyBestEffort()
|
||||
return gatewayKeyLimit(out, args[1:])
|
||||
}
|
||||
if (action == "list" && len(args) != 1) || (action != "list" && len(args) != 2) {
|
||||
return fmt.Errorf("usage: magpie gateway-key list | add <name> | rotate <id> | remove <id>")
|
||||
return fmt.Errorf("usage: magpie gateway-key list | add <name> | rotate <id> | remove <id> | limit <id> [off | day|week|month [--tokens N] [--cost USD] [--cache-reads]]")
|
||||
}
|
||||
switch action {
|
||||
case "list", "add", "rotate", "remove":
|
||||
@@ -33,18 +41,19 @@ func gatewayKeysTo(out io.Writer, args []string) error {
|
||||
return err
|
||||
}
|
||||
w := tabwriter.NewWriter(out, 0, 4, 2, ' ', 0)
|
||||
fmt.Fprintln(w, "ID\tNAME\tSTATE\tGATEWAY KEY")
|
||||
fmt.Fprintln(w, "ID\tNAME\tSTATE\tGATEWAY KEY\tLIMIT")
|
||||
now := time.Now()
|
||||
for _, k := range keys {
|
||||
state := "enabled"
|
||||
if k.Off {
|
||||
state = "disabled"
|
||||
}
|
||||
fmt.Fprintf(w, "%s\t%s\t%s\t%s\n", k.ID, strings.Map(func(r rune) rune {
|
||||
fmt.Fprintf(w, "%s\t%s\t%s\t%s\t%s\n", k.ID, strings.Map(func(r rune) rune {
|
||||
if r < 32 || r == 127 {
|
||||
return ' '
|
||||
}
|
||||
return r
|
||||
}, k.Name), state, k.Masked)
|
||||
}, k.Name), state, k.Masked, limitWords(k, now))
|
||||
}
|
||||
return w.Flush()
|
||||
}
|
||||
@@ -63,3 +72,133 @@ func gatewayKeysTo(out io.Writer, args []string) error {
|
||||
}
|
||||
return err
|
||||
}
|
||||
|
||||
// gatewayKeyLimit shows or sets a key's limit (#585):
|
||||
//
|
||||
// magpie gateway-key limit <id> what it has used of its limit
|
||||
// magpie gateway-key limit <id> off no limit
|
||||
// magpie gateway-key limit <id> day|week|month [--tokens N] [--cost USD] [--cache-reads]
|
||||
func gatewayKeyLimit(out io.Writer, args []string) error {
|
||||
usage := fmt.Errorf("usage: magpie gateway-key limit <id> [off | day|week|month [--tokens N] [--cost USD] [--cache-reads]]")
|
||||
if len(args) == 0 {
|
||||
return usage
|
||||
}
|
||||
id := args[0]
|
||||
if len(args) > 1 {
|
||||
var lim *access.Limit
|
||||
if args[1] != "off" {
|
||||
lim = &access.Limit{Period: args[1]}
|
||||
if !slices.Contains(access.Periods, lim.Period) {
|
||||
return fmt.Errorf("a limit's period is day, week or month, not %q", lim.Period)
|
||||
}
|
||||
rest := args[2:]
|
||||
for i := 0; i < len(rest); i++ {
|
||||
flag, val, has := strings.Cut(rest[i], "=")
|
||||
next := func() (string, error) {
|
||||
if has {
|
||||
return val, nil
|
||||
}
|
||||
if i+1 >= len(rest) {
|
||||
return "", fmt.Errorf("%s needs a value", flag)
|
||||
}
|
||||
i++
|
||||
return rest[i], nil
|
||||
}
|
||||
switch flag {
|
||||
case "--tokens":
|
||||
v, err := next()
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
n, err := parseTokens(v)
|
||||
if err != nil || n < 0 {
|
||||
return fmt.Errorf("--tokens is a count, like 1000000 or 2m: %q", v)
|
||||
}
|
||||
lim.Tokens = int64(n)
|
||||
case "--cost":
|
||||
v, err := next()
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if lim.Cost, err = strconv.ParseFloat(strings.TrimPrefix(v, "$"), 64); err != nil {
|
||||
return fmt.Errorf("--cost is US dollars, like 5 or 2.50: %q", v)
|
||||
}
|
||||
case "--cache-reads":
|
||||
lim.CacheReads = true
|
||||
default:
|
||||
return usage
|
||||
}
|
||||
}
|
||||
if !lim.Limited() {
|
||||
return fmt.Errorf("say --tokens, --cost or both, or off for no limit")
|
||||
}
|
||||
} else if len(args) > 2 {
|
||||
return usage
|
||||
}
|
||||
if _, err := access.Update("limit-key", access.Change{Key: id, Limit: lim}); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
keys, err := access.List()
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
i := slices.IndexFunc(keys, func(k access.Key) bool { return k.ID == id })
|
||||
if i < 0 {
|
||||
return fmt.Errorf("Key not found")
|
||||
}
|
||||
k, now := keys[i], time.Now()
|
||||
st := budget.Of(k, now)
|
||||
if st == nil {
|
||||
_, err = fmt.Fprintf(out, "%s: no limit\n", k.Name)
|
||||
return err
|
||||
}
|
||||
w := tabwriter.NewWriter(out, 0, 4, 2, ' ', 0)
|
||||
fmt.Fprintf(w, "Key\t%s (%s)\n", k.Name, k.ID)
|
||||
fmt.Fprintf(w, "Window\tthis %s, %s – %s (local time)\n", st.Period, st.Start.Format("2006-01-02 15:04"), st.Reset.Format("2006-01-02 15:04"))
|
||||
if st.TokenLimit > 0 {
|
||||
counted := "input, output and cache writes"
|
||||
if st.CacheReads {
|
||||
counted += ", and cache reads"
|
||||
}
|
||||
fmt.Fprintf(w, "Tokens\t%d used of %d, %d left (%s)\n", st.Tokens, st.TokenLimit, st.TokensLeft, counted)
|
||||
} else {
|
||||
fmt.Fprintf(w, "Tokens\t%d used, no cap\n", st.Tokens)
|
||||
}
|
||||
unpriced := ""
|
||||
if st.Unpriced > 0 {
|
||||
unpriced = fmt.Sprintf(", %d calls without a known price not in it", st.Unpriced)
|
||||
}
|
||||
if st.CostLimit > 0 {
|
||||
fmt.Fprintf(w, "Cost\t$%.2f used of $%.2f, $%.2f left (estimate at Usage prices%s)\n", st.Cost, st.CostLimit, st.CostLeft, unpriced)
|
||||
} else {
|
||||
fmt.Fprintf(w, "Cost\t$%.2f used, no cap (estimate%s)\n", st.Cost, unpriced)
|
||||
}
|
||||
state := "open"
|
||||
if st.Spent {
|
||||
state = "spent: requests are refused until " + st.Reset.Format("2006-01-02 15:04")
|
||||
}
|
||||
fmt.Fprintf(w, "State\t%s\n", state)
|
||||
return w.Flush()
|
||||
}
|
||||
|
||||
// limitWords is a key's limit in the list: "-" for none, else what it has
|
||||
// used of it, "1200/1000000 tokens/day".
|
||||
func limitWords(k access.Key, now time.Time) string {
|
||||
st := budget.Of(k, now)
|
||||
if st == nil {
|
||||
return "-"
|
||||
}
|
||||
var parts []string
|
||||
if st.TokenLimit > 0 {
|
||||
parts = append(parts, fmt.Sprintf("%d/%d tokens", st.Tokens, st.TokenLimit))
|
||||
}
|
||||
if st.CostLimit > 0 {
|
||||
parts = append(parts, fmt.Sprintf("$%.2f/$%.2f", st.Cost, st.CostLimit))
|
||||
}
|
||||
s := strings.Join(parts, " ") + " per " + st.Period
|
||||
if st.Spent {
|
||||
s += " (spent until " + st.Reset.Format("01-02 15:04") + ")"
|
||||
}
|
||||
return s
|
||||
}
|
||||
@@ -79,3 +79,55 @@ func TestGatewayKeyListWithReadOnlyMigration(t *testing.T) {
|
||||
t.Fatal("read-only migration blocked CLI list", out.String(), err)
|
||||
}
|
||||
}
|
||||
|
||||
// magpie gateway-key limit sets, shows and takes off a key's limit, and
|
||||
// list says what each key has used of its own (#585).
|
||||
func TestGatewayKeyLimitCLI(t *testing.T) {
|
||||
t.Setenv("HOME", t.TempDir())
|
||||
t.Setenv("USERPROFILE", os.Getenv("HOME"))
|
||||
t.Setenv("XDG_CONFIG_HOME", t.TempDir())
|
||||
if err := settings.Save(settings.Settings{LAN: true}); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
call := func(args ...string) (string, error) {
|
||||
var out bytes.Buffer
|
||||
err := gatewayKeysTo(&out, append([]string{"gateway-key"}, args...))
|
||||
return out.String(), err
|
||||
}
|
||||
if _, err := call("add", "Phone"); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
keys, _ := access.List()
|
||||
id := keys[0].ID
|
||||
if got, err := call("limit", id); err != nil || !strings.Contains(got, "no limit") {
|
||||
t.Fatal(got, err)
|
||||
}
|
||||
got, err := call("limit", id, "week", "--tokens", "2m", "--cost=5", "--cache-reads")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
keys, _ = access.List()
|
||||
if l := keys[0].Limit; l == nil || l.Period != "week" || l.Tokens != 2_000_000 || l.Cost != 5 || !l.CacheReads {
|
||||
t.Fatalf("limit kept as %+v", keys[0].Limit)
|
||||
}
|
||||
for _, want := range []string{"0 used of 2000000, 2000000 left", "cache reads", "$0.00 used of $5.00", "State", "open"} {
|
||||
if !strings.Contains(got, want) {
|
||||
t.Fatalf("limit shows %q, without %q", got, want)
|
||||
}
|
||||
}
|
||||
if got, _ := call("list"); !strings.Contains(got, "LIMIT") || !strings.Contains(got, "0/2000000 tokens $0.00/$5.00 per week") {
|
||||
t.Fatal(got)
|
||||
}
|
||||
for _, args := range [][]string{{"limit"}, {"limit", id, "hour", "--tokens", "5"}, {"limit", id, "day"}, {"limit", id, "day", "--tokens", "lots"},
|
||||
{"limit", id, "day", "--cost"}, {"limit", id, "day", "--what"}, {"limit", "missing", "day", "--tokens", "5"}, {"limit", id, "off", "extra"}} {
|
||||
if _, err := call(args...); err == nil {
|
||||
t.Error("accepted", args)
|
||||
}
|
||||
}
|
||||
if got, err := call("limit", id, "off"); err != nil || !strings.Contains(got, "no limit") {
|
||||
t.Fatal(got, err)
|
||||
}
|
||||
if keys, _ = access.List(); keys[0].Limit != nil {
|
||||
t.Fatal("off kept", keys[0].Limit)
|
||||
}
|
||||
}
|
||||
@@ -37,9 +37,15 @@ type Key struct {
|
||||
LAN bool `json:"lan,omitempty"` // default key and durable migration marker
|
||||
Secret string `json:"secret,omitempty"`
|
||||
Masked string `json:"masked,omitempty"`
|
||||
// Limit is the key's own budget (#585); nil for none.
|
||||
Limit *Limit `json:"limit,omitempty"`
|
||||
}
|
||||
|
||||
type Identity struct{ KeyID, KeyName string }
|
||||
// Identity is the gateway key a request came with, and its budget.
|
||||
type Identity struct {
|
||||
KeyID, KeyName string
|
||||
Limit *Limit
|
||||
}
|
||||
type contextKey struct{}
|
||||
|
||||
func WithIdentity(ctx context.Context, who Identity) context.Context {
|
||||
@@ -111,6 +117,8 @@ func random(n int) (string, error) {
|
||||
type Change struct {
|
||||
Key string `json:"key"`
|
||||
Name string `json:"name"`
|
||||
// Limit is what "limit-key" sets; nil or Unlimited takes the limit off.
|
||||
Limit *Limit `json:"limit,omitempty"`
|
||||
}
|
||||
|
||||
// Update writes the named key store atomically.
|
||||
@@ -159,6 +167,12 @@ func Update(action string, in Change) (string, error) {
|
||||
keys[i].Secret = secret
|
||||
case "rename-key":
|
||||
keys[i].Name = name
|
||||
case "limit-key":
|
||||
lim, err := in.Limit.Valid()
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
keys[i].Limit = lim
|
||||
case "on-key", "off-key":
|
||||
keys[i].Off = action == "off-key"
|
||||
case "remove-key":
|
||||
@@ -173,7 +187,7 @@ func Update(action string, in Change) (string, error) {
|
||||
default:
|
||||
return "", fmt.Errorf("unknown key action %q", action)
|
||||
}
|
||||
if defaultKey && action != "remove-key" && action != "rename-key" {
|
||||
if defaultKey && action != "remove-key" && action != "rename-key" && action != "limit-key" {
|
||||
mirror = &keys[i]
|
||||
}
|
||||
}
|
||||
@@ -220,7 +234,7 @@ func Authenticate(secret string) (Identity, bool) {
|
||||
if k.Off {
|
||||
return Identity{}, false
|
||||
}
|
||||
return Identity{k.ID, k.Name}, true
|
||||
return Identity{k.ID, k.Name, k.Limit}, true
|
||||
}
|
||||
}
|
||||
s := settings.Load()
|
||||
@@ -228,7 +242,7 @@ func Authenticate(secret string) (Identity, bool) {
|
||||
// fall back after either durable marker exists, or to a revoked mirror.
|
||||
if s.LAN && s.LANKeyID == "" && s.LANKey != "" && !strings.HasPrefix(s.LANKey, revokedLANPrefix) && !slices.ContainsFunc(keys, func(k Key) bool { return k.LAN }) && subtle.ConstantTimeCompare([]byte(secret), []byte(s.LANKey)) == 1 {
|
||||
k := legacyLANKey(s.LANKey)
|
||||
return Identity{k.ID, k.Name}, true
|
||||
return Identity{KeyID: k.ID, KeyName: k.Name}, true
|
||||
}
|
||||
return Identity{}, false
|
||||
}
|
||||
|
||||
@@ -0,0 +1,69 @@
|
||||
package access
|
||||
|
||||
import (
|
||||
"errors"
|
||||
"math"
|
||||
"time"
|
||||
)
|
||||
|
||||
// Limit is a gateway key's budget for a calendar window in local time
|
||||
// (#585): a day from midnight, a week from Monday's midnight, a month
|
||||
// from the 1st's. Tokens counts each call's uncached input, output
|
||||
// (reasoning included) and cache writes, and its cache reads too when
|
||||
// CacheReads is set; Cost is an estimate in US dollars at the prices
|
||||
// magpie shows on the Usage page, not a vendor's bill. A cap of 0 is no
|
||||
// cap; a limit with neither is no limit.
|
||||
type Limit struct {
|
||||
Period string `json:"period"`
|
||||
Tokens int64 `json:"tokens,omitempty"`
|
||||
Cost float64 `json:"cost,omitempty"`
|
||||
CacheReads bool `json:"cacheReads,omitempty"`
|
||||
}
|
||||
|
||||
// Periods are the windows a limit can be set for.
|
||||
var Periods = []string{"day", "week", "month"}
|
||||
|
||||
// Limited: l caps something.
|
||||
func (l *Limit) Limited() bool { return l != nil && (l.Tokens > 0 || l.Cost > 0) }
|
||||
|
||||
// Valid is l as it is kept: nil for no limit, an error for one that
|
||||
// can't be.
|
||||
func (l *Limit) Valid() (*Limit, error) {
|
||||
if l == nil {
|
||||
return nil, nil
|
||||
}
|
||||
if l.Tokens < 0 || l.Cost < 0 || math.IsNaN(l.Cost) || math.IsInf(l.Cost, 0) {
|
||||
return nil, errors.New("A limit can't be negative")
|
||||
}
|
||||
if !l.Limited() {
|
||||
return nil, nil
|
||||
}
|
||||
v := *l
|
||||
if v.Period == "" {
|
||||
v.Period = "day"
|
||||
}
|
||||
switch v.Period {
|
||||
case "day", "week", "month":
|
||||
default:
|
||||
return nil, errors.New("A limit's period is day, week or month")
|
||||
}
|
||||
v.Cost = math.Round(v.Cost*1e4) / 1e4
|
||||
return &v, nil
|
||||
}
|
||||
|
||||
// Window is the calendar window of period that now is in, in now's
|
||||
// location: its start and when the next begins.
|
||||
func Window(period string, now time.Time) (start, reset time.Time) {
|
||||
y, m, d := now.Date()
|
||||
loc := now.Location()
|
||||
switch period {
|
||||
case "week":
|
||||
back := (int(now.Weekday()) + 6) % 7 // days since Monday
|
||||
start = time.Date(y, m, d-back, 0, 0, 0, 0, loc)
|
||||
return start, time.Date(y, m, d-back+7, 0, 0, 0, 0, loc)
|
||||
case "month":
|
||||
return time.Date(y, m, 1, 0, 0, 0, 0, loc), time.Date(y, m+1, 1, 0, 0, 0, 0, loc)
|
||||
default:
|
||||
return time.Date(y, m, d, 0, 0, 0, 0, loc), time.Date(y, m, d+1, 0, 0, 0, 0, loc)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,48 @@
|
||||
package access
|
||||
|
||||
import (
|
||||
"testing"
|
||||
"time"
|
||||
)
|
||||
|
||||
func TestLimitWindow(t *testing.T) {
|
||||
loc := time.FixedZone("X", 8*3600)
|
||||
at := func(s string) time.Time {
|
||||
v, err := time.ParseInLocation("2006-01-02 15:04", s, loc)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
return v
|
||||
}
|
||||
for _, c := range []struct{ period, now, start, reset string }{
|
||||
{"day", "2026-10-03 23:59", "2026-10-03 00:00", "2026-10-04 00:00"},
|
||||
{"day", "2026-12-31 00:00", "2026-12-31 00:00", "2027-01-01 00:00"},
|
||||
{"week", "2026-10-03 12:00", "2026-09-28 00:00", "2026-10-05 00:00"}, // a Saturday: from Monday
|
||||
{"week", "2026-10-04 23:00", "2026-09-28 00:00", "2026-10-05 00:00"}, // Sunday is the week's last day
|
||||
{"week", "2026-10-05 00:00", "2026-10-05 00:00", "2026-10-12 00:00"},
|
||||
{"month", "2026-10-31 18:00", "2026-10-01 00:00", "2026-11-01 00:00"},
|
||||
{"month", "2026-12-15 08:00", "2026-12-01 00:00", "2027-01-01 00:00"},
|
||||
} {
|
||||
start, reset := Window(c.period, at(c.now))
|
||||
if !start.Equal(at(c.start)) || !reset.Equal(at(c.reset)) {
|
||||
t.Errorf("%s at %s: %s – %s", c.period, c.now, start, reset)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestLimitValid(t *testing.T) {
|
||||
if l, err := (&Limit{Period: "day"}).Valid(); l != nil || err != nil {
|
||||
t.Fatal("no cap is no limit", l, err)
|
||||
}
|
||||
if l, err := (*Limit)(nil).Valid(); l != nil || err != nil {
|
||||
t.Fatal(l, err)
|
||||
}
|
||||
if l, err := (&Limit{Tokens: 5}).Valid(); err != nil || l.Period != "day" {
|
||||
t.Fatal("a period is a day unless said", l, err)
|
||||
}
|
||||
for _, bad := range []Limit{{Period: "hour", Tokens: 1}, {Period: "day", Tokens: -1}, {Period: "day", Cost: -2}} {
|
||||
if _, err := bad.Valid(); err == nil {
|
||||
t.Error("accepted", bad)
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,302 @@
|
||||
// Package budget holds each gateway key to its own limit (#585): what a
|
||||
// key has used in its window is summed from usage.jsonl, the records the
|
||||
// Usage page shows, so it survives a restart; a request is let in only
|
||||
// while the key has something left after what its requests in flight are
|
||||
// expected to use, and those are settled with what the call recorded.
|
||||
//
|
||||
// The rules:
|
||||
// - A call counts in the window it started in (its record's time), and
|
||||
// windows are calendar ones in local time (access.Window).
|
||||
// - Tokens are a call's uncached input, output and cache writes, and
|
||||
// its cache reads when the limit says so; cost is an estimate at the
|
||||
// Usage page's prices, and a call with no known price adds none.
|
||||
// - A request refused before any upstream (no model, a refusal like
|
||||
// this one) uses nothing; a failed call counts the tokens it recorded.
|
||||
// - A streamed reply is settled when it ends, with the usage its vendor
|
||||
// reported; one cut off before that counts what was recorded.
|
||||
// - While a request is in flight it holds a reservation: its body's
|
||||
// size in tokens (4 bytes a token) plus the key's mean output per
|
||||
// call this window. A request is let in while used + reserved is
|
||||
// under the cap, so parallel requests overshoot by about one call.
|
||||
// - Only requests that came with a gateway key are limited: from
|
||||
// another computer one is required; on this one a request without one
|
||||
// (or with a token that is no key) isn't limited.
|
||||
package budget
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"strings"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"github.com/yetone/magpie/internal/access"
|
||||
"github.com/yetone/magpie/internal/provider"
|
||||
"github.com/yetone/magpie/internal/settings"
|
||||
"github.com/yetone/magpie/internal/usage"
|
||||
)
|
||||
|
||||
// sums is what a key's calls came to in a window.
|
||||
type sums struct {
|
||||
start time.Time
|
||||
calls int
|
||||
input, output, cacheRead, cacheWrite int64
|
||||
cost float64
|
||||
unpriced int
|
||||
}
|
||||
|
||||
func (s *sums) add(r usage.Record, cost float64, priced bool) {
|
||||
s.calls++
|
||||
s.input += int64(r.Input)
|
||||
s.output += int64(r.Output)
|
||||
s.cacheRead += int64(r.CacheRead)
|
||||
s.cacheWrite += int64(r.CacheWrite)
|
||||
s.cost += cost
|
||||
if !priced {
|
||||
s.unpriced++
|
||||
}
|
||||
}
|
||||
|
||||
func (s *sums) tokens(l *access.Limit) int64 {
|
||||
n := s.input + s.output + s.cacheWrite
|
||||
if l != nil && l.CacheReads {
|
||||
n += s.cacheRead
|
||||
}
|
||||
return n
|
||||
}
|
||||
|
||||
type held struct {
|
||||
n int
|
||||
tokens int64
|
||||
cost float64
|
||||
}
|
||||
|
||||
var state = struct {
|
||||
sync.Mutex
|
||||
sums map[string]*sums // by usage.Path() and key
|
||||
held map[string]*held
|
||||
}{sums: map[string]*sums{}, held: map[string]*held{}}
|
||||
|
||||
func slot(keyID string) string { return usage.Path() + "\x00" + keyID }
|
||||
|
||||
// counts reports whether r is a call a key's budget counts.
|
||||
func counts(r usage.Record) bool { return !r.IsRejected() && r.Computer == "" }
|
||||
|
||||
// sumsFor is keyID's sums for the window starting at start, read from
|
||||
// usage.jsonl the first time. Called with state held.
|
||||
func sumsFor(keyID string, start time.Time) *sums {
|
||||
k := slot(keyID)
|
||||
if s := state.sums[k]; s != nil && s.start.Equal(start) {
|
||||
return s
|
||||
}
|
||||
s := read(keyID, start)
|
||||
state.sums[k] = s
|
||||
return s
|
||||
}
|
||||
|
||||
func read(keyID string, start time.Time) *sums {
|
||||
s := &sums{start: start}
|
||||
cost := usage.NewCoster()
|
||||
usage.Visit(start, func(r usage.Record) {
|
||||
if r.CallerKeyID == keyID && counts(r) && !r.Time.Before(start) {
|
||||
c, ok := cost(r)
|
||||
s.add(r, c, ok)
|
||||
}
|
||||
})
|
||||
return s
|
||||
}
|
||||
|
||||
// Forget drops what is kept in memory, as a restart would.
|
||||
func Forget() {
|
||||
state.Lock()
|
||||
defer state.Unlock()
|
||||
state.sums = map[string]*sums{}
|
||||
state.held = map[string]*held{}
|
||||
}
|
||||
|
||||
// Append writes rec to usage.jsonl and, for a key's call, counts it
|
||||
// against the key's window. The two happen together so a window read from
|
||||
// the file meanwhile counts the call once.
|
||||
func Append(rec usage.Record) {
|
||||
if rec.CallerKeyID == "" {
|
||||
usage.Append(rec)
|
||||
return
|
||||
}
|
||||
state.Lock()
|
||||
defer state.Unlock()
|
||||
if rec.Time.IsZero() {
|
||||
rec.Time = time.Now()
|
||||
}
|
||||
usage.Append(rec)
|
||||
s := state.sums[slot(rec.CallerKeyID)]
|
||||
if s == nil || !counts(rec) || rec.Time.Before(s.start) {
|
||||
return
|
||||
}
|
||||
c, ok := usage.NewCoster()(rec)
|
||||
s.add(rec, c, ok)
|
||||
}
|
||||
|
||||
// Status is a key's limit and what it has used of it.
|
||||
type Status struct {
|
||||
Period string `json:"period"`
|
||||
Start time.Time `json:"start"`
|
||||
Reset time.Time `json:"reset"`
|
||||
Calls int `json:"calls"`
|
||||
Tokens int64 `json:"tokens"`
|
||||
Cost float64 `json:"cost"`
|
||||
Unpriced int `json:"unpriced,omitempty"`
|
||||
// InFlight requests, and the tokens and cost they are reserved at
|
||||
InFlight int `json:"inFlight,omitempty"`
|
||||
Reserved int64 `json:"reserved,omitempty"`
|
||||
ReservedCost float64 `json:"reservedCost,omitempty"`
|
||||
TokenLimit int64 `json:"tokenLimit,omitempty"`
|
||||
CostLimit float64 `json:"costLimit,omitempty"`
|
||||
CacheReads bool `json:"cacheReads,omitempty"`
|
||||
// TokensLeft and CostLeft are what is left of each cap, 0 at most
|
||||
TokensLeft int64 `json:"tokensLeft"`
|
||||
CostLeft float64 `json:"costLeft"`
|
||||
// Spent: the key is refused until Reset
|
||||
Spent bool `json:"spent"`
|
||||
}
|
||||
|
||||
func status(keyID string, l *access.Limit, s *sums, h *held, reset time.Time) Status {
|
||||
st := Status{Period: l.Period, Start: s.start, Reset: reset, Calls: s.calls, Tokens: s.tokens(l), Cost: s.cost, Unpriced: s.unpriced,
|
||||
TokenLimit: l.Tokens, CostLimit: l.Cost, CacheReads: l.CacheReads}
|
||||
if h != nil {
|
||||
st.InFlight, st.Reserved, st.ReservedCost = h.n, h.tokens, h.cost
|
||||
}
|
||||
if l.Tokens > 0 {
|
||||
st.TokensLeft = max(l.Tokens-st.Tokens, 0)
|
||||
st.Spent = st.Tokens >= l.Tokens
|
||||
}
|
||||
if l.Cost > 0 {
|
||||
st.CostLeft = max(l.Cost-st.Cost, 0)
|
||||
st.Spent = st.Spent || st.Cost >= l.Cost
|
||||
}
|
||||
return st
|
||||
}
|
||||
|
||||
// Of is key k's status now, nil when it has no limit. It reads the file
|
||||
// afresh rather than the gateway's sums, so it is right for a gateway
|
||||
// served by another process too.
|
||||
func Of(k access.Key, now time.Time) *Status {
|
||||
if !k.Limit.Limited() {
|
||||
return nil
|
||||
}
|
||||
start, reset := access.Window(k.Limit.Period, now)
|
||||
s := read(k.ID, start)
|
||||
state.Lock()
|
||||
h := state.held[slot(k.ID)]
|
||||
st := status(k.ID, k.Limit, s, h, reset)
|
||||
state.Unlock()
|
||||
return &st
|
||||
}
|
||||
|
||||
// Refusal is why a key's request was turned away.
|
||||
type Refusal struct {
|
||||
Status
|
||||
Key string
|
||||
}
|
||||
|
||||
// Error says it as a client is told: the key, what it used of which
|
||||
// cap, and when it resets.
|
||||
func (r *Refusal) Error() string {
|
||||
what := ""
|
||||
if r.TokenLimit > 0 && r.Tokens+r.Reserved >= r.TokenLimit {
|
||||
counted := "input, output and cache-write tokens"
|
||||
if r.CacheReads {
|
||||
counted = "input, output and cache tokens"
|
||||
}
|
||||
what = fmt.Sprintf("%d of its %d %s", r.Tokens, r.TokenLimit, counted)
|
||||
} else {
|
||||
what = fmt.Sprintf("$%.2f of its $%.2f estimated cost", r.Cost, r.CostLimit)
|
||||
}
|
||||
if !r.Spent {
|
||||
what += fmt.Sprintf(", with %d requests in flight holding the rest", r.InFlight)
|
||||
}
|
||||
return fmt.Sprintf("magpie gateway key %q has used %s for this %s; it resets at %s",
|
||||
r.Key, what, r.Period, r.Reset.Format("2006-01-02 15:04 MST"))
|
||||
}
|
||||
|
||||
// RetryAfter is how long until the window resets, in whole seconds.
|
||||
func (r *Refusal) RetryAfter(now time.Time) int {
|
||||
return max(int(r.Reset.Sub(now).Seconds()+0.999), 1)
|
||||
}
|
||||
|
||||
// Reserve lets a request of who's in, holding a reservation for it until
|
||||
// the release it returns is called, or refuses it. body is the request's
|
||||
// size in bytes (-1 unknown) and model the model it named, for a cost cap
|
||||
// before the key has a price of its own this window.
|
||||
func Reserve(who access.Identity, body int64, model string, now time.Time) (release func(), refused *Refusal) {
|
||||
l := who.Limit
|
||||
if who.KeyID == "" || !l.Limited() {
|
||||
return func() {}, nil
|
||||
}
|
||||
start, reset := access.Window(l.Period, now)
|
||||
state.Lock()
|
||||
defer state.Unlock()
|
||||
k := slot(who.KeyID)
|
||||
s := sumsFor(who.KeyID, start)
|
||||
h := state.held[k]
|
||||
if h == nil {
|
||||
h = &held{}
|
||||
state.held[k] = h
|
||||
}
|
||||
used := s.tokens(l)
|
||||
if (l.Tokens > 0 && used+h.tokens >= l.Tokens) || (l.Cost > 0 && s.cost+h.cost >= l.Cost) {
|
||||
return nil, &Refusal{Status: status(who.KeyID, l, s, h, reset), Key: who.KeyName}
|
||||
}
|
||||
est := max(body, 0) / 4
|
||||
if s.calls > 0 {
|
||||
est += s.output / int64(s.calls)
|
||||
}
|
||||
var cost float64
|
||||
if l.Cost > 0 {
|
||||
if used > 0 && s.cost > 0 {
|
||||
cost = float64(est) * s.cost / float64(used)
|
||||
} else if p, ok := priceOf(model); ok {
|
||||
cost = float64(est) * p
|
||||
}
|
||||
}
|
||||
h.n++
|
||||
h.tokens += est
|
||||
h.cost += cost
|
||||
once := sync.Once{}
|
||||
return func() {
|
||||
once.Do(func() {
|
||||
state.Lock()
|
||||
defer state.Unlock()
|
||||
h.n--
|
||||
h.tokens -= est
|
||||
h.cost -= cost
|
||||
if h.n == 0 && state.held[k] == h {
|
||||
delete(state.held, k)
|
||||
}
|
||||
})
|
||||
}, nil
|
||||
}
|
||||
|
||||
// priceOf is the input price per token of model when it names its
|
||||
// provider ("provider/model").
|
||||
func priceOf(model string) (float64, bool) {
|
||||
p, m, ok := strings.Cut(model, "/")
|
||||
if !ok || p == "" || m == "" {
|
||||
return 0, false
|
||||
}
|
||||
pr, ok := provider.EffectivePriceIn(settings.Load(), p, m)
|
||||
if !ok {
|
||||
return 0, false
|
||||
}
|
||||
return pr.Input / 1e6, true
|
||||
}
|
||||
|
||||
// ModelOf is the model a JSON request body names, "" for none.
|
||||
func ModelOf(body []byte) string {
|
||||
var v struct {
|
||||
Model string `json:"model"`
|
||||
}
|
||||
if json.Unmarshal(body, &v) != nil {
|
||||
return ""
|
||||
}
|
||||
return v.Model
|
||||
}
|
||||
@@ -443,6 +443,7 @@ func (s *Server) Handler() http.Handler {
|
||||
mux.HandleFunc("GET /v1/magpie/quotas", s.quotas)
|
||||
mux.HandleFunc("GET /v1/magpie/route", s.sessionRoute)
|
||||
mux.HandleFunc("GET /v1/magpie/concurrency", s.concurrency)
|
||||
mux.HandleFunc("GET /v1/magpie/limit", s.keyLimit)
|
||||
mux.HandleFunc("POST /v1/chat/completions", s.handle(provider.Chat))
|
||||
mux.HandleFunc("POST /chat/completions", s.handle(provider.Chat))
|
||||
mux.HandleFunc("POST /v1/responses", s.handle(provider.Responses))
|
||||
@@ -468,7 +469,7 @@ func (s *Server) Handler() http.Handler {
|
||||
mux.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
|
||||
writeError(w, provider.Chat, http.StatusNotFound, "magpie serves /v1/chat/completions, /v1/responses, /v1/messages, /v1/systemone, /v1/images/generations, /v1/images/edits, /v1/videos and /v1beta/models/*")
|
||||
})
|
||||
return s.counted(callerGuard(withCaller(mux)))
|
||||
return s.counted(callerGuard(withCaller(keyLimited(mux))))
|
||||
}
|
||||
|
||||
func (s *Server) info(w http.ResponseWriter, r *http.Request) {
|
||||
|
||||
@@ -0,0 +1,83 @@
|
||||
package gateway
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"io"
|
||||
"net/http"
|
||||
"strconv"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/yetone/magpie/internal/access"
|
||||
"github.com/yetone/magpie/internal/budget"
|
||||
"github.com/yetone/magpie/internal/provider"
|
||||
"github.com/yetone/magpie/internal/usage"
|
||||
)
|
||||
|
||||
// keyLimited holds a gateway key's requests to its limit (#585, see
|
||||
// package budget for the rules): one that would go over is refused with a
|
||||
// 429 before any provider is asked, saying the key, the limit and when it
|
||||
// resets; one let in holds a reservation until it ends. Reads, the
|
||||
// bridge's tool callbacks and counting tokens cost nothing and pass.
|
||||
func keyLimited(next http.Handler) http.Handler {
|
||||
return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
who := access.Caller(r.Context())
|
||||
if !who.Limit.Limited() || r.Method != http.MethodPost || strings.HasPrefix(r.URL.Path, "/_magpie/") || strings.HasSuffix(r.URL.Path, "/count_tokens") {
|
||||
next.ServeHTTP(w, r)
|
||||
return
|
||||
}
|
||||
size, model := r.ContentLength, ""
|
||||
if who.Limit.Cost > 0 || size < 0 {
|
||||
b, err := io.ReadAll(r.Body)
|
||||
r.Body.Close()
|
||||
if err != nil {
|
||||
writeError(w, protoOfPath(r.URL.Path), http.StatusBadRequest, err.Error())
|
||||
return
|
||||
}
|
||||
r.Body = io.NopCloser(bytes.NewReader(b))
|
||||
size, model = int64(len(b)), budget.ModelOf(b)
|
||||
}
|
||||
now := time.Now()
|
||||
release, refused := budget.Reserve(who, size, model, now)
|
||||
if refused != nil {
|
||||
msg := refused.Error()
|
||||
w.Header().Set("Retry-After", strconv.Itoa(refused.RetryAfter(now)))
|
||||
// the Anthropic and OpenAI SDKs would otherwise retry a 429
|
||||
w.Header().Set("X-Should-Retry", "false")
|
||||
w.Header().Set("X-Magpie-Limit-Reset", refused.Reset.Format(time.RFC3339))
|
||||
writeError(w, protoOfPath(r.URL.Path), http.StatusTooManyRequests, msg)
|
||||
c := callerOf(r)
|
||||
appendUsage(r, usage.Record{Time: now, Agent: c.agent, Via: c.via, Requested: model, Status: http.StatusTooManyRequests, Rejected: true,
|
||||
Error: msg, ErrType: "gateway_key_limit", Endpoint: r.URL.Path, Session: sessionOf(r.Header), NativeSession: nativeSessionOf(r.Header)})
|
||||
return
|
||||
}
|
||||
defer release()
|
||||
next.ServeHTTP(w, r)
|
||||
})
|
||||
}
|
||||
|
||||
// protoOfPath is the API a path is of, for an error said its way.
|
||||
func protoOfPath(path string) provider.Protocol {
|
||||
switch {
|
||||
case strings.HasSuffix(path, "/messages"):
|
||||
return provider.Anthropic
|
||||
case strings.HasPrefix(path, "/v1beta/"):
|
||||
return provider.Gemini
|
||||
case strings.HasSuffix(path, "/responses"):
|
||||
return provider.Responses
|
||||
}
|
||||
return provider.Chat
|
||||
}
|
||||
|
||||
// keyLimit answers GET /v1/magpie/limit: the calling gateway key's limit,
|
||||
// what it has used and when it resets; {"limited":false} for a key with
|
||||
// none or a request without one.
|
||||
func (s *Server) keyLimit(w http.ResponseWriter, r *http.Request) {
|
||||
who := access.Caller(r.Context())
|
||||
if who.KeyID == "" || !who.Limit.Limited() {
|
||||
writeJSON(w, 200, map[string]any{"limited": false, "key": who.KeyName})
|
||||
return
|
||||
}
|
||||
st := budget.Of(access.Key{ID: who.KeyID, Name: who.KeyName, Limit: who.Limit}, time.Now())
|
||||
writeJSON(w, 200, map[string]any{"limited": true, "key": who.KeyName, "limit": st})
|
||||
}
|
||||
@@ -0,0 +1,228 @@
|
||||
package gateway
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"io"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"strconv"
|
||||
"strings"
|
||||
"sync"
|
||||
"sync/atomic"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/yetone/magpie/internal/access"
|
||||
"github.com/yetone/magpie/internal/budget"
|
||||
"github.com/yetone/magpie/internal/provider"
|
||||
"github.com/yetone/magpie/internal/usage"
|
||||
)
|
||||
|
||||
// limitedUpstream is a provider whose every call uses 300 input and 50
|
||||
// output tokens, counting the calls it is sent.
|
||||
func limitedUpstream(t *testing.T, delay time.Duration) *atomic.Int64 {
|
||||
t.Helper()
|
||||
var calls atomic.Int64
|
||||
up := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
calls.Add(1)
|
||||
time.Sleep(delay)
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
io.WriteString(w, `{"choices":[{"message":{"role":"assistant","content":"ok"},"finish_reason":"stop"}],"usage":{"prompt_tokens":300,"completion_tokens":50}}`)
|
||||
}))
|
||||
t.Cleanup(up.Close)
|
||||
if err := provider.Save(provider.Provider{ID: "plan", Name: "Plan", Key: "upstream-secret", Chat: up.URL + "/v1", Models: []string{"m1"}}); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
return &calls
|
||||
}
|
||||
|
||||
func setLimit(t *testing.T, id string, l *access.Limit) {
|
||||
t.Helper()
|
||||
if _, err := access.Update("limit-key", access.Change{Key: id, Limit: l}); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
}
|
||||
|
||||
// A key past its limit is refused before any provider is asked, with a
|
||||
// 429 that names the key, the limit and when it resets; another key and a
|
||||
// keyless request from this computer go on; what was used survives a
|
||||
// restart; raising the limit lets it in again (#585).
|
||||
func TestKeyLimitRefusesBeforeUpstream(t *testing.T) {
|
||||
fresh(t)
|
||||
budget.Forget()
|
||||
calls := limitedUpstream(t, 0)
|
||||
keys, secrets := newCaller(t, "Capped", "Other")
|
||||
setLimit(t, keys[0].ID, &access.Limit{Period: "day", Tokens: 700})
|
||||
h := New().Handler()
|
||||
call := func(secret, path, body string) *httptest.ResponseRecorder {
|
||||
r := httptest.NewRequest("POST", path, strings.NewReader(body))
|
||||
if secret != "" {
|
||||
r.Header.Set("Authorization", "Bearer "+secret)
|
||||
} else {
|
||||
r.RemoteAddr = "127.0.0.1:5000"
|
||||
}
|
||||
w := httptest.NewRecorder()
|
||||
h.ServeHTTP(w, r)
|
||||
return w
|
||||
}
|
||||
for i := range 2 {
|
||||
if w := call(secrets[0], "/v1/chat/completions", chatReq); w.Code != 200 {
|
||||
t.Fatal("within the limit", i, w.Code, w.Body.String())
|
||||
}
|
||||
}
|
||||
if calls.Load() != 2 {
|
||||
t.Fatal(calls.Load())
|
||||
}
|
||||
w := call(secrets[0], "/v1/chat/completions", chatReq)
|
||||
if w.Code != http.StatusTooManyRequests {
|
||||
t.Fatal("spent key let in", w.Code, w.Body.String())
|
||||
}
|
||||
if calls.Load() != 2 {
|
||||
t.Fatal("a spent key's request reached the provider")
|
||||
}
|
||||
var e struct {
|
||||
Error struct{ Message, Type string }
|
||||
}
|
||||
json.Unmarshal(w.Body.Bytes(), &e)
|
||||
if e.Error.Type != "rate_limit_error" || !strings.Contains(e.Error.Message, `"Capped"`) || !strings.Contains(e.Error.Message, "700 of its 700") || !strings.Contains(e.Error.Message, "resets at") {
|
||||
t.Fatal("refusal says", w.Body.String())
|
||||
}
|
||||
if ra, _ := strconv.Atoi(w.Header().Get("Retry-After")); ra <= 0 || ra > 86400+3600 || w.Header().Get("X-Should-Retry") != "false" {
|
||||
t.Fatal("headers", w.Header())
|
||||
}
|
||||
rec := lastUsage(t)
|
||||
if !rec.Rejected || rec.CallerKeyID != keys[0].ID || rec.ErrType != "gateway_key_limit" || rec.Status != 429 {
|
||||
t.Fatalf("refusal recorded as %+v", rec)
|
||||
}
|
||||
// said in Anthropic's shape on its API
|
||||
w = call(secrets[0], "/v1/messages", `{"model":"plan/m1","max_tokens":10,"messages":[{"role":"user","content":"hi"}]}`)
|
||||
var a struct {
|
||||
Type string
|
||||
Error struct{ Type string }
|
||||
}
|
||||
json.Unmarshal(w.Body.Bytes(), &a)
|
||||
if w.Code != 429 || a.Type != "error" || a.Error.Type != "rate_limit_error" || calls.Load() != 2 {
|
||||
t.Fatal("anthropic refusal", w.Code, w.Body.String())
|
||||
}
|
||||
if w := call(secrets[1], "/v1/chat/completions", chatReq); w.Code != 200 {
|
||||
t.Fatal("another key held to the spent one's limit", w.Code)
|
||||
}
|
||||
if w := call("", "/v1/chat/completions", chatReq); w.Code != 200 {
|
||||
t.Fatal("a keyless loopback request was limited", w.Code)
|
||||
}
|
||||
// the gateway restarts: the window is read back from usage.jsonl
|
||||
budget.Forget()
|
||||
if w := call(secrets[0], "/v1/chat/completions", chatReq); w.Code != 429 {
|
||||
t.Fatal("a restart forgot what the key used", w.Code)
|
||||
}
|
||||
st := budget.Of(access.Key{ID: keys[0].ID, Limit: &access.Limit{Period: "day", Tokens: 700}}, time.Now())
|
||||
if st == nil || st.Tokens != 700 || st.Calls != 2 || !st.Spent || st.TokensLeft != 0 {
|
||||
t.Fatalf("status %+v", st)
|
||||
}
|
||||
setLimit(t, keys[0].ID, &access.Limit{Period: "day", Tokens: 1050})
|
||||
if w := call(secrets[0], "/v1/chat/completions", chatReq); w.Code != 200 {
|
||||
t.Fatal("a raised limit still refused", w.Code)
|
||||
}
|
||||
setLimit(t, keys[0].ID, nil)
|
||||
if w := call(secrets[0], "/v1/chat/completions", chatReq); w.Code != 200 {
|
||||
t.Fatal("no limit still refused", w.Code)
|
||||
}
|
||||
// the key asks after its own limit
|
||||
setLimit(t, keys[0].ID, &access.Limit{Period: "week", Tokens: 5000})
|
||||
r := httptest.NewRequest("GET", "/v1/magpie/limit", nil)
|
||||
r.Header.Set("Authorization", "Bearer "+secrets[0])
|
||||
lw := httptest.NewRecorder()
|
||||
h.ServeHTTP(lw, r)
|
||||
var got struct {
|
||||
Limited bool
|
||||
Limit budget.Status
|
||||
}
|
||||
json.Unmarshal(lw.Body.Bytes(), &got)
|
||||
if !got.Limited || got.Limit.Tokens != 1400 || got.Limit.TokensLeft != 3600 || got.Limit.Period != "week" {
|
||||
t.Fatal("limit endpoint", lw.Body.String())
|
||||
}
|
||||
}
|
||||
|
||||
// N requests at once can't all pass a check made before any is counted:
|
||||
// those in flight hold a reservation, so the key goes over by about one
|
||||
// call, not by N.
|
||||
func TestKeyLimitHoldsParallelRequests(t *testing.T) {
|
||||
fresh(t)
|
||||
budget.Forget()
|
||||
calls := limitedUpstream(t, 80*time.Millisecond)
|
||||
keys, secrets := newCaller(t, "Busy")
|
||||
setLimit(t, keys[0].ID, &access.Limit{Period: "month", Tokens: 2000})
|
||||
h := New().Handler()
|
||||
// a body of about the 300 tokens the call uses
|
||||
body := `{"model":"plan/m1","messages":[{"role":"user","content":"` + strings.Repeat("x", 1130) + `"}]}`
|
||||
call := func() int {
|
||||
r := httptest.NewRequest("POST", "/v1/chat/completions", strings.NewReader(body))
|
||||
r.Header.Set("Authorization", "Bearer "+secrets[0])
|
||||
w := httptest.NewRecorder()
|
||||
h.ServeHTTP(w, r)
|
||||
return w.Code
|
||||
}
|
||||
if call() != 200 {
|
||||
t.Fatal("first call")
|
||||
}
|
||||
var wg sync.WaitGroup
|
||||
var ok, refused atomic.Int64
|
||||
for range 20 {
|
||||
wg.Go(func() {
|
||||
switch call() {
|
||||
case 200:
|
||||
ok.Add(1)
|
||||
case 429:
|
||||
refused.Add(1)
|
||||
}
|
||||
})
|
||||
}
|
||||
wg.Wait()
|
||||
st := budget.Of(access.Key{ID: keys[0].ID, Limit: &access.Limit{Period: "month", Tokens: 2000}}, time.Now())
|
||||
if st.Tokens > 2000+350 || refused.Load() == 0 || ok.Load()+refused.Load() != 20 {
|
||||
t.Fatalf("20 at once: %d let in, %d refused, %d tokens used of 2000, %d upstream calls", ok.Load(), refused.Load(), st.Tokens, calls.Load())
|
||||
}
|
||||
if st.InFlight != 0 || st.Reserved != 0 {
|
||||
t.Fatal("reservations left behind", st)
|
||||
}
|
||||
if call() != 429 {
|
||||
t.Fatal("spent key let in")
|
||||
}
|
||||
}
|
||||
|
||||
// A window ends at the calendar's: what was used today counts nothing
|
||||
// tomorrow, and a call before the window started never counts.
|
||||
func TestKeyLimitWindowResets(t *testing.T) {
|
||||
fresh(t)
|
||||
budget.Forget()
|
||||
limitedUpstream(t, 0)
|
||||
keys, secrets := newCaller(t, "Daily")
|
||||
lim := &access.Limit{Period: "day", Tokens: 350}
|
||||
setLimit(t, keys[0].ID, lim)
|
||||
// yesterday's call
|
||||
budget.Append(usage.Record{Time: time.Now().AddDate(0, 0, -1), Provider: "plan", Model: "m1", Input: 5000, Status: 200, CallerKeyID: keys[0].ID})
|
||||
h := New().Handler()
|
||||
call := func() int {
|
||||
r := httptest.NewRequest("POST", "/v1/chat/completions", strings.NewReader(chatReq))
|
||||
r.Header.Set("Authorization", "Bearer "+secrets[0])
|
||||
w := httptest.NewRecorder()
|
||||
h.ServeHTTP(w, r)
|
||||
return w.Code
|
||||
}
|
||||
if call() != 200 {
|
||||
t.Fatal("yesterday's use counted today")
|
||||
}
|
||||
if call() != 429 {
|
||||
t.Fatal("spent key let in")
|
||||
}
|
||||
who := access.Identity{KeyID: keys[0].ID, KeyName: "Daily", Limit: lim}
|
||||
if _, refused := budget.Reserve(who, 10, "", time.Now()); refused == nil {
|
||||
t.Fatal("spent today")
|
||||
}
|
||||
_, reset := access.Window("day", time.Now())
|
||||
release, refused := budget.Reserve(who, 10, "", reset.Add(time.Minute))
|
||||
if refused != nil {
|
||||
t.Fatal("tomorrow's window still spent", refused)
|
||||
}
|
||||
release()
|
||||
}
|
||||
@@ -13,6 +13,7 @@ import (
|
||||
"sync"
|
||||
|
||||
"github.com/yetone/magpie/internal/access"
|
||||
"github.com/yetone/magpie/internal/budget"
|
||||
"github.com/yetone/magpie/internal/provider"
|
||||
"github.com/yetone/magpie/internal/settings"
|
||||
"github.com/yetone/magpie/internal/usage"
|
||||
@@ -257,7 +258,7 @@ func accountOf(p provider.Provider) string {
|
||||
func appendUsage(r *http.Request, rec usage.Record) {
|
||||
who := access.Caller(r.Context())
|
||||
rec.CallerKeyID, rec.CallerKeyName = who.KeyID, who.KeyName
|
||||
usage.Append(rec)
|
||||
budget.Append(rec)
|
||||
}
|
||||
|
||||
// lanKeyed marks a request from another machine that carried the key.
|
||||
|
||||
@@ -1654,6 +1654,20 @@ body.window .ag-sk-icon { width: 22px; height: 22px; border-radius: 6px; }
|
||||
.acc .plan { flex: none; }
|
||||
#gatewayKeys .acc { min-width: 0; }
|
||||
#gatewayKeys .acc .plan { flex: 0 1 auto; min-width: 0; overflow: hidden; text-overflow: ellipsis; }
|
||||
/* a gateway key's own limit (#585): the badge, what it used under the
|
||||
row, and its editor there */
|
||||
#gatewayKeys .acc.with-limit { flex-wrap: wrap; row-gap: 3px; }
|
||||
#gatewayKeys .amodels.limit.spent { color: var(--red); background: var(--red-soft); }
|
||||
#gatewayKeys .key-limit { flex-basis: 100%; min-width: 0; padding: 0 0 2px 25px; font-size: 11px; color: var(--muted); font-variant-numeric: tabular-nums; overflow-wrap: anywhere; }
|
||||
#gatewayKeys .key-limit.spent { color: var(--red); }
|
||||
#gatewayKeys .key-limit-ed .key-limit { padding-left: 0; }
|
||||
#gatewayKeys .klf { display: grid; grid-template-columns: minmax(0, 1fr) minmax(0, 1fr) minmax(0, 1fr); gap: 6px; }
|
||||
#gatewayKeys .klf-field { display: grid; gap: 3px; min-width: 0; }
|
||||
#gatewayKeys .klf-name { font-size: 11px; color: var(--muted); }
|
||||
#gatewayKeys .klf input, #gatewayKeys .klf select { width: 100%; min-width: 0; height: 28px; padding: 0 8px; border: 1px solid var(--line); border-radius: 7px; background: var(--card); color: var(--fg); font: inherit; font-size: 12px; }
|
||||
#gatewayKeys .klf-check { display: inline-flex; align-items: center; gap: 6px; font-size: 12px; color: var(--fg-2); }
|
||||
#gatewayKeys .klf-err { color: var(--red); }
|
||||
@media (max-width: 520px) { #gatewayKeys .klf { grid-template-columns: 1fr; } }
|
||||
/* the Settings row's picker: the key row's badge, always shown */
|
||||
.row.pref .proto.pick { flex: none; margin-left: auto; display: inline-flex; align-items: center; gap: 3px; height: 24px; padding: 0 6px 0 9px; font: inherit; font-size: 12px; color: var(--fg-2); border: 0; border-radius: 7px; background: var(--pill); white-space: nowrap; cursor: default; transition: background .12s, color .12s; }
|
||||
.row.pref .proto.pick svg { opacity: .7; transition: transform .15s var(--ease); }
|
||||
|
||||
@@ -3309,7 +3309,7 @@ function renderGatewayView() {
|
||||
renderGateway();
|
||||
$("#gatewayKeysBlock").hidden = !providers.gateway.lan;
|
||||
if (!providers.gateway.lan) gatewayKeyDraft = null;
|
||||
if (gatewayKeyDraft === null && !$("#gatewayKeys .rename-in")) renderGatewayKeys();
|
||||
if (gatewayKeyDraft === null && gatewayLimit === null && !$("#gatewayKeys .rename-in")) renderGatewayKeys();
|
||||
renderConnect();
|
||||
renderGatewayModels();
|
||||
renderActivity();
|
||||
|
||||
@@ -1,6 +1,9 @@
|
||||
// Caller keys use the same account list and edit controls as provider keys.
|
||||
let gatewayKeys = null, gatewayKeyDraft = null, gatewayKeysBusy = false;
|
||||
let gatewayKeysRead = 0;
|
||||
// gatewayLimit is the key whose limit editor is open, and its draft
|
||||
// (#585): what is picked in it is sent only by its Save.
|
||||
let gatewayLimit = null;
|
||||
let connectURL = "", connectKeyID = "", connectSecret = null;
|
||||
|
||||
function connectPick(id, label, text, options, value, choose) {
|
||||
@@ -56,6 +59,7 @@ async function gatewayKeyAction(action, body) {
|
||||
connectSecret = null;
|
||||
gatewayKeys = out.keys || [];
|
||||
gatewayKeyDraft = null;
|
||||
if (action === "limit-key" || !gatewayKeys.some((k) => k.id === gatewayLimit?.id)) gatewayLimit = null;
|
||||
renderGatewayKeys();
|
||||
if (view === "gateway" && providers) selectConnectKey(connectKeyID);
|
||||
return out;
|
||||
@@ -150,7 +154,9 @@ function renderGatewayKeys() {
|
||||
rm.onclick = () => askGatewayKey(k, false);
|
||||
const rotate = el("button", "text quiet", t("Rotate key"));
|
||||
rotate.onclick = () => askGatewayKey(k, true);
|
||||
row.append(tick, gatewayRename(k), el("span", "plan mono", k.masked), el("span", "grow"), copyCallerKeyBtn(k), rotate, rm);
|
||||
row.append(tick, gatewayRename(k), el("span", "plan mono", k.masked), el("span", "grow"), gatewayLimitBadge(k), copyCallerKeyBtn(k), rotate, rm);
|
||||
const under = gatewayLimit?.id === k.id ? gatewayLimitEditor(k) : k.used ? gatewayLimitLine(k) : null;
|
||||
if (under) row.classList.add("with-limit"), row.append(under);
|
||||
list.append(row);
|
||||
}
|
||||
if (gatewayKeyDraft !== null) list.append(gatewayKeyForm());
|
||||
@@ -185,3 +191,155 @@ function askGatewayKey(k, rotate) {
|
||||
$("#modal").classList.add("lib");
|
||||
cancel.focus({ preventScroll: true });
|
||||
}
|
||||
|
||||
// ---- a gateway key's own limit (#585) ----
|
||||
|
||||
const limitPeriods = ["day", "week", "month"];
|
||||
const limitPer = { day: "per day", week: "per week", month: "per month" };
|
||||
|
||||
// limitTokens reads a count as typed: 2000000, 2,000,000, 500k, 2M, 1.5m,
|
||||
// 100万, 1亿; "" for none, NaN for what isn't one.
|
||||
function limitTokens(v) {
|
||||
const s = String(v ?? "").trim().toLowerCase().replace(/[,\s_]/g, "");
|
||||
if (!s) return 0;
|
||||
const m = s.match(/^(\d+(?:\.\d+)?)(k|m|b|万|亿)?$/);
|
||||
if (!m) return NaN;
|
||||
const mult = { k: 1e3, m: 1e6, b: 1e9, "万": 1e4, "亿": 1e8 }[m[2]] || 1;
|
||||
return Math.round(parseFloat(m[1]) * mult);
|
||||
}
|
||||
function limitCost(v) {
|
||||
const s = String(v ?? "").trim().replace(/^[$＄]/, "");
|
||||
if (!s) return 0;
|
||||
return /^\d+(\.\d+)?$/.test(s) ? parseFloat(s) : NaN;
|
||||
}
|
||||
const usd = (n) => "$" + (n >= 100 ? n.toFixed(0) : n.toFixed(2));
|
||||
|
||||
// gatewayLimitBadge opens a key's limit editor: "Limit" on hover when it
|
||||
// has none, else what it has used, "1.2M / 2M · day", always shown.
|
||||
function gatewayLimitBadge(k) {
|
||||
const u = k.used;
|
||||
const b = el("button", "amodels limit" + (u ? " set" : "") + (u?.spent ? " spent" : "") + (gatewayLimit?.id === k.id ? " open" : ""));
|
||||
b.type = "button";
|
||||
if (!u) b.textContent = t("Limit");
|
||||
else {
|
||||
const parts = [];
|
||||
if (u.tokenLimit) parts.push(fmtN(u.tokens) + " / " + fmtN(u.tokenLimit));
|
||||
if (u.costLimit) parts.push(usd(u.cost) + " / " + usd(u.costLimit));
|
||||
b.textContent = (u.spent ? t("Spent") + " · " : "") + parts.join(" · ") + " · " + t(limitPer[u.period] || u.period);
|
||||
}
|
||||
b.title = t(u ? "Change this key's limit" : "Set a limit for this key");
|
||||
b.setAttribute("aria-label", b.title);
|
||||
b.setAttribute("aria-expanded", gatewayLimit?.id === k.id ? "true" : "false");
|
||||
b.onclick = (e) => {
|
||||
e.preventDefault();
|
||||
if (gatewayLimit?.id === k.id) gatewayLimit = null;
|
||||
else {
|
||||
const l = k.limit || {};
|
||||
gatewayLimit = { id: k.id, period: l.period || "day", tokens: l.tokens ? String(l.tokens) : "", cost: l.cost ? String(l.cost) : "", cacheReads: !!l.cacheReads };
|
||||
}
|
||||
renderGatewayKeys();
|
||||
};
|
||||
return b;
|
||||
}
|
||||
|
||||
// gatewayLimitLine says what a limited key has used, has left and when
|
||||
// it resets, under its row.
|
||||
function gatewayLimitLine(k) {
|
||||
const u = k.used;
|
||||
const line = el("div", "key-limit" + (u.spent ? " spent" : ""));
|
||||
const parts = [];
|
||||
if (u.tokenLimit) parts.push(t("{used} of {limit} tokens · {left} left", { used: fmtN(u.tokens), limit: fmtN(u.tokenLimit), left: fmtN(u.tokensLeft) }));
|
||||
if (u.costLimit) parts.push(t("{used} of {limit} estimated · {left} left", { used: usd(u.cost), limit: usd(u.costLimit), left: usd(u.costLeft) }));
|
||||
const at = new Date(u.reset);
|
||||
parts.push(t(u.spent ? "refused until {when}" : "resets {when}", { when: resetClock(at) }));
|
||||
if (u.inFlight) parts.push(t(u.inFlight === 1 ? "1 in flight" : "{n} in flight", { n: u.inFlight }));
|
||||
line.textContent = parts.join(" · ");
|
||||
line.title = t("Counted from {start}; resets {when}", { start: new Date(u.start).toLocaleString(), when: at.toLocaleString() })
|
||||
+ (u.unpriced ? "\n" + t("{n} calls without a known price aren't in the cost", { n: u.unpriced }) : "");
|
||||
return line;
|
||||
}
|
||||
|
||||
// gatewayLimitEditor stages a key's limit: nothing is sent until Save.
|
||||
function gatewayLimitEditor(k) {
|
||||
const d = gatewayLimit;
|
||||
const box = el("div", "acct-models key-limit-ed");
|
||||
box.onclick = (e) => e.stopPropagation();
|
||||
const kept = k.limit || {};
|
||||
const changed = () => d.period !== (kept.period || "day") || limitTokens(d.tokens) !== (kept.tokens || 0)
|
||||
|| limitCost(d.cost) !== (kept.cost || 0) || d.cacheReads !== !!kept.cacheReads;
|
||||
const fields = el("div", "klf");
|
||||
const period = el("select", "klf-period");
|
||||
period.setAttribute("aria-label", t("Limit window"));
|
||||
for (const p of limitPeriods) {
|
||||
const o = el("option", "", t(limitPer[p]));
|
||||
o.value = p;
|
||||
period.append(o);
|
||||
}
|
||||
period.value = d.period;
|
||||
const tokens = input(d.tokens, t("No token cap"));
|
||||
tokens.setAttribute("aria-label", t("Token limit"));
|
||||
tokens.inputMode = "decimal";
|
||||
const cost = input(d.cost, t("No cost cap"));
|
||||
cost.setAttribute("aria-label", t("Estimated cost limit, US$"));
|
||||
cost.inputMode = "decimal";
|
||||
const tl = el("label", "klf-field");
|
||||
tl.append(el("span", "klf-name", t("Tokens")), tokens);
|
||||
const cl = el("label", "klf-field");
|
||||
cl.append(el("span", "klf-name", t("Estimated cost, US$")), cost);
|
||||
const pl = el("label", "klf-field");
|
||||
pl.append(el("span", "klf-name", t("Window")), period);
|
||||
fields.append(pl, tl, cl);
|
||||
const cache = el("label", "klf-check");
|
||||
const cb = el("input");
|
||||
cb.type = "checkbox";
|
||||
cb.checked = d.cacheReads;
|
||||
cache.append(cb, el("span", "", t("Count cache reads too")));
|
||||
const hint = el("div", "hint", t("Tokens are each call's uncached input, output and cache writes. The cost is an estimate at the Usage page's prices, not a bill. Windows are calendar ones in this computer's time (a week from Monday). Requests from this computer without a key aren't limited."));
|
||||
const err = el("div", "hint klf-err");
|
||||
const unsaved = el("span", "hint munsaved", t("unsaved"));
|
||||
unsaved.title = t("Made when you Save; Cancel drops it");
|
||||
const save = el("button", "text primary", t("Save"));
|
||||
const update = () => {
|
||||
const bad = Number.isNaN(limitTokens(d.tokens)) ? t("Tokens is a count, like 2000000 or 2M")
|
||||
: Number.isNaN(limitCost(d.cost)) ? t("The cost is US dollars, like 5 or 2.50") : "";
|
||||
err.textContent = bad;
|
||||
err.hidden = !bad;
|
||||
unsaved.hidden = !changed();
|
||||
save.disabled = !!bad || !changed();
|
||||
};
|
||||
period.onchange = () => { d.period = period.value; update(); };
|
||||
tokens.oninput = () => { d.tokens = tokens.value; update(); };
|
||||
cost.oninput = () => { d.cost = cost.value; update(); };
|
||||
cb.onchange = () => { d.cacheReads = cb.checked; update(); };
|
||||
const close = () => { gatewayLimit = null; renderGatewayKeys(); };
|
||||
for (const i of [tokens, cost]) i.onkeydown = (e) => {
|
||||
e.stopPropagation();
|
||||
if (e.key === "Escape") close();
|
||||
else if (e.key === "Enter" && !save.disabled) save.click();
|
||||
};
|
||||
const cancel = el("button", "text", t("Cancel"));
|
||||
cancel.onclick = close;
|
||||
save.onclick = async () => {
|
||||
if (save.disabled) return;
|
||||
save.disabled = true;
|
||||
const tk = limitTokens(d.tokens), c = limitCost(d.cost);
|
||||
const limit = tk || c ? { period: d.period, tokens: tk, cost: c, cacheReads: d.cacheReads } : null;
|
||||
const out = await gatewayKeyAction("limit-key", { key: k.id, limit });
|
||||
if (!out) { update(); return; }
|
||||
status(t(limit ? "Limit saved for {name}" : "Limit removed for {name}", { name: k.name }), "ok");
|
||||
};
|
||||
const foot = el("div", "acm-foot");
|
||||
foot.append(cache, el("span", "grow"), unsaved);
|
||||
if (k.limit) {
|
||||
const off = el("button", "text", t("No limit"));
|
||||
off.title = t("Clears the caps; Save makes it");
|
||||
off.onclick = () => { d.tokens = ""; d.cost = ""; tokens.value = ""; cost.value = ""; update(); };
|
||||
foot.append(off);
|
||||
}
|
||||
foot.append(cancel, save);
|
||||
box.append(fields, hint, err);
|
||||
if (k.used) box.append(gatewayLimitLine(k));
|
||||
box.append(foot);
|
||||
update();
|
||||
return box;
|
||||
}
|
||||
@@ -6,6 +6,35 @@ const I18N = {
|
||||
zh: {
|
||||
"Gateway keys": "网关密钥",
|
||||
"Gateway key": "网关密钥",
|
||||
"Limit": "限额",
|
||||
"Spent": "已用完",
|
||||
"per day": "每日",
|
||||
"per week": "每周",
|
||||
"per month": "每月",
|
||||
"Change this key's limit": "修改此密钥的限额",
|
||||
"Set a limit for this key": "为此密钥设置限额",
|
||||
"{used} of {limit} tokens · {left} left": "已用 {used} / {limit} Token · 剩余 {left}",
|
||||
"{used} of {limit} estimated · {left} left": "估算费用 {used} / {limit} · 剩余 {left}",
|
||||
"refused until {when}": "{when} 前拒绝请求",
|
||||
"resets {when}": "{when} 重置",
|
||||
"1 in flight": "1 个请求进行中",
|
||||
"Counted from {start}; resets {when}": "自 {start} 起计；{when} 重置",
|
||||
"{n} calls without a known price aren't in the cost": "{n} 次调用没有已知价格，未计入费用",
|
||||
"Limit window": "限额周期",
|
||||
"Window": "周期",
|
||||
"No token cap": "不限 Token",
|
||||
"No cost cap": "不限费用",
|
||||
"Token limit": "Token 限额",
|
||||
"Estimated cost limit, US$": "估算费用限额（美元）",
|
||||
"Estimated cost, US$": "估算费用（美元）",
|
||||
"Count cache reads too": "缓存读取也计入",
|
||||
"Tokens are each call's uncached input, output and cache writes. The cost is an estimate at the Usage page's prices, not a bill. Windows are calendar ones in this computer's time (a week from Monday). Requests from this computer without a key aren't limited.": "Token 按每次调用未命中缓存的输入、输出和缓存写入计。费用按用量页的价格估算，并非账单。周期按本机时间的自然日、周、月计算（每周从周一开始）。本机不带密钥的请求不受限额。",
|
||||
"Tokens is a count, like 2000000 or 2M": "Token 应为数量，例如 2000000、2M 或 200万",
|
||||
"The cost is US dollars, like 5 or 2.50": "费用以美元计，例如 5 或 2.50",
|
||||
"Made when you Save; Cancel drops it": "点击保存后生效，取消则放弃",
|
||||
"Limit saved for {name}": "已保存 {name} 的限额",
|
||||
"Limit removed for {name}": "已取消 {name} 的限额",
|
||||
"Clears the caps; Save makes it": "清空上限，点击保存后生效",
|
||||
"All gateway keys": "全部网关密钥",
|
||||
"All accounts": "全部账号",
|
||||
"account not recorded": "未记录账号",
|
||||
|
||||
@@ -3,10 +3,28 @@ package gui
|
||||
import (
|
||||
"encoding/json"
|
||||
"net/http"
|
||||
"time"
|
||||
|
||||
"github.com/yetone/magpie/internal/access"
|
||||
"github.com/yetone/magpie/internal/budget"
|
||||
)
|
||||
|
||||
// keyRow is a gateway key as the Gateway page lists it, with what it has
|
||||
// used of its limit (#585).
|
||||
type keyRow struct {
|
||||
access.Key
|
||||
Used *budget.Status `json:"used,omitempty"`
|
||||
}
|
||||
|
||||
func withLimits(keys []access.Key) []keyRow {
|
||||
now := time.Now()
|
||||
out := make([]keyRow, 0, len(keys))
|
||||
for _, k := range keys {
|
||||
out = append(out, keyRow{Key: k, Used: budget.Of(k, now)})
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
func callerKeyRoutes(mux *http.ServeMux) {
|
||||
mux.HandleFunc("GET /api/caller-keys", func(w http.ResponseWriter, r *http.Request) {
|
||||
w.Header().Set("Cache-Control", "no-store")
|
||||
@@ -16,7 +34,7 @@ func callerKeyRoutes(mux *http.ServeMux) {
|
||||
fail(w, err)
|
||||
return
|
||||
}
|
||||
writeJSON(w, map[string]any{"keys": keys})
|
||||
writeJSON(w, map[string]any{"keys": withLimits(keys)})
|
||||
})
|
||||
mux.HandleFunc("POST /api/caller-keys/{action}", func(w http.ResponseWriter, r *http.Request) {
|
||||
var in access.Change
|
||||
@@ -36,6 +54,6 @@ func callerKeyRoutes(mux *http.ServeMux) {
|
||||
return
|
||||
}
|
||||
w.Header().Set("Cache-Control", "no-store")
|
||||
writeJSON(w, map[string]any{"keys": keys, "secret": secret})
|
||||
writeJSON(w, map[string]any{"keys": withLimits(keys), "secret": secret})
|
||||
})
|
||||
}
|
||||
@@ -21,6 +21,14 @@ literal custom names and narrow layouts in Chinese and English on both engines.
|
||||
Sharing off keeps the original API-key field and hides the gateway-key picker;
|
||||
sharing on names the arbitrary local option separately from the Magpie key.
|
||||
|
||||
`gateway-key-limit.test.cjs` checks a gateway key's own limit (#585): each
|
||||
limited key's badge and line (used, limit, left, reset, requests in flight,
|
||||
calls without a price), a spent key's "refused until", and the limit editor
|
||||
under a key's row. Window, tokens, cost and cache reads are staged, marked
|
||||
unsaved and sent only by Save; Cancel drops them, a bad amount disables Save,
|
||||
and No limit + Save sends `null`. No click scrolls the page, nothing has a
|
||||
left border, and the open editor fits a 560px window, in Chinese and English.
|
||||
|
||||
`api-key-usage.test.cjs` exercises the existing provider key list: adding,
|
||||
enabling and disabling, choosing the first key, renaming and removing.
|
||||
The Usage page's Gateway key rows, request filter and CSV exports identify
|
||||
@@ -42,7 +50,7 @@ The settings theme picker is also switched and reloaded in Chromium and
|
||||
WebKit. No separate colours are defined for gateway keys.
|
||||
|
||||
```sh
|
||||
node --test --test-concurrency=1 internal/gui/tests/gateway-caller-keys.test.cjs internal/gui/tests/gateway-connect-keys.test.cjs internal/gui/tests/api-key-usage.test.cjs internal/gui/tests/api-key-theme.test.cjs
|
||||
node --test --test-concurrency=1 internal/gui/tests/gateway-caller-keys.test.cjs internal/gui/tests/gateway-connect-keys.test.cjs internal/gui/tests/gateway-key-limit.test.cjs internal/gui/tests/api-key-usage.test.cjs internal/gui/tests/api-key-theme.test.cjs
|
||||
```
|
||||
|
||||
## Other Browser Regressions
|
||||
|
||||
+7
@@ -14,6 +14,7 @@ function fixture(lang, theme, events, options = {}) {
|
||||
{ id: "server", name: "Server", masked: "sk-magpie-key-…222222" },
|
||||
{ id: "work", name: "Work", masked: "sk-magpie-key-…333333" },
|
||||
];
|
||||
for (const k of keys) if (options.limits?.[k.id]) Object.assign(k, options.limits[k.id]);
|
||||
let serial = 0;
|
||||
let lan = !!options.lan;
|
||||
let lanKeyID = "", rotations = 0, lanSecret = "";
|
||||
@@ -102,6 +103,12 @@ function fixture(lang, theme, events, options = {}) {
|
||||
}
|
||||
if (action === "remove-key") keys = keys.filter((v) => v !== k);
|
||||
if (action === "on-key" || action === "off-key") k.off = action === "off-key";
|
||||
if (action === "limit-key") {
|
||||
// what a key has used is the fixture's; a new limit starts unused
|
||||
k.limit = body.limit || undefined;
|
||||
k.used = body.limit ? { period: body.limit.period, start: new Date().toISOString(), reset: new Date(Date.now() + 864e5).toISOString(),
|
||||
calls: 0, tokens: 0, cost: 0, tokenLimit: body.limit.tokens, costLimit: body.limit.cost, tokensLeft: body.limit.tokens, costLeft: body.limit.cost, spent: false } : undefined;
|
||||
}
|
||||
if (action === "copy-key") secret = k.id === lanKeyID ? lanSecret : secrets.get(k.id);
|
||||
return json({ keys, secret });
|
||||
}
|
||||
|
||||
@@ -0,0 +1,114 @@
|
||||
// A gateway key's own limit (#585): what each key has used of it, and
|
||||
// its editor under the key's row, staged until Save.
|
||||
const assert = require("node:assert/strict");
|
||||
const { test } = require("node:test");
|
||||
const { chromium, webkit } = require("playwright");
|
||||
const { fixture } = require("./fixtures/caller-keys.cjs");
|
||||
|
||||
const day = 864e5;
|
||||
const limits = () => ({
|
||||
laptop: {
|
||||
limit: { period: "day", tokens: 1000000 },
|
||||
used: { period: "day", start: new Date(Date.now() - day / 2).toISOString(), reset: new Date(Date.now() + day / 2).toISOString(), calls: 12,
|
||||
tokens: 250000, cost: 0.8, tokenLimit: 1000000, tokensLeft: 750000, costLeft: 0, inFlight: 2, spent: false },
|
||||
},
|
||||
server: {
|
||||
limit: { period: "month", cost: 5 },
|
||||
used: { period: "month", start: new Date(Date.now() - 3 * day).toISOString(), reset: new Date(Date.now() + 3 * day).toISOString(), calls: 40,
|
||||
tokens: 900000, cost: 5.2, costLimit: 5, tokensLeft: 0, costLeft: 0, unpriced: 3, spent: true },
|
||||
},
|
||||
});
|
||||
|
||||
for (const engine of (process.env.BROWSER ? [process.env.BROWSER] : ["chromium", "webkit"])) {
|
||||
for (const lang of ["en", "zh"]) {
|
||||
test(`${engine} ${lang}: a gateway key's limit is shown and set with Save`, async (t) => {
|
||||
const browser = await (engine === "webkit" ? webkit.launch() : chromium.launch({ channel: "chromium" }));
|
||||
t.after(() => browser.close());
|
||||
const page = await (await browser.newContext({ viewport: { width: 1000, height: 420 }, reducedMotion: "reduce" })).newPage();
|
||||
page.setDefaultTimeout(6000);
|
||||
const events = [], errors = [];
|
||||
page.on("pageerror", (e) => errors.push(e.message));
|
||||
await page.route("**/*", fixture(lang, "light", events, { lan: true, limits: limits() }));
|
||||
const zh = lang === "zh";
|
||||
const w = zh
|
||||
? { limit: "为此密钥设置限额", change: "修改此密钥的限额", left: /剩余/, refused: /前拒绝请求/, tokens: "Token 限额", cost: "估算费用限额（美元）", window: "限额周期", cache: "缓存读取也计入", save: "保存", cancel: "取消", none: "不限制", unsaved: "未保存", week: "每周", bad: /Token 应为数量/ }
|
||||
: { limit: "Set a limit for this key", change: "Change this key's limit", left: /tokens · .* left/, refused: /refused until/, tokens: "Token limit", cost: "Estimated cost limit, US$", window: "Limit window", cache: "Count cache reads too", save: "Save", cancel: "Cancel", none: "No limit", unsaved: "unsaved", week: "per week", bad: /Tokens is a count/ };
|
||||
await page.goto("http://magpie.test/?view=gateway");
|
||||
await page.locator("#gatewayKeys .acc[data-key]").last().waitFor();
|
||||
const row = (id) => page.locator(`#gatewayKeys .acc[data-key="${id}"]`);
|
||||
// what each limited key has used, has left and when it resets
|
||||
assert.match(await row("laptop").locator(".key-limit").textContent(), w.left);
|
||||
assert.match(await row("laptop").locator(".key-limit").textContent(), zh ? /2 个请求进行中/ : /2 in flight/);
|
||||
assert.match(await row("laptop").locator(".amodels.limit").textContent(), zh ? /每日/ : /per day/);
|
||||
assert.match(await row("server").locator(".key-limit").textContent(), w.refused);
|
||||
assert.match(await row("server").locator(".key-limit").textContent(), /\$5\.20/);
|
||||
assert.equal(await row("server").locator(".amodels.limit.spent").count(), 1);
|
||||
assert.match(await row("server").locator(".key-limit").getAttribute("title"), zh ? /3 次调用没有已知价格/ : /3 calls without a known price/);
|
||||
assert.equal(await row("work").locator(".key-limit").count(), 0, "a key without a limit says nothing under its row");
|
||||
// nothing about it is drawn with a coloured left border
|
||||
for (const sel of [".key-limit", ".amodels.limit"]) {
|
||||
const styles = await page.locator("#gatewayKeys " + sel).evaluateAll((els) => els.map((e) => getComputedStyle(e).borderLeftStyle));
|
||||
assert(styles.every((s) => s === "none"), sel + " has a left border");
|
||||
}
|
||||
const view = page.locator("#view-gateway");
|
||||
const scrolled = async () => [await view.evaluate((v) => v.scrollTop), await page.evaluate(() => window.scrollY)];
|
||||
await page.mouse.move(500, 300);
|
||||
await page.mouse.wheel(0, 120);
|
||||
await page.waitForTimeout(100);
|
||||
const at = await scrolled();
|
||||
// open Work's editor: staged, nothing sent
|
||||
await row("work").getByRole("button", { name: w.limit, exact: true }).click();
|
||||
const ed = row("work").locator(".key-limit-ed");
|
||||
await ed.waitFor();
|
||||
assert.deepEqual(await scrolled(), at, "opening the editor moved the page");
|
||||
assert.equal(await ed.evaluate((e) => getComputedStyle(e).borderLeftStyle), "none");
|
||||
await ed.getByRole("textbox", { name: w.tokens, exact: true }).fill("2M");
|
||||
await ed.getByRole("combobox", { name: w.window, exact: true }).selectOption("week");
|
||||
await ed.getByLabel(w.cache).check();
|
||||
assert.equal(await ed.locator(".munsaved").isVisible(), true);
|
||||
assert.equal(await ed.locator(".munsaved").textContent(), w.unsaved);
|
||||
await ed.getByRole("button", { name: w.cancel, exact: true }).click();
|
||||
await ed.waitFor({ state: "detached" });
|
||||
assert.equal(events.filter((e) => e.action === "limit-key").length, 0, "Cancel sent the limit");
|
||||
assert.deepEqual(await scrolled(), at, "Cancel moved the page");
|
||||
// opened again, the draft is gone
|
||||
await row("work").getByRole("button", { name: w.limit, exact: true }).click();
|
||||
await ed.waitFor();
|
||||
assert.equal(await ed.getByRole("textbox", { name: w.tokens, exact: true }).inputValue(), "");
|
||||
assert.equal(await ed.locator(".munsaved").isVisible(), false);
|
||||
await ed.getByRole("textbox", { name: w.tokens, exact: true }).fill("lots");
|
||||
assert.match(await ed.locator(".klf-err").textContent(), w.bad);
|
||||
assert.equal(await ed.getByRole("button", { name: w.save, exact: true }).isDisabled(), true);
|
||||
await ed.getByRole("textbox", { name: w.tokens, exact: true }).fill(zh ? "200万" : "2M");
|
||||
await ed.getByRole("textbox", { name: w.cost, exact: true }).fill("5");
|
||||
await ed.getByRole("combobox", { name: w.window, exact: true }).selectOption("week");
|
||||
await ed.getByLabel(w.cache).check();
|
||||
assert.equal(events.filter((e) => e.action === "limit-key").length, 0, "a pick was sent before Save");
|
||||
await ed.getByRole("button", { name: w.save, exact: true }).click();
|
||||
await ed.waitFor({ state: "detached" });
|
||||
const sent = events.filter((e) => e.action === "limit-key");
|
||||
assert.deepEqual(sent.map((e) => e.body), [{ key: "work", limit: { period: "week", tokens: 2000000, cost: 5, cacheReads: true } }]);
|
||||
assert.match(await row("work").locator(".amodels.limit").textContent(), new RegExp(w.week));
|
||||
assert.match(await row("work").locator(".key-limit").textContent(), w.left);
|
||||
assert.deepEqual(await scrolled(), at, "Save moved the page");
|
||||
// the limit taken off
|
||||
await row("laptop").getByRole("button", { name: w.change, exact: true }).click();
|
||||
const led = row("laptop").locator(".key-limit-ed");
|
||||
await led.waitFor();
|
||||
assert.equal(await led.getByRole("textbox", { name: w.tokens, exact: true }).inputValue(), "1000000");
|
||||
await led.getByRole("button", { name: w.none, exact: true }).click();
|
||||
assert.equal(events.filter((e) => e.action === "limit-key").length, 1);
|
||||
await led.getByRole("button", { name: w.save, exact: true }).click();
|
||||
await led.waitFor({ state: "detached" });
|
||||
assert.deepEqual(events.filter((e) => e.action === "limit-key").at(-1).body, { key: "laptop", limit: null });
|
||||
assert.equal(await row("laptop").locator(".key-limit").count(), 0);
|
||||
assert.deepEqual(await scrolled(), at, "no click moved the page");
|
||||
// narrow: the editor fits
|
||||
await row("server").getByRole("button", { name: w.change, exact: true }).click();
|
||||
await row("server").locator(".key-limit-ed").waitFor();
|
||||
await page.setViewportSize({ width: 560, height: 740 });
|
||||
assert.equal(await page.evaluate(() => document.documentElement.scrollWidth > window.innerWidth), false);
|
||||
assert.deepEqual(errors, []);
|
||||
});
|
||||
}
|
||||
}
|
||||
@@ -646,3 +646,22 @@ func LedgerSeries(p Period, rows []Row) (bucket string, pts []SeriesPoint) {
|
||||
}
|
||||
return bucket, pts
|
||||
}
|
||||
|
||||
// Visit calls fn with each record from since on, as the ledger's index
|
||||
// reads usage.jsonl: only its blocks that reach since are decoded.
|
||||
func Visit(since time.Time, fn func(Record)) { readLogSnapshot().visit(since, fn) }
|
||||
|
||||
// NewCoster prices one record at a time at the effective prices the
|
||||
// ledger uses, read once: its cost in USD, and false for a call with
|
||||
// tokens but no known price.
|
||||
func NewCoster() func(Record) (float64, bool) {
|
||||
priceOf, renamed := pricer(), provider.Renamed()
|
||||
return func(r Record) (float64, bool) {
|
||||
if id, ok := renamed[r.Provider]; ok {
|
||||
r.Provider = id
|
||||
}
|
||||
var t Totals
|
||||
t.add(r, priceOf(r))
|
||||
return t.Cost, t.Unpriced == 0
|
||||
}
|
||||
}
|
||||
@@ -86,6 +86,7 @@ const usage = `magpie — one place to pick every agent's model
|
||||
magpie serve run the gateway alone (the app runs it too)
|
||||
magpie healthcheck exit 0 when the gateway answers (a container's HEALTHCHECK)
|
||||
magpie gateway-key list|add <name>|rotate <id>|remove <id> manage the keys clients use to call a shared gateway
|
||||
magpie gateway-key limit <id> [off|day|week|month --tokens N --cost USD --cache-reads] a key's own limit, and what it used
|
||||
magpie mcp image the image and video generation MCP server an agent is given from the library (stdio)
|
||||
magpie usage [today|7d|30d|all] tokens and cost per agent, model and subscription account (30d)
|
||||
magpie usage --csv [--account <name>] [today|7d|30d|all] every request as CSV (or one account's): the model asked for, sent and served, tokens, cost, time, status, account
|
||||
|
||||
Reference in new issue
Block a user