Files
magpie/gateway_keys_cli_test.go
T
yetone 0879894af3 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).
2026-10-03 00:18:16 +08:00

134 lines
4.4 KiB
Go

package main
import (
"bytes"
"os"
"strings"
"testing"
"github.com/yetone/magpie/internal/access"
"github.com/yetone/magpie/internal/settings"
)
func TestGatewayKeyCLI(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 {
t.Helper()
var out bytes.Buffer
if err := gatewayKeysTo(&out, append([]string{"gateway-key"}, args...)); err != nil {
t.Fatal(err)
}
return strings.TrimSpace(out.String())
}
if got := call("list"); !strings.Contains(got, "GATEWAY KEY") {
t.Fatal(got)
}
secret := call("add", "Remote laptop")
keys, _ := access.List()
if len(keys) != 1 {
t.Fatal(keys)
}
id := keys[0].ID
if who, ok := access.Authenticate(secret); !ok || who.KeyID != id || who.KeyName != "Remote laptop" {
t.Fatal(who, ok)
}
if got := call(); strings.Contains(got, secret) || !strings.Contains(got, id) || !strings.Contains(got, "Remote laptop") {
t.Fatal("list must mask credentials", got)
}
next := call("rotate", id)
if next == secret {
t.Fatal("rotation reused the secret")
}
if _, ok := access.Authenticate(secret); ok {
t.Fatal("old credential works after rotation")
}
if who, ok := access.Authenticate(next); !ok || who.KeyID != id {
t.Fatal("rotation changed identity", who, ok)
}
call("remove", id)
if _, ok := access.Authenticate(next); ok {
t.Fatal("removed credential still works")
}
for _, args := range [][]string{{"add"}, {"add", ""}, {"list", "extra"}, {"rotate", "missing"}, {"remove", "missing"}, {"unknown", "x"}} {
if err := gatewayKeysTo(&bytes.Buffer{}, append([]string{"gateway-key"}, args...)); err == nil {
t.Error("accepted", args)
}
}
}
func TestGatewayKeyListWithReadOnlyMigration(t *testing.T) {
t.Setenv("XDG_CONFIG_HOME", t.TempDir())
s := settings.Settings{LAN: true, LANKey: "sk-magpie-fixture-cli"}
if err := settings.Save(s); err != nil {
t.Fatal(err)
}
if err := os.Chmod(settings.Path(), 0o400); err != nil {
t.Fatal(err)
}
t.Cleanup(func() { os.Chmod(settings.Path(), 0o600) })
if err := settings.Save(s); err == nil {
t.Skip("settings.json remains writable")
}
var out bytes.Buffer
if err := gatewayKeysTo(&out, []string{"gateway-key", "list"}); err != nil || !strings.Contains(out.String(), "Magpie") {
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)
}
}