ModelRouter/internal/lua/plugins/billing.lua

-- billing.lua — usage accounting plugin for ModelRouter.
--
-- Computes what each request cost, from three configurable dimensions:
--
--   source  a flat per-request price for an upstream source
--   model   a per-token price for a model id (prompt / completion separately)
--   key     an override price for one gateway key
--
-- It then keeps running totals for the whole gateway, per source, per model
-- and per key, and publishes them in `plugin.state` so the kernel can serve
-- them at GET /api/plugins/billing/state — which is what its own dashboard
-- component reads.
--
-- ACCOUNTING BOUNDARY (important, and deliberate):
-- this plugin REPORTS; it does not ENFORCE. The gateway's own quota accounting
-- (internal/gateway/stats.go, enforced at request admission) stays
-- authoritative for limits. Two independent accounting paths that disagree are
-- worse than one that is slightly less featureful, so nothing here feeds back
-- into routing or quota decisions.
--
-- PRICE CONFIGURATION
-- Prices are supplied as a Lua table assigned to `billing.prices` before the
-- plugin is loaded, OR at runtime through PUT /api/plugins/billing/state. The
-- shape is:
--
--   billing.prices = {
--     currency = "USD",              -- display only, no conversion happens
--     default  = { prompt = 0, completion = 0, per_request = 0 },
--     sources  = {
--       ["localzen"]  = { per_request = 0.0 },
--       ["trae"]      = { per_request = 0.01 },
--     },
--     models   = {
--       ["gpt-5.4"]   = { prompt = 1.25e-6, completion = 1e-5 },  -- USD per TOKEN
--       ["kimi-k3"]   = { prompt = 6e-7,  completion = 2.5e-6 },
--       ["kolors"]    = { per_request = 0.04 },                    -- image: flat
--     },
--     keys     = {
--       -- by gateway key (the same value the audit log masks to ***xxxxxx)
--       ["***a1b2c3"] = { prompt = 1.1e-6, completion = 9e-6 },
--     },
--   }
--
-- Precedence for a token price: keys > models > default. A flat per_request
-- price, when present at any level, is ADDED on top of the token cost, so an
-- image model can carry both (e.g. tokens billed plus a fixed fee).
--
-- Numbers are USD per single token, which is how providers publish prices. That
-- makes a typical entry look like 1.25e-6; the plugin multiplies by the token
-- count, so no unit conversion happens anywhere.

local plugin = {
  name = "billing",
  version = "1.0.0",
  description = "Per-source / per-model / per-key cost accounting with a dashboard",
  author = "ModelRouter",
}

-- ---------- prices ----------

-- plugin.prices can be pre-seeded by embedding this file (an operator edits the
-- table below) or replaced at runtime through the state API. It is a SEPARATE
-- field from plugin.state on purpose: PUT /api/plugins/billing/state replaces
-- `state` wholesale, and prices must not live there or a price update would
-- wipe the accumulated totals. See docs/plugins.md.
local DEFAULT_PRICES = {
  currency = "USD",
  default  = { prompt = 0, completion = 0, per_request = 0 },
  sources  = {},
  models   = {},
  keys     = {},
}
plugin.prices = DEFAULT_PRICES

-- ---------- accumulated totals ----------

-- state is what the kernel serves at GET /api/plugins/billing/state. It holds
-- ACCUMULATED TOTALS ONLY — prices live in plugin.prices (see above), so
-- replacing state never destroys a price table and updating prices never
-- destroys history.
--
-- Structure:
--   total          { cost, requests, prompt_tokens, completion_tokens }
--   by_source      { <name> = { cost, requests, ...tokens } }
--   by_model       { <model> = { cost, ... } }
--   by_key         { <masked key id> = { cost, ... } }
--   by_day         { "YYYY-MM-DD" = { cost, ... } }
--   top_sources    [ {name, cost, requests}, ... ]   sorted, capped
--   top_models     [ ... ]
--   top_keys       [ ... ]
--
-- Sorted top-N lists are maintained incrementally rather than re-sorted on
-- every request: this hook runs once per request on the hot path, so it does
-- map updates only. The sort happens when state is READ.
local function emptyBucket()
  return { cost = 0, requests = 0, prompt_tokens = 0, completion_tokens = 0, failures = 0 }
end

plugin.state = {
  total     = emptyBucket(),
  by_source = {},
  by_model  = {},
  by_key    = {},
  by_day    = {},
  started   = os.time and 0 or 0,
}

local function bucket(tbl, k)
  local b = tbl[k]
  if b == nil then
    b = emptyBucket()
    tbl[k] = b
  end
  return b
end

local function add(b, cost, prompt, completion, ok)
  b.cost = b.cost + cost
  b.requests = b.requests + 1
  b.prompt_tokens = b.prompt_tokens + prompt
  b.completion_tokens = b.completion_tokens + completion
  if not ok then b.failures = b.failures + 1 end
end

-- ---------- pricing ----------

-- lookup walks keys > models > default and returns a price triple plus whether
-- a flat per_request component applies.
local function priceFor(payload)
  local p = plugin.prices or DEFAULT_PRICES
  local d = p.default or {}
  local out = { prompt = d.prompt or 0, completion = d.completion or 0, per_request = 0 }

  -- model dimension (a token price overrides the default's token prices)
  local mp = p.models and p.models[payload.model]
  if mp then
    if mp.prompt ~= nil then out.prompt = mp.prompt end
    if mp.completion ~= nil then out.completion = mp.completion end
    if mp.per_request ~= nil then out.per_request = out.per_request + mp.per_request end
  end

  -- source dimension: usually a flat fee, but may also carry token prices
  local sp = p.sources and p.sources[payload.source]
  if sp then
    if sp.prompt ~= nil then out.prompt = sp.prompt end
    if sp.completion ~= nil then out.completion = sp.completion end
    if sp.per_request ~= nil then out.per_request = out.per_request + sp.per_request end
  end

  -- key dimension wins over the others (an operator pricing one customer
  -- specially must be able to override both the model and the source price)
  local kp = p.keys and p.keys[payload.key]
  if kp then
    if kp.prompt ~= nil then out.prompt = kp.prompt end
    if kp.completion ~= nil then out.completion = kp.completion end
    if kp.per_request ~= nil then out.per_request = out.per_request + kp.per_request end
  end
  return out
end

-- costFor computes one request's price.
--
-- A FAILED request still costs money whenever the upstream billed for it, which
-- the kernel cannot know; the conservative and useful default is to charge
-- failed requests their token cost (a 500 after generation still consumed
-- tokens) but NOT a flat per_request fee that was never actually charged. That
-- is what the `ok` flag selects below, and it is the single most debatable
-- policy in this file — it is a config toggle so an operator can flip it.
local function costFor(payload)
  local price = priceFor(payload)
  local prompt = tonumber(payload.prompt_tokens) or 0
  local completion = tonumber(payload.completion_tokens) or 0
  local cost = prompt * price.prompt + completion * price.completion

  local flat = price.per_request
  if not payload.ok and not plugin.count_failures then
    flat = 0
  end
  return cost + flat
end

-- ---------- day bucket ----------

local function dayKey(epoch_seconds)
  -- os.date is available in LuaJIT; fall back to a UTC-ish arithmetic stamp if
  -- the host build has no os.date (keeps the plugin from erroring out on a
  -- stripped runtime, which would otherwise look like a plugin failure).
  if os and os.date then
    return os.date("!%Y-%m-%d", epoch_seconds)
  end
  return tostring(math.floor(epoch_seconds / 86400))
end

-- ---------- hooks ----------

plugin.hooks = {
  -- chain_step gives the per-tier walk; request_end gives the final accounting.
  -- Subscribing to chain_step is OPTIONAL here: the totals are driven by
  -- request_end alone, and the degradation counters below are pure observation.
  -- A gateway with thousands of requests can drop this hook to save the
  -- per-step Lua call without losing a single billed request.
  chain_step  = "on_chain_step",
  request_end = "on_request_end",
}

-- Tracks how often a request had to drop below the top tier, and which tier
-- actually served it. Without this, "tier 1 was cooling" and "tier 1 served it"
-- are indistinguishable in the accounts, and a quietly degraded gateway looks
-- exactly like a healthy one.
plugin.state.degraded_reqs = 0
plugin.state.by_tier_served = {}
plugin.state.skip_reasons = {}

function plugin.on_chain_step(payload)
  if payload == nil then return nil end
  local s = plugin.state
  if s == nil then return nil end
  if s.by_tier_served == nil then s.by_tier_served = {} end
  if s.skip_reasons == nil then s.skip_reasons = {} end

  if payload.kind == "selected" then
    local t = tostring(payload.tier or "?")
    s.by_tier_served[t] = (s.by_tier_served[t] or 0) + 1
  elseif payload.kind == "tier_skip" or payload.kind == "tier_busy" then
    -- reason text is the ACTIONABLE part; normalise the volatile bits so the
    -- same cause aggregates instead of creating a new row per request.
    local r = tostring(payload.reason or payload.kind or "unknown")
    r = string.gsub(r, "within [%d%.%a]+", "within <wait>")
    s.skip_reasons[r] = (s.skip_reasons[r] or 0) + 1
  end
  return nil
end

function plugin.on_request_end(payload)
  if payload == nil then return nil end
  local prompt = tonumber(payload.prompt_tokens) or 0
  local completion = tonumber(payload.completion_tokens) or 0
  local ok = payload.ok and true or false
  local cost = costFor(payload)

  local s = plugin.state
  -- Rebuild any missing container. This is reached in two real situations:
  -- a fresh plugin, and an admin who PUT a partial state (e.g. only "prices"),
  -- which legitimately replaces `state` with a sparse table. Checking only the
  -- outer table would leave `s.total` nil and crash the hook on the next call.
  if s == nil then s = {} plugin.state = s end
  if s.total == nil then s.total = emptyBucket() end
  if s.by_source == nil then s.by_source = {} end
  if s.by_model == nil then s.by_model = {} end
  if s.by_key == nil then s.by_key = {} end
  if s.by_day == nil then s.by_day = {} end
  if s.started == nil then s.started = payload.time or 0 end
  if s.degraded_reqs == nil then s.degraded_reqs = 0 end
  -- Degradation is counted here rather than in the chain_step hook because
  -- request_end sees the whole walk at once: one degraded request must count
  -- once, whereas the walk may contain several skipped tiers.
  if payload.degraded then s.degraded_reqs = s.degraded_reqs + 1 end

  add(s.total, cost, prompt, completion, ok)
  if payload.source ~= nil and payload.source ~= "" then
    add(bucket(s.by_source, payload.source), cost, prompt, completion, ok)
  end
  if payload.model ~= nil and payload.model ~= "" then
    add(bucket(s.by_model, payload.model), cost, prompt, completion, ok)
  end
  if payload.key ~= nil and payload.key ~= "" then
    add(bucket(s.by_key, payload.key), cost, prompt, completion, ok)
  end

  -- Daily rollup, so the dashboard can draw a trend without the browser
  -- re-deriving it. Keyed off the request's own timestamp, not os.time(), so a
  -- replayed or imported record lands on the right day.
  local ts = payload.time
  if ts ~= nil and ts > 0 then
    if ts > 1000000000000 then ts = ts / 1000 end  -- kernel sends unix MILLIseconds
    add(bucket(s.by_day, dayKey(ts)), cost, prompt, completion, ok)
  end
  return nil  -- last stage: nobody downstream would read a return value
end

-- ---------- dashboard UI ----------

-- A whole page. The kernel injects this HTML and evaluates the <script> after
-- the DOM exists, and exposes `pluginAPI` for talking to the gateway.
plugin.ui = {
  page = {
    page_id = "billing",
    title = "Billing",
    icon = "💰",
    order = 40,
    mount = [==[
<div id="billing-root" style="padding:16px">
  <div class="kpis" id="billing-kpis" style="display:grid;grid-template-columns:repeat(auto-fit,minmax(170px,1fr));gap:12px;margin-bottom:18px"></div>
  <div style="display:grid;grid-template-columns:repeat(auto-fit,minmax(320px,1fr));gap:16px">
    <div class="card" style="padding:14px">
      <h3 style="margin:0 0 10px;font-size:14px">Per source</h3>
      <div id="billing-by-source"></div>
    </div>
    <div class="card" style="padding:14px">
      <h3 style="margin:0 0 10px;font-size:14px">Per model</h3>
      <div id="billing-by-model"></div>
    </div>
    <div class="card" style="padding:14px">
      <h3 style="margin:0 0 10px;font-size:14px">Per gateway key</h3>
      <div id="billing-by-key"></div>
    </div>
  </div>
  <div class="card" style="padding:14px;margin-top:16px">
    <h3 style="margin:0 0 10px;font-size:14px">Daily</h3>
    <div id="billing-by-day"></div>
  </div>
</div>
<script>
(function () {
  var ROOT = "billing";
  function fmt(n) {
    if (n === null || n === undefined) return "-";
    n = Number(n);
    if (!isFinite(n)) return "-";
    if (n === 0) return "0";
    if (Math.abs(n) < 0.000001) return n.toExponential(2);
    return n.toFixed(Math.abs(n) < 1 ? 6 : 4);
  }
  function money(v, cur) { return (cur || "USD") + " " + fmt(v); }
  function esc(s) {
    return String(s == null ? "" : s).replace(/[&<>"]/g, function (c) {
      return { "&": "&amp;", "<": "&lt;", ">": "&gt;", '"': "&quot;" }[c];
    });
  }
  function row(name, b, cur) {
    return "<tr><td><b>" + esc(name) + "</b></td><td>" + money(b.cost, cur) +
      "</td><td>" + (b.requests || 0) + "</td><td>" + (b.prompt_tokens || 0) +
      "</td><td>" + (b.completion_tokens || 0) + "</td></tr>";
  }
  function tableFor(el, obj, cur, empty) {
    var keys = Object.keys(obj || {});
    if (!keys.length) { el.innerHTML = '<div class="muted">' + empty + "</div>"; return; }
    keys.sort(function (a, b) { return (obj[b].cost || 0) - (obj[a].cost || 0); });
    var h = "<table style='width:100%;border-collapse:collapse;font-size:13px'>" +
      "<tr style='text-align:left;opacity:.65'><th>name</th><th>cost</th><th>reqs</th>" +
      "<th>prompt</th><th>completion</th></tr>";
    for (var i = 0; i < keys.length; i++) {
      var k = keys[i];
      h += "<tr style='border-top:1px solid rgba(120,90,150,.14)'>" + row(k, obj[k], cur) + "</tr>";
    }
    el.innerHTML = h + "</table>";
  }
  function render(st) {
    if (!st) return;
    var cur = (st.currency || "USD");
    var t = st.total || {};
    document.getElementById("billing-kpis").innerHTML = [
      ["Total", money(t.cost, cur)],
      ["Requests", t.requests || 0],
      ["Degraded", s.degraded_reqs || 0],
      ["Prompt tokens", t.prompt_tokens || 0],
      ["Completion tokens", t.completion_tokens || 0],
      ["Failures", t.failures || 0]
    ].map(function (kv) {
      return "<div class='card' style='padding:12px'><div style='font-size:11px;opacity:.65'>" +
        kv[0] + "</div><div style='font-size:19px;font-weight:600;margin-top:4px'>" +
        esc(kv[1]) + "</div></div>";
    }).join("");
    tableFor(document.getElementById("billing-by-source"), st.by_source, cur, "no per-source data yet");
    tableFor(document.getElementById("billing-by-model"), st.by_model, cur, "no per-model data yet");
    tableFor(document.getElementById("billing-by-key"), st.by_key, cur, "no per-key data yet");
    tableFor(document.getElementById("billing-by-day"), st.by_day, cur, "no daily data yet");
  }
  async function refresh() {
    try {
      var r = await fetch("/api/plugins/" + ROOT + "/state", { credentials: "same-origin" });
      if (!r.ok) return;
      var j = await r.json();
      render(j.state);
    } catch (e) { /* the pane is optional decoration; never break the page */ }
  }
  window.__billingRefresh = refresh;
  refresh();
  if (window.pluginAPI && pluginAPI.onTabShown) pluginAPI.onTabShown(refresh);
})();
</script>
]==],
  },
  -- Two elements on the EXISTING status page: a headline tile and a
  -- per-source cost breakdown, so the number is visible without opening the
  -- Billing tab.
  elements = {
    {
      target = "status",
      anchor = "top",
      order = 5,
      mount = [==[
<div class="card" id="billing-status-tile" style="padding:12px;margin-bottom:12px">
  <div style="font-size:11px;opacity:.65">Total spend (billing plugin)</div>
  <div id="billing-status-total" style="font-size:22px;font-weight:600;margin-top:4px">—</div>
  <div id="billing-status-sub" style="font-size:12px;opacity:.65;margin-top:2px"></div>
</div>
<script>
(function () {
  function fmt(n) {
    n = Number(n || 0);
    if (n === 0) return "0";
    if (Math.abs(n) < 0.000001) return n.toExponential(2);
    return n.toFixed(Math.abs(n) < 1 ? 6 : 4);
  }
  async function tick() {
    try {
      var r = await fetch("/api/plugins/billing/state", { credentials: "same-origin" });
      if (!r.ok) return;
      var j = await r.json();
      var st = j.state;
      if (!st || !st.total) return;
      var cur = st.currency || "USD";
      document.getElementById("billing-status-total").textContent = cur + " " + fmt(st.total.cost);
      var parts = [];
      var srcs = st.by_source || {};
      var names = Object.keys(srcs).sort(function (a, b) {
        return (srcs[b].cost || 0) - (srcs[a].cost || 0);
      });
      for (var i = 0; i < Math.min(3, names.length); i++) {
        parts.push(names[i] + " " + fmt(srcs[names[i]].cost));
      }
      document.getElementById("billing-status-sub").textContent =
        (st.total.requests || 0) + " requests" + (parts.length ? " · top: " + parts.join(" · ") : "");
    } catch (e) { /* decoration only */ }
  }
  if (window.pluginAPI && pluginAPI.onTabShown) pluginAPI.onTabShown(tick);
  tick();
})();
</script>
]==],
    },
  },
}

return plugin