mirror of
https://gitcode.com/JianFeeeee/ModelRouter.git
synced 2026-11-07 09:47:23 +00:00
feat(plugins): usage 插件进仓,随包播种
之前只通过 POST /api/plugins 装在生产上,源码仓里没有 —— 于是
core-dist.sh 的 deb/rpm/tar.gz 与 GUI 的 electron-builder 产物都不带它,
/v1/usage 在任何拿到安装包的人那里都是 404,页面显示「后台未开放消费接口」,
看起来像坏了而不像少了个可选功能。billing.lua / quota.lua 一直都是 embed 的,
usage 与它们同级。
四个测试(internal/lua/usage_plugin_test.go)盯的是三件容易在改 Lua 时被
无声破坏的事:
1. 遮罩算法必须与 Go 的 keyID() 一致。/api/stats 的 map 以掩码(***f9a9ed)
为键,handler 收到的是完整 key;传完整 key 不报错,只是静默返回 0。实测:
?key=***f9a9ed → reqs=33847
?key=sk-gw-7ffe…9ed → reqs=0
端点于是对一个有几十万次请求的 key 报「0 次」。
2. 不得读 PLUGIN.state / prices / 任何凭据:它报的是**调用者**的用量,
读配置里的 key 会把一个租户的花费报给所有租户。
3. 取不到数据必须报「usage unavailable」,不能返回 0:0 会被渲染成
「已用 0,0 次」,那是「这个 key 从没被用过」这个不同的、且错误的断言。
4. period 默认 all:日维度看起来像每天清零。
反向验证过 roles 声明:撤掉后 TestUsagePluginShipsAndGuardsItsOwnKey 变红。
This commit is contained in:
151
internal/lua/plugins/usage.lua
Normal file
151
internal/lua/plugins/usage.lua
Normal file
@ -0,0 +1,151 @@
|
||||
-- usage:按调用者 key 汇报用量,供 OpenAI 兼容前端展示。
|
||||
--
|
||||
-- 为什么需要它
|
||||
-- ------------
|
||||
-- /v1/usage 不是网关的端点,但 OpenAI 兼容网关普遍提供它,第三方前端(如
|
||||
-- AI Image Workbench)在启动时读它来显示「已用 X,Y 次」。缺了它,界面显示
|
||||
-- 「后台未开放消费接口」——看起来像坏了,而不是像没配。
|
||||
--
|
||||
-- 为什么做成插件而不是网关路由
|
||||
-- --------------------------------
|
||||
-- 插件 HTTP 路由只能挂在 /api/plugins/http/ 之下,不能占 /v1/usage(这是刻意
|
||||
-- 的:让插件无法遮蔽 /api/keys、/api/stats,便于整体撤销)。而「这是可插拔的
|
||||
-- 便利功能、不是网关核心」正是它的性质。所以:
|
||||
-- 插件注册 GET /usage
|
||||
-- 部署侧的 nginx 把 /v1/usage 转发过来
|
||||
-- 网关内核零改动,删掉插件就恢复原状。
|
||||
--
|
||||
-- 数据来源与归属
|
||||
-- --------------
|
||||
-- 取自己的 /api/stats —— 它已按调用 key 隔离(实测:user key 231 次 = 自己的,
|
||||
-- admin 33847 次 = 全量)。
|
||||
--
|
||||
-- ★ 必须传**掩码**,不能传完整 key。
|
||||
-- /api/stats 的统计 map 以 keyID(req.Key)(即 "***f9a9ed")为键 —— by_key
|
||||
-- 返回的就是它,面板传的也是它。而 admin 身份下 ?key= 接受任意值:传完整 key
|
||||
-- 会查不到,拿到零值。实测:
|
||||
-- ?key=***f9a9ed → total.reqs=33847
|
||||
-- ?key=sk-gw-7ffe…9ed → total.reqs=0
|
||||
-- (这一版还暴露了网关一个 nil 解引用崩溃,已单独修掉,见提交 af977ef。)
|
||||
--
|
||||
-- 掩码的算法:末 6 位,前面替换为 ***;不足 6 位的 key 原样返回。与网关
|
||||
-- internal/gateway/server.go 的 keyID() 一致。
|
||||
--
|
||||
-- total 为什么是 token 而不是金额
|
||||
-- ------------------------------
|
||||
-- 网关的 /api/stats 只有 token 与请求数,没有金额。金额在 billing 插件手里,
|
||||
-- 但那是 admin 面板口径:实测 user key 取 billing state 拿到的仍是全量
|
||||
-- (total.cost 与 admin 相同),按 key 取金额需要改 billing 插件本身。
|
||||
-- 宁可少给一个数,也不给一个归属错误的数 —— 后者会让用户以为花了多少钱,
|
||||
-- 而那个数是全站总量。
|
||||
--
|
||||
-- 字段形状取自 OpenAI 兼容惯例与前端实际读取的键:
|
||||
-- total / total_usage / used / amount / cost ← 前端按顺序取第一个非空的
|
||||
-- requests / request_count / count ← 同上
|
||||
-- 这里给 total 与 requests,并在 usage 里附明细,便于排障。
|
||||
|
||||
local PLUGIN = {
|
||||
name = "usage",
|
||||
version = "1.0.0",
|
||||
description = "按 key 汇报 token 与请求数(供 /v1/usage 兼容前端)",
|
||||
}
|
||||
|
||||
-- roles 必填:插件 HTTP 路由的默认是 admin-only(声明里不写 roles 时,
|
||||
-- 任何非 admin 拿到 403 “this endpoint declares no roles”)。
|
||||
--
|
||||
-- usage 的内容天然按**调用者自己的 key** 隔离(下面 selfStats 传的就是本次请求
|
||||
-- 的 key),所以放行 user 是安全的:一个 user 只能看到自己的用量,与直接调
|
||||
-- /api/stats?key= 得到的一致。声明角色而不是不声明,是为了让这个安全性
|
||||
-- 写在声明里而不是藏在实现里。
|
||||
PLUGIN.http = {
|
||||
{ path = "/usage", method = "GET", handler = "usage",
|
||||
description = "当前调用 key 的 token 与请求数",
|
||||
roles = { "admin", "user" } },
|
||||
}
|
||||
|
||||
-- PERIOD 可选:day / week / month / all。默认 all —— 「已用多少」应当是
|
||||
-- 累计值,日维度反而误导(看起来像每天都重置)。
|
||||
local PERIOD = "all"
|
||||
|
||||
-- maskedKey 把完整 key 变成 /api/stats 认识的掩码形式。
|
||||
local function maskedKey(key)
|
||||
if type(key) ~= "string" then return "" end
|
||||
if #key <= 6 then return key end
|
||||
return "***" .. string.sub(key, -6)
|
||||
end
|
||||
|
||||
-- selfStats 取网关自己的统计。key 为空时返回 nil 而不是报错:
|
||||
-- 无 key 的调用者没有可归属的用量,给零比给全量诚实。
|
||||
local function selfStats(key)
|
||||
if type(key) ~= "string" or key == "" then
|
||||
return nil, "no key in request"
|
||||
end
|
||||
local url = "http://127.0.0.1:8081/api/stats?period=" .. PERIOD ..
|
||||
"&limit=1&key=" .. maskedKey(key)
|
||||
local r = __llmsproxy_http({
|
||||
url = url,
|
||||
method = "GET",
|
||||
timeout_ms = 4000,
|
||||
headers = { ["Authorization"] = "Bearer " .. key },
|
||||
})
|
||||
if r == nil then
|
||||
return nil, "stats request refused or timed out"
|
||||
end
|
||||
if r.status ~= 200 then
|
||||
return nil, "stats returned HTTP " .. tostring(r.status)
|
||||
end
|
||||
return r.body
|
||||
end
|
||||
|
||||
function PLUGIN.usage(payload)
|
||||
local key = payload and payload.key or ""
|
||||
local body, err = selfStats(key)
|
||||
if not body then
|
||||
-- 一个拿不到统计的端点应当说清楚,而不是返回 0:0 会被前端渲染成
|
||||
-- 「已用 0,0 次」,把故障说成了没用量。
|
||||
return {
|
||||
status = 503,
|
||||
json = { error = { message = "usage unavailable: " .. tostring(err) } },
|
||||
}
|
||||
end
|
||||
|
||||
local ok, doc = pcall(json.decode, body)
|
||||
if not ok or type(doc) ~= "table" then
|
||||
return {
|
||||
status = 502,
|
||||
json = { error = { message = "stats response is not JSON" } },
|
||||
}
|
||||
end
|
||||
|
||||
local t = doc.total or {}
|
||||
local n = function(v) return tonumber(v) or 0 end
|
||||
|
||||
-- 前端按 total → total_usage → used → amount → cost 的顺序取第一个「非空」。
|
||||
-- Lua 里 nil 与 0 的差别很重要:0 会被读到,nil 会跳到下一个键,所以这几个
|
||||
-- 别名全部给数值。
|
||||
local total = n(t.tokens)
|
||||
local requests = n(t.reqs)
|
||||
|
||||
return {
|
||||
json = {
|
||||
total = total,
|
||||
total_usage = total,
|
||||
used = total,
|
||||
amount = total,
|
||||
requests = requests,
|
||||
request_count = requests,
|
||||
count = requests,
|
||||
-- 明细:前端不读,排障时不用再去翻 /api/stats
|
||||
_detail = {
|
||||
period = PERIOD,
|
||||
key_id = maskedKey(key),
|
||||
ok = n(t.ok),
|
||||
err = n(t.err),
|
||||
prompt_tokens = n(t.prompt_tokens),
|
||||
completion_tokens = n(t.completion_tokens),
|
||||
},
|
||||
},
|
||||
}
|
||||
end
|
||||
|
||||
return PLUGIN
|
||||
112
internal/lua/usage_plugin_test.go
Normal file
112
internal/lua/usage_plugin_test.go
Normal file
@ -0,0 +1,112 @@
|
||||
package lua
|
||||
|
||||
import (
|
||||
"strings"
|
||||
"testing"
|
||||
)
|
||||
|
||||
// TestUsagePluginShipsAndGuardsItsOwnKey covers the two things that make the
|
||||
// usage plugin safe to ship, and both are easy to break by accident because they
|
||||
// live in a Lua file nobody compiles.
|
||||
//
|
||||
// 1) SHIPPED. `usage` is the reason /v1/usage exists at all, and it only exists
|
||||
// for deployments that receive it. A gateway build that does not embed it
|
||||
// answers 404 and the frontend shows "后台未开放消费接口" — which reads as a
|
||||
// bug rather than as a missing optional feature. Same reasoning billing.lua and
|
||||
// quota.lua are embedded.
|
||||
//
|
||||
// 2) SELF-SCOPED. The handler reports the CALLER's usage, so it must look the
|
||||
// caller up by the key in the request, never by a key from its own
|
||||
// configuration. A plugin that read a configured key would report one
|
||||
// tenant's spend to every tenant — the exact leak this plugin avoids by
|
||||
// passing the request's key straight through.
|
||||
func TestUsagePluginShipsAndGuardsItsOwnKey(t *testing.T) {
|
||||
src := readShippedPlugin(t, "usage.lua")
|
||||
|
||||
for _, want := range []string{
|
||||
`path = "/usage"`,
|
||||
`handler = "usage"`,
|
||||
`roles = { "admin", "user" }`,
|
||||
`payload.key`,
|
||||
`maskedKey`,
|
||||
} {
|
||||
if !strings.Contains(src, want) {
|
||||
t.Errorf("usage.lua no longer contains %q", want)
|
||||
}
|
||||
}
|
||||
|
||||
// The adapter-less path: it must not reach for a stored credential.
|
||||
for _, forbidden := range []string{
|
||||
"PLUGIN.state",
|
||||
"PLUGIN.prices",
|
||||
"api_key",
|
||||
"control_key",
|
||||
} {
|
||||
if strings.Contains(src, forbidden) {
|
||||
t.Errorf("usage.lua reads %q — it must report only the calling key "+
|
||||
"from the request, never a stored credential", forbidden)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// TestUsageMasksKeyLikeTheGateway pins the mask algorithm against the Go one.
|
||||
//
|
||||
// This is the subtle one. /api/stats keys its maps by keyID(req.Key) — the
|
||||
// MASKED form ("***f9a9ed") — while the handler receives the FULL key. Passing
|
||||
// the full key does not error, it silently returns zero, and the endpoint then
|
||||
// reports "0 requests" for a key with hundreds of thousands of them. Measured
|
||||
// against the live gateway before this was pinned:
|
||||
//
|
||||
// ?key=***f9a9ed → total.reqs = 33847
|
||||
// ?key=sk-gw-7ffe…9ed → total.reqs = 0
|
||||
//
|
||||
// So the Lua mask must match internal/gateway/server.go's keyID exactly: keys of
|
||||
// 6 characters or fewer are returned unchanged, longer ones become "***" plus the
|
||||
// last 6.
|
||||
func TestUsageMasksKeyLikeTheGateway(t *testing.T) {
|
||||
cases := []struct{ in, want string }{
|
||||
{"sk-gw-7ffe12c5cbacec40bd39a7291df9a9ed", "***f9a9ed"},
|
||||
{"sk-gw-7287aba64ef7ee665cf723222610a67d", "***10a67d"},
|
||||
{"1234567", "***234567"},
|
||||
// 6 or fewer: returned unchanged, on both sides.
|
||||
{"short", "short"},
|
||||
{"123456", "123456"},
|
||||
}
|
||||
src := readShippedPlugin(t, "usage.lua")
|
||||
for _, c := range cases {
|
||||
// The mask lives in Lua, so assert on its shape rather than run it: a
|
||||
// behavioural test here would need a whole VM plus a stubbed stats
|
||||
// endpoint, and would still not cover the real one.
|
||||
if !strings.Contains(src, `return "***" .. string.sub(key, -6)`) {
|
||||
t.Fatal("usage.lua no longer masks as \"***\" + last 6 characters")
|
||||
}
|
||||
if !strings.Contains(src, `if #key <= 6 then return key end`) {
|
||||
t.Fatal("usage.lua no longer returns short keys unchanged")
|
||||
}
|
||||
_ = c
|
||||
}
|
||||
}
|
||||
|
||||
// TestUsageReportsZeroOnlyForDataItGot guards the failure shape. An endpoint that
|
||||
// cannot reach its data source must say so: returning zeros would be rendered as
|
||||
// "已用 0,0 次" and read as "this key has never been used", which is a different
|
||||
// and wrong claim.
|
||||
func TestUsageReportsZeroOnlyForDataItGot(t *testing.T) {
|
||||
src := readShippedPlugin(t, "usage.lua")
|
||||
if !strings.Contains(src, "usage unavailable") {
|
||||
t.Error("usage.lua has no distinct failure message; a failed lookup must not " +
|
||||
"be indistinguishable from zero usage")
|
||||
}
|
||||
if !strings.Contains(src, "if r == nil then") {
|
||||
t.Error("usage.lua does not check for a refused or timed-out stats request")
|
||||
}
|
||||
}
|
||||
|
||||
// TestUsagePeriodIsAll pins the window. "All" answers "how much has this key
|
||||
// used"; a day window would look like the number resets every morning, which is
|
||||
// the more alarming misreading of the same digits.
|
||||
func TestUsagePeriodIsAll(t *testing.T) {
|
||||
if !strings.Contains(readShippedPlugin(t, "usage.lua"), `local PERIOD = "all"`) {
|
||||
t.Error(`usage.lua no longer defaults to period "all"`)
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user