jianf/ModelRouter
1
0
Fork 0
mirror of https://gitcode.com/JianFeeeee/ModelRouter.git synced 2026-10-15 12:14:54 +00:00
Code Issues Packages Projects Releases Wiki Activity

refactor(quota): 配额改为按模型，删除整钥总配额

Browse Source 
用户明确要求：配额应当是密钥对应的**每个模型的单独配额**，而非整体配额。

## 语义变更

删除 GWKey.TokenQuota / ReqQuota / Period / Hours（整钥总额）。
ModelScope 新增 ReqQuota —— 请求数配额下沉到每条模型范围。

现在：每条 models[] 各自带 token 配额 + 请求数配额 + 重置周期，
彼此独立。一个模型用满只影响该模型。

★ 为什么不保留整钥总额：它会让「把 A 模型的额度挪给 B」变成一次全局
重分配；按模型独立计费则每个模型各自可控，运维能直接看出哪个模型在吃预算。

## 连带改动

- checkQuota 合并 key 级与 scope 级判定；checkKeyQuotaRetry 整体删除
  （顺带修掉上轮遗留的双重判定：入口不再先判空再重算）
- core：CreateKeyWithQuota / UpdateKeyWithQuota / ApplyQuota 全部删除，
  改由 ValidateScopeQuotas 校验每条 scope 的配额
- admin key：scope 上的配额不强制（admin 的 scope 仍限制模型范围，
  但不强制配额）—— 否则管理员会把自己锁在门外
- /api/v1/keys 不再回显 key 级配额字段（scope 里已含）
- WebUI：删除整钥配额徽标 / 「配额」按钮 / 创建表单的配额组 /
  putScope 的整钥回传；模型砖块与范围编辑器新增「请求数配额」输入，
  徽标显示 `1.0K 77×·1h`（未设配额显示 ∞）

## 判据

- TestOneModelsQuotaDoesNotBlockAnother 是本次核心保证。
  ★ 它第一版是**假判据**：m2 从不消耗，key-wide 计数器与 m1 自己的计数器
  读数恰好相同，退回 key-wide 仍通过。变异测试抓到后改为「先用 m2 花掉
  远超 m1 配额的量，再验证 m1 仍可用」—— 这样两种设计才可区分。
- TestUncappedModelNeverBlocked / TestAdminKeyScopesAreNotEnforced 新增
- UI 契约判据重写：整钥配额界面必须彻底消失（13 个符号）、
  scope 编辑器必须往返 req_quota、putScope 只发 scope 列表
- 错误消息点名具体模型（TestKeyAPIRejectionNamesTheModel）
- 3/3 变异全被抓

实测（真实进程 + 浏览器）：m2 配额 500000 连打 25 次全成功，
m1 配额 1000 立即 429「token quota exceeded for "m1" (4315/1000)」，
此后 m2/m3 仍 200。UI：整钥配额元素全为 0，砖块各显配额，
编辑器预填/保存正确，零 JS 异常。
This commit is contained in:
JianFeeeee
2026-09-27 19:02:13 +08:00
parent 5530912d32
commit c51066f0b6
11 changed files with 515 additions and 645 deletions
Download Patch File Download Diff File Expand all files Collapse all files

145
internal/gateway/key_quota_wiring_test.go
View File

@ -82,8 +82,8 @@ func chatAs(t *testing.T, g *Gateway, key, model string) (*httptest.ResponseReco
// 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-a", Role: "user",
Models: []config.ModelScope{{Model: "m1", TokenQuota: 8, Period: "hour"}}},
config.GWKey{Key: "sk-b", Role: "user", Name: "b"},
)
// 4 tokens per call, budget 8 -> the third call crosses it
@ -114,10 +114,10 @@ func TestKeyTokenQuotaBlocksWithRetryAfter(t *testing.T) {
// 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"}}},
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: 1000, Period: "hour"}}},
)
rr, _ := chatAs(t, g, "sk-a", "m1")
if rr.Code != 200 {
@ -162,7 +162,8 @@ func TestScopeModelTokenQuotaIsolatedPerKey(t *testing.T) {
// 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-admin", Role: "admin",
Models: []config.ModelScope{{Model: "m1", TokenQuota: 1, ReqQuota: 1, Period: "hour"}}},
config.GWKey{Key: "sk-b", Role: "user"},
)
for i := 1; i <= 3; i++ {
@ -176,8 +177,8 @@ func TestAdminKeyIsNeverQuotaCapped(t *testing.T) {
// 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-a", Role: "user",
Models: []config.ModelScope{{Model: "m1", ReqQuota: 2, Period: "hour"}}},
config.GWKey{Key: "sk-b", Role: "user"},
)
for i := 1; i <= 2; i++ {
@ -197,8 +198,8 @@ func TestKeyRequestQuotaBlocks(t *testing.T) {
// 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-a", Role: "user",
Models: []config.ModelScope{{Model: "other-model", TokenQuota: 1000, Period: "hour"}}},
config.GWKey{Key: "sk-b", Role: "user"},
)
rr, code := chatAs(t, g, "sk-a", "m1")
@ -256,8 +257,8 @@ 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"}}})
config.GWKey{Key: "sk-a", Role: "user",
Models: []config.ModelScope{{Model: "AUTO", TokenQuota: 4, Period: "hour"}}})
ctx := quotaCtx(t, g, "sk-a")
// exhaust the key first
@ -304,7 +305,10 @@ func newQuotaGW(t *testing.T, upURL string, keys ...config.GWKey) *Gateway {
Keys: keys,
Sources: []config.Source{{
Name: "up", BaseURL: upURL, Adapter: "openai",
Models: []config.Model{{ID: "m1", Priority: 100}},
Models: []config.Model{
{ID: "m1", Priority: 100},
{ID: "m2", Priority: 90},
},
}},
}
if err := cfg.ApplyDefaults(); err != nil {
@ -325,3 +329,116 @@ func newQuotaGW(t *testing.T, upURL string, keys ...config.GWKey) *Gateway {
}
return g
}
// The core of per-model quotas: a model that runs out of budget must stop
// being served on its own, while every other model on the SAME key keeps
// working. A key-wide total would fail this — it would block m2 because m1 was
// capped, which is exactly the coupling this design removes.
func TestOneModelsQuotaDoesNotBlockAnother(t *testing.T) {
up := upstream(t, &upstreamCtrl{})
defer 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: filepath.Join(td, "adapters"),
RuntimeFile: filepath.Join(td, "runtime.json"),
Keys: []config.GWKey{{
Key: "sk-a", Role: "user", Name: "two-models",
Models: []config.ModelScope{
{Model: "m1", TokenQuota: 4, Period: "hour"},
{Model: "m2", TokenQuota: 1000, Period: "hour"},
},
}},
Sources: []config.Source{{
Name: "up", BaseURL: up.URL, Adapter: "openai",
Models: []config.Model{
{ID: "m1", Priority: 100},
{ID: "m2", Priority: 90},
},
}},
}
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-a"})
if err != nil {
t.Fatalf("gateway: %v", err)
}
// Spend on m2 FIRST. Without this the two designs are
// indistinguishable: a key-wide counter and m1's own counter would both
// read 0 before m1 is used, so the test would pass either way (it did —
// see the commit that rewrote it).
for i := 1; i <= 3; i++ {
if rr, _ := chatAs(t, g, "sk-a", "m2"); rr.Code != 200 {
t.Fatalf("m2 priming call %d: want 200, got %d", i, rr.Code)
}
}
// m1's budget is 4 and each call costs 4, so the key-wide total (m2+m1) is
// already 12 when m1 starts: a key-wide cap would refuse m1 immediately.
if rr, _ := chatAs(t, g, "sk-a", "m1"); rr.Code != 200 {
t.Fatalf("m1 must be served on its own budget (m2's spend is not its problem), got %d (%s)",
rr.Code, rr.Body.String())
}
if rr, code := chatAs(t, g, "sk-a", "m1"); rr.Code != http.StatusTooManyRequests {
t.Fatalf("m1 second call: want 429, got %d (%s)", rr.Code, code)
}
// m2 must keep working
for i := 1; i <= 3; i++ {
if rr, _ := chatAs(t, g, "sk-a", "m2"); rr.Code != 200 {
t.Fatalf("m2 call %d must be served while m1 is capped, got %d (%s)", i, rr.Code, rr.Body.String())
}
}
// and the message must name m1, not the key
rr, _ := chatAs(t, g, "sk-a", "m1")
if !strings.Contains(rr.Body.String(), "m1") {
t.Errorf("rejection should name the capped model: %s", rr.Body.String())
}
}
// A model with no quota on it is never blocked by a sibling's cap, and an
// uncapped key is never blocked at all.
func TestUncappedModelNeverBlocked(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: "m1", TokenQuota: 1, Period: "hour"},
{Model: "m2"},
}})
if rr, _ := chatAs(t, g, "sk-a", "m1"); rr.Code != 200 {
t.Fatalf("m1 first: want 200, got %d", rr.Code)
}
if rr, _ := chatAs(t, g, "sk-a", "m1"); rr.Code != http.StatusTooManyRequests {
t.Fatalf("m1 second: want 429, got %d", rr.Code)
}
for i := 1; i <= 4; i++ {
if rr, _ := chatAs(t, g, "sk-a", "m2"); rr.Code != 200 {
t.Fatalf("m2 (uncapped) call %d: want 200, got %d", i, rr.Code)
}
}
}
// An admin key is never capped even when its scopes carry budgets: a cap that
// locked the operator out would be unrecoverable through the UI.
func TestAdminKeyScopesAreNotEnforced(t *testing.T) {
up := upstream(t, &upstreamCtrl{})
defer up.Close()
g := newQuotaGW(t, up.URL,
config.GWKey{Key: "sk-admin", Role: "admin", Models: []config.ModelScope{
{Model: "m1", TokenQuota: 1, ReqQuota: 1, Period: "hour"},
}})
for i := 1; i <= 4; i++ {
if rr, _ := chatAs(t, g, "sk-admin", "m1"); rr.Code != 200 {
t.Fatalf("admin call %d: want 200 (admin scopes are not enforced), got %d", i, rr.Code)
}
}
}