ModelRouter/internal/lua/plugins/usage.lua

-- 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 = modelrouter.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