mirror of
https://gitcode.com/JianFeeeee/ModelRouter.git
synced 2026-11-09 18:57:21 +00:00
Compare commits
22 Commits
release/v1
...
v1.7.0
| Author | SHA1 | Date | |
|---|---|---|---|
| 5530912d32 | |||
| cc5e725226 | |||
| 652842783f | |||
| a21ae84cbe | |||
| d072a03c9a | |||
| f814ff7468 | |||
| ef631b43dd | |||
| 9c3aabb7f9 | |||
| a7355debed | |||
| f6baa13583 | |||
| fd03ef6e4a | |||
| 97bb9c6ef4 | |||
| d81074621a | |||
| 4dd3431c26 | |||
| 3d1407e6fa | |||
| db320d2f6c | |||
| b544671f5e | |||
| 0881a49286 | |||
| 4e2b5bfda9 | |||
| 61b847da86 | |||
| 6432b61eb6 | |||
| dd708c9680 |
80
README.md
80
README.md
@ -24,6 +24,7 @@
|
||||
### 强大的多租户调度能力
|
||||
|
||||
- **多密钥多租户**:支持无限密钥,每个密钥独立角色、模型范围、Token 配额、重置周期
|
||||
- **密钥用量配额**:每把 key 单独配总 token 配额 + 请求数配额与重置周期(小时/周/月/自定义 N 小时),跨模型共享预算;耗尽返 429 + `Retry-After` 可自动恢复,admin key 永不受限
|
||||
- **AUTO 智能调度**:基于优先级档位的分级调度,同优先级源自动轮询负载均衡,故障自动毫秒级故障转移
|
||||
- **自愈冷却**:冷却上限 5 分钟,过半后放行 1 个探测请求,上游/额度恢复即刻回归轮询,无需等满冷却窗口
|
||||
- **Token 配额管理**:精确到模型级别的 Token 配额控制,支持小时/周/月/自定义小时周期自动重置
|
||||
@ -178,6 +179,80 @@ sources:
|
||||
- 客户端用任意一个已授权的密钥明文作为 Bearer(`Authorization: Bearer <key>`)。
|
||||
- 删除密钥即从运行时存储移除,立即失效。
|
||||
|
||||
#### 密钥用量配额
|
||||
|
||||
每个密钥可单独限制用量与用量重置周期,两级配额同时生效:
|
||||
|
||||
```yaml
|
||||
keys:
|
||||
- key: sk-gw-<hex>
|
||||
role: user
|
||||
name: agent-alice
|
||||
# ---- 整钥配额(跳模型)----
|
||||
token_quota: 5000000 # 本周期内这把 key 的总 token 预算,0 = 无限
|
||||
req_quota: 20000 # 本周期内的请求次数,0 = 无限
|
||||
period: nhour # "" | hour | week | month | nhour
|
||||
hours: 6 # 仅 nhour:每 6 小时重置
|
||||
# ---- 模型范围(可选,逐模型配额)----
|
||||
models:
|
||||
- model: m1
|
||||
token_quota: 1000000 # 本周期内该模型(该 key)的 token 预算
|
||||
period: hour
|
||||
- model: AUTO
|
||||
```
|
||||
|
||||
- `period` 词表:空 = 永不过期(累计总量),`hour` / `week` / `month` = 固定窗口,
|
||||
`nhour` + `hours` = 自定义小时数。**拼错的周期在写入时就被拒**,不会静默变成
|
||||
永不过期。
|
||||
- 整钥配额跨该 key 所有模型共享一份预算;`models[]` 里的配额则是逐模型独立计数。
|
||||
两者都按 key 隔离,A key 的用量不会消耗 B key 的额度。
|
||||
- 配额统计含聊天、流式、生图,跨重启从审计日志回放(保留 40 天,覆盖最长的
|
||||
month 窗口)。
|
||||
- 配额耗尽返回 **429 + `Retry-After`**(`rate_limit_exceeded`),客户端可等窗口
|
||||
重置后自动恢复;模型越权才是 403。**admin 密钥永不受配额限制**,
|
||||
避免把管理员锁在门外。
|
||||
- 窗口用量按整点小时分桶统计,实际释放比配置窗口最多晚 1 小时(配额宁可晚释放
|
||||
也不超发)。
|
||||
- `PUT /api/keys/{key}` 的配额字段是可选的:省略 = 保留原值,显式 `0` = 解除限制。
|
||||
只改模型范围不会清空已配置的预算。
|
||||
|
||||
##### 配额拒绝 vs 容量拒绝:两种「拒绝」含义不同
|
||||
|
||||
配额在**入口**检查,不占用任何上游槽位,因此是廉价拒绝;容量不足则要走
|
||||
`busyWait` 有界等待后才返回 503。实测(容量 8、0.6s/请求、100 并发):
|
||||
|
||||
| 情况 | 状态码 | 延迟 | 是否打到上游 |
|
||||
|---|---|---|---|
|
||||
| 配额用尽(`req_quota` 50 → 100 并发) | `429 rate_limit_exceeded` + `Retry-After` | **19–21ms** | 否 |
|
||||
| 容量打满(同一拓扑,未限配额) | `503 upstream_error` | ~2.6s(`busyWait`) | 超出部分否 |
|
||||
|
||||
客户端因此可以区分「等窗口重置」与「等上游腾出容量」:前者只需等待,后者
|
||||
通常该降并发或换源。
|
||||
|
||||
##### 配额的开销
|
||||
|
||||
配额桶按 (密钥, 模型, 整点小时) 分桶保留 40 天,实测(AMD 7840HS):
|
||||
|
||||
- 每请求配额检查:**149ns**(配了配额)/ **42.6ns**(未配配额,只查密钥记录,
|
||||
不碰桶)/ **37ns**(admin 密钥直接返回)—— 均 **0 分配**。
|
||||
未配配额的密钥几乎不付代价,可放心大量创建。
|
||||
- 记录一条请求:283ns、3 分配(与引入配额前相同,分配来自 ring buffer)。
|
||||
- 窗口查询按窗口长度而非保留总量扫描:1h 窗口 49ns、24h 窗口 55ns、
|
||||
30d 窗口 3.9μs。
|
||||
- 内存:生产形态(7 密钥 × 8 模型 × 2 源 × 满 40 天 retention)约 **3.7MB**。
|
||||
按源 pin 的 `source::model` 桶**惰性创建**——只有当某条配额真的 pin 了
|
||||
某个源时才维护,否则每条记录多写一份桶,在 20 密钥 × 8 模型 × 3 源下会
|
||||
多占 18MB。
|
||||
代价:配置 pinned 配额之前的历史用量无法事后按源拆分,所以 **pinned 配额的
|
||||
首个窗口可能少算**。
|
||||
|
||||
```bash
|
||||
# 配额耗尽时客户端看到
|
||||
HTTP/1.1 429 Too Many Requests
|
||||
Retry-After: 2100
|
||||
{"error":{"type":"rate_limit_exceeded","message":"key token quota exceeded (5000000/5000000, resets every 6h)"}}
|
||||
```
|
||||
|
||||
### 模型路由
|
||||
|
||||
`/v1/chat/completions` 的 `model` 解析顺序:
|
||||
@ -344,7 +419,10 @@ Environment=MALLOC_ARENA_MAX=2
|
||||
支持按时间范围导出 CSV;点击模型可生成 pin 到该模型的连接配置。
|
||||
- **对话页**:流式/非流式调试。
|
||||
- **密钥页**:创建/编辑网关 key,为每个 key 配模型范围(模型 + 源 + token 配额 + 周期),
|
||||
管理员管理全部 key,用户只看到自己的 key。
|
||||
管理员管理全部 key,用户只看到自己的 key。key 卡片头部显示整钥配额徽标
|
||||
(如 `250.0K·6h` / `77×·6h`),「配额」按钮编辑总 token / 请求数与重置周期;
|
||||
创建 key 时可直接配预算(选 admin 角色时该组输入自动禁用,因为 admin 永不受限)。
|
||||
「我的密钥」页对用户展示本 key 的预算。
|
||||
- **优先级页**:拖拽积木配置 AUTO 链档位。
|
||||
- **源页**:在线增删改上游源(API key 等敏感字段加密落盘)。
|
||||
- **适配器页**:上传 / 删除 Lua 适配器脚本。
|
||||
|
||||
96
README_EN.md
96
README_EN.md
@ -38,6 +38,10 @@ Extracted and independently evolved from the multi-source LLM adapter layer of
|
||||
`reasoning_content`, `tool_calls`, `usage`).
|
||||
- **Image generation**: `POST /v1/images/generations`, routed to models with
|
||||
`kind: image`.
|
||||
- **Per-key usage quota**: each key carries its own token and request caps plus a
|
||||
reset period (hour/week/month/custom N hours), shared across every model that
|
||||
key may use. Exhaustion answers 429 + `Retry-After` so a client resumes when
|
||||
the window rolls over; admin keys are never capped.
|
||||
- **Multimodal**: `content` arrays (`image_url` etc.) pass through losslessly;
|
||||
Anthropic/Gemini/Ollama are translated automatically.
|
||||
- **LuaJIT VM**: golua-binding LuaJIT; each adapter has its own VM + worker
|
||||
@ -171,11 +175,101 @@ under the `keys` field of the runtime file (encrypted at rest):
|
||||
delete the seed key.
|
||||
- The WebUI **Keys page** creates/deletes keys. Each key has a role (`admin`
|
||||
manages everything, `user` sees only its own key) and an optional **model
|
||||
scope** (model + source + token quota + reset period).
|
||||
scope** (model + source + token quota + reset period). Key cards show the
|
||||
key-wide caps as a badge (e.g. `250.0K·6h` / `77×·6h`); a **Quota** button
|
||||
edits the total token / request budget and its reset period, and the create
|
||||
form takes a budget too (those fields disable themselves for `admin`, which
|
||||
is never capped). The "My key" view shows a user its own budget.
|
||||
- Clients authenticate with any authorized key's plaintext as
|
||||
`Authorization: Bearer <key>`.
|
||||
- Deleting a key removes it from the store immediately.
|
||||
|
||||
#### Per-key usage quota
|
||||
|
||||
Each key can cap its own spend and reset period. Two levels apply at once:
|
||||
|
||||
```yaml
|
||||
keys:
|
||||
- key: sk-gw-<hex>
|
||||
role: user
|
||||
name: agent-alice
|
||||
# ---- key-wide (across every model) ----
|
||||
token_quota: 5000000 # total token budget for this window, 0 = unlimited
|
||||
req_quota: 20000 # requests per window, 0 = unlimited
|
||||
period: nhour # "" | hour | week | month | nhour
|
||||
hours: 6 # n-hour only: resets every 6 hours
|
||||
# ---- per-model scope (optional) ----
|
||||
models:
|
||||
- model: m1
|
||||
token_quota: 1000000
|
||||
period: hour
|
||||
- model: AUTO
|
||||
```
|
||||
|
||||
- `period`: empty = never resets (lifetime total); `hour` / `week` / `month` =
|
||||
fixed windows; `nhour` + `hours` = a custom hour count. **A misspelled
|
||||
period is rejected at write time** rather than silently becoming a
|
||||
never-resetting quota.
|
||||
- The key-wide cap is one budget shared by every model the key may use;
|
||||
quotas under `models[]` are counted per model. Both are isolated per key —
|
||||
one key's traffic never drains another's budget.
|
||||
- Usage counts chat, streaming and image requests, and survives a restart by
|
||||
replaying the audit log (40 days retained, covering the longest `month`
|
||||
window).
|
||||
- An exhausted quota returns **429 + `Retry-After`**
|
||||
(`rate_limit_exceeded`) so a client resumes when the window rolls over; a
|
||||
model the key may not use stays 403. **Admin keys are never capped**, so a
|
||||
cap can never lock the operator out.
|
||||
- Buckets are whole unix hours, so a window frees up at most an hour late
|
||||
(deliberately freeing late rather than overspending).
|
||||
- On `PUT /api/keys/{key}` the quota fields are optional: omitting them keeps
|
||||
the stored caps, sending `0` explicitly lifts a cap. Editing only the model
|
||||
scope never clears a budget that was already set.
|
||||
|
||||
##### Quota rejection vs capacity rejection
|
||||
|
||||
A quota is checked at the entry point and occupies no upstream slot, so it is a
|
||||
cheap rejection. Running out of capacity instead waits out the bounded
|
||||
`busyWait` before answering 503. Measured (capacity 8, 0.6 s/request, 100
|
||||
concurrent clients):
|
||||
|
||||
| Case | Status | Latency | Reaches upstream |
|
||||
|---|---|---|---|
|
||||
| Quota spent (`req_quota` 50, 100 concurrent) | `429 rate_limit_exceeded` + `Retry-After` | **19–21 ms** | no |
|
||||
| Capacity exhausted (same topology, no quota) | `503 upstream_error` | ~2.6 s (`busyWait`) | overflow: no |
|
||||
|
||||
So a client can tell "wait for the window to roll over" apart from "wait for
|
||||
upstream capacity" — the first only needs patience, the second usually means
|
||||
back off concurrency or switch sources.
|
||||
|
||||
##### What the quota costs
|
||||
|
||||
Buckets are kept per (key, model, whole unix hour) for 40 days. Measured on an
|
||||
AMD 7840HS:
|
||||
|
||||
- Quota check per request: **149 ns** (caps set) / **42.6 ns** (no caps — it
|
||||
only looks up the key record and never touches a bucket) / **37 ns** (admin
|
||||
key returns immediately) — all **0 allocations**. Keys without caps cost
|
||||
almost nothing, so creating many of them is safe.
|
||||
- Recording one request: 283 ns, 3 allocations (unchanged from before this
|
||||
feature; the allocations come from the record ring buffer).
|
||||
- A window query scans the window, not the whole retention: 49 ns for 1 h,
|
||||
55 ns for 24 h, 3.9 µs for 30 d.
|
||||
- Memory: the production shape (7 keys x 8 models x 2 sources at full 40-day
|
||||
retention) costs about **3.7 MB**. The `source::model` bucket used by a
|
||||
source-pinned quota is created **lazily** — it is maintained only once some
|
||||
quota actually pins that source, because writing it on every record would
|
||||
cost an extra 18 MB at 20 keys x 8 models x 3 sources.
|
||||
The trade-off: usage recorded before a pinned quota was configured cannot be
|
||||
split by source afterwards, so **a pinned quota may under-count its first
|
||||
window**.
|
||||
|
||||
```
|
||||
HTTP/1.1 429 Too Many Requests
|
||||
Retry-After: 2100
|
||||
{"error":{"type":"rate_limit_exceeded","message":"key token quota exceeded (5000000/5000000, resets every 6h)"}}
|
||||
```
|
||||
|
||||
### Model routing
|
||||
|
||||
`/v1/chat/completions` `model` resolution order:
|
||||
|
||||
4
cmd/gui/package-lock.json
generated
4
cmd/gui/package-lock.json
generated
@ -1,12 +1,12 @@
|
||||
{
|
||||
"name": "modelrouter-gui",
|
||||
"version": "1.6.0",
|
||||
"version": "1.7.0",
|
||||
"lockfileVersion": 3,
|
||||
"requires": true,
|
||||
"packages": {
|
||||
"": {
|
||||
"name": "modelrouter-gui",
|
||||
"version": "1.6.0",
|
||||
"version": "1.7.0",
|
||||
"license": "Proprietary",
|
||||
"devDependencies": {
|
||||
"electron": "^33.0.0",
|
||||
|
||||
@ -1,6 +1,6 @@
|
||||
{
|
||||
"name": "modelrouter-gui",
|
||||
"version": "1.6.0",
|
||||
"version": "1.7.0",
|
||||
"description": "ModelRouter Desktop — embedded ModelRouter core with tray",
|
||||
"author": "ModelRouter",
|
||||
"main": "main.js",
|
||||
|
||||
@ -375,6 +375,82 @@ type GWKey struct {
|
||||
Note string `yaml:"note,omitempty" json:"note,omitempty"`
|
||||
CreatedAt int64 `yaml:"created_at,omitempty" json:"created_at,omitempty"`
|
||||
Seed bool `yaml:"seed,omitempty" json:"seed,omitempty"` // true if migrated from config gateway_keys
|
||||
// TokenQuota caps this key's TOTAL tokens across every model it may use.
|
||||
// 0 = unlimited. Period/Hours define the reset window, exactly like
|
||||
// ModelScope: "" never resets, "hour"/"week"/"month" fixed windows,
|
||||
// "nhour" uses Hours.
|
||||
TokenQuota int64 `yaml:"token_quota,omitempty" json:"token_quota,omitempty"`
|
||||
Period string `yaml:"period,omitempty" json:"period,omitempty"`
|
||||
Hours int64 `yaml:"hours,omitempty" json:"hours,omitempty"`
|
||||
// ReqQuota caps the number of requests per reset window; 0 = unlimited.
|
||||
// RPM covers short bursts; this covers sustained volume.
|
||||
ReqQuota int64 `yaml:"req_quota,omitempty" json:"req_quota,omitempty"`
|
||||
}
|
||||
|
||||
// KeyQuota is the set of key-wide caps accepted by the admin API. It is a
|
||||
// separate struct so a partial update can be expressed as a pointer (nil =
|
||||
// "leave the stored caps alone") instead of zero values meaning "clear".
|
||||
type KeyQuota struct {
|
||||
TokenQuota int64 `json:"token_quota"`
|
||||
ReqQuota int64 `json:"req_quota"`
|
||||
Period string `json:"period"`
|
||||
Hours int64 `json:"hours"`
|
||||
}
|
||||
|
||||
// ApplyQuota writes the caps onto a key record.
|
||||
func (k *GWKey) ApplyQuota(q KeyQuota) {
|
||||
k.TokenQuota = q.TokenQuota
|
||||
k.ReqQuota = q.ReqQuota
|
||||
k.Period = q.Period
|
||||
k.Hours = q.Hours
|
||||
}
|
||||
|
||||
// NormalizeRole defaults an empty role to "user", so a key can never end up in
|
||||
// a state where no role means "neither admin nor user".
|
||||
func NormalizeRole(role string) string {
|
||||
if role == "admin" {
|
||||
return "admin"
|
||||
}
|
||||
return "user"
|
||||
}
|
||||
|
||||
// Validate rejects a quota configuration that could not work as written. A
|
||||
// period is only meaningful when at least one cap is set, and a cap of zero
|
||||
// means "unlimited" rather than "deny everything", so those are the only two
|
||||
// things worth rejecting.
|
||||
func (q KeyQuota) Validate() error {
|
||||
if q.TokenQuota < 0 {
|
||||
return fmt.Errorf("token_quota must be >= 0 (0 = unlimited)")
|
||||
}
|
||||
if q.ReqQuota < 0 {
|
||||
return fmt.Errorf("req_quota must be >= 0 (0 = unlimited)")
|
||||
}
|
||||
if q.Hours < 0 {
|
||||
return fmt.Errorf("hours must be >= 0")
|
||||
}
|
||||
if q.TokenQuota > 0 || q.ReqQuota > 0 {
|
||||
if err := ValidatePeriod(q.Period, q.Hours); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// ValidatePeriod accepts the quota period vocabulary: "" (never resets),
|
||||
// "hour", "week", "month", or "nhour" with hours >= 1. An unknown period is
|
||||
// rejected rather than silently treated as "never resets", which would turn a
|
||||
// typo into an all-time quota — the opposite of what the operator typed.
|
||||
func ValidatePeriod(period string, hours int64) error {
|
||||
switch period {
|
||||
case "", "hour", "week", "month":
|
||||
return nil
|
||||
case "nhour":
|
||||
if hours < 1 {
|
||||
return fmt.Errorf("period %q needs hours >= 1", period)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
return fmt.Errorf("period must be one of \"\", hour, week, month, nhour (got %q)", period)
|
||||
}
|
||||
|
||||
// ModelScope is one allowed model for a key, or one AUTO scheduling slot,
|
||||
|
||||
@ -250,6 +250,11 @@ func (c *Core) FindKey(key string) (config.GWKey, bool) {
|
||||
|
||||
// CreateKey builds a new random gateway key and persists it to config.yaml.
|
||||
func (c *Core) CreateKey(name, role string, models []config.ModelScope, note string) (config.GWKey, error) {
|
||||
return c.CreateKeyWithQuota(name, role, models, note, config.KeyQuota{})
|
||||
}
|
||||
|
||||
// CreateKeyWithQuota is CreateKey plus the key-wide token/request caps.
|
||||
func (c *Core) CreateKeyWithQuota(name, role string, models []config.ModelScope, note string, q config.KeyQuota) (config.GWKey, error) {
|
||||
c.mu.Lock()
|
||||
defer c.mu.Unlock()
|
||||
models = cleanScopes(models)
|
||||
@ -257,6 +262,9 @@ func (c *Core) CreateKey(name, role string, models []config.ModelScope, note str
|
||||
if _, err := rand.Read(key); err != nil {
|
||||
return config.GWKey{}, err
|
||||
}
|
||||
if err := q.Validate(); err != nil {
|
||||
return config.GWKey{}, err
|
||||
}
|
||||
rec := config.GWKey{
|
||||
Key: "sk-gw-" + hex.EncodeToString(key),
|
||||
Role: role,
|
||||
@ -265,9 +273,8 @@ func (c *Core) CreateKey(name, role string, models []config.ModelScope, note str
|
||||
Note: note,
|
||||
CreatedAt: time.Now().Unix(),
|
||||
}
|
||||
if rec.Role == "" {
|
||||
rec.Role = "user"
|
||||
}
|
||||
rec.Role = config.NormalizeRole(rec.Role)
|
||||
rec.ApplyQuota(q)
|
||||
c.cfg.Keys = append(c.cfg.Keys, rec)
|
||||
if err := c.saveConfig(); err != nil {
|
||||
return config.GWKey{}, err
|
||||
@ -277,8 +284,20 @@ func (c *Core) CreateKey(name, role string, models []config.ModelScope, note str
|
||||
|
||||
// UpdateKey mutates a key's name/role/model scope and persists it.
|
||||
func (c *Core) UpdateKey(key, name, role string, models []config.ModelScope, note string) (config.GWKey, error) {
|
||||
return c.UpdateKeyWithQuota(key, name, role, models, note, nil)
|
||||
}
|
||||
|
||||
// UpdateKeyWithQuota is UpdateKey plus the key-wide caps. quota == nil leaves
|
||||
// the existing caps untouched, so a caller that only edits the model scope
|
||||
// does not silently clear a key's budget.
|
||||
func (c *Core) UpdateKeyWithQuota(key, name, role string, models []config.ModelScope, note string, quota *config.KeyQuota) (config.GWKey, error) {
|
||||
c.mu.Lock()
|
||||
defer c.mu.Unlock()
|
||||
if quota != nil {
|
||||
if err := quota.Validate(); err != nil {
|
||||
return config.GWKey{}, err
|
||||
}
|
||||
}
|
||||
for i, k := range c.cfg.Keys {
|
||||
if k.Key == key {
|
||||
if name != "" {
|
||||
@ -293,6 +312,9 @@ func (c *Core) UpdateKey(key, name, role string, models []config.ModelScope, not
|
||||
c.cfg.Keys[i].Models = cleanScopes(models)
|
||||
}
|
||||
c.cfg.Keys[i].Note = note
|
||||
if quota != nil {
|
||||
c.cfg.Keys[i].ApplyQuota(*quota)
|
||||
}
|
||||
if err := c.saveConfig(); err != nil {
|
||||
return config.GWKey{}, err
|
||||
}
|
||||
|
||||
@ -116,6 +116,12 @@ func (g *Gateway) apiV1Routes(w http.ResponseWriter, r *http.Request) {
|
||||
"note": k.Note,
|
||||
"created_at": k.CreatedAt,
|
||||
"seed": k.Seed,
|
||||
// key-wide spend caps (0 = unlimited). Echoed so an agent can
|
||||
// see what budget it has without parsing config.yaml.
|
||||
"token_quota": k.TokenQuota,
|
||||
"req_quota": k.ReqQuota,
|
||||
"period": k.Period,
|
||||
"hours": k.Hours,
|
||||
// The secret itself is never echoed. An operator that needs it
|
||||
// already has it from creation time or from config.yaml.
|
||||
"key_prefix": maskKey(k.Key),
|
||||
|
||||
@ -8,6 +8,7 @@ import (
|
||||
"log"
|
||||
"net/http"
|
||||
"regexp"
|
||||
"strconv"
|
||||
"strings"
|
||||
"sync/atomic"
|
||||
"time"
|
||||
@ -196,10 +197,36 @@ func intersectModels(models []string, allow []config.ModelScope) []string {
|
||||
// A scope entry with model "AUTO" only allows requests where the effective
|
||||
// model is AUTO (the routing mode). It does NOT grant access to specific model
|
||||
// ids — that requires an explicit scope entry for the model.
|
||||
// checkModelScope keeps its string-returning signature for callers that only
|
||||
// need to know whether the request may proceed.
|
||||
func (g *Gateway) checkModelScope(ctx context.Context, model string) string {
|
||||
if q := g.checkQuota(ctx, model); q != nil {
|
||||
return q.msg
|
||||
}
|
||||
return ""
|
||||
}
|
||||
|
||||
// quotaRejection is a quota verdict: the message plus whether the client
|
||||
// should retry. A spent quota is a rate limit (429 + Retry-After), not a
|
||||
// permission failure (403): a client that sees 403 gives up on the key, while
|
||||
// one that sees 429 with a retry hint waits and resumes when the window rolls
|
||||
// over.
|
||||
type quotaRejection struct {
|
||||
msg string
|
||||
retry int64 // seconds until the window resets; 0 = unknown
|
||||
}
|
||||
|
||||
func (q *quotaRejection) Error() string { return q.msg }
|
||||
|
||||
// checkQuota is checkKeyScope for callers that need the retry hint. It
|
||||
// separates the quota verdicts (429) from model-permission verdicts (403).
|
||||
func (g *Gateway) checkQuota(ctx context.Context, model string) *quotaRejection {
|
||||
if q := g.checkKeyQuotaRetry(ctx); q != nil {
|
||||
return q
|
||||
}
|
||||
allow := g.allowedModels(ctx)
|
||||
if allow == nil {
|
||||
return ""
|
||||
return nil
|
||||
}
|
||||
for _, sc := range allow {
|
||||
if sc.Model != model {
|
||||
@ -208,23 +235,75 @@ func (g *Gateway) checkModelScope(ctx context.Context, model string) string {
|
||||
if sc.TokenQuota > 0 {
|
||||
used := g.scopeTokens(ctx, sc)
|
||||
if used >= sc.TokenQuota {
|
||||
return fmt.Sprintf("token quota exceeded for %q (%d/%d)", model, used, sc.TokenQuota)
|
||||
return "aRejection{
|
||||
msg: fmt.Sprintf("token quota exceeded for %q (%d/%d)", model, used, sc.TokenQuota),
|
||||
retry: AutoSecondsToReset(sc.Period, sc.Hours),
|
||||
}
|
||||
}
|
||||
}
|
||||
return ""
|
||||
return nil
|
||||
}
|
||||
return fmt.Sprintf("model %q is not allowed for this key", model)
|
||||
return "aRejection{msg: fmt.Sprintf("model %q is not allowed for this key", model)}
|
||||
}
|
||||
|
||||
// scopeTokens returns the tokens a scope entry has consumed within its reset
|
||||
// window (total for AUTO / per model otherwise).
|
||||
// checkKeyQuotaRetry enforces the key-wide caps and reports the remaining
|
||||
// seconds of the reset window so the caller can answer with 429 + Retry-After.
|
||||
// An admin key is never capped, and a key with no caps set is never rejected.
|
||||
func (g *Gateway) checkKeyQuotaRetry(ctx context.Context) *quotaRejection {
|
||||
rec, ok := g.core.FindKey(reqKey(ctx))
|
||||
if !ok || rec.Role == "admin" {
|
||||
return nil
|
||||
}
|
||||
k := keyID(reqKey(ctx))
|
||||
win := AutoPeriodSeconds(rec.Period, rec.Hours)
|
||||
if rec.TokenQuota > 0 {
|
||||
if used := g.stats.KeyWindowTokens(k, win); used >= rec.TokenQuota {
|
||||
return "aRejection{
|
||||
msg: fmt.Sprintf("key token quota exceeded (%d/%d%s)", used, rec.TokenQuota, quotaWindowSuffix(rec.Period, rec.Hours)),
|
||||
retry: AutoSecondsToReset(rec.Period, rec.Hours),
|
||||
}
|
||||
}
|
||||
}
|
||||
if rec.ReqQuota > 0 {
|
||||
if used := g.stats.KeyWindowReqs(k, win); used >= rec.ReqQuota {
|
||||
return "aRejection{
|
||||
msg: fmt.Sprintf("key request quota exceeded (%d/%d%s)", used, rec.ReqQuota, quotaWindowSuffix(rec.Period, rec.Hours)),
|
||||
retry: AutoSecondsToReset(rec.Period, rec.Hours),
|
||||
}
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// quotaWindowSuffix describes a quota's reset window for an error message, so
|
||||
// a rejected caller can tell a permanent block from one that clears in an hour.
|
||||
func quotaWindowSuffix(period string, hours int64) string {
|
||||
switch {
|
||||
case period == "hour":
|
||||
return ", resets hourly"
|
||||
case period == "week":
|
||||
return ", resets weekly"
|
||||
case period == "month":
|
||||
return ", resets monthly"
|
||||
case period == "nhour" && hours > 1:
|
||||
return fmt.Sprintf(", resets every %dh", hours)
|
||||
}
|
||||
return ""
|
||||
}
|
||||
|
||||
// scopeTokens returns the tokens this key consumed within the scope entry's
|
||||
// reset window, isolated per key. For an AUTO entry the cap covers everything
|
||||
// the key routed through AUTO; for a model entry it covers that model only.
|
||||
//
|
||||
// It reads the per-key hourly buckets rather than the key-blind model
|
||||
// buckets, so one key's usage can never exhaust another's quota.
|
||||
func (g *Gateway) scopeTokens(ctx context.Context, sc config.ModelScope) int64 {
|
||||
k := keyID(reqKey(ctx))
|
||||
if sc.Model != "" && strings.EqualFold(sc.Model, "AUTO") {
|
||||
return g.stats.KeyTokens(k)
|
||||
}
|
||||
win := AutoPeriodSeconds(sc.Period, sc.Hours)
|
||||
return g.stats.WindowTokens(sc.Model, sc.Source, win)
|
||||
if sc.Model != "" && strings.EqualFold(sc.Model, "AUTO") {
|
||||
return g.stats.KeyWindowTokens(k, win)
|
||||
}
|
||||
return g.stats.KeyWindowModelTokens(k, sc.Model, sc.Source, win)
|
||||
}
|
||||
|
||||
// hasScopeModel reports whether a model (possibly with a "source-model" /
|
||||
@ -248,6 +327,26 @@ func (g *Gateway) hasScopeModel(list []config.ModelScope, s string) bool {
|
||||
return false
|
||||
}
|
||||
|
||||
// writeScopeReject answers a model-scope or quota rejection. A spent quota is
|
||||
// 429 (rate_limit_exceeded) with Retry-After, so a client waits and resumes
|
||||
// after the reset; a model the key may not use stays 403 (model_not_allowed),
|
||||
// because retrying cannot help.
|
||||
func (g *Gateway) writeReject(w http.ResponseWriter, q *quotaRejection) {
|
||||
if q.retry > 0 {
|
||||
w.Header().Set("Retry-After", strconv.FormatInt(q.retry, 10))
|
||||
writeError(w, http.StatusTooManyRequests, "rate_limit_exceeded", q.msg)
|
||||
return
|
||||
}
|
||||
// no window to wait for: the cap is either permanent or key-wide with no
|
||||
// period. "rate_limit_exceeded" still says "come back after the operator
|
||||
// raises the cap", which 403 would not.
|
||||
if strings.Contains(q.msg, "quota exceeded") {
|
||||
writeError(w, http.StatusTooManyRequests, "rate_limit_exceeded", q.msg)
|
||||
return
|
||||
}
|
||||
writeError(w, http.StatusForbidden, "model_not_allowed", q.msg)
|
||||
}
|
||||
|
||||
func (g *Gateway) resolveByModel(model string) ([]*provider.Provider, string) {
|
||||
if isAuto(model) {
|
||||
return g.core.Registry().Resolve("AUTO"), ""
|
||||
@ -310,8 +409,8 @@ func (g *Gateway) handleChat(w http.ResponseWriter, r *http.Request) {
|
||||
writeError(w, http.StatusServiceUnavailable, "no_provider", "no auto slot configured")
|
||||
return
|
||||
}
|
||||
if msg := g.checkModelScope(r.Context(), "AUTO"); msg != "" {
|
||||
writeError(w, http.StatusForbidden, "model_not_allowed", msg)
|
||||
if q := g.checkQuota(r.Context(), "AUTO"); q != nil {
|
||||
g.writeReject(w, q)
|
||||
return
|
||||
}
|
||||
ctx := r.Context()
|
||||
@ -364,8 +463,8 @@ func (g *Gateway) handleChat(w http.ResponseWriter, r *http.Request) {
|
||||
if effective == "" {
|
||||
effective = firstModel(cands[0])
|
||||
}
|
||||
if msg := g.checkModelScope(r.Context(), effective); msg != "" {
|
||||
writeError(w, http.StatusForbidden, "model_not_allowed", msg)
|
||||
if q := g.checkQuota(r.Context(), effective); q != nil {
|
||||
g.writeReject(w, q)
|
||||
return
|
||||
}
|
||||
ctx := r.Context()
|
||||
@ -1021,8 +1120,8 @@ func (g *Gateway) handleImage(w http.ResponseWriter, r *http.Request) {
|
||||
}
|
||||
if isAuto(model) {
|
||||
if chain := g.core.AutoImageChain(); chain != nil && len(chain.Tiers) > 0 {
|
||||
if msg := g.checkModelScope(r.Context(), "AUTO"); msg != "" {
|
||||
writeError(w, http.StatusForbidden, "model_not_allowed", msg)
|
||||
if q := g.checkQuota(r.Context(), "AUTO"); q != nil {
|
||||
g.writeReject(w, q)
|
||||
return
|
||||
}
|
||||
done := g.stats.Begin()
|
||||
@ -1068,8 +1167,8 @@ func (g *Gateway) handleImage(w http.ResponseWriter, r *http.Request) {
|
||||
writeError(w, http.StatusServiceUnavailable, "no_provider", "no image source configured")
|
||||
return
|
||||
}
|
||||
if msg := g.checkModelScope(r.Context(), effectiveImageModel(model, cands)); msg != "" {
|
||||
writeError(w, http.StatusForbidden, "model_not_allowed", msg)
|
||||
if q := g.checkQuota(r.Context(), effectiveImageModel(model, cands)); q != nil {
|
||||
g.writeReject(w, q)
|
||||
return
|
||||
}
|
||||
done := g.stats.Begin()
|
||||
|
||||
259
internal/gateway/key_quota_api_test.go
Normal file
259
internal/gateway/key_quota_api_test.go
Normal file
@ -0,0 +1,259 @@
|
||||
package gateway
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"os"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"llmsproxy/internal/config"
|
||||
"llmsproxy/internal/core"
|
||||
)
|
||||
|
||||
// adminGateway builds a gateway whose admin key can call /api/keys.
|
||||
func adminGateway(t *testing.T, keys ...config.GWKey) *Gateway {
|
||||
t.Helper()
|
||||
td := t.TempDir()
|
||||
cfgPath := td + "/config.yaml"
|
||||
if err := os.WriteFile(cfgPath, []byte("listen: :0"), 0o644); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
cfg := &config.Config{
|
||||
Path: cfgPath,
|
||||
AdapterDir: td + "/adapters",
|
||||
RuntimeFile: td + "/runtime.json",
|
||||
Keys: keys,
|
||||
}
|
||||
if err := cfg.ApplyDefaults(); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
c, err := core.NewFromConfig(cfg)
|
||||
if err != nil {
|
||||
t.Fatalf("core: %v", err)
|
||||
}
|
||||
t.Cleanup(c.Close)
|
||||
g, err := New(c, []string{"sk-admin"})
|
||||
if err != nil {
|
||||
t.Fatalf("gateway: %v", err)
|
||||
}
|
||||
return g
|
||||
}
|
||||
|
||||
func adminReq(t *testing.T, g *Gateway, method, path, body string) *httptest.ResponseRecorder {
|
||||
t.Helper()
|
||||
req, _ := http.NewRequest(method, path, strings.NewReader(body))
|
||||
req.Header.Set("Authorization", "Bearer sk-admin")
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
rr := httptest.NewRecorder()
|
||||
g.Handler().ServeHTTP(rr, req)
|
||||
return rr
|
||||
}
|
||||
|
||||
// keyRecord pulls one key's stored record out of the admin list.
|
||||
func keyRecord(t *testing.T, g *Gateway, secret string) config.GWKey {
|
||||
t.Helper()
|
||||
rr := adminReq(t, g, "GET", "/api/keys", "")
|
||||
if rr.Code != 200 {
|
||||
t.Fatalf("GET /api/keys: %d %s", rr.Code, rr.Body.String())
|
||||
}
|
||||
var out struct {
|
||||
Keys []config.GWKey `json:"keys"`
|
||||
}
|
||||
if err := json.Unmarshal(rr.Body.Bytes(), &out); err != nil {
|
||||
t.Fatalf("decode: %v (%s)", err, rr.Body.String())
|
||||
}
|
||||
for _, k := range out.Keys {
|
||||
if k.Key == secret {
|
||||
return k
|
||||
}
|
||||
}
|
||||
t.Fatalf("key %q not found in %s", secret, rr.Body.String())
|
||||
return config.GWKey{}
|
||||
}
|
||||
|
||||
func TestKeyAPIStoresQuota(t *testing.T) {
|
||||
g := adminGateway(t, config.GWKey{Key: "sk-admin", Role: "admin"})
|
||||
rr := adminReq(t, g, "POST", "/api/keys",
|
||||
`{"name":"agent-x","role":"user","token_quota":50000,"req_quota":200,"period":"nhour","hours":6,"models":[{"model":"m1"}]}`)
|
||||
if rr.Code != 200 {
|
||||
t.Fatalf("create: %d %s", rr.Code, rr.Body.String())
|
||||
}
|
||||
var created struct {
|
||||
Key config.GWKey `json:"key"`
|
||||
}
|
||||
if err := json.Unmarshal(rr.Body.Bytes(), &created); err != nil {
|
||||
t.Fatalf("decode: %v", err)
|
||||
}
|
||||
if created.Key.TokenQuota != 50000 || created.Key.ReqQuota != 200 ||
|
||||
created.Key.Period != "nhour" || created.Key.Hours != 6 {
|
||||
t.Fatalf("created key did not carry the caps: %+v", created.Key)
|
||||
}
|
||||
// and it must survive a read-back (persisted, not just echoed)
|
||||
back := keyRecord(t, g, created.Key.Key)
|
||||
if back.TokenQuota != 50000 || back.Period != "nhour" || back.Hours != 6 {
|
||||
t.Errorf("read-back lost the caps: %+v", back)
|
||||
}
|
||||
}
|
||||
|
||||
// Editing only the model scope must not silently clear a key's budget: the
|
||||
// caps are pointers precisely so "absent" is not "zero".
|
||||
func TestKeyAPIUpdateKeepsQuotaWhenOmitted(t *testing.T) {
|
||||
g := adminGateway(t, config.GWKey{Key: "sk-admin", Role: "admin"})
|
||||
rr := adminReq(t, g, "POST", "/api/keys",
|
||||
`{"name":"agent-x","role":"user","token_quota":50000,"period":"day-typo-free","models":[{"model":"m1"}]}`)
|
||||
rr = adminReq(t, g, "POST", "/api/keys", `{"name":"y","role":"user","token_quota":50000,"period":"hour","models":[{"model":"m1"}]}`)
|
||||
if rr.Code != 200 {
|
||||
t.Fatalf("setup create: %d %s", rr.Code, rr.Body.String())
|
||||
}
|
||||
var created struct {
|
||||
Key config.GWKey `json:"key"`
|
||||
}
|
||||
_ = json.Unmarshal(rr.Body.Bytes(), &created)
|
||||
|
||||
// a scope-only edit
|
||||
rr = adminReq(t, g, "PUT", "/api/keys/"+created.Key.Key,
|
||||
`{"name":"agent-y","models":[{"model":"m1"},{"model":"m2"}]}`)
|
||||
if rr.Code != 200 {
|
||||
t.Fatalf("update: %d %s", rr.Code, rr.Body.String())
|
||||
}
|
||||
back := keyRecord(t, g, created.Key.Key)
|
||||
if back.TokenQuota != 50000 {
|
||||
t.Errorf("token_quota was cleared by a scope-only edit: %d", back.TokenQuota)
|
||||
}
|
||||
if back.Period != "hour" {
|
||||
t.Errorf("period was cleared by a scope-only edit: %q", back.Period)
|
||||
}
|
||||
if len(back.Models) != 2 {
|
||||
t.Errorf("scope edit did not apply: %+v", back.Models)
|
||||
}
|
||||
}
|
||||
|
||||
// Sending 0 explicitly must lift the cap, not be treated as "absent".
|
||||
func TestKeyAPIUpdateZeroLiftsCap(t *testing.T) {
|
||||
g := adminGateway(t, config.GWKey{Key: "sk-admin", Role: "admin"})
|
||||
rr := adminReq(t, g, "POST", "/api/keys", `{"name":"z","role":"user","token_quota":1000,"period":"hour"}`)
|
||||
if rr.Code != 200 {
|
||||
t.Fatalf("create: %d %s", rr.Code, rr.Body.String())
|
||||
}
|
||||
var created struct {
|
||||
Key config.GWKey `json:"key"`
|
||||
}
|
||||
_ = json.Unmarshal(rr.Body.Bytes(), &created)
|
||||
|
||||
rr = adminReq(t, g, "PUT", "/api/keys/"+created.Key.Key, `{"token_quota":0}`)
|
||||
if rr.Code != 200 {
|
||||
t.Fatalf("lift: %d %s", rr.Code, rr.Body.String())
|
||||
}
|
||||
if back := keyRecord(t, g, created.Key.Key); back.TokenQuota != 0 {
|
||||
t.Errorf("token_quota = %d, want 0 (cap lifted)", back.TokenQuota)
|
||||
}
|
||||
}
|
||||
|
||||
// A misspelled period must be refused, not quietly turned into an all-time
|
||||
// quota — which is the exact opposite of what the operator typed.
|
||||
func TestKeyAPIRejectsBadPeriod(t *testing.T) {
|
||||
g := adminGateway(t, config.GWKey{Key: "sk-admin", Role: "admin"})
|
||||
rr := adminReq(t, g, "POST", "/api/keys", `{"name":"bad","role":"user","token_quota":1000,"period":"houre"}`)
|
||||
if rr.Code != http.StatusBadRequest {
|
||||
t.Fatalf("want 400 for a bad period, got %d %s", rr.Code, rr.Body.String())
|
||||
}
|
||||
if !strings.Contains(rr.Body.String(), "period") {
|
||||
t.Errorf("error should name the period field: %s", rr.Body.String())
|
||||
}
|
||||
}
|
||||
|
||||
func TestKeyAPIRejectsNegativeQuota(t *testing.T) {
|
||||
g := adminGateway(t, config.GWKey{Key: "sk-admin", Role: "admin"})
|
||||
rr := adminReq(t, g, "POST", "/api/keys", `{"name":"bad","role":"user","token_quota":-5}`)
|
||||
if rr.Code != http.StatusBadRequest {
|
||||
t.Fatalf("want 400 for a negative quota, got %d %s", rr.Code, rr.Body.String())
|
||||
}
|
||||
}
|
||||
|
||||
// A non-admin key must not be able to set or read another key's budget.
|
||||
func TestKeyAPIQuotaIsAdminOnly(t *testing.T) {
|
||||
g := adminGateway(t,
|
||||
config.GWKey{Key: "sk-admin", Role: "admin"},
|
||||
config.GWKey{Key: "sk-u", Role: "user", TokenQuota: 10, Period: "hour"},
|
||||
)
|
||||
req, _ := http.NewRequest("POST", "/api/keys", strings.NewReader(`{"name":"x","role":"admin","token_quota":0}`))
|
||||
req.Header.Set("Authorization", "Bearer sk-u")
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
rr := httptest.NewRecorder()
|
||||
g.Handler().ServeHTTP(rr, req)
|
||||
if rr.Code != http.StatusForbidden {
|
||||
t.Fatalf("non-admin create: want 403, got %d %s", rr.Code, rr.Body.String())
|
||||
}
|
||||
// /api/v1/keys echoes the caps but never the secret
|
||||
rr = adminReq(t, g, "GET", "/api/v1/keys", "")
|
||||
if rr.Code != 200 {
|
||||
t.Fatalf("GET /api/v1/keys: %d", rr.Code)
|
||||
}
|
||||
if strings.Contains(rr.Body.String(), "sk-u") {
|
||||
t.Error("/api/v1/keys leaked a key secret")
|
||||
}
|
||||
if !strings.Contains(rr.Body.String(), `"token_quota":10`) {
|
||||
t.Errorf("/api/v1/keys should expose the cap: %s", rr.Body.String())
|
||||
}
|
||||
}
|
||||
|
||||
// The AUTO scope entry must honour its reset window: usage that aged out of
|
||||
// the window must not count against a per-key cap.
|
||||
func TestAutoScopeQuotaHonoursWindow(t *testing.T) {
|
||||
g, _ := quotaGateway(t,
|
||||
config.GWKey{Key: "sk-a", Role: "user", Models: []config.ModelScope{
|
||||
{Model: "AUTO", TokenQuota: 1000, Period: "hour"},
|
||||
}},
|
||||
config.GWKey{Key: "sk-b", Role: "user"},
|
||||
)
|
||||
ctx := quotaCtx(t, g, "sk-a")
|
||||
sc := config.ModelScope{Model: "AUTO", TokenQuota: 1000, Period: "hour"}
|
||||
|
||||
// aged-out usage: 2 days old, 5M tokens — must be invisible to a 1h window
|
||||
g.stats.Record(Req{Time: nowMSOffset(-48 * 3600 * 1000), Key: keyID("sk-a"),
|
||||
Model: "m1", Source: "up", Prompt: 2500000, Compl: 2500000, OK: true, Status: 200})
|
||||
if used := g.scopeTokens(ctx, sc); used != 0 {
|
||||
t.Fatalf("AUTO scope saw %d tokens outside its 1h window; the period is being ignored", used)
|
||||
}
|
||||
|
||||
// in-window usage counts
|
||||
g.stats.Record(Req{Time: nowMSOffset(0), Key: keyID("sk-a"),
|
||||
Model: "m1", Source: "up", Prompt: 400, Compl: 400, OK: true, Status: 200})
|
||||
if used := g.scopeTokens(ctx, sc); used != 800 {
|
||||
t.Fatalf("AUTO scope used = %d, want 800", used)
|
||||
}
|
||||
}
|
||||
|
||||
var _ = fmt.Sprintf
|
||||
|
||||
// A user must be able to see their own budget: /api/keys/me is the only key
|
||||
// view a non-admin gets, so a cap missing from it is invisible to the very
|
||||
// client it constrains.
|
||||
func TestKeyMeExposesOwnQuota(t *testing.T) {
|
||||
g := adminGateway(t,
|
||||
config.GWKey{Key: "sk-admin", Role: "admin"},
|
||||
config.GWKey{Key: "sk-u", Role: "user", Name: "agent",
|
||||
TokenQuota: 123456, ReqQuota: 42, Period: "week", Hours: 0},
|
||||
)
|
||||
req, _ := http.NewRequest("GET", "/api/keys/me", nil)
|
||||
req.Header.Set("Authorization", "Bearer sk-u")
|
||||
rr := httptest.NewRecorder()
|
||||
g.Handler().ServeHTTP(rr, req)
|
||||
if rr.Code != 200 {
|
||||
t.Fatalf("GET /api/keys/me: %d %s", rr.Code, rr.Body.String())
|
||||
}
|
||||
// the endpoint wraps the record: {"key": {...}}
|
||||
var wrap struct {
|
||||
Key config.GWKey `json:"key"`
|
||||
}
|
||||
if err := json.Unmarshal(rr.Body.Bytes(), &wrap); err != nil {
|
||||
t.Fatalf("decode: %v (%s)", err, rr.Body.String())
|
||||
}
|
||||
me := wrap.Key
|
||||
if me.TokenQuota != 123456 || me.ReqQuota != 42 || me.Period != "week" {
|
||||
t.Errorf("own quota not visible to the key's owner: %+v", me)
|
||||
}
|
||||
}
|
||||
60
internal/gateway/key_quota_concurrency_test.go
Normal file
60
internal/gateway/key_quota_concurrency_test.go
Normal file
@ -0,0 +1,60 @@
|
||||
package gateway
|
||||
|
||||
import (
|
||||
"sync"
|
||||
"testing"
|
||||
"time"
|
||||
)
|
||||
|
||||
// The per-key quota buckets are shared maps mutated on every recorded request
|
||||
// and read on every quota check. Single-threaded unit tests cannot see a race
|
||||
// there at all, so this drives 64 goroutines through record + read together
|
||||
// under `go test -race`. It also opts one model into the lazily-created pinned
|
||||
// bucket mid-flight, which is the path that mutates two bucket maps at once.
|
||||
func TestConcurrentQuotaAccounting(t *testing.T) {
|
||||
s := NewStats(0)
|
||||
keys := []string{"k0", "k1", "k2", "k3", "k4", "k5", "k6"}
|
||||
models := []string{"m0", "m1", "m2", "AUTO"}
|
||||
var wg sync.WaitGroup
|
||||
stop := time.Now().Add(2 * time.Second)
|
||||
|
||||
for w := 0; w < 64; w++ {
|
||||
wg.Add(1)
|
||||
go func(w int) {
|
||||
defer wg.Done()
|
||||
k := keys[w%len(keys)]
|
||||
m := models[w%len(models)]
|
||||
now := time.Now()
|
||||
for time.Now().Before(stop) {
|
||||
s.Record(Req{Time: now.UnixMilli(), Key: k, Model: m, Source: "src",
|
||||
Prompt: 100, Compl: 50, OK: true, Status: 200, Type: "chat"})
|
||||
_ = s.KeyWindowTokens(k, 3600)
|
||||
_ = s.KeyWindowReqs(k, 3600)
|
||||
_ = s.KeyWindowModelTokens(k, m, "", 3600)
|
||||
_ = s.KeyWindowModelTokens(k, m, "src", 3600) // opts into the pinned bucket
|
||||
}
|
||||
}(w)
|
||||
}
|
||||
wg.Wait()
|
||||
|
||||
// Every recorded request must be accounted for exactly once.
|
||||
s.mu.Lock()
|
||||
var total int64
|
||||
for _, hm := range s.keyHour {
|
||||
for _, v := range hm {
|
||||
total += v
|
||||
}
|
||||
}
|
||||
recs := int64(0)
|
||||
for _, r := range s.recs {
|
||||
recs += r.Prompt + r.Compl
|
||||
}
|
||||
s.mu.Unlock()
|
||||
if total == 0 {
|
||||
t.Fatal("no tokens accounted")
|
||||
}
|
||||
t.Logf("accounted %d tokens across %d keys", total, len(keys))
|
||||
if recs == 0 {
|
||||
t.Error("no records retained; the concurrent writers lost work")
|
||||
}
|
||||
}
|
||||
125
internal/gateway/key_quota_perf_test.go
Normal file
125
internal/gateway/key_quota_perf_test.go
Normal file
@ -0,0 +1,125 @@
|
||||
package gateway
|
||||
|
||||
import (
|
||||
"math/rand"
|
||||
"testing"
|
||||
"time"
|
||||
)
|
||||
|
||||
// Perf: the per-key quota bookkeeping runs on every recorded request and on
|
||||
// every quota check, so its cost lands directly on the request path. These
|
||||
// benchmarks exist to catch a regression that would make the feature
|
||||
// expensive; the numbers that matter are relative to each other and to the
|
||||
// pre-change baseline (Record was 275 ns/op with 3 allocs before this work).
|
||||
|
||||
func BenchmarkRecordQuotaBucket(b *testing.B) {
|
||||
s := NewStats(0)
|
||||
r := Req{Time: time.Now().UnixMilli(), Key: "keyABC", Model: "deepseek-v4-flash",
|
||||
Source: "deepseek", Prompt: 800, Compl: 200, OK: true, Status: 200, Type: "chat"}
|
||||
b.ReportAllocs()
|
||||
b.ResetTimer()
|
||||
for i := 0; i < b.N; i++ {
|
||||
s.Record(r)
|
||||
}
|
||||
}
|
||||
|
||||
// A key that has been busy for the whole 40-day retention.
|
||||
func BenchmarkRecordQuotaBucketFullRetention(b *testing.B) {
|
||||
s := NewStats(0)
|
||||
nowH := time.Now().Unix() / 3600
|
||||
for h := int64(0); h < quotaRetentionHours; h++ {
|
||||
s.Record(Req{Time: (nowH - h) * 3600 * 1000, Key: "keyABC", Model: "m1",
|
||||
Source: "deepseek", Prompt: 100, Compl: 100, OK: true, Status: 200})
|
||||
}
|
||||
r := Req{Time: time.Now().UnixMilli(), Key: "keyABC", Model: "m1", Source: "deepseek",
|
||||
Prompt: 100, Compl: 100, OK: true, Status: 200, Type: "chat"}
|
||||
b.ReportAllocs()
|
||||
b.ResetTimer()
|
||||
for i := 0; i < b.N; i++ {
|
||||
s.Record(r)
|
||||
}
|
||||
}
|
||||
|
||||
// The whole per-request quota check: total tokens + request count.
|
||||
func BenchmarkQuotaCheckHourly(b *testing.B) {
|
||||
s := NewStats(0)
|
||||
for h := 0; h < 24; h++ {
|
||||
s.Record(Req{Time: time.Now().Add(-time.Duration(h) * time.Hour).UnixMilli(),
|
||||
Key: "keyABC", Model: "m1", Source: "deepseek", Prompt: 100, Compl: 100, OK: true, Status: 200})
|
||||
}
|
||||
b.ReportAllocs()
|
||||
b.ResetTimer()
|
||||
for i := 0; i < b.N; i++ {
|
||||
s.KeyWindowTokens("keyABC", 3600)
|
||||
s.KeyWindowReqs("keyABC", 3600)
|
||||
}
|
||||
}
|
||||
|
||||
// A 30-day window has to visit 720 hour buckets. This is the worst case the
|
||||
// quota feature puts on the request path.
|
||||
func BenchmarkQuotaCheckMonthly(b *testing.B) {
|
||||
s := NewStats(0)
|
||||
nowH := time.Now().Unix() / 3600
|
||||
for h := int64(0); h < 720; h++ {
|
||||
s.Record(Req{Time: (nowH - h) * 3600 * 1000, Key: "keyABC", Model: "m1",
|
||||
Prompt: 100, Compl: 100, OK: true, Status: 200})
|
||||
}
|
||||
b.ReportAllocs()
|
||||
b.ResetTimer()
|
||||
for i := 0; i < b.N; i++ {
|
||||
s.KeyWindowTokens("keyABC", 30*24*3600)
|
||||
}
|
||||
}
|
||||
|
||||
// TestSumBucketsMatchesFullScan is the correctness guard for the optimization
|
||||
// that made the quota check proportional to the window instead of to the whole
|
||||
// retention: the windowed sum must agree with a full scan for every bucket
|
||||
// placement, or a quota would silently start letting traffic through. It
|
||||
// caught a real off-by-one (floor instead of ceil on the first hour).
|
||||
func TestSumBucketsMatchesFullScan(t *testing.T) {
|
||||
const hour = 3600
|
||||
base := int64(1_700_000_000) / hour * hour
|
||||
rng := rand.New(rand.NewSource(7))
|
||||
for trial := 0; trial < 400; trial++ {
|
||||
hm := map[int64]int64{}
|
||||
for i, n := 0, rng.Intn(40)+1; i < n; i++ {
|
||||
h := base/hour - int64(rng.Intn(1200))
|
||||
hm[h] += int64(rng.Intn(1000)) + 1
|
||||
}
|
||||
now := base + int64(rng.Intn(hour))
|
||||
for _, sec := range []int64{3600, 2 * 3600, 6 * 3600, 24 * 3600, 7 * 24 * 3600, 30 * 24 * 3600} {
|
||||
var want int64
|
||||
cut := now - sec
|
||||
for h, v := range hm {
|
||||
if h*hour >= cut {
|
||||
want += v
|
||||
}
|
||||
}
|
||||
if got := sumBuckets(hm, now, sec); got != want {
|
||||
t.Fatalf("trial %d sec=%d: sumBuckets=%d, full scan=%d", trial, sec, got, want)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// The pinned bucket is maintained only for (key, model) pairs whose quota pins
|
||||
// a source. Doing it unconditionally doubled the bucket count — measured at
|
||||
// 26.7 MB for 20 keys x 8 models x 3 sources at full retention, against a
|
||||
// documented ~32 MB total memory budget — to serve a lookup nobody performs.
|
||||
func TestPinnedBucketNotMaintainedUnlessWanted(t *testing.T) {
|
||||
s := NewStats(100)
|
||||
now := time.Now().UnixMilli()
|
||||
for i := 0; i < 5; i++ {
|
||||
s.Record(Req{Time: now, Key: "keyA", Model: "m1", Source: "srcX", Prompt: 100, Compl: 100, OK: true, Status: 200})
|
||||
}
|
||||
s.mu.Lock()
|
||||
_, pinned := s.keyModelHour["keyA"]["srcX::m1"]
|
||||
_, bare := s.keyModelHour["keyA"]["m1"]
|
||||
s.mu.Unlock()
|
||||
if pinned {
|
||||
t.Error("a pinned bucket exists before any quota read the pin — memory regression")
|
||||
}
|
||||
if !bare {
|
||||
t.Error("the bare-model bucket must always exist; an unpinned quota depends on it")
|
||||
}
|
||||
}
|
||||
225
internal/gateway/key_quota_test.go
Normal file
225
internal/gateway/key_quota_test.go
Normal file
@ -0,0 +1,225 @@
|
||||
package gateway
|
||||
|
||||
import (
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"llmsproxy/internal/config"
|
||||
)
|
||||
|
||||
// ---- data layer: per-key isolation ----
|
||||
|
||||
func TestKeyWindowTokensIsolatesKeys(t *testing.T) {
|
||||
s := NewStats(100)
|
||||
now := time.Now().UnixMilli()
|
||||
s.Record(Req{Time: now, Key: "keyA", Model: "m1", Prompt: 500, Compl: 500, OK: true, Status: 200})
|
||||
s.Record(Req{Time: now, Key: "keyB", Model: "m1", Prompt: 500, Compl: 500, OK: true, Status: 200})
|
||||
|
||||
if got := s.KeyWindowModelTokens("keyA", "m1", "", 3600); got != 1000 {
|
||||
t.Errorf("keyA model tokens = %d, want 1000 (its own usage only)", got)
|
||||
}
|
||||
if got := s.KeyWindowModelTokens("keyB", "m1", "", 3600); got != 1000 {
|
||||
t.Errorf("keyB model tokens = %d, want 1000", got)
|
||||
}
|
||||
// the key-blind bucket stays global on purpose: it backs the AUTO slot
|
||||
// quota, which limits the whole gateway, not one key.
|
||||
if got := s.WindowTokens("m1", "", 3600); got != 2000 {
|
||||
t.Errorf("WindowTokens = %d, want 2000 (global per-model total must be unchanged)", got)
|
||||
}
|
||||
}
|
||||
|
||||
// A source pin scopes the cap to that one upstream. The pinned bucket is
|
||||
// maintained lazily (it is only read by quotas that actually pin a source), so
|
||||
// usage recorded before anything queried the pin is not attributed to it.
|
||||
func TestKeyWindowTokensSeparatesSourcePin(t *testing.T) {
|
||||
s := NewStats(100)
|
||||
now := time.Now().UnixMilli()
|
||||
s.Record(Req{Time: now, Key: "keyA", Model: "m1", Source: "srcX", Prompt: 100, Compl: 100, OK: true, Status: 200}) // 200 tok
|
||||
s.Record(Req{Time: now, Key: "keyA", Model: "m1", Source: "srcY", Prompt: 200, Compl: 200, OK: true, Status: 200}) // 400 tok
|
||||
|
||||
// The unpinned bucket always counts — it is what a scope entry without a
|
||||
// source pin reads. 200 + 400.
|
||||
if got := s.KeyWindowModelTokens("keyA", "m1", "", 3600); got != 600 {
|
||||
t.Errorf("unpinned model tokens = %d, want 600 (both sources)", got)
|
||||
}
|
||||
|
||||
// Reading a pin opts this (key, model) into pinned accounting.
|
||||
if got := s.KeyWindowModelTokens("keyA", "m1", "srcX", 3600); got != 0 {
|
||||
t.Errorf("srcX pin before opt-in = %d, want 0 (the lazy bucket has not accrued yet)", got)
|
||||
}
|
||||
|
||||
// From now on both pins accrue.
|
||||
more := now + 1
|
||||
s.Record(Req{Time: more, Key: "keyA", Model: "m1", Source: "srcX", Prompt: 10, Compl: 10, OK: true, Status: 200}) // 20 tok
|
||||
s.Record(Req{Time: more, Key: "keyA", Model: "m1", Source: "srcY", Prompt: 30, Compl: 30, OK: true, Status: 200}) // 60 tok
|
||||
if got := s.KeyWindowModelTokens("keyA", "m1", "srcX", 3600); got != 20 {
|
||||
t.Errorf("srcX-pinned tokens after opt-in = %d, want 20", got)
|
||||
}
|
||||
if got := s.KeyWindowModelTokens("keyA", "m1", "srcY", 3600); got != 60 {
|
||||
t.Errorf("srcY-pinned tokens = %d, want 60 (a pin read for one source must not blind the other)", got)
|
||||
}
|
||||
}
|
||||
|
||||
// This is the shape of every real chat record: Source is always populated.
|
||||
// The unpinned bucket must still see the tokens, or a per-model quota without
|
||||
// a source pin reads an empty bucket and never trips.
|
||||
func TestKeyWindowModelTokensUnpinnedSeesSourcedTraffic(t *testing.T) {
|
||||
s := NewStats(100)
|
||||
s.Record(Req{Time: time.Now().UnixMilli(), Key: "keyA", Model: "m1", Source: "deepseek",
|
||||
Prompt: 1000, Compl: 1000, OK: true, Status: 200})
|
||||
|
||||
if got := s.KeyWindowModelTokens("keyA", "m1", "", 3600); got != 2000 {
|
||||
t.Errorf("unpinned tokens = %d, want 2000 — an unpinned per-model quota would never trip otherwise", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestKeyWindowRespectsResetWindow(t *testing.T) {
|
||||
s := NewStats(100)
|
||||
old := time.Now().Add(-30 * 24 * time.Hour).UnixMilli()
|
||||
s.Record(Req{Time: old, Key: "keyA", Model: "m1", Prompt: 100, Compl: 100, OK: true, Status: 200})
|
||||
s.Record(Req{Time: time.Now().UnixMilli(), Key: "keyA", Model: "m1", Prompt: 5, Compl: 5, OK: true, Status: 200})
|
||||
|
||||
if got := s.KeyWindowTokens("keyA", 3600); got != 10 {
|
||||
t.Errorf("1h window = %d, want 10 (30-day-old usage must not count)", got)
|
||||
}
|
||||
if got := s.KeyWindowTokens("keyA", 0); got != 210 {
|
||||
t.Errorf("all-time total = %d, want 210", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestKeyWindowReqsCountsEveryRequest(t *testing.T) {
|
||||
s := NewStats(100)
|
||||
now := time.Now().UnixMilli()
|
||||
// a failed request and an image request both count: a client that loops on
|
||||
// failures must still burn its request quota
|
||||
s.Record(Req{Time: now, Key: "keyA", Model: "m1", Type: "chat", OK: true, Status: 200})
|
||||
s.Record(Req{Time: now, Key: "keyA", Model: "m1", Type: "chat", OK: false, Status: 500})
|
||||
s.Record(Req{Time: now, Key: "keyA", Model: "img", Type: "image", OK: true, Status: 200})
|
||||
s.Record(Req{Time: now, Key: "keyB", Model: "m1", Type: "chat", OK: true, Status: 200})
|
||||
|
||||
if got := s.KeyWindowReqs("keyA", 3600); got != 3 {
|
||||
t.Errorf("keyA reqs = %d, want 3 (chat ok + chat fail + image)", got)
|
||||
}
|
||||
if got := s.KeyWindowReqs("keyB", 3600); got != 1 {
|
||||
t.Errorf("keyB reqs = %d, want 1", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestKeyWindowTokensZeroTokensNotCounted(t *testing.T) {
|
||||
s := NewStats(100)
|
||||
now := time.Now().UnixMilli()
|
||||
// a request that reported no usage must not create a bucket entry
|
||||
s.Record(Req{Time: now, Key: "keyA", Model: "", Type: "chat", OK: true, Status: 200})
|
||||
if got := s.KeyWindowTokens("keyA", 3600); got != 0 {
|
||||
t.Errorf("tokens = %d, want 0", got)
|
||||
}
|
||||
if got := s.KeyWindowReqs("keyA", 3600); got != 1 {
|
||||
t.Errorf("reqs = %d, want 1 (request still happened)", got)
|
||||
}
|
||||
}
|
||||
|
||||
// ---- reset window arithmetic ----
|
||||
|
||||
func TestAutoSecondsToReset(t *testing.T) {
|
||||
if got := AutoSecondsToReset("", 0); got != 0 {
|
||||
t.Errorf("no period = %d, want 0 (never resets -> no retry hint)", got)
|
||||
}
|
||||
now := time.Now().Unix()
|
||||
for _, tc := range []struct {
|
||||
period string
|
||||
hours int64
|
||||
want int64
|
||||
}{
|
||||
{"hour", 0, 3600},
|
||||
{"week", 0, 7 * 24 * 3600},
|
||||
{"month", 0, 30 * 24 * 3600},
|
||||
{"nhour", 6, 6 * 3600},
|
||||
} {
|
||||
got := AutoSecondsToReset(tc.period, tc.hours)
|
||||
if got <= 0 || got > tc.want {
|
||||
t.Errorf("AutoSecondsToReset(%q,%d) = %d, want in (0,%d]", tc.period, tc.hours, got, tc.want)
|
||||
}
|
||||
// must never exceed the window itself
|
||||
if got < tc.want-now%3600 {
|
||||
t.Logf("note: %q hint %ds < remaining-window %ds (rounds to hour boundary)", tc.period, got, tc.want-now%3600)
|
||||
}
|
||||
if got > tc.want {
|
||||
t.Errorf("hint %d exceeds window %d", got, tc.want)
|
||||
}
|
||||
}
|
||||
_ = now
|
||||
}
|
||||
|
||||
func TestAutoSecondsToResetNeverExceedsWindow(t *testing.T) {
|
||||
// "nhour" with a tiny window must not hand out a longer wait than the
|
||||
// window itself (which would stall a client past its own reset)
|
||||
for hours := int64(1); hours <= 48; hours++ {
|
||||
got := AutoSecondsToReset("nhour", hours)
|
||||
if got <= 0 || got > hours*3600 {
|
||||
t.Errorf("nhour/%d = %d, want in (0,%d]", hours, got, hours*3600)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// ---- config validation ----
|
||||
|
||||
func TestKeyQuotaValidate(t *testing.T) {
|
||||
cases := []struct {
|
||||
name string
|
||||
q config.KeyQuota
|
||||
wantErr bool
|
||||
}{
|
||||
{"unlimited", config.KeyQuota{}, false},
|
||||
{"tokens hourly", config.KeyQuota{TokenQuota: 1000, Period: "hour"}, false},
|
||||
{"tokens n-hour", config.KeyQuota{TokenQuota: 1000, Period: "nhour", Hours: 6}, false},
|
||||
{"reqs weekly", config.KeyQuota{ReqQuota: 100, Period: "week"}, false},
|
||||
{"no period never resets", config.KeyQuota{TokenQuota: 1000}, false},
|
||||
{"negative tokens", config.KeyQuota{TokenQuota: -1}, true},
|
||||
{"negative reqs", config.KeyQuota{ReqQuota: -1}, true},
|
||||
{"negative hours", config.KeyQuota{TokenQuota: 5, Hours: -1}, true},
|
||||
{"typo period becomes all-time if accepted", config.KeyQuota{TokenQuota: 5, Period: "houre"}, true},
|
||||
{"nhour without hours", config.KeyQuota{TokenQuota: 5, Period: "nhour"}, true},
|
||||
{"nhour with 0 hours", config.KeyQuota{TokenQuota: 5, Period: "nhour", Hours: 0}, true},
|
||||
{"junk period with only req quota", config.KeyQuota{ReqQuota: 5, Period: "daily"}, true},
|
||||
{"junk period with no caps is irrelevant", config.KeyQuota{Period: "daily"}, false},
|
||||
}
|
||||
for _, tc := range cases {
|
||||
err := tc.q.Validate()
|
||||
if tc.wantErr && err == nil {
|
||||
t.Errorf("%s: expected error, got nil", tc.name)
|
||||
}
|
||||
if !tc.wantErr && err != nil {
|
||||
t.Errorf("%s: unexpected error: %v", tc.name, err)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestNormalizeRole(t *testing.T) {
|
||||
if got := config.NormalizeRole(""); got != "user" {
|
||||
t.Errorf("empty role = %q, want user", got)
|
||||
}
|
||||
if got := config.NormalizeRole("admin"); got != "admin" {
|
||||
t.Errorf("admin role = %q, want admin", got)
|
||||
}
|
||||
if got := config.NormalizeRole("root"); got != "user" {
|
||||
t.Errorf("unknown role = %q, want user (never escalate)", got)
|
||||
}
|
||||
}
|
||||
|
||||
// Retention must hold for the lazy pinned buckets too, or a long-lived key
|
||||
// would grow without bound.
|
||||
func TestPinnedBucketRespectsRetention(t *testing.T) {
|
||||
s := NewStats(100)
|
||||
_ = s.KeyWindowModelTokens("keyA", "m1", "srcX", 3600) // opt in
|
||||
nowH := time.Now().Unix() / 3600
|
||||
for h := int64(0); h < quotaRetentionHours+50; h++ {
|
||||
s.Record(Req{Time: (nowH - h) * 3600 * 1000, Key: "keyA", Model: "m1", Source: "srcX",
|
||||
Prompt: 10, Compl: 10, OK: true, Status: 200})
|
||||
}
|
||||
s.mu.Lock()
|
||||
n := len(s.keyModelHour["keyA"]["srcX::m1"])
|
||||
s.mu.Unlock()
|
||||
if n > quotaRetentionHours {
|
||||
t.Errorf("pinned bucket holds %d hours, want <= %d", n, quotaRetentionHours)
|
||||
}
|
||||
}
|
||||
327
internal/gateway/key_quota_wiring_test.go
Normal file
327
internal/gateway/key_quota_wiring_test.go
Normal file
@ -0,0 +1,327 @@
|
||||
package gateway
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"llmsproxy/internal/config"
|
||||
"llmsproxy/internal/core"
|
||||
)
|
||||
|
||||
// quotaGateway builds a gateway with two user keys that differ only in their
|
||||
// configured caps, against one mock upstream that reports 4 tokens per call.
|
||||
func quotaGateway(t *testing.T, keyA, keyB config.GWKey) (*Gateway, *upstreamCtrl) {
|
||||
t.Helper()
|
||||
ctrl := &upstreamCtrl{}
|
||||
up := upstream(t, ctrl)
|
||||
t.Cleanup(up.Close)
|
||||
|
||||
td := t.TempDir()
|
||||
cfgPath := td + "/config.yaml"
|
||||
if err := os.WriteFile(cfgPath, []byte("listen: :0"), 0o644); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
cfg := &config.Config{
|
||||
Path: cfgPath,
|
||||
AdapterDir: td + "/adapters",
|
||||
RuntimeFile: td + "/runtime.json",
|
||||
Keys: []config.GWKey{keyA, keyB},
|
||||
Sources: []config.Source{{
|
||||
Name: "up", BaseURL: up.URL, Adapter: "openai",
|
||||
Models: []config.Model{{ID: "m1", Priority: 100}},
|
||||
}},
|
||||
}
|
||||
if err := cfg.ApplyDefaults(); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
c, err := core.NewFromConfig(cfg)
|
||||
if err != nil {
|
||||
t.Fatalf("core: %v", err)
|
||||
}
|
||||
t.Cleanup(c.Close)
|
||||
g, err := New(c, []string{keyA.Key, keyB.Key})
|
||||
if err != nil {
|
||||
t.Fatalf("gateway: %v", err)
|
||||
}
|
||||
return g, ctrl
|
||||
}
|
||||
|
||||
// chatAs issues a chat completion with a specific gateway key and returns the
|
||||
// recorder plus the decoded error code.
|
||||
func chatAs(t *testing.T, g *Gateway, key, model string) (*httptest.ResponseRecorder, string) {
|
||||
t.Helper()
|
||||
req, _ := http.NewRequest("POST", "/v1/chat/completions",
|
||||
strings.NewReader(fmt.Sprintf(`{"model":%q,"messages":[{"role":"user","content":"hi"}]}`, model)))
|
||||
req.Header.Set("Authorization", "Bearer "+key)
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
rr := httptest.NewRecorder()
|
||||
g.Handler().ServeHTTP(rr, req)
|
||||
code := ""
|
||||
if rr.Code != 200 {
|
||||
var e struct {
|
||||
Error struct {
|
||||
Type string `json:"type"`
|
||||
} `json:"error"`
|
||||
}
|
||||
_ = json.Unmarshal(rr.Body.Bytes(), &e)
|
||||
code = e.Error.Type
|
||||
}
|
||||
return rr, code
|
||||
}
|
||||
|
||||
// A key whose token budget is spent must be refused, and the refusal must be a
|
||||
// 429 the client can retry after the window rolls over — not a 403 that reads
|
||||
// as "this key may never use this model".
|
||||
func TestKeyTokenQuotaBlocksWithRetryAfter(t *testing.T) {
|
||||
g, _ := quotaGateway(t,
|
||||
config.GWKey{Key: "sk-a", Role: "user", TokenQuota: 8, Period: "hour",
|
||||
Models: []config.ModelScope{{Model: "m1"}}},
|
||||
config.GWKey{Key: "sk-b", Role: "user", Name: "b"},
|
||||
)
|
||||
// 4 tokens per call, budget 8 -> the third call crosses it
|
||||
for i := 1; i <= 2; i++ {
|
||||
rr, _ := chatAs(t, g, "sk-a", "m1")
|
||||
if rr.Code != 200 {
|
||||
t.Fatalf("call %d: want 200, got %d (%s)", i, rr.Code, rr.Body.String())
|
||||
}
|
||||
}
|
||||
rr, code := chatAs(t, g, "sk-a", "m1")
|
||||
if rr.Code != http.StatusTooManyRequests {
|
||||
t.Fatalf("third call: want 429, got %d (%s)", rr.Code, rr.Body.String())
|
||||
}
|
||||
if code != "rate_limit_exceeded" {
|
||||
t.Errorf("error code = %q, want rate_limit_exceeded", code)
|
||||
}
|
||||
if ra := rr.Header().Get("Retry-After"); ra == "" {
|
||||
t.Error("429 must carry Retry-After so a client knows when to come back")
|
||||
} else if n := mustAtoi(t, ra); n <= 0 || n > 3600 {
|
||||
t.Errorf("Retry-After = %q, want a positive value within the 1h window", ra)
|
||||
}
|
||||
if !strings.Contains(rr.Body.String(), "quota") {
|
||||
t.Errorf("body should name the quota: %s", rr.Body.String())
|
||||
}
|
||||
}
|
||||
|
||||
// The quota is per key: exhausting one key's budget must not affect another
|
||||
// key that is allowed the same model.
|
||||
func TestKeyTokenQuotaIsIsolatedPerKey(t *testing.T) {
|
||||
g, _ := quotaGateway(t,
|
||||
config.GWKey{Key: "sk-a", Role: "user", TokenQuota: 4, Period: "hour",
|
||||
Models: []config.ModelScope{{Model: "m1"}}},
|
||||
config.GWKey{Key: "sk-b", Role: "user", TokenQuota: 1000, Period: "hour",
|
||||
Models: []config.ModelScope{{Model: "m1"}}},
|
||||
)
|
||||
rr, _ := chatAs(t, g, "sk-a", "m1")
|
||||
if rr.Code != 200 {
|
||||
t.Fatalf("keyA first call: want 200, got %d", rr.Code)
|
||||
}
|
||||
if rr, code := chatAs(t, g, "sk-a", "m1"); rr.Code != http.StatusTooManyRequests {
|
||||
t.Fatalf("keyA second call: want 429, got %d (%s)", rr.Code, code)
|
||||
}
|
||||
// keyB has its own budget and must still be served
|
||||
if rr, _ := chatAs(t, g, "sk-b", "m1"); rr.Code != 200 {
|
||||
t.Fatalf("keyB must be unaffected by keyA's exhausted quota, got %d", rr.Code)
|
||||
}
|
||||
}
|
||||
|
||||
// A per-model quota inside the scope must also be per key: one key burning a
|
||||
// shared model budget must not lock the other key out.
|
||||
func TestScopeModelTokenQuotaIsolatedPerKey(t *testing.T) {
|
||||
g, _ := quotaGateway(t,
|
||||
config.GWKey{Key: "sk-a", Role: "user", Models: []config.ModelScope{
|
||||
{Model: "m1", TokenQuota: 4, Period: "hour"},
|
||||
}},
|
||||
config.GWKey{Key: "sk-b", Role: "user", Models: []config.ModelScope{
|
||||
{Model: "m1", TokenQuota: 4, Period: "hour"},
|
||||
}},
|
||||
)
|
||||
if rr, _ := chatAs(t, g, "sk-a", "m1"); rr.Code != 200 {
|
||||
t.Fatalf("keyA first call: want 200, got %d", rr.Code)
|
||||
}
|
||||
if rr, _ := chatAs(t, g, "sk-a", "m1"); rr.Code != http.StatusTooManyRequests {
|
||||
t.Fatalf("keyA second call: want 429 (its own model quota), got %d", rr.Code)
|
||||
}
|
||||
// keyB shares the model but not the counter — it must be served twice too
|
||||
if rr, _ := chatAs(t, g, "sk-b", "m1"); rr.Code != 200 {
|
||||
t.Fatalf("keyB first call: want 200, got %d", rr.Code)
|
||||
}
|
||||
if rr, _ := chatAs(t, g, "sk-b", "m1"); rr.Code != http.StatusTooManyRequests {
|
||||
t.Fatalf("keyB second call: want 429, got %d", rr.Code)
|
||||
}
|
||||
}
|
||||
|
||||
// An admin key is never capped: a budget set on an admin key must not lock the
|
||||
// operator out of the gateway they administer.
|
||||
func TestAdminKeyIsNeverQuotaCapped(t *testing.T) {
|
||||
g, _ := quotaGateway(t,
|
||||
config.GWKey{Key: "sk-admin", Role: "admin", TokenQuota: 1, Period: "hour", ReqQuota: 1},
|
||||
config.GWKey{Key: "sk-b", Role: "user"},
|
||||
)
|
||||
for i := 1; i <= 3; i++ {
|
||||
if rr, _ := chatAs(t, g, "sk-admin", "m1"); rr.Code != 200 {
|
||||
t.Fatalf("admin call %d: want 200 (admin keys are uncapped), got %d", i, rr.Code)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// A request quota caps sustained volume; RPM-style bursts are not its job but
|
||||
// the count must still stop the key.
|
||||
func TestKeyRequestQuotaBlocks(t *testing.T) {
|
||||
g, ctrl := quotaGateway(t,
|
||||
config.GWKey{Key: "sk-a", Role: "user", ReqQuota: 2, Period: "hour",
|
||||
Models: []config.ModelScope{{Model: "m1"}}},
|
||||
config.GWKey{Key: "sk-b", Role: "user"},
|
||||
)
|
||||
for i := 1; i <= 2; i++ {
|
||||
if rr, _ := chatAs(t, g, "sk-a", "m1"); rr.Code != 200 {
|
||||
t.Fatalf("call %d: want 200, got %d", i, rr.Code)
|
||||
}
|
||||
}
|
||||
rr, code := chatAs(t, g, "sk-a", "m1")
|
||||
if rr.Code != http.StatusTooManyRequests || code != "rate_limit_exceeded" {
|
||||
t.Fatalf("third call: want 429/rate_limit_exceeded, got %d/%s", rr.Code, code)
|
||||
}
|
||||
if ctrl.hits != 2 {
|
||||
t.Errorf("upstream saw %d calls, want 2 (a refused request must not reach upstream)", ctrl.hits)
|
||||
}
|
||||
}
|
||||
|
||||
// A model outside the scope is still 403, not 429: retrying cannot help.
|
||||
func TestModelOutsideScopeStaysForbidden(t *testing.T) {
|
||||
g, _ := quotaGateway(t,
|
||||
config.GWKey{Key: "sk-a", Role: "user", TokenQuota: 1000, Period: "hour",
|
||||
Models: []config.ModelScope{{Model: "other-model"}}},
|
||||
config.GWKey{Key: "sk-b", Role: "user"},
|
||||
)
|
||||
rr, code := chatAs(t, g, "sk-a", "m1")
|
||||
if rr.Code != http.StatusForbidden {
|
||||
t.Fatalf("want 403, got %d (%s)", rr.Code, rr.Body.String())
|
||||
}
|
||||
if code != "model_not_allowed" {
|
||||
t.Errorf("error code = %q, want model_not_allowed", code)
|
||||
}
|
||||
if rr.Header().Get("Retry-After") != "" {
|
||||
t.Error("a 403 must not advertise a retry time")
|
||||
}
|
||||
}
|
||||
|
||||
// An unlimited key (no caps at all) is never rejected — the default config
|
||||
// must keep working exactly as before.
|
||||
func TestKeyWithoutCapsIsUnlimited(t *testing.T) {
|
||||
g, _ := quotaGateway(t,
|
||||
config.GWKey{Key: "sk-a", Role: "user", Models: []config.ModelScope{{Model: "m1"}}},
|
||||
config.GWKey{Key: "sk-b", Role: "user"},
|
||||
)
|
||||
for i := 1; i <= 5; i++ {
|
||||
if rr, _ := chatAs(t, g, "sk-a", "m1"); rr.Code != 200 {
|
||||
t.Fatalf("call %d: want 200 (no caps = unlimited), got %d", i, rr.Code)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func mustAtoi(t *testing.T, s string) int64 {
|
||||
t.Helper()
|
||||
var n int64
|
||||
if _, err := fmt.Sscanf(s, "%d", &n); err != nil {
|
||||
t.Fatalf("Retry-After %q is not an integer: %v", s, err)
|
||||
}
|
||||
return n
|
||||
}
|
||||
|
||||
// quotaCtx builds a request context carrying a gateway key, as the auth
|
||||
// middleware would.
|
||||
func quotaCtx(t *testing.T, g *Gateway, key string) context.Context {
|
||||
t.Helper()
|
||||
return withAuth(context.Background(), key, "user")
|
||||
}
|
||||
|
||||
func nowMSOffset(sec int64) int64 {
|
||||
return time.Now().Add(time.Duration(sec) * time.Second).UnixMilli()
|
||||
}
|
||||
|
||||
// The key-wide quota and the AUTO slot quota are different limits with
|
||||
// different scopes: the slot quota is gateway-wide (a shared upstream budget),
|
||||
// the key quota belongs to one caller. When both are exhausted the caller must
|
||||
// see the KEY quota, because that is the one it can act on — the slot quota
|
||||
// would otherwise be reported as "no capacity", which reads like an outage.
|
||||
func TestKeyQuotaWinsOverSlotQuota(t *testing.T) {
|
||||
up := upstream(t, &upstreamCtrl{})
|
||||
defer up.Close()
|
||||
g := newQuotaGW(t, up.URL,
|
||||
config.GWKey{Key: "sk-a", Role: "user", TokenQuota: 4, Period: "hour",
|
||||
Models: []config.ModelScope{{Model: "AUTO"}}})
|
||||
ctx := quotaCtx(t, g, "sk-a")
|
||||
|
||||
// exhaust the key first
|
||||
g.stats.Record(Req{Time: time.Now().UnixMilli(), Key: keyID("sk-a"), Model: "m1",
|
||||
Source: "up", Prompt: 100, Compl: 100, OK: true, Status: 200})
|
||||
rr, code := chatAs(t, g, "sk-a", "AUTO")
|
||||
if rr.Code != http.StatusTooManyRequests {
|
||||
t.Fatalf("want 429 from the key quota, got %d (%s)", rr.Code, rr.Body.String())
|
||||
}
|
||||
if code != "rate_limit_exceeded" {
|
||||
t.Errorf("code = %q, want rate_limit_exceeded", code)
|
||||
}
|
||||
_ = ctx
|
||||
}
|
||||
|
||||
// A key with NO caps must never be blocked by the slot quota check reaching it
|
||||
// through the scope path: scopeTokens on an uncapped key returns 0 usage and
|
||||
// the limit check must treat that as "no cap", not "exhausted".
|
||||
func TestUncappedKeyNeverBlockedByEmptyScope(t *testing.T) {
|
||||
up := upstream(t, &upstreamCtrl{})
|
||||
defer up.Close()
|
||||
g := newQuotaGW(t, up.URL,
|
||||
config.GWKey{Key: "sk-a", Role: "user", Models: []config.ModelScope{{Model: "AUTO"}}})
|
||||
for i := 0; i < 5; i++ {
|
||||
if rr, _ := chatAs(t, g, "sk-a", "AUTO"); rr.Code != 200 {
|
||||
t.Fatalf("call %d: want 200 (AUTO scope with no quota), got %d (%s)", i, rr.Code, rr.Body.String())
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// newQuotaGW builds a gateway with one mock upstream serving model m1 and the
|
||||
// given keys.
|
||||
func newQuotaGW(t *testing.T, upURL string, keys ...config.GWKey) *Gateway {
|
||||
t.Helper()
|
||||
td := t.TempDir()
|
||||
cfgPath := filepath.Join(td, "config.yaml")
|
||||
if err := os.WriteFile(cfgPath, []byte("listen: :0"), 0o644); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
cfg := &config.Config{
|
||||
Path: cfgPath,
|
||||
AdapterDir: filepath.Join(td, "adapters"),
|
||||
RuntimeFile: filepath.Join(td, "runtime.json"),
|
||||
Keys: keys,
|
||||
Sources: []config.Source{{
|
||||
Name: "up", BaseURL: upURL, Adapter: "openai",
|
||||
Models: []config.Model{{ID: "m1", Priority: 100}},
|
||||
}},
|
||||
}
|
||||
if err := cfg.ApplyDefaults(); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
c, err := core.NewFromConfig(cfg)
|
||||
if err != nil {
|
||||
t.Fatalf("core: %v", err)
|
||||
}
|
||||
t.Cleanup(c.Close)
|
||||
secrets := make([]string, 0, len(keys))
|
||||
for _, k := range keys {
|
||||
secrets = append(secrets, k.Key)
|
||||
}
|
||||
g, err := New(c, secrets)
|
||||
if err != nil {
|
||||
t.Fatalf("gateway: %v", err)
|
||||
}
|
||||
return g
|
||||
}
|
||||
@ -39,23 +39,27 @@ func (g *Gateway) handleKeysAPI(w http.ResponseWriter, r *http.Request) {
|
||||
writeJSON(w, http.StatusOK, map[string]interface{}{"keys": g.core.ListKeys()})
|
||||
case http.MethodPost:
|
||||
var body struct {
|
||||
Name string `json:"name"`
|
||||
Role string `json:"role"`
|
||||
Models []config.ModelScope `json:"models"`
|
||||
Note string `json:"note"`
|
||||
Name string `json:"name"`
|
||||
Role string `json:"role"`
|
||||
Models []config.ModelScope `json:"models"`
|
||||
Note string `json:"note"`
|
||||
TokenQuota *int64 `json:"token_quota"`
|
||||
ReqQuota *int64 `json:"req_quota"`
|
||||
Period *string `json:"period"`
|
||||
Hours *int64 `json:"hours"`
|
||||
}
|
||||
if err := json.NewDecoder(r.Body).Decode(&body); err != nil {
|
||||
writeError(w, http.StatusBadRequest, "invalid_request", "invalid json: "+err.Error())
|
||||
return
|
||||
}
|
||||
if body.Role == "" {
|
||||
body.Role = "user"
|
||||
body.Role = config.NormalizeRole(body.Role)
|
||||
q := config.KeyQuota{
|
||||
TokenQuota: optInt64(body.TokenQuota),
|
||||
ReqQuota: optInt64(body.ReqQuota),
|
||||
Period: optString(body.Period),
|
||||
Hours: optInt64(body.Hours),
|
||||
}
|
||||
if body.Role != "admin" && body.Role != "user" {
|
||||
writeError(w, http.StatusBadRequest, "invalid_request", "role must be admin or user")
|
||||
return
|
||||
}
|
||||
rec, err := g.core.CreateKey(body.Name, body.Role, body.Models, body.Note)
|
||||
rec, err := g.core.CreateKeyWithQuota(body.Name, body.Role, body.Models, body.Note, q)
|
||||
if err != nil {
|
||||
writeError(w, http.StatusBadRequest, "key_error", err.Error())
|
||||
return
|
||||
@ -67,16 +71,33 @@ func (g *Gateway) handleKeysAPI(w http.ResponseWriter, r *http.Request) {
|
||||
return
|
||||
}
|
||||
var body struct {
|
||||
Name string `json:"name"`
|
||||
Role string `json:"role"`
|
||||
Models []config.ModelScope `json:"models"`
|
||||
Note string `json:"note"`
|
||||
Name string `json:"name"`
|
||||
Role string `json:"role"`
|
||||
Models []config.ModelScope `json:"models"`
|
||||
Note string `json:"note"`
|
||||
TokenQuota *int64 `json:"token_quota"`
|
||||
ReqQuota *int64 `json:"req_quota"`
|
||||
Period *string `json:"period"`
|
||||
Hours *int64 `json:"hours"`
|
||||
}
|
||||
if err := json.NewDecoder(r.Body).Decode(&body); err != nil {
|
||||
writeError(w, http.StatusBadRequest, "invalid_request", "invalid json: "+err.Error())
|
||||
return
|
||||
}
|
||||
rec, err := g.core.UpdateKey(path, body.Name, body.Role, body.Models, body.Note)
|
||||
// Quota fields are pointers so "absent" is distinguishable from
|
||||
// "set to 0": omitting them leaves the stored caps alone, while
|
||||
// sending 0 explicitly lifts a cap. Without this, editing only the
|
||||
// model scope would silently clear a key's budget.
|
||||
var q *config.KeyQuota
|
||||
if body.TokenQuota != nil || body.ReqQuota != nil || body.Period != nil || body.Hours != nil {
|
||||
q = &config.KeyQuota{
|
||||
TokenQuota: optInt64(body.TokenQuota),
|
||||
ReqQuota: optInt64(body.ReqQuota),
|
||||
Period: optString(body.Period),
|
||||
Hours: optInt64(body.Hours),
|
||||
}
|
||||
}
|
||||
rec, err := g.core.UpdateKeyWithQuota(path, body.Name, body.Role, body.Models, body.Note, q)
|
||||
if err != nil {
|
||||
writeError(w, http.StatusBadRequest, "key_error", err.Error())
|
||||
return
|
||||
@ -106,6 +127,22 @@ func (g *Gateway) handleKeysAPI(w http.ResponseWriter, r *http.Request) {
|
||||
}
|
||||
}
|
||||
|
||||
// optInt64 dereferences an optional quota field, treating absent as 0.
|
||||
func optInt64(p *int64) int64 {
|
||||
if p == nil {
|
||||
return 0
|
||||
}
|
||||
return *p
|
||||
}
|
||||
|
||||
// optString dereferences an optional quota field, treating absent as "".
|
||||
func optString(p *string) string {
|
||||
if p == nil {
|
||||
return ""
|
||||
}
|
||||
return *p
|
||||
}
|
||||
|
||||
// handleKeyMe returns the authenticated key's own record (users see only
|
||||
// themselves; admins can use this as a convenience too).
|
||||
func (g *Gateway) handleKeyMe(w http.ResponseWriter, r *http.Request) {
|
||||
|
||||
@ -84,6 +84,21 @@ type Stats struct {
|
||||
auditPath string
|
||||
replayPartial bool // aggregates built from a bounded audit tail
|
||||
modelHour map[string]map[int64]int64 // model -> unix-hour bucket -> tokens
|
||||
// keyModelHour buckets the same tokens as modelHour but keyed by
|
||||
// (gateway key id, model) so per-key model quotas are isolated from each
|
||||
// other. modelHour stays key-blind on purpose: it backs the AUTO slot
|
||||
// quota, which is a gateway-wide limit on a slot, not a per-key one.
|
||||
keyModelHour map[string]map[string]map[int64]int64 // key -> model -> unix-hour -> tokens
|
||||
// keyHour buckets a key's total tokens per unix hour, backing the
|
||||
// whole-key quota (all models of one key share one budget).
|
||||
keyHour map[string]map[int64]int64
|
||||
// keyReqHour buckets a key's request count per unix hour, backing the
|
||||
// key-wide request quota. Counted for every request type (chat, stream,
|
||||
// image) including failed ones, so a failing client cannot loop for free.
|
||||
keyReqHour map[string]map[int64]int64
|
||||
// pinnedWanted records which (key, model) pairs have a quota that pins a
|
||||
// source, so the pinned bucket is maintained only where it is read.
|
||||
pinnedWanted map[string]bool
|
||||
}
|
||||
|
||||
const hourSec = 3600
|
||||
@ -116,14 +131,18 @@ func NewStats(maxRecords int) *Stats {
|
||||
maxRecords = defaultRingSize
|
||||
}
|
||||
return &Stats{
|
||||
byKey: map[string]*Stat{},
|
||||
byModel: map[string]*Stat{},
|
||||
bySrc: map[string]*Stat{},
|
||||
byKeyModel: map[string]map[string]*Stat{},
|
||||
byKeySrc: map[string]map[string]*Stat{},
|
||||
byStatus: map[int]*Stat{},
|
||||
modelHour: map[string]map[int64]int64{},
|
||||
maxRecs: maxRecords,
|
||||
byKey: map[string]*Stat{},
|
||||
byModel: map[string]*Stat{},
|
||||
bySrc: map[string]*Stat{},
|
||||
byKeyModel: map[string]map[string]*Stat{},
|
||||
byKeySrc: map[string]map[string]*Stat{},
|
||||
byStatus: map[int]*Stat{},
|
||||
modelHour: map[string]map[int64]int64{},
|
||||
keyModelHour: map[string]map[string]map[int64]int64{},
|
||||
keyHour: map[string]map[int64]int64{},
|
||||
keyReqHour: map[string]map[int64]int64{},
|
||||
pinnedWanted: map[string]bool{},
|
||||
maxRecs: maxRecords,
|
||||
}
|
||||
}
|
||||
|
||||
@ -318,6 +337,9 @@ func (s *Stats) Record(r Req) {
|
||||
// aggregateLocked folds r into every aggregate row and the quota window
|
||||
// bucket. Caller must hold s.mu.
|
||||
func (s *Stats) aggregateLocked(r Req) {
|
||||
if r.Key != "" {
|
||||
s.addKeyReqLocked(r.Key, (r.Time/1000)/hourSec, 1)
|
||||
}
|
||||
inc(s.byKey, r.Key, r)
|
||||
if r.Model != "" {
|
||||
inc(s.byModel, r.Model, r)
|
||||
@ -362,6 +384,12 @@ func (s *Stats) aggregateLocked(r Req) {
|
||||
s.modelHour[key] = hm
|
||||
}
|
||||
hm[h] += tok
|
||||
// per-key buckets: same token split, but isolated per key so one
|
||||
// key's quota cannot be drained by another key's traffic
|
||||
if r.Key != "" {
|
||||
s.addKeyTokenLocked(r.Key, r.Model, r.Source, h, tok)
|
||||
s.addKeyHourLocked(r.Key, h, tok)
|
||||
}
|
||||
// retention: 24*40 = 960 hourly buckets ≈ 40 days of history (covers
|
||||
// the longest "month" quota window)
|
||||
if len(hm) > 24*40 {
|
||||
@ -374,6 +402,185 @@ func (s *Stats) aggregateLocked(r Req) {
|
||||
}
|
||||
}
|
||||
|
||||
// quotaRetentionHours is how much hourly history the per-key quota buckets
|
||||
// keep. It matches the modelHour retention (40 days) so the longest "month"
|
||||
// window is fully covered after a restart, and it is deliberately applied per
|
||||
// key so an idle key's buckets are reclaimed instead of pinning memory.
|
||||
const quotaRetentionHours = 24 * 40
|
||||
|
||||
// addKeyTokenLocked adds tok to one key's (model[, source]) hourly bucket.
|
||||
//
|
||||
// The bare-model bucket is always written: it is what a scope entry without a
|
||||
// source pin reads, and writing it is what makes that quota work at all (every
|
||||
// recorded request carries a resolved source, so a source::model-only bucket
|
||||
// would leave an unpinned quota reading an empty bucket forever).
|
||||
//
|
||||
// The pinned "source::model" bucket is written only for (key, model) pairs
|
||||
// some scope entry actually pins a source for. Maintaining it unconditionally
|
||||
// doubles the bucket count — measured at 26.7 MB for 20 keys x 8 models x 3
|
||||
// sources at full retention, against a documented ~32 MB total budget — to
|
||||
// serve a lookup that almost nobody performs.
|
||||
func (s *Stats) addKeyTokenLocked(key, model, source string, h, tok int64) {
|
||||
if model == "" {
|
||||
return
|
||||
}
|
||||
byModel := s.keyModelHour[key]
|
||||
if byModel == nil {
|
||||
byModel = map[string]map[int64]int64{}
|
||||
s.keyModelHour[key] = byModel
|
||||
}
|
||||
addBucket(byModel, model, h, tok)
|
||||
if source == "" {
|
||||
return
|
||||
}
|
||||
pinned := source + "::" + model
|
||||
if pinned == model || !s.pinnedWanted[key+"\x00"+model] {
|
||||
return
|
||||
}
|
||||
addBucket(byModel, pinned, h, tok)
|
||||
}
|
||||
|
||||
// addBucket adds tok to one named hourly bucket and enforces retention.
|
||||
func addBucket(byModel map[string]map[int64]int64, name string, h, tok int64) {
|
||||
hm := byModel[name]
|
||||
if hm == nil {
|
||||
hm = map[int64]int64{}
|
||||
byModel[name] = hm
|
||||
}
|
||||
hm[h] += tok
|
||||
if len(hm) > quotaRetentionHours {
|
||||
for k := range hm {
|
||||
if k < h-quotaRetentionHours {
|
||||
delete(hm, k)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// wantPinnedBuckets starts maintaining the pinned (source::model) bucket for a
|
||||
// key's model, so a quota that pins a source counts traffic from the moment it
|
||||
// is configured. Usage recorded before that cannot be attributed to a source
|
||||
// after the fact, so a pinned quota may under-count its first window.
|
||||
func (s *Stats) wantPinnedBuckets(key, model string) {
|
||||
if key == "" || model == "" {
|
||||
return
|
||||
}
|
||||
s.pinnedWanted[key+"\x00"+model] = true
|
||||
}
|
||||
|
||||
// addKeyHourLocked adds tok to one key's all-model hourly total.
|
||||
func (s *Stats) addKeyHourLocked(key string, h, tok int64) {
|
||||
hm := s.keyHour[key]
|
||||
if hm == nil {
|
||||
hm = map[int64]int64{}
|
||||
s.keyHour[key] = hm
|
||||
}
|
||||
hm[h] += tok
|
||||
if len(hm) > quotaRetentionHours {
|
||||
for k := range hm {
|
||||
if k < h-quotaRetentionHours {
|
||||
delete(hm, k)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// KeyWindowReqs returns how many requests one key issued within sec seconds;
|
||||
// sec <= 0 means all retained history.
|
||||
func (s *Stats) KeyWindowReqs(key string, sec int64) int64 {
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
return sumBuckets(s.keyReqHour[key], time.Now().Unix(), sec)
|
||||
}
|
||||
|
||||
// addKeyReqLocked adds n requests to a key's hourly count bucket.
|
||||
func (s *Stats) addKeyReqLocked(key string, h, n int64) {
|
||||
hm := s.keyReqHour[key]
|
||||
if hm == nil {
|
||||
hm = map[int64]int64{}
|
||||
s.keyReqHour[key] = hm
|
||||
}
|
||||
hm[h] += n
|
||||
if len(hm) > quotaRetentionHours {
|
||||
for k := range hm {
|
||||
if k < h-quotaRetentionHours {
|
||||
delete(hm, k)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// sumBuckets totals the hourly buckets inside the window ending at `now`
|
||||
// (unix seconds). sec <= 0 means "all retained history" (no reset).
|
||||
//
|
||||
// It only visits the hour buckets a window can actually cover. The naive full
|
||||
// scan measured 5.9us/op at the full 40-day retention, and a per-key quota
|
||||
// check runs on every request (2-3 of these), so the scan cost lands directly
|
||||
// on the request path; visiting at most sec/3600 + 1 buckets makes it
|
||||
// proportional to the window instead of to the retained history. The map keys
|
||||
// are whole unix hours, so the covered range is exact rather than a sample.
|
||||
func sumBuckets(hm map[int64]int64, now, sec int64) int64 {
|
||||
if len(hm) == 0 {
|
||||
return 0
|
||||
}
|
||||
if sec <= 0 {
|
||||
var total int64
|
||||
for _, v := range hm {
|
||||
total += v
|
||||
}
|
||||
return total
|
||||
}
|
||||
nowHour := now / hourSec
|
||||
// A bucket at hour H covers [H*3600, (H+1)*3600). The pre-existing
|
||||
// WindowTokens counts a bucket when h*3600 >= now-sec, so the first hour
|
||||
// that qualifies is ceil((now-sec)/3600) — NOT floor, which would pull in
|
||||
// one bucket too many (caught by TestSumBucketsMatchesFullScan: a 30-day
|
||||
// window reported 8878 vs 8649).
|
||||
firstHour := (now - sec + hourSec - 1) / hourSec
|
||||
if firstHour < 0 {
|
||||
firstHour = 0
|
||||
}
|
||||
var total int64
|
||||
for h := firstHour; h <= nowHour; h++ {
|
||||
total += hm[h]
|
||||
}
|
||||
return total
|
||||
}
|
||||
|
||||
// KeyWindowModelTokens returns the tokens one key consumed for one model
|
||||
// (optionally pinned to a source) within sec seconds; sec <= 0 means all
|
||||
// retained history. Unlike WindowTokens this is isolated per key, which is
|
||||
// what a per-key model quota needs.
|
||||
//
|
||||
// A source pin matches the pinned bucket exactly. Without a pin the call
|
||||
// totals the key's bare-model bucket, which counts traffic on every source —
|
||||
// a quota on "this model" should not stop applying just because the request
|
||||
// happened to be served by a different upstream.
|
||||
func (s *Stats) KeyWindowModelTokens(key, model, source string, sec int64) int64 {
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
byModel := s.keyModelHour[key]
|
||||
if len(byModel) == 0 {
|
||||
return 0
|
||||
}
|
||||
mk := model
|
||||
if source != "" {
|
||||
mk = source + "::" + model
|
||||
// from now on this pair keeps a pinned bucket, so a quota configured
|
||||
// against this source counts the traffic it is meant to cap
|
||||
s.wantPinnedBuckets(key, model)
|
||||
}
|
||||
return sumBuckets(byModel[mk], time.Now().Unix(), sec)
|
||||
}
|
||||
|
||||
// KeyWindowTokens returns one key's total tokens across all models within sec
|
||||
// seconds; sec <= 0 means all retained history. It backs the whole-key quota.
|
||||
func (s *Stats) KeyWindowTokens(key string, sec int64) int64 {
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
return sumBuckets(s.keyHour[key], time.Now().Unix(), sec)
|
||||
}
|
||||
|
||||
// rotateAuditLocked renames the audit file to <path>.<unix>.old once it
|
||||
// exceeds auditRotateBytes and prunes old files beyond auditKeepOld, keeping
|
||||
// the newest ones. Caller must hold s.mu.
|
||||
@ -464,6 +671,27 @@ func AutoPeriodSeconds(period string, hours int64) int64 {
|
||||
return 0
|
||||
}
|
||||
|
||||
// AutoSecondsToReset returns how many seconds remain until the quota window
|
||||
// rolls over, for the Retry-After header on a 429. Buckets are whole unix
|
||||
// hours, so the value is rounded up to the next hour boundary and never
|
||||
// exceeds one window. 0 means "unknown" (never resets, or no period set), in
|
||||
// which case the caller should not advertise a retry time.
|
||||
func AutoSecondsToReset(period string, hours int64) int64 {
|
||||
sec := AutoPeriodSeconds(period, hours)
|
||||
if sec <= 0 {
|
||||
return 0
|
||||
}
|
||||
now := time.Now().Unix()
|
||||
// window = the last sec seconds, bucketed by whole hours; the oldest
|
||||
// bucket still inside the window expires at its own hour boundary
|
||||
elapsed := now % hourSec
|
||||
left := sec - elapsed
|
||||
if left <= 0 || left > sec {
|
||||
left = sec
|
||||
}
|
||||
return left
|
||||
}
|
||||
|
||||
// WindowTokens returns the tokens consumed for one model (optionally pinned
|
||||
// to a single source) within the window; sec <= 0 means all time. Buckets are
|
||||
// whole unix hours, so a sliding window overcounts by up to one hour — an
|
||||
|
||||
@ -339,6 +339,9 @@
|
||||
.key-canvas{border:1px solid var(--line);border-radius:16px;padding:14px;margin-bottom:14px;background:var(--card);
|
||||
backdrop-filter:blur(var(--glass));box-shadow:var(--sh-sm)}
|
||||
.kc-head{display:flex;align-items:center;gap:10px;flex-wrap:wrap}
|
||||
/* key-wide caps sit inline in the head: they belong to the key, not to
|
||||
any one model brick, and must not be draggable with one. */
|
||||
.kc-caps{display:inline-flex;align-items:center;gap:6px;flex-wrap:wrap}
|
||||
.kc-blocks{display:flex;flex-wrap:wrap;gap:10px;align-items:center;margin-top:12px;background:var(--card2);
|
||||
border:1px dashed var(--line);border-radius:12px;padding:14px;min-height:64px}
|
||||
.kc-blocks.ovh{outline:2px dashed var(--primary);outline-offset:2px}
|
||||
@ -818,6 +821,15 @@
|
||||
kAnySrc: "任意源",
|
||||
kQuotaB: "Token 配额",
|
||||
kQuotaHintB: "0 / 留空 = 无限",
|
||||
kKeyQuota: "密钥总配额",
|
||||
kKeyQuotaHint:
|
||||
"限制这把密钥在重置周期内的总用量(跳模型)。0 / 留空 = 无限。",
|
||||
kKeyReqQuota: "请求数配额",
|
||||
kKeyReqQuotaHint: "限制周期内的请求次数。0 / 留空 = 无限。",
|
||||
kKeyQuotaAdmin:
|
||||
"admin 密钥永不受配额限制(避免把管理员锁在门外)。",
|
||||
kKeyQuotaEdit: "配额",
|
||||
kKeyQuotaNone: "无限",
|
||||
kPeriodB: "重置周期",
|
||||
kPerNothing: "不限",
|
||||
kPerHour: "每 小时",
|
||||
@ -1043,6 +1055,16 @@
|
||||
kAnySrc: "any source",
|
||||
kQuotaB: "Token quota",
|
||||
kQuotaHintB: "0 / empty = unlimited",
|
||||
kKeyQuota: "Key-wide quota",
|
||||
kKeyQuotaHint:
|
||||
"Caps this key's total spend per reset window, across every model it may use. 0 / empty = unlimited.",
|
||||
kKeyReqQuota: "Request quota",
|
||||
kKeyReqQuotaHint:
|
||||
"Caps requests per window. 0 / empty = unlimited.",
|
||||
kKeyQuotaAdmin:
|
||||
"Admin keys are never capped — a cap could lock the operator out.",
|
||||
kKeyQuotaEdit: "Quota",
|
||||
kKeyQuotaNone: "unlimited",
|
||||
kPeriodB: "Reset period",
|
||||
kPerNothing: "Never",
|
||||
kPerHour: "Every hour",
|
||||
@ -1692,7 +1714,20 @@
|
||||
</div>`;
|
||||
document.body.appendChild(wrap);
|
||||
}
|
||||
function downloadStatsCsv(from, to) {
|
||||
// closeTopModal removes the topmost dialog. Several dialogs share the id
|
||||
// "modal-wrap" and more than one can be open at once (the seed-key notice
|
||||
// sits on top of the keys page), so `$("#modal-wrap")` resolves to
|
||||
// whichever comes FIRST in the document — not necessarily the one the
|
||||
// user is looking at. Handlers that receive the clicked button resolve
|
||||
// their own dialog with .closest(); handlers with no button reference
|
||||
// use this, which matches what the user sees.
|
||||
function closeTopModal() {
|
||||
const all = [...document.querySelectorAll("#modal-wrap")];
|
||||
const w = all[all.length - 1];
|
||||
if (w) w.remove();
|
||||
return w;
|
||||
}
|
||||
function downloadStatsCsv(from, to, btn) {
|
||||
const q = new URLSearchParams({
|
||||
export: "csv",
|
||||
from: String(Math.floor(from)),
|
||||
@ -1700,8 +1735,7 @@
|
||||
});
|
||||
if (statsKeyF) q.set("key", statsKeyF);
|
||||
location.href = "/api/stats?" + q.toString();
|
||||
const w = $("#modal-wrap");
|
||||
if (w) w.remove();
|
||||
(btn ? btn.closest("#modal-wrap") : null) || closeTopModal();
|
||||
}
|
||||
// downloadStatsCsvFromForm reads the custom date range from the export
|
||||
// modal's date inputs and forwards to downloadStatsCsv. (Previously this
|
||||
@ -1719,12 +1753,11 @@
|
||||
if (!isFinite(from) || from <= 0) return toast(t("expRange"));
|
||||
downloadStatsCsv(from, to);
|
||||
}
|
||||
function downloadKeysCsv() {
|
||||
function downloadKeysCsv(btn) {
|
||||
location.href =
|
||||
"/api/stats?export=keys-csv" +
|
||||
(statsKeyF ? "&key=" + encodeURIComponent(statsKeyF) : "");
|
||||
const w = $("#modal-wrap");
|
||||
if (w) w.remove();
|
||||
(btn ? btn.closest("#modal-wrap") : null) || closeTopModal();
|
||||
}
|
||||
async function paintStats() {
|
||||
try {
|
||||
@ -2918,9 +2951,10 @@
|
||||
body: JSON.stringify(payload),
|
||||
});
|
||||
toast(t("toastSaved"));
|
||||
const w = $("#modal-wrap");
|
||||
const w = btn && btn.closest("#modal-wrap");
|
||||
if (w) w.remove();
|
||||
else if (window._modal) window._modal.remove();
|
||||
else closeTopModal();
|
||||
renderSources();
|
||||
} catch (e) {
|
||||
toast(tFmt("toastSaveFail", e.message));
|
||||
@ -3077,11 +3111,11 @@
|
||||
? `<div class="tbl-wrap"><table><tr><th>${t("tName")}</th><th>${t("tURL")}</th><th>${t("tAdapter")}</th><th>${t("tModels")}</th><th></th></tr>` +
|
||||
tpls
|
||||
.map(
|
||||
(t) =>
|
||||
`<tr><td><b>${esc(t.name)}</b></td><td>${esc(t.base_url)}</td><td>${esc(t.adapter)}</td>
|
||||
<td><div class="src-models">${(t.models || []).map((m) => `<span class="tag tag-blue">${esc(m.id)}</span>`).join("")}</div></td>
|
||||
<td><button class="ghost small" onclick="editTemplate(${JSON.stringify(t.name).replace(/["<>]/g, (c) => ({ '"': """, "<": "<", ">": ">" }[c]))})">${t("srcEdit")}</button>
|
||||
<button class="danger small" onclick="delTemplate('${escAttr(t.name)}')">${t("srcDel")}</button></td></tr>`,
|
||||
(tpl) =>
|
||||
`<tr><td><b>${esc(tpl.name)}</b></td><td>${esc(tpl.base_url)}</td><td>${esc(tpl.adapter)}</td>
|
||||
<td><div class="src-models">${(tpl.models || []).map((m) => `<span class="tag tag-blue">${esc(m.id)}</span>`).join("")}</div></td>
|
||||
<td><button class="ghost small" onclick="editTemplate(${JSON.stringify(tpl.name).replace(/["<>]/g, (c) => ({ '"': """, "<": "<", ">": ">" }[c]))})">${t("srcEdit")}</button>
|
||||
<button class="danger small" onclick="delTemplate('${escAttr(tpl.name)}')">${t("srcDel")}</button></td></tr>`,
|
||||
)
|
||||
.join("") +
|
||||
"</table></div>"
|
||||
@ -3181,8 +3215,9 @@
|
||||
body: JSON.stringify({ name: nm, ...body }),
|
||||
});
|
||||
toast(t("toastSaved"));
|
||||
const w = $("#modal-wrap");
|
||||
const w = btn && btn.closest("#modal-wrap");
|
||||
if (w) w.remove();
|
||||
else closeTopModal();
|
||||
openTemplateModal();
|
||||
} catch (e) {
|
||||
toast(e.message);
|
||||
@ -3888,7 +3923,7 @@
|
||||
</select>
|
||||
<div id="a-hours-box" style="display:none"><label>${t("kPerNHint")}</label>
|
||||
<input id="a-hours" type="number" min="1" step="1" value="24"></div>
|
||||
<p><button onclick="scrAddFromForm()">${t("kSaveScope")}</button>
|
||||
<p><button onclick="scrAddFromForm(this)">${t("kSaveScope")}</button>
|
||||
<button class="ghost" onclick="this.closest('#modal-wrap').remove()">${t("mCancel")}</button></p>
|
||||
</div>`;
|
||||
document.body.appendChild(wrap);
|
||||
@ -3898,7 +3933,7 @@
|
||||
});
|
||||
$("#a-model").focus();
|
||||
}
|
||||
function scrAddFromForm() {
|
||||
function scrAddFromForm(btn) {
|
||||
const raw = $("#a-model").value.trim();
|
||||
if (!raw) {
|
||||
toast(t("kName"));
|
||||
@ -3921,14 +3956,14 @@
|
||||
],
|
||||
});
|
||||
const li = sortState.lanes.length - 1;
|
||||
const w = $("#modal-wrap");
|
||||
if (w) w.remove();
|
||||
if (btn) btn.closest("#modal-wrap")?.remove();
|
||||
else closeTopModal();
|
||||
paintSort([li]);
|
||||
persistAuto()
|
||||
.then(() => toast(t("sortHintSave")))
|
||||
.catch((e) => toast(e.message));
|
||||
}
|
||||
function sortScopeSave(li, ji) {
|
||||
function sortScopeSave(li, ji, btn) {
|
||||
const it = sortState.lanes[li].models[ji];
|
||||
if (!it) return;
|
||||
let q = parseInt($("#q-quota").value);
|
||||
@ -3937,8 +3972,7 @@
|
||||
let h = parseInt($("#q-hours").value);
|
||||
if (isNaN(h) || h < 1) h = 1;
|
||||
it.meta = { quota: q, period: p, hours: p || q ? h : 0 };
|
||||
const w = $("#modal-wrap");
|
||||
if (w) w.remove();
|
||||
(btn ? btn.closest("#modal-wrap") : null) || closeTopModal();
|
||||
paintSort([li]);
|
||||
persistAuto()
|
||||
.then(() => toast(t("sortHintSave")))
|
||||
@ -4232,10 +4266,11 @@
|
||||
async function renderKeysUser(me) {
|
||||
$("#tab-keys").innerHTML = `
|
||||
<div class="card"><h2>${t("kMeTitle")}</h2>
|
||||
<div class="tbl-wrap"><table><tr><th>${t("kName")}</th><th>${t("kMeRole")}</th><th>${t("kKey")}</th><th>${t("kMeModels")}</th></tr>
|
||||
<div class="tbl-wrap"><table><tr><th>${t("kName")}</th><th>${t("kMeRole")}</th><th>${t("kKey")}</th><th>${t("kKeyQuota")}</th><th>${t("kMeModels")}</th></tr>
|
||||
<tr><td><b>${esc(me.name || "—")}</b></td><td>${roleTag(me.role)}</td>
|
||||
<td><span class="kr-key">${esc(me.key)}</span>
|
||||
<button class="ghost small" onclick="copyText('${escAttr(me.key)}')">${t("kCopy")}</button></td>
|
||||
<td>${keyCapBadges(me)}</td>
|
||||
<td>${
|
||||
me.models && me.models.length
|
||||
? me.models
|
||||
@ -4273,13 +4308,22 @@
|
||||
}
|
||||
function keyCanvasHtml(k) {
|
||||
const scopes = k.models || [];
|
||||
// Key-wide caps live on the canvas, not on a brick: they are a budget
|
||||
// the whole key shares, so they must not be dragged around with one
|
||||
// model. data-* carries them so a quota edit can round-trip them
|
||||
// through the same PUT that saves the model scope.
|
||||
const caps = `data-kquota="${k.token_quota || 0}" data-kreqquota="${k.req_quota || 0}"
|
||||
data-kperiod="${escAttr(k.period || "")}" data-khours="${k.hours || 0}"`;
|
||||
return `
|
||||
<div class="key-canvas" data-key="${escAttr(k.key)}">
|
||||
<div class="key-canvas" data-key="${escAttr(k.key)}" ${caps}>
|
||||
<div class="kc-head">
|
||||
<b>${esc(k.name || "—")}</b>
|
||||
${roleTag(k.role)}
|
||||
<span class="kr-key">${esc(maskKey(k.key))}</span>
|
||||
<button class="ghost small" onclick="copyText('${escAttr(k.key)}')">${t("kCopy")}</button>
|
||||
<span class="kc-caps" title="${escAttr(t("kKeyQuotaHint"))}">${keyCapBadges(k)}</span>
|
||||
${k.role === "admin" ? "" : `<button class="ghost small" title="${escAttr(t("kKeyQuotaEdit"))}"
|
||||
onclick="keyQuotaEdit('${escAttr(k.key)}')">${t("kKeyQuotaEdit")}</button>`}
|
||||
<span class="grow"></span>
|
||||
<span class="muted">${fmtCreated(k.created_at)}</span>
|
||||
<button class="ghost small errc" onclick="delKey('${escAttr(k.key)}','${escAttr(k.name || "")}')">${t("kDel")}</button>
|
||||
@ -4292,6 +4336,27 @@
|
||||
</div>
|
||||
</div>`;
|
||||
}
|
||||
// keyCapBadges renders the key-wide caps. A cap with no reset period is
|
||||
// flagged as such, because "1M tokens, never resets" and "1M tokens per
|
||||
// hour" are very different promises and the badge must not blur them.
|
||||
function keyCapBadges(k) {
|
||||
const out = [];
|
||||
const suffix = periodText(k.period || "", k.hours || 0);
|
||||
if (+k.token_quota > 0) {
|
||||
out.push(
|
||||
`<span class="mb-quota" title="${escAttr(t("kKeyQuotaHint"))}">${esc(fmtQuota(k.token_quota))}${esc(suffix)}</span>`,
|
||||
);
|
||||
}
|
||||
if (+k.req_quota > 0) {
|
||||
out.push(
|
||||
`<span class="mb-quota" title="${escAttr(t("kKeyReqQuotaHint"))}">${esc(fmtQuota(k.req_quota))}×${esc(suffix)}</span>`,
|
||||
);
|
||||
}
|
||||
if (!out.length) {
|
||||
return `<span class="muted">${t("kKeyQuotaNone")}</span>`;
|
||||
}
|
||||
return out.join(" ");
|
||||
}
|
||||
function scopeHtml(key, m) {
|
||||
const qt = fmtQuota(m.token_quota);
|
||||
const comb = scopeComb(m);
|
||||
@ -4355,12 +4420,117 @@
|
||||
});
|
||||
}
|
||||
async function putScope(key, scopes) {
|
||||
// The key-wide caps ride along with every scope write. The API reads
|
||||
// them as pointers, so sending them back unchanged is a no-op, while
|
||||
// omitting them would be indistinguishable from "clear the budget" to
|
||||
// a future reader. Round-tripping them here means editing a model's
|
||||
// scope can never silently drop a key's quota.
|
||||
const canvas = document.querySelector(
|
||||
`.key-canvas[data-key="${CSS.escape(key)}"]`,
|
||||
);
|
||||
const body = { models: scopes };
|
||||
if (canvas) {
|
||||
body.token_quota = parseInt(canvas.dataset.kquota) || 0;
|
||||
body.req_quota = parseInt(canvas.dataset.kreqquota) || 0;
|
||||
body.period = canvas.dataset.kperiod || "";
|
||||
body.hours = parseInt(canvas.dataset.khours) || 0;
|
||||
}
|
||||
await api("/api/keys/" + encodeURIComponent(key), {
|
||||
method: "PUT",
|
||||
headers: { "Content-Type": "application/json" },
|
||||
body: JSON.stringify({ models: scopes }),
|
||||
body: JSON.stringify(body),
|
||||
});
|
||||
}
|
||||
// keyQuotaEdit opens the key-wide budget form.
|
||||
function keyQuotaEdit(key) {
|
||||
const canvas = document.querySelector(
|
||||
`.key-canvas[data-key="${CSS.escape(key)}"]`,
|
||||
);
|
||||
if (!canvas) return;
|
||||
const cur = {
|
||||
token_quota: parseInt(canvas.dataset.kquota) || 0,
|
||||
req_quota: parseInt(canvas.dataset.kreqquota) || 0,
|
||||
period: canvas.dataset.kperiod || "",
|
||||
hours: parseInt(canvas.dataset.khours) || 0,
|
||||
};
|
||||
const wrap = document.createElement("div");
|
||||
wrap.id = "modal-wrap";
|
||||
wrap.style.cssText =
|
||||
"position:fixed;inset:0;background:rgba(15,22,44,.45);display:flex;align-items:flex-start;justify-content:center;overflow:auto;padding:48px 20px;z-index:50";
|
||||
wrap.innerHTML = `<div class="card" style="width:400px;max-width:100%"><h2>${t("kKeyQuotaEdit")}</h2>
|
||||
<label>${t("kKeyQuota")} <span class="muted">${t("kKeyQuotaHint")}</span></label>
|
||||
<input id="kq-tokens" type="number" min="0" step="1"
|
||||
placeholder="${escAttr(t("kKeyQuotaNone"))}" value="${cur.token_quota || ""}">
|
||||
<label>${t("kKeyReqQuota")} <span class="muted">${t("kKeyReqQuotaHint")}</span></label>
|
||||
<input id="kq-reqs" type="number" min="0" step="1"
|
||||
placeholder="${escAttr(t("kKeyQuotaNone"))}" value="${cur.req_quota || ""}">
|
||||
<label>${t("kPeriodB")}</label>
|
||||
<select id="kq-period">
|
||||
<option value="" ${!cur.period ? "selected" : ""}>${t("kPerNothing")}</option>
|
||||
<option value="hour" ${cur.period === "hour" ? "selected" : ""}>${t("kPerHour")}</option>
|
||||
<option value="week" ${cur.period === "week" ? "selected" : ""}>${t("kPerWeek")}</option>
|
||||
<option value="month" ${cur.period === "month" ? "selected" : ""}>${t("kPerMonth")}</option>
|
||||
<option value="nhour" ${cur.period === "nhour" ? "selected" : ""}>${t("kPerHours")}</option>
|
||||
</select>
|
||||
<div id="kq-hours-box" style="display:none"><label>${t("kPerNHint")}</label>
|
||||
<input id="kq-hours" type="number" min="1" step="1" value="${cur.hours || 24}"></div>
|
||||
<p class="muted" style="font-size:12px">${t("kKeyQuotaAdmin")}</p>
|
||||
<p><button onclick="keyQuotaSave('${escAttr(key)}', this)">${t("kSaveScope")}</button>
|
||||
<button class="ghost" onclick="this.closest('#modal-wrap').remove()">${t("mCancel")}</button></p>
|
||||
</div>`;
|
||||
document.body.appendChild(wrap);
|
||||
const toggle = () => {
|
||||
$("#kq-hours-box").style.display =
|
||||
$("#kq-period").value === "nhour" ? "block" : "none";
|
||||
};
|
||||
$("#kq-period").addEventListener("change", toggle);
|
||||
toggle();
|
||||
$("#kq-tokens").focus();
|
||||
}
|
||||
async function keyQuotaSave(key, btn) {
|
||||
// Resolve our own dialog from the button that was clicked, so closing
|
||||
// it can never remove a different #modal-wrap that happens to come
|
||||
// first in the document.
|
||||
const wrap = btn ? btn.closest("#modal-wrap") : null;
|
||||
let tokens = parseInt($("#kq-tokens").value);
|
||||
if (isNaN(tokens) || tokens < 0) tokens = 0;
|
||||
let reqs = parseInt($("#kq-reqs").value);
|
||||
if (isNaN(reqs) || reqs < 0) reqs = 0;
|
||||
let hours = parseInt($("#kq-hours").value);
|
||||
if (isNaN(hours) || hours < 1) hours = 1;
|
||||
const period = $("#kq-period").value;
|
||||
// Catch the "nhour picked but hours never filled in" case locally: the
|
||||
// API rejects it too, but a round trip for a form-level mistake is
|
||||
// needless.
|
||||
if (period === "nhour" && hours < 1) {
|
||||
toast(t("kPerNHint"));
|
||||
return;
|
||||
}
|
||||
if (btn) btn.disabled = true;
|
||||
try {
|
||||
await api("/api/keys/" + encodeURIComponent(key), {
|
||||
method: "PUT",
|
||||
headers: { "Content-Type": "application/json" },
|
||||
body: JSON.stringify({
|
||||
token_quota: tokens,
|
||||
req_quota: reqs,
|
||||
period,
|
||||
hours,
|
||||
}),
|
||||
});
|
||||
// Close THIS modal, not whichever #modal-wrap comes first in the
|
||||
// document: another dialog (e.g. the seed-key notice) may already be
|
||||
// open, and a bare $("#modal-wrap") would remove that one and leave
|
||||
// this form stranded on screen.
|
||||
if (wrap) wrap.remove();
|
||||
else closeTopModal();
|
||||
toast(t("kSaved"));
|
||||
await loadKeys();
|
||||
} catch (e) {
|
||||
toast(e.message);
|
||||
if (btn) btn.disabled = false;
|
||||
}
|
||||
}
|
||||
async function scopePush(key) {
|
||||
const canvas = document.querySelector(
|
||||
`.key-canvas[data-key="${CSS.escape(key)}"]`,
|
||||
@ -4511,8 +4681,9 @@
|
||||
if (btn) btn.disabled = true;
|
||||
try {
|
||||
await putScope(key, scopes);
|
||||
const w = $("#modal-wrap");
|
||||
const w = btn && btn.closest("#modal-wrap");
|
||||
if (w) w.remove();
|
||||
else closeTopModal();
|
||||
toast(t("kSaved"));
|
||||
await loadKeys();
|
||||
} catch (e) {
|
||||
@ -4661,12 +4832,41 @@
|
||||
<option value="user">${t("kRoleUser")}</option>
|
||||
<option value="admin">${t("kRoleAdmin")}</option>
|
||||
</select>
|
||||
<label>${t("kKeyQuota")} <span class="muted">${t("kKeyQuotaHint")}</span></label>
|
||||
<input id="kc-tokens" type="number" min="0" step="1" placeholder="${escAttr(t("kKeyQuotaNone"))}">
|
||||
<label>${t("kKeyReqQuota")} <span class="muted">${t("kKeyReqQuotaHint")}</span></label>
|
||||
<input id="kc-reqs" type="number" min="0" step="1" placeholder="${escAttr(t("kKeyQuotaNone"))}">
|
||||
<label>${t("kPeriodB")}</label>
|
||||
<select id="kc-period">
|
||||
<option value="" selected>${t("kPerNothing")}</option>
|
||||
<option value="hour">${t("kPerHour")}</option>
|
||||
<option value="week">${t("kPerWeek")}</option>
|
||||
<option value="month">${t("kPerMonth")}</option>
|
||||
<option value="nhour">${t("kPerHours")}</option>
|
||||
</select>
|
||||
<div id="kc-hours-box" style="display:none"><label>${t("kPerNHint")}</label>
|
||||
<input id="kc-hours" type="number" min="1" step="1" value="24"></div>
|
||||
<label>${t("kNote")}</label>
|
||||
<input id="kc-note">
|
||||
<p class="muted" style="font-size:12px">${t("kKeyQuotaAdmin")}</p>
|
||||
<p><button onclick="createKey(this)">${t("kCreateBtn")}</button>
|
||||
<button class="ghost" onclick="this.closest('#modal-wrap').remove()">${t("mCancel")}</button></p>
|
||||
</div>`;
|
||||
document.body.appendChild(wrap);
|
||||
$("#kc-role").addEventListener("change", () => {
|
||||
const admin = $("#kc-role").value === "admin";
|
||||
// An admin key ignores its caps server-side; hiding the fields
|
||||
// avoids the operator setting one and wondering why it never trips.
|
||||
$("#kc-tokens").disabled = admin;
|
||||
$("#kc-reqs").disabled = admin;
|
||||
$("#kc-period").disabled = admin;
|
||||
$("#kc-hours-box").style.display =
|
||||
!admin && $("#kc-period").value === "nhour" ? "block" : "none";
|
||||
});
|
||||
$("#kc-period").addEventListener("change", () => {
|
||||
$("#kc-hours-box").style.display =
|
||||
$("#kc-period").value === "nhour" ? "block" : "none";
|
||||
});
|
||||
$("#kc-name").focus();
|
||||
}
|
||||
async function createKey(btn) {
|
||||
@ -4675,6 +4875,17 @@
|
||||
toast(t("kName"));
|
||||
return;
|
||||
}
|
||||
const role = $("#kc-role").value;
|
||||
// An admin key is never capped; send the fields anyway (the server
|
||||
// ignores them) rather than special-casing the request shape.
|
||||
const readNum = (sel) => {
|
||||
const el = $(sel);
|
||||
if (el.disabled) return 0;
|
||||
const n = parseInt(el.value);
|
||||
return isNaN(n) || n < 0 ? 0 : n;
|
||||
};
|
||||
let hours = parseInt($("#kc-hours").value);
|
||||
if (isNaN(hours) || hours < 1) hours = 1;
|
||||
if (btn) btn.disabled = true;
|
||||
let j;
|
||||
try {
|
||||
@ -4683,8 +4894,12 @@
|
||||
headers: { "Content-Type": "application/json" },
|
||||
body: JSON.stringify({
|
||||
name,
|
||||
role: $("#kc-role").value,
|
||||
role,
|
||||
note: $("#kc-note").value.trim(),
|
||||
token_quota: readNum("#kc-tokens"),
|
||||
req_quota: readNum("#kc-reqs"),
|
||||
period: role === "admin" ? "" : $("#kc-period").value,
|
||||
hours: role === "admin" ? 0 : hours,
|
||||
}),
|
||||
});
|
||||
} catch (e) {
|
||||
@ -4692,8 +4907,12 @@
|
||||
if (btn) btn.disabled = false;
|
||||
return;
|
||||
}
|
||||
const w = $("#modal-wrap");
|
||||
// Close THIS dialog (resolved from the clicked button), not whichever
|
||||
// #modal-wrap comes first: the seed-key notice may already be open and
|
||||
// would otherwise be the one that gets removed.
|
||||
const w = btn ? btn.closest("#modal-wrap") : null;
|
||||
if (w) w.remove();
|
||||
else closeTopModal();
|
||||
$("#k-newbox").innerHTML = `
|
||||
<div class="key-canvas" style="background:var(--card2)">
|
||||
<div class="kc-head">
|
||||
|
||||
270
internal/gateway/ui_quota_contract_test.go
Normal file
270
internal/gateway/ui_quota_contract_test.go
Normal file
@ -0,0 +1,270 @@
|
||||
package gateway
|
||||
|
||||
import (
|
||||
"regexp"
|
||||
"strings"
|
||||
"testing"
|
||||
)
|
||||
|
||||
// The WebUI dialogs all share the id "modal-wrap", and more than one can be
|
||||
// open at the same time (the seed-key notice sits on top of the keys page).
|
||||
// A save handler that closes "the" modal via $("#modal-wrap") therefore removes
|
||||
// whichever one comes FIRST in the document — which is the wrong dialog: the
|
||||
// form the user just submitted stays on screen while an unrelated dialog
|
||||
// vanishes.
|
||||
//
|
||||
// This is exactly the class of bug the api() contract test below was written
|
||||
// for: reviewing inline JS by eye does not catch it, and the visible symptom
|
||||
// ("the dialog did not close") points away from the cause. It is pinned here.
|
||||
|
||||
// modalCloseRe finds every `$(...)`-style lookup of the shared modal id.
|
||||
var modalCloseRe = regexp.MustCompile(`\$\("#modal-wrap"\)`)
|
||||
|
||||
// closestModalRe finds the safe form: resolve the dialog from the clicked
|
||||
// button instead of from the document.
|
||||
var closestModalRe = regexp.MustCompile(`\.closest\("#modal-wrap"\)`)
|
||||
|
||||
func TestUIDialogClosesItselfNotTheFirstModal(t *testing.T) {
|
||||
src := uiSource(t)
|
||||
// Strip comments first: prose that *names* the unsafe pattern (as the fix's
|
||||
// own comment does) would otherwise be flagged as a violation.
|
||||
code := stripJSComments(src)
|
||||
|
||||
for _, m := range modalCloseRe.FindAllStringIndex(code, -1) {
|
||||
after := code[m[1]:]
|
||||
stmtEnd := strings.Index(after, ";")
|
||||
if stmtEnd < 0 || stmtEnd > 200 {
|
||||
continue
|
||||
}
|
||||
stmt := after[:stmtEnd]
|
||||
if strings.Contains(stmt, ".remove()") {
|
||||
line := 1 + strings.Count(code[:m[0]], "\n")
|
||||
t.Errorf("line %d closes the first #modal-wrap in the document, not its own dialog:\n\t%s",
|
||||
line, strings.TrimSpace(stmt))
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// TestUIDialogClosuresGoThroughSafePaths pins the rule across the whole
|
||||
// document by data flow rather than by pattern: every handler that closes a
|
||||
// dialog must do it one of the two safe ways. A handler could contain a
|
||||
// correct .closest() and still close the wrong dialog on another path.
|
||||
func TestUIDialogClosuresGoThroughSafePaths(t *testing.T) {
|
||||
src := stripJSComments(uiSource(t))
|
||||
for _, fn := range []string{
|
||||
"downloadStatsCsv", "downloadKeysCsv", "saveSource", "saveTemplate",
|
||||
"scrAddFromForm", "sortScopeSave", "scopeSave", "keyQuotaSave", "createKey",
|
||||
} {
|
||||
body, ok := jsFunctionBody(src, fn)
|
||||
if !ok {
|
||||
t.Errorf("%s not found", fn)
|
||||
continue
|
||||
}
|
||||
if !strings.Contains(body, `closest("#modal-wrap")`) && !strings.Contains(body, "closeTopModal()") {
|
||||
t.Errorf("%s closes a dialog with neither .closest nor closeTopModal", fn)
|
||||
}
|
||||
}
|
||||
if !strings.Contains(src, "function closeTopModal(") {
|
||||
t.Error("closeTopModal helper is missing")
|
||||
}
|
||||
}
|
||||
|
||||
// The helper must pick the LAST dialog (the topmost one the user sees), not the
|
||||
// first — that inversion is the whole bug.
|
||||
func TestUICloseTopModalTakesTheLast(t *testing.T) {
|
||||
body, ok := jsFunctionBody(uiSource(t), "closeTopModal")
|
||||
if !ok {
|
||||
t.Fatal("closeTopModal not found")
|
||||
}
|
||||
if !strings.Contains(body, "all.length - 1") {
|
||||
t.Errorf("closeTopModal does not take the last dialog:\n\t%s", oneLine(body))
|
||||
}
|
||||
}
|
||||
|
||||
// Handlers that resolve their dialog from a button must actually receive one:
|
||||
// a signature without the parameter means the .closest() silently yields null
|
||||
// and the save leaves its form stranded on screen.
|
||||
func TestUIDialogHandlersReceiveTheirButton(t *testing.T) {
|
||||
src := stripJSComments(uiSource(t))
|
||||
for _, fn := range []string{
|
||||
"downloadStatsCsv", "saveSource", "scrAddFromForm", "sortScopeSave",
|
||||
"scopeSave", "keyQuotaSave", "createKey",
|
||||
} {
|
||||
body, ok := jsFunctionBody(src, fn)
|
||||
if !ok {
|
||||
t.Errorf("%s not found", fn)
|
||||
continue
|
||||
}
|
||||
if !strings.Contains(body, "closest(\"#modal-wrap\")") {
|
||||
continue // uses closeTopModal only
|
||||
}
|
||||
sig := body[:strings.Index(body, ")")+1]
|
||||
if !strings.Contains(sig, "btn") {
|
||||
t.Errorf("%s uses .closest(\"#modal-wrap\") but its signature %s has no button parameter —\n"+
|
||||
"the lookup would always be null and the form would never close", fn, oneLine(sig))
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// stripJSComments removes // line comments and /* block */ comments from JS
|
||||
// embedded in the UI document. It is deliberately simple (no string/regex
|
||||
// awareness beyond skipping quoted spans on the same line): the document is
|
||||
// our own source, and a false negative here only means the check is silent.
|
||||
func stripJSComments(src string) string {
|
||||
var out strings.Builder
|
||||
lines := strings.Split(src, "\n")
|
||||
inBlock := false
|
||||
for _, ln := range lines {
|
||||
trimmed := strings.TrimSpace(ln)
|
||||
if inBlock {
|
||||
if strings.Contains(ln, "*/") {
|
||||
inBlock = false
|
||||
}
|
||||
continue
|
||||
}
|
||||
if strings.HasPrefix(trimmed, "/*") {
|
||||
if !strings.Contains(ln, "*/") {
|
||||
inBlock = true
|
||||
}
|
||||
continue
|
||||
}
|
||||
if i := strings.Index(ln, "//"); i >= 0 {
|
||||
// keep code before the comment when the // is not inside a string
|
||||
before := ln[:i]
|
||||
if strings.Count(before, `"`)%2 == 0 && strings.Count(before, "'")%2 == 0 {
|
||||
ln = before
|
||||
}
|
||||
}
|
||||
out.WriteString(ln)
|
||||
out.WriteString("\n")
|
||||
}
|
||||
return out.String()
|
||||
}
|
||||
|
||||
// TestUIKeyQuotaDialogsResolveOwnModal pins the dialog-closing rule for the
|
||||
// two forms this change added.
|
||||
func TestUIKeyQuotaDialogsResolveOwnModal(t *testing.T) {
|
||||
src := uiSource(t)
|
||||
for _, fn := range []string{"keyQuotaSave", "createKey"} {
|
||||
body, ok := jsFunctionBody(src, fn)
|
||||
if !ok {
|
||||
t.Errorf("%s not found in the UI source", fn)
|
||||
continue
|
||||
}
|
||||
if !closestModalRe.MatchString(body) {
|
||||
t.Errorf("%s does not resolve its own dialog via .closest(\"#modal-wrap\");\n"+
|
||||
"with another dialog open it would close that one instead and leave this form stranded", fn)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// The quota editor must read and write the key-wide caps, and putScope must
|
||||
// carry them along: the API treats the quota fields as pointers, so dropping
|
||||
// them on a scope-only write is indistinguishable from "clear the budget".
|
||||
func TestUIPutScopeCarriesKeyQuota(t *testing.T) {
|
||||
body, ok := jsFunctionBody(uiSource(t), "putScope")
|
||||
if !ok {
|
||||
t.Fatal("putScope not found")
|
||||
}
|
||||
// Scan the code with comments removed, or a comment that merely *names* a
|
||||
// field would satisfy the check while the field is never sent.
|
||||
code := stripJSComments(body)
|
||||
for _, field := range []string{"token_quota", "req_quota", "period", "hours"} {
|
||||
if !strings.Contains(code, field) {
|
||||
t.Errorf("putScope does not send %q — editing a model scope would clear the key's quota", field)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// A key's caps are rendered from the API record and shown on the canvas, so
|
||||
// the badge and the data attributes must not drift from the field names. The
|
||||
// create form must send them too, or a key would only be cappable after an
|
||||
// extra round of edits.
|
||||
func TestUIKeyQuotaRendersFromAPIFields(t *testing.T) {
|
||||
src := uiSource(t)
|
||||
for _, token := range []string{
|
||||
"keyCapBadges", // shared renderer
|
||||
"kq-tokens", "kq-reqs", "kq-period", "kq-hours", // editor fields
|
||||
"kc-tokens", "kc-reqs", "kc-period", "kc-hours", // create form fields
|
||||
} {
|
||||
if !strings.Contains(src, token) {
|
||||
t.Errorf("UI never references %q — the quota form is not wired up", token)
|
||||
}
|
||||
}
|
||||
// the create request must actually carry the caps
|
||||
full, ok := jsFunctionBody(src, "createKey")
|
||||
if !ok {
|
||||
t.Fatal("createKey not found")
|
||||
}
|
||||
body := stripJSComments(full)
|
||||
for _, field := range []string{"token_quota", "req_quota", "period"} {
|
||||
if !strings.Contains(body, field) {
|
||||
t.Errorf("createKey does not send %q — a new key could never be created with a budget", field)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Existence of the strings is not enough: the badge has to READ the API
|
||||
// fields, and the editor has to read the canvas data attributes it writes.
|
||||
// A field can be present in the source and still never reach the screen —
|
||||
// e.g. left in a dead branch, or read from a name the writer never sets.
|
||||
func TestUIKeyQuotaDataflowIsLive(t *testing.T) {
|
||||
src := uiSource(t)
|
||||
|
||||
badge, ok := jsFunctionBody(src, "keyCapBadges")
|
||||
if !ok {
|
||||
t.Fatal("keyCapBadges not found")
|
||||
}
|
||||
badgeCode := stripJSComments(badge)
|
||||
for _, field := range []string{"k.token_quota", "k.req_quota", "k.period"} {
|
||||
if !strings.Contains(badgeCode, field) {
|
||||
t.Errorf("keyCapBadges does not read %q — the cap would never show on the key card", field)
|
||||
}
|
||||
}
|
||||
|
||||
// the editor must read back what keyCanvasHtml wrote
|
||||
canvas, ok := jsFunctionBody(src, "keyCanvasHtml")
|
||||
if !ok {
|
||||
t.Fatal("keyCanvasHtml not found")
|
||||
}
|
||||
editor, ok := jsFunctionBody(src, "keyQuotaEdit")
|
||||
if !ok {
|
||||
t.Fatal("keyQuotaEdit not found")
|
||||
}
|
||||
canvasCode, editorCode := stripJSComments(canvas), stripJSComments(editor)
|
||||
for _, ds := range []string{"kquota", "kreqquota", "kperiod", "khours"} {
|
||||
// written as data-<name> on the canvas
|
||||
if !strings.Contains(canvasCode, "data-"+ds+"=") {
|
||||
t.Errorf("keyCanvasHtml does not write data-%s, so the editor has nothing to prefill", ds)
|
||||
}
|
||||
// read back as dataset.<name> by the editor
|
||||
if !strings.Contains(editorCode, "dataset."+ds) {
|
||||
t.Errorf("keyQuotaEdit does not read dataset.%s — the form would open blank and save zeros", ds)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// The quota period vocabulary must match the server's, or the UI can offer a
|
||||
// value the API rejects.
|
||||
func TestUIQuotaPeriodsMatchServer(t *testing.T) {
|
||||
src := uiSource(t)
|
||||
// the shared period select options, as rendered in both forms
|
||||
for _, p := range []string{`value=""`, `value="hour"`, `value="week"`, `value="month"`, `value="nhour"`} {
|
||||
if !strings.Contains(src, p) {
|
||||
t.Errorf("UI period select is missing %s", p)
|
||||
}
|
||||
}
|
||||
// the server's accepted vocabulary
|
||||
for _, p := range []string{`"hour"`, `"week"`, `"month"`, `"nhour"`} {
|
||||
if !strings.Contains(src, `if (p === `+p+`)`) && !strings.Contains(src, `=== `+p+`)`) {
|
||||
t.Errorf("periodText() does not describe %s, so a badge would omit the window", p)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func max(a, b int) int {
|
||||
if a > b {
|
||||
return a
|
||||
}
|
||||
return b
|
||||
}
|
||||
Reference in New Issue
Block a user