From d3203d88fc517b733e46ca9ae3bcffecc902c83a Mon Sep 17 00:00:00 2001 From: JianFeeeee Date: Wed, 7 Oct 2026 16:32:58 +0800 Subject: [PATCH] =?UTF-8?q?feat(plugins):=20usage=20=E6=8F=92=E4=BB=B6?= =?UTF-8?q?=E8=BF=9B=E4=BB=93=EF=BC=8C=E9=9A=8F=E5=8C=85=E6=92=AD=E7=A7=8D?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 之前只通过 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 变红。 (cherry picked from commit ae2b1219599da6a960f453ca2d648a8d1a3b95f8) --- internal/lua/plugins/usage.lua | 151 ++++++++++++++++++++++++++++++ internal/lua/usage_plugin_test.go | 112 ++++++++++++++++++++++ 2 files changed, 263 insertions(+) create mode 100644 internal/lua/plugins/usage.lua create mode 100644 internal/lua/usage_plugin_test.go diff --git a/internal/lua/plugins/usage.lua b/internal/lua/plugins/usage.lua new file mode 100644 index 0000000..6193004 --- /dev/null +++ b/internal/lua/plugins/usage.lua @@ -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 \ No newline at end of file diff --git a/internal/lua/usage_plugin_test.go b/internal/lua/usage_plugin_test.go new file mode 100644 index 0000000..745298b --- /dev/null +++ b/internal/lua/usage_plugin_test.go @@ -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"`) + } +} \ No newline at end of file