mirror of
https://gitcode.com/JianFeeeee/ModelRouter.git
synced 2026-10-03 23:54:06 +00:00
feat(plugin): Lua 插件机制 + 计费插件 + 插件文档
插件 = plugin_dir 下的单个 .lua 文件，做两件事：挂请求流水线的钩子、在启动时
贡献 WebUI 界面（整页或往现有页面追加组件）。两者独立。
## 流水线 stage（三个）
request_start 已解析鉴权、未选源
routed 已选定 (source, model)、未发往上游
request_end 每请求恰好一次，带最终计量
request_end 挂在 gateway.writeRec——四条入口路径（直连/AUTO × 流式/非流式）的
唯一汇合点：既不漏（流式 token 只有流结束才知道）也不重。
## 计费插件（plugins/billing.lua，默认 seed，开箱可用）
源 / 模型 / 密钥三个维度定价。token 价优先级 keys > models > default；per_request
固定价是**叠加**的（生图模型可以既算 token 又收固定费）。单位是 USD/单 token，
即各家 provider 的公布口径。累计 total / by_source / by_model / by_key / by_day。
失败请求保留 token 费用、丢弃固定费（可经 count_failures 翻转）。
界面 = 一个独立页 + 状态页顶部一块总开销 tile。
## 一个明确的设计边界
计费插件**只报表，不执法**。网关自己的配额会计（stats.go，入口强制）才是限额
权威，插件不参与任何路由/配额决策。两套独立会计若对不上，比一套功能略少的
更糟。
## ★ 中途改掉的一个根本设计错误
最初让插件复用适配器的**弹性 worker 池**（多状态）。这对适配器是对的（它们无
状态），对插件是错的：计费插件往 plugin.state 累加，多状态意味着总量被劈成
几份；而 SetState 写价格只写进其中一个 worker，钩子恰好跑到另一个时**所有请求
按 0 计费**。改为**单状态 + 互斥锁**。代价写进文档：钩子必须短、同步、不阻塞，
卡住的钩子会卡住所有插件的钩子。
这个 bug 是测试逼出来的——先写了 SetState+Fire 的用例，数字全是 0 才挖出来。
另一个连带缺陷：只带 prices 的 PUT 会整体替换 state，把累计量清零。改为
prices/state 分离——prices 是配置、state 是历史，改价不动账。
## 撞到的三个 Lua 绑定的坑（都写进注释）
- SetGlobal **会 pop 栈**：连着调两次，第二次从空栈取，赋成 nil
- GetField 索引越界是 **SIGABRT 整个进程**，不是 panic，recover 救不了
- Call(nargs, n) **不接受函数索引**，它调的是 nargs 个参数正下方那个；
传索引会调到参数上（"attempt to call a table value"）
另外 GetField/SetField 用绝对索引，SetTop(0) 之后必须重取。
## 错误隔离
钩子 error() 不影响转发：捕获 → 记进 hook_errors → 跳下一个插件。适配器出错
会让源进冷却，插件出错**零惩罚**——插件是可选功能。/api/plugins 的 hook_errors
让"坏掉的插件"可见而不是静默消失。
## 界面注入
GET /api/ui-inject 一次返回所有插件的扩展（侧栏需要全部 page 才能建好）。
WebUI 在首次 render **之前** await 注入：先插 HTML 再重建 <script> 让它执行
（innerHTML/template 插入的 script 不会执行，这正是要的效果——避免脚本跑在
自己 DOM 之前）。注入失败不影响仪表盘。
browser 侧 pluginAPI 暴露 fetchState / postState / onTabShown。
## 文档
docs/plugins.md —— 快速上手、加载与热更新、三个 stage 的完整字段表、界面扩展、
状态与 HTTP API、运行时约束（单状态/异常隔离/内置函数）、计费插件的定价与
计费策略、排错表、与适配器的对比表。
## 判据（328 个测试全绿，插件相关 33 个）
- 计费断言的是**具体金额**（0.00625 / 0.0402 / 0.0075…），不是"能加载"
- 4 个变异都红：钩子异常不隔离 / prices 清空累计 / 忽略 key 优先级 /
毫秒时间戳不换算
- UI 侧 6 个判据把注入顺序、script 执行时机、pluginAPI 名称、tab 路由、
anchor 四种形式、失败非致命全钉住
- 鉴权：state 读任意角色、写仅 admin
This commit is contained in:
JianFeeeee
408
docs/plugins.md
Normal file
408
docs/plugins.md
Normal file
@ -0,0 +1,408 @@
|
||||
# 插件系统（Plugin System）
|
||||
|
||||
> **English**: this document is the reference for writing ModelRouter plugins.
|
||||
> The Chinese version is the primary one; section titles map 1:1.
|
||||
|
||||
ModelRouter 的插件是**单个 `.lua` 文件**，放在 `config.yaml` 的 `plugin_dir` 目录里。
|
||||
插件能做两件事：
|
||||
|
||||
1. **挂钩子**：在请求流水线的若干 stage 上注册回调，看到每个请求的完整信息，
|
||||
并可以把结果累加进自己的状态。
|
||||
2. **贡献界面**：在启动时返回 HTML / CSS / JS，由内核注入 WebUI——可以是一整个
|
||||
新页面，也可以是往现有页面里追加一个组件。
|
||||
|
||||
两者互相独立：只想统计请求数的插件不必碰界面；只想加个仪表盘的插件不必碰钩子。
|
||||
|
||||
---
|
||||
|
||||
## 1. 快速上手
|
||||
|
||||
一个最小的插件：
|
||||
|
||||
```lua
|
||||
-- plugins/hello.lua
|
||||
local plugin = {
|
||||
name = "hello",
|
||||
version = "1.0.0",
|
||||
description = "示例插件",
|
||||
author = "you",
|
||||
}
|
||||
|
||||
-- 声明钩子
|
||||
plugin.hooks = {
|
||||
request_end = "on_request_end",
|
||||
}
|
||||
|
||||
-- 钩子实现
|
||||
function plugin.on_request_end(payload)
|
||||
-- payload 是解码后的 table，不是 JSON 字符串
|
||||
log("info", string.format("%s via %s: %d prompt tokens",
|
||||
payload.model, payload.source, payload.prompt_tokens or 0))
|
||||
return nil -- 最后一个 stage 没有下游，return 无意义
|
||||
end
|
||||
|
||||
-- 贡献界面
|
||||
plugin.ui = {
|
||||
page = {
|
||||
page_id = "hello", -- kebab-case
|
||||
title = "Hello",
|
||||
icon = "👋",
|
||||
order = 90, -- 侧栏排序
|
||||
mount = [[<div id="hello">hello</div>]],
|
||||
},
|
||||
}
|
||||
|
||||
return plugin -- 必须返回一个 table
|
||||
```
|
||||
|
||||
放进 `plugin_dir` 后重启即生效。`GET /api/plugins` 确认它被加载了。
|
||||
|
||||
---
|
||||
|
||||
## 2. 加载与生命周期
|
||||
|
||||
```
|
||||
core.New
|
||||
└─ lua.NewVM(adapter_dir).Start() 适配器状态
|
||||
└─ lua.NewPlugins(vm, plugin_dir)
|
||||
├─ SeedBundled() 仅当目录不存在时写入内置插件（目前是 billing）
|
||||
└─ LoadDir() 按文件名字典序逐个加载
|
||||
```
|
||||
|
||||
**加载失败不影响网关启动。** 一个语法错误的插件会被记录在 `GET /api/plugins`
|
||||
的 `error` 字段里，永远不会被调用。这与适配器一致，但理由更强：插件是可选的
|
||||
第三方扩展，因为一个 `.lua` 打错字就让网关起不来是错误的取舍。
|
||||
|
||||
**目录一旦存在就是权威的。** 与适配器同规则：首启会 seed 内置插件，之后目录里
|
||||
的文件说了算，删除或编辑内置插件都是真实生效的操作。
|
||||
|
||||
### 2.1 热更新
|
||||
|
||||
| 方式 | 效果 |
|
||||
|---|---|
|
||||
| `POST /api/plugins {name, code}` | 写文件 + 立即加载新版本（旧的 Lua 状态被关闭重建，**累计量清零**） |
|
||||
| `DELETE /api/plugins/{name}` | 删文件 + 卸载 |
|
||||
| 改文件后 `POST` 同名 | 同上 |
|
||||
|
||||
改文件但**不** POST，需要重启才生效。
|
||||
|
||||
---
|
||||
|
||||
## 3. 流水线 stage
|
||||
|
||||
一个请求依次经过三个 stage。插件可以为任意 stage 注册钩子；未注册的 stage
|
||||
被忽略，所以插件不会因为网关将来新增 stage 而报错。
|
||||
|
||||
```
|
||||
客户端请求
|
||||
│
|
||||
┌─────────▼──────────┐
|
||||
│ request_start │ 已解析、已鉴权，尚未选源
|
||||
│ · type │ "chat" | "stream" | "image"
|
||||
│ · model │ 客户端请求的原始 model（"AUTO" 也在这里）
|
||||
│ · key / role │ 掩码后的网关 key id（"***a1b2c3"）与角色
|
||||
│ · source │ 空（还没选源）
|
||||
│ · stream │
|
||||
│ · messages_count │
|
||||
│ · tools_count │
|
||||
│ · ts │ unix 秒
|
||||
└─────────┬──────────┘
|
||||
│ （调度：tier 遍历 → 槽位轮转 → 冷却/配额过滤）
|
||||
┌─────────▼──────────┐
|
||||
│ routed │ 已选定 (source, model)，尚未发往上游
|
||||
│ · source / model │ 实际选中的
|
||||
│ · tier │ AUTO 链的档位；直连 = -1；AUTO = -2
|
||||
│ · stream / key │
|
||||
└─────────┬──────────┘
|
||||
│ （HTTP 往返 / SSE 流）
|
||||
┌─────────▼──────────┐
|
||||
│ request_end │ 每个请求恰好一次，成功失败都触发
|
||||
│ · ok / status │
|
||||
│ · latency_ms │
|
||||
│ · first_byte_ms │ 流式的首字节时间
|
||||
│ · prompt_tokens │ 上游真实 usage，缺失时为字节估算
|
||||
│ · completion_tokens
|
||||
│ · cache_hit_tokens / cache_miss_tokens
|
||||
│ · image_count │ 生图数量（图片不计 token）
|
||||
│ · error │ 失败原因，成功时为 ""
|
||||
│ · time │ unix **毫秒**
|
||||
└─────────┬──────────┘
|
||||
│
|
||||
写审计 + 聚合统计
|
||||
```
|
||||
|
||||
### 3.1 触发点在哪
|
||||
|
||||
| stage | 代码位置 | 说明 |
|
||||
|---|---|---|
|
||||
| `request_start` | `gateway/chat.go` `handleChat` | 每个 chat 请求一次 |
|
||||
| `routed` | `singleChat` / `streamChat` / `singleChatAuto` / `streamChatAuto` | 成功选定源之后，各一次 |
|
||||
| `request_end` | `gateway/chat.go` `writeRec` | **所有出口的唯一汇合点**，每个请求一次 |
|
||||
|
||||
`request_end` 放在 `writeRec` 是因为四条入口路径（直连/AUTO × 流式/非流式）都
|
||||
经过它，既不会漏（流式的 token 数只有流结束才知道），也不会重复。
|
||||
|
||||
**hot path 注意事项**：没有插件注册某 stage 时，`Fire` 立刻返回（一次 `RLock`
|
||||
加一次 map 查找）。装了插件之后，每个请求会在该 stage 上多一次 Lua 调用——
|
||||
这是同步的，在关键路径上。计费插件那种"每请求一次"是正常的；把重活放进钩子是
|
||||
反模式。
|
||||
|
||||
### 3.2 钩子的返回值
|
||||
|
||||
- 返回 `nil` 或不返回 = **没有意见**，payload 原样传给下一个插件
|
||||
- 返回 table = 其中的键会**合并进 payload**，并作为 `Fire` 的返回值
|
||||
|
||||
前三个 stage 的返回值目前没有内部消费者（最后一个 stage 之后就是写审计），
|
||||
所以计费插件改用 `plugin.state` + `/state` 端点来暴露数据。
|
||||
|
||||
---
|
||||
|
||||
## 4. 界面扩展
|
||||
|
||||
### 4.1 整页
|
||||
|
||||
```lua
|
||||
plugin.ui = {
|
||||
page = {
|
||||
page_id = "billing", -- 必填,kebab-case。侧栏 data-tab 与 #tab-billing
|
||||
title = "Billing", -- 必填,侧栏文字
|
||||
icon = "💰", -- 可选
|
||||
order = 40, -- 侧栏排序,默认 100
|
||||
mount = [[...HTML...]],
|
||||
},
|
||||
}
|
||||
```
|
||||
|
||||
### 4.2 往现有页面追加元素
|
||||
|
||||
```lua
|
||||
plugin.ui = {
|
||||
elements = {
|
||||
{
|
||||
target = "status", -- status | chat | keys | sort | sources | adapters
|
||||
anchor = "top", -- "top" | "bottom" | "before:<sel>" | "after:<sel>"
|
||||
order = 5,
|
||||
mount = [[...HTML...]],
|
||||
},
|
||||
},
|
||||
}
|
||||
```
|
||||
|
||||
### 4.3 `mount` 里可以带 `<script>` 和 `<style>`
|
||||
|
||||
内核的注入顺序是：**先插 HTML，再执行 `<script>`**，所以脚本里访问
|
||||
`document.getElementById` 一定能拿到已渲染的节点。
|
||||
|
||||
### 4.4 插件可用的浏览器端 API
|
||||
|
||||
| API | 作用 |
|
||||
|---|---|
|
||||
| `window.pluginAPI.fetchState(name)` | 等价于 `GET /api/plugins/<name>/state` |
|
||||
| `window.pluginAPI.onTabShown(fn)` | 注册"页面切到可见时"的回调（轮询类组件用） |
|
||||
| `window.pluginAPI.postState(name, obj)` | `PUT` 自己的 state（写操作，需 admin） |
|
||||
|
||||
### 4.5 一个完整例子
|
||||
|
||||
见 `internal/lua/plugins/billing.lua`——它同时用了两种 UI 形式、完整的价格配置、
|
||||
以及全部三个 stage。
|
||||
|
||||
---
|
||||
|
||||
## 5. 插件状态与 HTTP API
|
||||
|
||||
### 5.1 端点
|
||||
|
||||
```
|
||||
GET /api/plugins 已加载插件清单 + 钩子 + UI + 错误
|
||||
POST /api/plugins 上传/替换（admin）{name, code}
|
||||
DELETE /api/plugins/{name} 删除（admin）
|
||||
GET /api/plugins/{name}/state 读取插件自己发布的状态
|
||||
PUT /api/plugins/{name}/state 替换状态（admin）
|
||||
```
|
||||
|
||||
`GET .../state` 对**任意角色开放**：它是报表数据（开销、计数），用户自己的
|
||||
计费组件要能渲染。而 `PUT` 需要 admin。
|
||||
|
||||
### 5.2 约定：`prices` 与 `state` 分离
|
||||
|
||||
`PUT .../state` 的载荷里如果有 `prices` 键，它会被写进插件的 **`plugin.prices`**
|
||||
字段并从 `state` 里剔除。
|
||||
|
||||
**为什么**：`state` 是累计量（历史），`prices` 是配置。一次改价如果整体替换
|
||||
`state`，累计量就没了——账目会在改价那一刻清零。分离之后：
|
||||
|
||||
- 只带 `prices` → 改配置，**不动** `state`（累计量保留）
|
||||
- 带其它键 → 替换 `state`（这是显式的重置）
|
||||
|
||||
> 插件可以不遵守这个约定，把整个载荷放进 `state` 自行处理。约定只为内置的
|
||||
> 计费插件存在。
|
||||
|
||||
### 5.3 `GET /api/plugins` 的响应
|
||||
|
||||
```json
|
||||
{
|
||||
"plugins": [{
|
||||
"name": "billing", "version": "1.0.0",
|
||||
"description": "...", "author": "...",
|
||||
"hooks": ["request_end"],
|
||||
"ui": { "page": "billing", "elements": 1 },
|
||||
"loaded": true
|
||||
}],
|
||||
"hook_errors": { "request_end": { "count": 3, "last_error": "billing: ..." } },
|
||||
"plugin_dir": "/etc/llmsproxy/plugins"
|
||||
}
|
||||
```
|
||||
|
||||
`hook_errors` 是**排错入口**：一个坏掉的插件表现为"功能缺失"而不是报错，
|
||||
这个计数让它可见。
|
||||
|
||||
---
|
||||
|
||||
## 6. 运行时约束（重要）
|
||||
|
||||
### 6.1 一个插件 = 一个 Lua 状态
|
||||
|
||||
适配器可以有多个独立的 worker 状态（它们无状态，这是对的）。**插件不行**：
|
||||
钩子通常往 `plugin.state` 里累加，多个状态就意味着总额被劈成几份，而写价格只
|
||||
写进其中一个状态，钩子恰好跑到另一个时**所有请求按 0 计费**。
|
||||
|
||||
因此插件持有**唯一一个**状态，所有入口（`Fire` / `State` / `SetState`）用互斥锁
|
||||
串行化。代价：**一个卡住的钩子会卡住所有插件的钩子**。所以——
|
||||
|
||||
> **钩子必须短、同步、不阻塞。** 不要在里面做 HTTP 请求、睡眠或重计算。
|
||||
|
||||
### 6.2 异常被隔离
|
||||
|
||||
钩子里 `error()` 不会影响转发：异常被捕获、记进 `hook_errors`、跳到下一个插件。
|
||||
适配器变换出错会让源进入冷却，但插件出错**不产生任何惩罚**——插件是可选功能。
|
||||
|
||||
### 6.3 内置可用函数
|
||||
|
||||
与适配器相同：`json.encode` / `json.decode`、`log(level, msg)`、
|
||||
`hmac_sha256_hex` / `sha256_hex` / `base64_encode` / `tohex`。
|
||||
|
||||
⚠️ `os.date` 在极简运行时里可能不可用（计费插件因此有降级路径）。
|
||||
|
||||
### 6.4 命名
|
||||
|
||||
- 全局 `__llmsproxy_plugin` 存放插件返回的 table，**不要占用**
|
||||
- 插件间状态完全隔离，一个插件改不了另一个的全局
|
||||
|
||||
---
|
||||
|
||||
## 7. 内置计费插件
|
||||
|
||||
`plugins/billing.lua`，默认随首启 seed 进来，开箱可用。
|
||||
|
||||
### 7.1 计费维度
|
||||
|
||||
| 维度 | 用途 | 优先级 |
|
||||
|---|---|---|
|
||||
| `sources` | 源的固定价（如按请求计费的转发服务） | 中 |
|
||||
| `models` | 模型的 per-token 价 | 高 |
|
||||
| `keys` | 单个网关 key 的覆盖价 | **最高** |
|
||||
|
||||
**token 价**优先级：`keys` > `models` > `default`。
|
||||
**`per_request` 固定价是叠加的**（不覆盖），所以一个生图模型可以既算 token 又收固定费。
|
||||
|
||||
### 7.2 价格单位
|
||||
|
||||
**USD / 单个 token**。这是各家 provider 的公布口径，所以典型数值长这样：
|
||||
`1.25e-6`。插件内部乘以 token 数，全程不做单位换算。
|
||||
|
||||
```json
|
||||
{
|
||||
"prices": {
|
||||
"currency": "USD",
|
||||
"default": { "prompt": 0, "completion": 0, "per_request": 0 },
|
||||
"sources": { "trae": { "per_request": 0.01 } },
|
||||
"models": {
|
||||
"gpt-5.4": { "prompt": 1.25e-6, "completion": 1e-5 },
|
||||
"kolors": { "per_request": 0.04 }
|
||||
},
|
||||
"keys": { "***a1b2c3": { "prompt": 1.1e-6, "completion": 9e-6 } }
|
||||
}
|
||||
}
|
||||
```
|
||||
|
||||
配置价目：
|
||||
|
||||
```bash
|
||||
curl -X PUT http://127.0.0.1:8080/api/plugins/billing/state \
|
||||
-H "Authorization: Bearer $ADMIN_KEY" -H "Content-Type: application/json" \
|
||||
-d '{"prices":{"models":{"gpt-5.4":{"prompt":1.25e-6,"completion":1e-5}}}}'
|
||||
```
|
||||
|
||||
读回账目：
|
||||
|
||||
```bash
|
||||
curl -H "Authorization: Bearer $ADMIN_KEY" \
|
||||
http://127.0.0.1:8080/api/plugins/billing/state
|
||||
```
|
||||
|
||||
### 7.3 累计维度
|
||||
|
||||
`total` / `by_source` / `by_model` / `by_key` / `by_day`（`YYYY-MM-DD` UTC），
|
||||
每项含 `cost`、`requests`、`prompt_tokens`、`completion_tokens`、`failures`。
|
||||
|
||||
### 7.4 计费策略：失败的请求怎么算
|
||||
|
||||
**保留 token 费用，丢弃固定费用。** 理由：上游在生成后才 500，token 确实被消耗
|
||||
了；但那个从未真正发生的固定费不该收。
|
||||
|
||||
这是本插件里**最可争议的一条**。想改成"失败也收固定费"：
|
||||
|
||||
```json
|
||||
{ "prices": { ... }, "count_failures": true }
|
||||
```
|
||||
|
||||
### 7.5 ⚠️ 计费插件只报表，不执法
|
||||
|
||||
**网关自己的配额会计（`internal/gateway/stats.go`，入口处强制）才是限额权威。**
|
||||
本插件不参与任何路由或配额决策。
|
||||
|
||||
理由：两套独立的会计路径如果对不上，比一套功能略少的更糟。计费是**观察**，
|
||||
配额是**控制**，二者分开。
|
||||
|
||||
### 7.6 数据从哪来
|
||||
|
||||
`request_end` 的 `prompt_tokens` / `completion_tokens` 优先取上游真实的
|
||||
`usage`；上游没报时网关用字节估算（`len/3+1`）。流式请求在流结束后用上游真实
|
||||
数字**覆盖**估算值。所以计费数字的精度取决于上游是否报 usage。
|
||||
|
||||
---
|
||||
|
||||
## 8. 排错
|
||||
|
||||
| 现象 | 查什么 |
|
||||
|---|---|
|
||||
| 插件没出现在 `/api/plugins` | 目录对不对（响应里的 `plugin_dir`）；文件是不是 `.lua` |
|
||||
| `loaded: false` 且有 `error` | 语法错误或没 `return table`，错误信息在 `error` 字段 |
|
||||
| 功能"没反应"但无报错 | 查 `hook_errors`——坏钩子只记录不抛出 |
|
||||
| 钩子没被调用 | 该 stage 确实触发了吗（`routed` 只在**成功**选源后触发，调度全失败时不触发） |
|
||||
| 界面空白 | 内核的 `GET /api/ui-inject` 里有没有你的 `page_id`；脚本有没有报错（浏览器 console） |
|
||||
| `pluginAPI` 未定义 | 脚本在注入前执行了；确认用的是 `mount` 而不是别的方式插入 |
|
||||
|
||||
### 调试钩子
|
||||
|
||||
插件的 `log(level, msg)` 输出到网关的日志：
|
||||
|
||||
```
|
||||
journalctl -u llmsproxy -f | grep -i plugin
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## 9. 与适配器的区别
|
||||
|
||||
| | 适配器 | 插件 |
|
||||
|---|---|---|
|
||||
| 位置 | `adapter_dir` | `plugin_dir` |
|
||||
| 作用 | 上游协议转换 | 请求流水线 + 界面 |
|
||||
| 钩子 | `transform_request` / `transform_response` / `transform_stream_chunk` / `build_headers` / `transform_error` | `request_start` / `routed` / `request_end` |
|
||||
| 状态 | 每个源独立，多 worker | 单状态，见 §6.1 |
|
||||
| 出错后果 | 该 (源,模型) 进入冷却退避 | 仅该功能缺失 |
|
||||
| 鉴权 | 需要 `gateway_keys` 之外的独立凭据 | 网关 key |
|
||||
| 必需性 | 核心 | 可选，缺了网关照常跑 |
|
||||
|
||||
完整适配器协议见 [lua-adapters.md](lua-adapters.md)。
|
||||
@ -19,6 +19,7 @@ type Config struct {
|
||||
GatewayKeys []string `yaml:"gateway_keys"`
|
||||
DefaultModel string `yaml:"default_model"` // e.g. "AUTO" or a model id
|
||||
AdapterDir string `yaml:"adapter_dir"`
|
||||
PluginDir string `yaml:"plugin_dir,omitempty"` // Lua plugin directory (optional; hooks + WebUI extensions)
|
||||
RuntimeFile string `yaml:"runtime_file"`
|
||||
TLSCertFile string `yaml:"tls_cert_file,omitempty"` // PEM cert; when set together with tls_key_file, serve HTTPS
|
||||
TLSKeyFile string `yaml:"tls_key_file,omitempty"` // PEM private key
|
||||
@ -132,6 +133,7 @@ func writeDefaultConfig(path string) error {
|
||||
GatewayKeys: []string{key},
|
||||
DefaultModel: "AUTO",
|
||||
AdapterDir: filepath.Join(abs, "adapters"),
|
||||
PluginDir: filepath.Join(abs, "plugins"),
|
||||
RuntimeFile: filepath.Join(abs, "runtime.json"),
|
||||
}
|
||||
out, err := yaml.Marshal(&cfg)
|
||||
@ -288,6 +290,9 @@ func (c *Config) ApplyDefaults() error {
|
||||
if c.AdapterDir == "" {
|
||||
c.AdapterDir = "adapters"
|
||||
}
|
||||
// PluginDir is deliberately NOT defaulted: with no plugins a missing
|
||||
// directory would be created and seeded on every start for nothing. An
|
||||
// operator opts in by setting it (or by dropping files into it).
|
||||
if c.RuntimeFile == "" {
|
||||
c.RuntimeFile = "runtime.json"
|
||||
}
|
||||
|
||||
@ -30,6 +30,7 @@ type Core struct {
|
||||
mu sync.Mutex
|
||||
cfg *config.Config
|
||||
vm *lua.VM
|
||||
plugins *lua.Plugins
|
||||
store *config.Store
|
||||
scheduler *scheduler.Scheduler
|
||||
registry *provider.Registry
|
||||
@ -53,6 +54,27 @@ func NewFromConfig(cfg *config.Config) (*Core, error) {
|
||||
if err := c.vm.Start(); err != nil {
|
||||
return nil, fmt.Errorf("lua vm: %w", err)
|
||||
}
|
||||
// Plugins load AFTER the VM is up (a plugin is just another Lua state) but
|
||||
// BEFORE anything can serve traffic, so a plugin's request_end hook sees
|
||||
// the very first request. A plugin that fails to load is logged and skipped
|
||||
// rather than failing startup: plugins are optional extensions, and failing
|
||||
// to boot the gateway because a third-party .lua has a typo would be the
|
||||
// wrong trade.
|
||||
c.plugins = lua.NewPlugins(c.vm, cfg.PluginDir)
|
||||
if err := c.plugins.SeedBundled(); err != nil {
|
||||
log.Printf("[core] seed bundled plugins: %v", err)
|
||||
}
|
||||
if err := c.plugins.LoadDir(); err != nil {
|
||||
log.Printf("[core] plugin dir: %v", err)
|
||||
}
|
||||
for _, p := range c.plugins.List() {
|
||||
if msg, bad := p["error"]; bad {
|
||||
log.Printf("[core] plugin %v failed to load: %v", p["name"], msg)
|
||||
continue
|
||||
}
|
||||
log.Printf("[core] plugin %v %v loaded (hooks=%v ui=%v)",
|
||||
p["name"], p["version"], p["hooks"], p["ui"])
|
||||
}
|
||||
c.store = config.NewStore(cfg.RuntimeFile)
|
||||
// Share one box between the runtime store and config.yaml so a single
|
||||
// master.key seals both files. config.Load left the config holding
|
||||
@ -221,6 +243,11 @@ func buildRetries(cfg *config.Config) int {
|
||||
// VM exposes the Lua adapter runtime.
|
||||
func (c *Core) VM() *lua.VM { return c.vm }
|
||||
|
||||
// Plugins exposes the loaded plugin set. Never nil once New* has returned, so
|
||||
// the gateway can fire stages unconditionally; an unconfigured plugin
|
||||
// directory yields an empty registry whose Fire is a no-op.
|
||||
func (c *Core) Plugins() *lua.Plugins { return c.plugins }
|
||||
|
||||
func (c *Core) Scheduler() *scheduler.Scheduler { return c.scheduler }
|
||||
|
||||
func (c *Core) Registry() *provider.Registry { return c.registry }
|
||||
|
||||
@ -14,6 +14,7 @@ import (
|
||||
"time"
|
||||
|
||||
"llmsproxy/internal/config"
|
||||
"llmsproxy/internal/lua"
|
||||
"llmsproxy/internal/provider"
|
||||
"llmsproxy/internal/scheduler"
|
||||
"llmsproxy/internal/types"
|
||||
@ -414,6 +415,7 @@ func (g *Gateway) handleChat(w http.ResponseWriter, r *http.Request) {
|
||||
Type: "chat",
|
||||
OK: false,
|
||||
}
|
||||
g.fireStart(ctx, &req, "chat", "AUTO", len(req.Messages), len(req.Tools))
|
||||
// quotaExhausted reports a slot whose token window has been used up;
|
||||
// exhausted slots are dropped from scheduling without penalty.
|
||||
quotaExhausted := func(sl *scheduler.Slot) bool {
|
||||
@ -817,6 +819,7 @@ func (g *Gateway) singleChat(w http.ResponseWriter, ctx context.Context, cands [
|
||||
recordChatUsage(rec, req, resp)
|
||||
rec.Source = usedSrc
|
||||
rec.Model = usedModel
|
||||
g.fireRouted(ctx, "chat", usedSrc, usedModel, -1, false)
|
||||
// Non-streaming: the whole response arrives at once, so TTFB equals
|
||||
// the total latency.
|
||||
rec.FirstByteMs = rec.LatMs
|
||||
@ -824,7 +827,18 @@ func (g *Gateway) singleChat(w http.ResponseWriter, ctx context.Context, cands [
|
||||
writeChatCompletion(w, resp, effective)
|
||||
}
|
||||
|
||||
// writeRec records a finished request (audit + aggregates).
|
||||
// writeRec records a finished request (audit + aggregates) and fires the
|
||||
// plugin request_end stage.
|
||||
//
|
||||
// This is the ONE place every request passes through on its way out, which is
|
||||
// what makes it the right hook point: the four entry points (single/stream ×
|
||||
// direct/auto) all funnel here, so a plugin sees each request exactly once with
|
||||
// its final accounting. Firing earlier would miss the streamed ones (their
|
||||
// numbers are only known once the stream finishes), and firing in each entry
|
||||
// point would mean four call sites to keep in sync.
|
||||
//
|
||||
// Hooks run AFTER the record is written: a plugin must not be able to delay or
|
||||
// lose the audit trail, and a plugin that throws is contained by Fire.
|
||||
func (g *Gateway) writeRec(rec *Req) {
|
||||
if rec == nil {
|
||||
return
|
||||
@ -833,6 +847,81 @@ func (g *Gateway) writeRec(rec *Req) {
|
||||
rec.Time = time.Now().UnixMilli()
|
||||
}
|
||||
g.stats.Record(*rec)
|
||||
g.fireEnd(rec)
|
||||
}
|
||||
|
||||
// fireStart dispatches the plugin request_start stage: the request has been
|
||||
// parsed and authorized but no upstream slot has been chosen yet, so `source`
|
||||
// is empty. A plugin that only wants volume/acceptance counts can subscribe
|
||||
// here and stay out of the per-request hot path entirely.
|
||||
func (g *Gateway) fireStart(ctx context.Context, req *chatRequest, kind, model string, msgs, tools int) {
|
||||
ps := g.core.Plugins()
|
||||
if ps == nil || ps.Count() == 0 {
|
||||
return
|
||||
}
|
||||
ps.Fire(lua.StageRequestStart, map[string]interface{}{
|
||||
"stage": string(lua.StageRequestStart),
|
||||
"type": kind,
|
||||
"model": model,
|
||||
"key": keyID(reqKey(ctx)),
|
||||
"role": reqRole(ctx),
|
||||
"source": "",
|
||||
"stream": req.Stream,
|
||||
"messages_count": msgs,
|
||||
"tools_count": tools,
|
||||
"ts": time.Now().Unix(),
|
||||
})
|
||||
}
|
||||
|
||||
// fireRouted dispatches the plugin routed stage once a (source, model) slot has
|
||||
// been selected. tier is the AUTO tier index, or -1 on the direct path, so a
|
||||
// plugin can tell "this came from tier 1" from "this bypassed the chain".
|
||||
func (g *Gateway) fireRouted(ctx context.Context, kind, source, model string, tier int, stream bool) {
|
||||
ps := g.core.Plugins()
|
||||
if ps == nil || ps.Count() == 0 {
|
||||
return
|
||||
}
|
||||
ps.Fire(lua.StageRouted, map[string]interface{}{
|
||||
"stage": string(lua.StageRouted),
|
||||
"type": kind,
|
||||
"source": source,
|
||||
"model": model,
|
||||
"key": keyID(reqKey(ctx)),
|
||||
"tier": tier,
|
||||
"stream": stream,
|
||||
"ts": time.Now().Unix(),
|
||||
})
|
||||
}
|
||||
|
||||
// fireEnd dispatches the plugin request_end stage for one finished request.
|
||||
func (g *Gateway) fireEnd(rec *Req) {
|
||||
ps := g.core.Plugins()
|
||||
if ps == nil || ps.Count() == 0 {
|
||||
return
|
||||
}
|
||||
payload := map[string]interface{}{
|
||||
"stage": string(lua.StageRequestEnd),
|
||||
"type": rec.Type,
|
||||
"model": rec.Model,
|
||||
"source": rec.Source,
|
||||
"key": rec.Key,
|
||||
"ok": rec.OK,
|
||||
"status": rec.Status,
|
||||
"latency_ms": rec.LatMs,
|
||||
"first_byte_ms": rec.FirstByteMs,
|
||||
"prompt_tokens": rec.Prompt,
|
||||
"completion_tokens": rec.Compl,
|
||||
"cache_hit_tokens": rec.CacheHit,
|
||||
"cache_miss_tokens": rec.CacheMiss,
|
||||
"image_count": rec.ImageCount,
|
||||
"error": rec.Err,
|
||||
"time": rec.Time,
|
||||
}
|
||||
// The merged result is intentionally discarded: request_end is the last
|
||||
// stage, so there is nobody downstream to read a plugin's additions. Plugins
|
||||
// that need to publish derived numbers (the billing plugin) do it in their
|
||||
// OWN state and expose them through the /api/plugins/<name>/state endpoint.
|
||||
ps.Fire(lua.StageRequestEnd, payload)
|
||||
}
|
||||
|
||||
// mergeUsage combines token usage across stream chunks additively. Some
|
||||
@ -1039,6 +1128,7 @@ func (g *Gateway) streamChat(w http.ResponseWriter, ctx context.Context, cands [
|
||||
// failover it differs from the first candidate). Direct streams previously
|
||||
// discarded it.
|
||||
rec.Source = usedSrc
|
||||
g.fireRouted(ctx, "stream", usedSrc, rec.Model, -1, true)
|
||||
rec.Prompt = estimatePromptTokens(req)
|
||||
g.pumpStream(w, rec, chunks, effective, t0)
|
||||
}
|
||||
@ -1064,6 +1154,10 @@ func (g *Gateway) singleChatAuto(w http.ResponseWriter, ctx context.Context, cha
|
||||
recordChatUsage(rec, req, resp)
|
||||
rec.Source = usedSrc
|
||||
rec.Model = usedModel
|
||||
// AUTO has no single tier to report: the chain may have walked several
|
||||
// before this slot served the request, so -2 means "resolved by the chain"
|
||||
// and a plugin can tell that apart from the direct path's -1.
|
||||
g.fireRouted(ctx, "chat", usedSrc, usedModel, -2, false)
|
||||
rec.FirstByteMs = rec.LatMs
|
||||
g.writeRec(rec)
|
||||
writeChatCompletion(w, resp, usedModel)
|
||||
@ -1091,6 +1185,7 @@ func (g *Gateway) streamChatAuto(w http.ResponseWriter, ctx context.Context, cha
|
||||
rec.Model = usedModel
|
||||
}
|
||||
rec.Source = usedSrc
|
||||
g.fireRouted(ctx, "stream", usedSrc, rec.Model, -2, true)
|
||||
rec.Prompt = estimatePromptTokens(req)
|
||||
g.pumpStream(w, rec, chunks, usedModel, t0)
|
||||
}
|
||||
|
||||
203
internal/gateway/plugins_api.go
Normal file
203
internal/gateway/plugins_api.go
Normal file
@ -0,0 +1,203 @@
|
||||
package gateway
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"net/http"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
|
||||
"llmsproxy/internal/lua"
|
||||
)
|
||||
|
||||
var (
|
||||
errPluginName = errors.New("plugin name must be non-empty and contain no path separator or dot")
|
||||
errNoPluginDir = errors.New("no plugin_dir configured; set plugin_dir in config.yaml to enable plugins")
|
||||
)
|
||||
|
||||
// Plugin management API.
|
||||
//
|
||||
// GET /api/plugins list loaded plugins + their hooks/UI/errors
|
||||
// POST /api/plugins upload/replace one plugin (.lua), hot-applied
|
||||
// DELETE /api/plugins/{name} remove a plugin
|
||||
// GET /api/plugins/{name}/state the plugin's own published state
|
||||
// PUT /api/plugins/{name}/state replace that state (admin only)
|
||||
//
|
||||
// Why state is a first-class endpoint: request_end is the LAST stage, so a
|
||||
// hook's return value has no downstream consumer inside the gateway. A plugin
|
||||
// that accumulates numbers (the billing plugin does exactly this) therefore
|
||||
// keeps them in its own Lua state and serves them here, which is what its UI
|
||||
// component fetches. This keeps plugin data clearly separated from the
|
||||
// gateway's own stats — see the accounting note in docs/plugins.md: the
|
||||
// gateway's quota accounting stays authoritative, a plugin only reports.
|
||||
|
||||
// handlePluginUI serves the merged UI extensions the WebUI injects at boot.
|
||||
//
|
||||
// It is a single GET (not per-plugin) because the browser needs ALL extensions
|
||||
// before it can build the sidebar: a page contributed by one plugin and an
|
||||
// element contributed by another land in the same payload, and fetching them
|
||||
// separately would mean the sidebar has to be rebuilt as each arrives.
|
||||
//
|
||||
// Served to any authenticated role: the WebUI is authenticated before it asks,
|
||||
// and a plugin's own widgets need to render for the user whose key they show.
|
||||
func (g *Gateway) handlePluginUI(w http.ResponseWriter, r *http.Request) {
|
||||
if r.Method != http.MethodGet {
|
||||
writeError(w, http.StatusMethodNotAllowed, "method_not_allowed", "use GET")
|
||||
return
|
||||
}
|
||||
ps := g.core.Plugins()
|
||||
if ps == nil {
|
||||
writeJSON(w, http.StatusOK, map[string]interface{}{"ui": lua.UIExtension{}})
|
||||
return
|
||||
}
|
||||
writeJSON(w, http.StatusOK, map[string]interface{}{
|
||||
"ui": ps.UI(),
|
||||
"stages": []string{
|
||||
string(lua.StageRequestStart),
|
||||
string(lua.StageRouted),
|
||||
string(lua.StageRequestEnd),
|
||||
},
|
||||
})
|
||||
}
|
||||
|
||||
func (g *Gateway) handlePluginsAPI(w http.ResponseWriter, r *http.Request) {
|
||||
ps := g.core.Plugins()
|
||||
if ps == nil {
|
||||
writeJSON(w, http.StatusOK, map[string]interface{}{"plugins": []interface{}{}})
|
||||
return
|
||||
}
|
||||
path := strings.TrimPrefix(r.URL.Path, "/api/plugins")
|
||||
path = strings.Trim(path, "/")
|
||||
|
||||
// /api/plugins/{name}/state
|
||||
if strings.HasSuffix(path, "/state") {
|
||||
name := strings.TrimSuffix(path, "/state")
|
||||
if name == "" {
|
||||
writeError(w, http.StatusBadRequest, "invalid_request", "plugin name required")
|
||||
return
|
||||
}
|
||||
g.handlePluginState(w, r, name)
|
||||
return
|
||||
}
|
||||
|
||||
if reqRole(r.Context()) != "admin" {
|
||||
writeError(w, http.StatusForbidden, "forbidden", "admin role required")
|
||||
return
|
||||
}
|
||||
|
||||
if path == "" {
|
||||
switch r.Method {
|
||||
case http.MethodGet:
|
||||
writeJSON(w, http.StatusOK, map[string]interface{}{
|
||||
"plugins": ps.List(),
|
||||
"hook_errors": ps.HookErrors(),
|
||||
"plugin_dir": g.pluginDir(),
|
||||
})
|
||||
case http.MethodPost:
|
||||
var body struct {
|
||||
Name string `json:"name"`
|
||||
Code string `json:"code"`
|
||||
}
|
||||
if err := json.NewDecoder(r.Body).Decode(&body); err != nil {
|
||||
writeError(w, http.StatusBadRequest, "invalid_request", "invalid json: "+err.Error())
|
||||
return
|
||||
}
|
||||
if err := g.installPlugin(body.Name, body.Code); err != nil {
|
||||
writeError(w, http.StatusBadRequest, "plugin_error", err.Error())
|
||||
return
|
||||
}
|
||||
writeJSON(w, http.StatusOK, map[string]interface{}{"ok": true, "name": body.Name})
|
||||
default:
|
||||
writeError(w, http.StatusMethodNotAllowed, "method_not_allowed", "")
|
||||
}
|
||||
return
|
||||
}
|
||||
|
||||
switch r.Method {
|
||||
case http.MethodDelete:
|
||||
if err := g.removePlugin(path); err != nil {
|
||||
writeError(w, http.StatusBadRequest, "plugin_error", err.Error())
|
||||
return
|
||||
}
|
||||
writeJSON(w, http.StatusOK, map[string]interface{}{"ok": true})
|
||||
case http.MethodGet:
|
||||
// A GET on a specific plugin is almost always a client that meant to
|
||||
// delete it but let fetch default to GET; name the verb.
|
||||
writeError(w, http.StatusNotFound, "not_found",
|
||||
"plugin source not exposed; use DELETE /api/plugins/"+path+" to remove it, "+
|
||||
"or GET /api/plugins/"+path+"/state for its published state")
|
||||
default:
|
||||
writeError(w, http.StatusMethodNotAllowed, "method_not_allowed", "")
|
||||
}
|
||||
}
|
||||
|
||||
// handlePluginState serves GET (read state) and PUT (replace state).
|
||||
func (g *Gateway) handlePluginState(w http.ResponseWriter, r *http.Request, name string) {
|
||||
ps := g.core.Plugins()
|
||||
switch r.Method {
|
||||
case http.MethodGet:
|
||||
// Any role may read: plugin state is reporting data (cost, counts),
|
||||
// and the caller has already been authenticated. Admin-only would stop
|
||||
// a user key's own billing widget from rendering.
|
||||
writeJSON(w, http.StatusOK, map[string]interface{}{
|
||||
"plugin": name,
|
||||
"state": ps.State(name),
|
||||
})
|
||||
case http.MethodPut:
|
||||
if reqRole(r.Context()) != "admin" {
|
||||
writeError(w, http.StatusForbidden, "forbidden", "admin role required")
|
||||
return
|
||||
}
|
||||
var state interface{}
|
||||
if err := json.NewDecoder(r.Body).Decode(&state); err != nil {
|
||||
writeError(w, http.StatusBadRequest, "invalid_request", "invalid json: "+err.Error())
|
||||
return
|
||||
}
|
||||
if err := ps.SetState(name, state); err != nil {
|
||||
writeError(w, http.StatusBadRequest, "plugin_error", err.Error())
|
||||
return
|
||||
}
|
||||
writeJSON(w, http.StatusOK, map[string]interface{}{"ok": true})
|
||||
default:
|
||||
writeError(w, http.StatusMethodNotAllowed, "method_not_allowed", "")
|
||||
}
|
||||
}
|
||||
|
||||
func (g *Gateway) pluginDir() string {
|
||||
if c := g.core.Config(); c != nil {
|
||||
return c.PluginDir
|
||||
}
|
||||
return ""
|
||||
}
|
||||
|
||||
// installPlugin writes a plugin to disk and hot-loads it. A syntax error is
|
||||
// returned to the caller AND the file is left on disk so the operator can fix
|
||||
// it, matching how adapters behave (the file is authoritative once present).
|
||||
func (g *Gateway) installPlugin(name, code string) error {
|
||||
if name == "" {
|
||||
return errPluginName
|
||||
}
|
||||
if strings.ContainsAny(name, `/\.`) {
|
||||
return errPluginName
|
||||
}
|
||||
dir := g.pluginDir()
|
||||
if dir == "" {
|
||||
return errNoPluginDir
|
||||
}
|
||||
if err := os.MkdirAll(dir, 0755); err != nil {
|
||||
return err
|
||||
}
|
||||
if err := os.WriteFile(filepath.Join(dir, name+".lua"), []byte(code), 0644); err != nil {
|
||||
return err
|
||||
}
|
||||
return g.core.Plugins().LoadSource(name, code)
|
||||
}
|
||||
|
||||
func (g *Gateway) removePlugin(name string) error {
|
||||
if g.pluginDir() == "" {
|
||||
return errNoPluginDir
|
||||
}
|
||||
_ = os.Remove(filepath.Join(g.pluginDir(), name+".lua"))
|
||||
return g.core.Plugins().Unload(name)
|
||||
}
|
||||
187
internal/gateway/plugins_api_test.go
Normal file
187
internal/gateway/plugins_api_test.go
Normal file
@ -0,0 +1,187 @@
|
||||
package gateway
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"llmsproxy/internal/config"
|
||||
"llmsproxy/internal/core"
|
||||
"llmsproxy/internal/lua"
|
||||
)
|
||||
|
||||
// gatewayWithBilling boots a gateway with the bundled billing plugin loaded, so
|
||||
// the UI-injection endpoint is exercised against a real plugin rather than a
|
||||
// hand-written stub. The other plugin tests in this package assert on the
|
||||
// WebUI source; this one asserts on the HTTP contract the browser consumes.
|
||||
func gatewayWithBilling(t *testing.T) *Gateway {
|
||||
t.Helper()
|
||||
dir := t.TempDir()
|
||||
cfgPath := filepath.Join(dir, "config.yaml")
|
||||
body := "listen: :0\n" +
|
||||
"adapter_dir: " + filepath.Join(dir, "adapters") + "\n" +
|
||||
"plugin_dir: " + filepath.Join(dir, "plugins") + "\n" +
|
||||
"runtime_file: " + filepath.Join(dir, "runtime.json") + "\n" +
|
||||
"gateway_keys:\n - sk-test\n"
|
||||
if err := os.WriteFile(cfgPath, []byte(body), 0600); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
cfg, err := config.Load(cfgPath)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
c, err := core.NewFromConfig(cfg)
|
||||
if err != nil {
|
||||
t.Fatalf("core: %v", err)
|
||||
}
|
||||
t.Cleanup(c.Close)
|
||||
// Load the shipped plugin explicitly: seeding only runs for a directory that
|
||||
// does not exist yet, and this test wants a known plugin regardless.
|
||||
src, err := lua.ReadBundledPlugin("billing")
|
||||
if err != nil {
|
||||
t.Fatalf("read bundled billing: %v", err)
|
||||
}
|
||||
if err := c.Plugins().LoadSource("billing", src); err != nil {
|
||||
t.Fatalf("load billing: %v", err)
|
||||
}
|
||||
g, err := New(c)
|
||||
if err != nil {
|
||||
t.Fatalf("gateway: %v", err)
|
||||
}
|
||||
return g
|
||||
}
|
||||
|
||||
// TestUIInjectServesPluginUI: GET /api/ui-inject is the single call the WebUI
|
||||
// makes at boot, and it must carry BOTH a contributed page and contributed
|
||||
// elements — the browser builds the sidebar from the page and mounts the
|
||||
// elements into existing panes from the same payload, so a partial response
|
||||
// would produce a page with no body or a missing tile.
|
||||
func TestUIInjectServesPluginUI(t *testing.T) {
|
||||
g := gatewayWithBilling(t)
|
||||
rr := doReq(t, g, http.MethodGet, "/api/ui-inject", "")
|
||||
if rr.Code != http.StatusOK {
|
||||
t.Fatalf("status=%d body=%s", rr.Code, rr.Body.String())
|
||||
}
|
||||
var view struct {
|
||||
UI struct {
|
||||
Page *struct {
|
||||
PageID string `json:"page_id"`
|
||||
Title string `json:"title"`
|
||||
Mount string `json:"mount"`
|
||||
} `json:"page"`
|
||||
Elements []struct {
|
||||
Target string `json:"target"`
|
||||
Mount string `json:"mount"`
|
||||
} `json:"elements"`
|
||||
} `json:"ui"`
|
||||
Stages []string `json:"stages"`
|
||||
}
|
||||
if err := json.Unmarshal(rr.Body.Bytes(), &view); err != nil {
|
||||
t.Fatalf("decode: %v", err)
|
||||
}
|
||||
if view.UI.Page == nil || view.UI.Page.PageID != "billing" {
|
||||
t.Fatalf("no billing page in the inject payload")
|
||||
}
|
||||
if !strings.Contains(view.UI.Page.Mount, "billing-root") {
|
||||
t.Error("the page mount came back empty")
|
||||
}
|
||||
if len(view.UI.Elements) == 0 {
|
||||
t.Error("billing contributes an element to the status page but it is missing")
|
||||
}
|
||||
for _, e := range view.UI.Elements {
|
||||
if e.Target != "status" {
|
||||
t.Errorf("element target = %q, want \"status\"", e.Target)
|
||||
}
|
||||
}
|
||||
if len(view.Stages) != 3 {
|
||||
t.Errorf("stages = %v, want the three pipeline stages", view.Stages)
|
||||
}
|
||||
}
|
||||
|
||||
// TestUIInjectIsEmptyWithoutPlugins: a gateway with no plugins must still answer
|
||||
// 200 with an empty (not missing, not null) payload. The WebUI calls this
|
||||
// unconditionally at boot, so a 404 or a null `ui` would break every dashboard.
|
||||
func TestUIInjectIsEmptyWithoutPlugins(t *testing.T) {
|
||||
g := newTestGateway(t)
|
||||
rr := doReq(t, g, http.MethodGet, "/api/ui-inject", "")
|
||||
if rr.Code != http.StatusOK {
|
||||
t.Fatalf("status=%d body=%s", rr.Code, rr.Body.String())
|
||||
}
|
||||
if !strings.Contains(rr.Body.String(), `"ui"`) {
|
||||
t.Error("no ui key in the response; the WebUI would have nothing to read")
|
||||
}
|
||||
}
|
||||
|
||||
// newRecorderFor pushes a request through the full handler chain.
|
||||
func newRecorderFor(t *testing.T, g *Gateway, req *http.Request) *httptest.ResponseRecorder {
|
||||
t.Helper()
|
||||
rr := httptest.NewRecorder()
|
||||
g.Handler().ServeHTTP(rr, req)
|
||||
return rr
|
||||
}
|
||||
|
||||
// TestPluginsListAndStateAPI covers the management surface the plugin docs
|
||||
// promise: listing, and reading a plugin's own published state.
|
||||
func TestPluginsListAndStateAPI(t *testing.T) {
|
||||
g := gatewayWithBilling(t)
|
||||
rr := doReq(t, g, http.MethodGet, "/api/plugins", "")
|
||||
if rr.Code != http.StatusOK {
|
||||
t.Fatalf("GET /api/plugins = %d: %s", rr.Code, rr.Body.String())
|
||||
}
|
||||
for _, want := range []string{"billing", "hook_errors", "plugin_dir", "request_end"} {
|
||||
if !strings.Contains(rr.Body.String(), want) {
|
||||
t.Errorf("/api/plugins response lacks %q", want)
|
||||
}
|
||||
}
|
||||
// state read: the plugin published its (empty) state, so the key exists.
|
||||
rr = doReq(t, g, http.MethodGet, "/api/plugins/billing/state", "")
|
||||
if rr.Code != http.StatusOK {
|
||||
t.Fatalf("GET state = %d: %s", rr.Code, rr.Body.String())
|
||||
}
|
||||
if !strings.Contains(rr.Body.String(), `"total"`) {
|
||||
t.Errorf("billing state lacks the total bucket: %s", rr.Body.String())
|
||||
}
|
||||
}
|
||||
|
||||
// TestPluginStatePUTIsAdminOnly: only the WRITE side is gated. A user key must
|
||||
// be able to READ its own billing widget's data, but must not be able to
|
||||
// rewrite the price table.
|
||||
func TestPluginStatePUTIsAdminOnly(t *testing.T) {
|
||||
g := gatewayWithBilling(t)
|
||||
|
||||
// Build a user-role key and remember its secret.
|
||||
rec, err := g.core.CreateKey("viewer", "user", nil, "")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
userKey := rec.Key
|
||||
|
||||
// A user key may read the state.
|
||||
req, _ := http.NewRequest(http.MethodGet, "/api/plugins/billing/state", nil)
|
||||
req.Header.Set("Authorization", "Bearer "+userKey)
|
||||
rr := newRecorderFor(t, g, req)
|
||||
if rr.Code != http.StatusOK {
|
||||
t.Errorf("user GET state = %d, want 200 (the billing widget must render for users)", rr.Code)
|
||||
}
|
||||
|
||||
// A user key may NOT write it.
|
||||
put, _ := http.NewRequest(http.MethodPut, "/api/plugins/billing/state",
|
||||
strings.NewReader(`{"prices":{"default":{"prompt":0}}}`))
|
||||
put.Header.Set("Authorization", "Bearer "+userKey)
|
||||
put.Header.Set("Content-Type", "application/json")
|
||||
prr := newRecorderFor(t, g, put)
|
||||
if prr.Code != http.StatusForbidden {
|
||||
t.Errorf("user PUT state = %d, want 403 (a user must not rewrite the price table)", prr.Code)
|
||||
}
|
||||
|
||||
// An admin key may.
|
||||
adm := doReq(t, g, http.MethodPut, "/api/plugins/billing/state",
|
||||
`{"prices":{"default":{"prompt":1e-6}}}`)
|
||||
if adm.Code != http.StatusOK {
|
||||
t.Errorf("admin PUT state = %d: %s", adm.Code, adm.Body.String())
|
||||
}
|
||||
}
|
||||
@ -228,6 +228,10 @@ func (g *Gateway) routes(w http.ResponseWriter, r *http.Request) {
|
||||
g.handleSourcesAPI(w, r)
|
||||
case r.URL.Path == "/api/source_templates" || strings.HasPrefix(r.URL.Path, "/api/source_templates/"):
|
||||
g.handleSourceTemplatesAPI(w, r)
|
||||
case r.URL.Path == "/api/ui-inject" || strings.HasPrefix(r.URL.Path, "/api/ui-inject/"):
|
||||
g.handlePluginUI(w, r)
|
||||
case r.URL.Path == "/api/plugins" || strings.HasPrefix(r.URL.Path, "/api/plugins/"):
|
||||
g.handlePluginsAPI(w, r)
|
||||
case r.URL.Path == "/api/chat":
|
||||
g.handleChat(w, r)
|
||||
case r.URL.Path == "/api/status":
|
||||
|
||||
@ -4815,8 +4815,150 @@
|
||||
if (tab === "keys") return renderKeys();
|
||||
if (tab === "sort") return renderSort();
|
||||
if (tab === "sources") return renderSources();
|
||||
return renderAdapters();
|
||||
if (tab === "adapters") return renderAdapters();
|
||||
// A page contributed by a plugin has no renderer here: its <script>
|
||||
// already ran at injection time and owns its own DOM. We only fire the
|
||||
// "shown" callbacks so it can refresh when the user lands on it.
|
||||
if (PLUGIN_PAGES.has(tab)) return notifyPluginTab(tab);
|
||||
return undefined;
|
||||
}
|
||||
|
||||
// ---- plugin injection -------------------------------------------
|
||||
// Pages and elements contributed by Lua plugins (see docs/plugins.md).
|
||||
//
|
||||
// The server merges every plugin's extension into one payload at
|
||||
// GET /api/ui-inject, because the sidebar needs all of them before it can
|
||||
// be built. Injection happens once at boot, BEFORE the first goTab, so a
|
||||
// plugin page is a real tab rather than a special case in the router.
|
||||
const PLUGIN_PAGES = new Set();
|
||||
const PLUGIN_TAB_CBS = {};
|
||||
const PLUGIN_ELEMENTS = [];
|
||||
|
||||
// pluginAPI is the small surface a plugin's script may rely on. Kept
|
||||
// deliberately tiny: plugins are untrusted, and every convenience here is
|
||||
// one more thing to keep working across kernel changes.
|
||||
window.pluginAPI = {
|
||||
async fetchState(name) {
|
||||
const r = await fetch("/api/plugins/" + encodeURIComponent(name) + "/state", {
|
||||
credentials: "same-origin",
|
||||
});
|
||||
if (!r.ok) throw new Error("state " + r.status);
|
||||
return (await r.json()).state;
|
||||
},
|
||||
async postState(name, obj) {
|
||||
const r = await fetch("/api/plugins/" + encodeURIComponent(name) + "/state", {
|
||||
method: "PUT",
|
||||
credentials: "same-origin",
|
||||
headers: { "Content-Type": "application/json" },
|
||||
body: JSON.stringify(obj),
|
||||
});
|
||||
if (!r.ok) throw new Error((await r.json().catch(() => ({}))).error?.message || r.status);
|
||||
return true;
|
||||
},
|
||||
onTabShown(fn) {
|
||||
PLUGIN_TAB_CBS.__last = PLUGIN_TAB_CBS.__last || [];
|
||||
PLUGIN_TAB_CBS.__last.push(fn);
|
||||
},
|
||||
};
|
||||
|
||||
function notifyPluginTab(tab) {
|
||||
const fns = PLUGIN_TAB_CBS[tab] || PLUGIN_TAB_CBS.__last || [];
|
||||
fns.forEach((f) => {
|
||||
try {
|
||||
f();
|
||||
} catch (e) {
|
||||
console.warn("plugin tab callback failed", e);
|
||||
}
|
||||
});
|
||||
}
|
||||
|
||||
// injectPluginUI adds the sidebar button + pane for a plugin page and
|
||||
// mounts plugin elements into existing panes.
|
||||
async function injectPluginUI() {
|
||||
let payload;
|
||||
try {
|
||||
const r = await fetch("/api/ui-inject", { credentials: "same-origin" });
|
||||
if (!r.ok) return;
|
||||
payload = await r.json();
|
||||
} catch (e) {
|
||||
return; // plugins are optional; the UI must work without them
|
||||
}
|
||||
const ui = (payload && payload.ui) || {};
|
||||
const main = $("#main");
|
||||
const nav = $("#sb-nav");
|
||||
if (!main || !nav) return;
|
||||
|
||||
// --- page ---
|
||||
if (ui.page && ui.page.page_id && ui.page.mount) {
|
||||
const id = String(ui.page.page_id);
|
||||
if (!document.getElementById("tab-" + id)) {
|
||||
const pane = document.createElement("div");
|
||||
pane.id = "tab-" + id;
|
||||
pane.className = "tab-pane hidden";
|
||||
main.appendChild(pane);
|
||||
const btn = document.createElement("button");
|
||||
btn.className = "sb-i";
|
||||
btn.dataset.tab = id;
|
||||
btn.title = ui.page.title || id;
|
||||
btn.innerHTML =
|
||||
'<span style="font-size:18px;line-height:1">' +
|
||||
esc(ui.page.icon || "•") +
|
||||
"</span>";
|
||||
btn.onclick = () => goTab(id);
|
||||
nav.appendChild(btn);
|
||||
PLUGIN_PAGES.add(id);
|
||||
// The breadcrumb map is local to this file, so extend it here.
|
||||
if (typeof NAV_NAME === "object") NAV_NAME[id] = ui.page.title || id;
|
||||
}
|
||||
const pane = document.getElementById("tab-" + id);
|
||||
if (pane && !pane.dataset.pluginMounted) {
|
||||
pane.dataset.pluginMounted = "1";
|
||||
// Split the mount so <script>/<style> run only AFTER the markup is
|
||||
// in the document. Setting innerHTML with a <script> tag does not
|
||||
// execute it, which is exactly what we want to avoid the opposite
|
||||
// problem: running before its own DOM exists.
|
||||
const tpl = document.createElement("template");
|
||||
tpl.innerHTML = ui.page.mount;
|
||||
pane.appendChild(tpl.content);
|
||||
// Move each script into a fresh element so it executes.
|
||||
pane.querySelectorAll("script").forEach((old) => {
|
||||
const s = document.createElement("script");
|
||||
Array.from(old.attributes).forEach((a) => s.setAttribute(a.name, a.value));
|
||||
s.textContent = old.textContent;
|
||||
old.replaceWith(s);
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
// --- elements into existing pages ---
|
||||
(ui.elements || []).forEach((el, i) => {
|
||||
const target = document.getElementById("tab-" + el.target);
|
||||
if (!target || !el.mount) return;
|
||||
PLUGIN_ELEMENTS.push(el);
|
||||
const wrap = document.createElement("div");
|
||||
wrap.className = "plugin-el";
|
||||
wrap.dataset.target = el.target;
|
||||
wrap.dataset.idx = String(i);
|
||||
const tpl = document.createElement("template");
|
||||
tpl.innerHTML = el.mount;
|
||||
wrap.appendChild(tpl.content);
|
||||
const anchor = String(el.anchor || "bottom");
|
||||
if (anchor === "top") target.prepend(wrap);
|
||||
else if (anchor.startsWith("before:") || anchor.startsWith("after:")) {
|
||||
const [kind, sel] = anchor.split(/:(.+)/);
|
||||
const ref = target.querySelector(sel);
|
||||
if (ref) ref.parentNode.insertBefore(wrap, kind === "before" ? ref : ref.nextSibling);
|
||||
else target.appendChild(wrap);
|
||||
} else target.appendChild(wrap);
|
||||
wrap.querySelectorAll("script").forEach((old) => {
|
||||
const s = document.createElement("script");
|
||||
Array.from(old.attributes).forEach((a) => s.setAttribute(a.name, a.value));
|
||||
s.textContent = old.textContent;
|
||||
old.replaceWith(s);
|
||||
});
|
||||
});
|
||||
}
|
||||
|
||||
(async () => {
|
||||
try {
|
||||
const me = await api("/api/keys/me");
|
||||
@ -4836,7 +4978,15 @@
|
||||
window.addEventListener("pagehide", () => releaseRecords(false));
|
||||
window.addEventListener("beforeunload", () => releaseRecords(false));
|
||||
|
||||
refresh("status");
|
||||
// Plugin injection runs BEFORE the first render: a plugin page must exist
|
||||
// in #main and the sidebar before goTab runs, otherwise the sidebar shows
|
||||
// no entry and the pane is missing for a moment. Awaited (not fired and
|
||||
// forgotten) so a slow /api/ui-inject cannot race the first paint.
|
||||
injectPluginUI()
|
||||
.catch(() => {})
|
||||
.finally(() => {
|
||||
refresh("status");
|
||||
});
|
||||
</script>
|
||||
</body>
|
||||
</html>
|
||||
|
||||
158
internal/gateway/ui_plugin_test.go
Normal file
158
internal/gateway/ui_plugin_test.go
Normal file
@ -0,0 +1,158 @@
|
||||
package gateway
|
||||
|
||||
import (
|
||||
"strings"
|
||||
"testing"
|
||||
)
|
||||
|
||||
// The plugin UI injection is JavaScript inside the embedded index.html, and it
|
||||
// is the ONLY thing that turns a plugin's `ui` block into a visible page or
|
||||
// element. These tests pin the wiring on the JS side; the server side (what the
|
||||
// payload contains) is covered by TestUIInjectServesPluginUI and the lua
|
||||
// package's TestBillingPluginDeclaresUI.
|
||||
//
|
||||
// What makes this worth pinning: a missing hook here fails SILENTLY. The page
|
||||
// simply never appears, there is no error anywhere, and it looks like "the
|
||||
// plugin didn't declare a page" rather than "the UI forgot to inject it".
|
||||
|
||||
// uiSource returns the embedded WebUI document.
|
||||
func uiSourceX(t *testing.T) string {
|
||||
t.Helper()
|
||||
return uiSource(t)
|
||||
}
|
||||
|
||||
// TestUIFetchesPluginInjection: the boot sequence must ask the kernel what to
|
||||
// inject. Without this fetch the whole feature is inert.
|
||||
func TestUIFetchesPluginInjection(t *testing.T) {
|
||||
src := uiSourceX(t)
|
||||
if !strings.Contains(src, "/api/ui-inject") {
|
||||
t.Error("the WebUI never calls /api/ui-inject; plugin pages and elements can never appear")
|
||||
}
|
||||
}
|
||||
|
||||
// TestUIInjectsBeforeFirstRender: injection must be awaited before the first
|
||||
// refresh, otherwise the sidebar is built without the plugin entry and the
|
||||
// first paint races the fetch. This is an ordering contract, so it is asserted
|
||||
// on the source order rather than trusted.
|
||||
func TestUIInjectsBeforeFirstRender(t *testing.T) {
|
||||
src := uiSourceX(t)
|
||||
iInject := strings.Index(src, "injectPluginUI()")
|
||||
iRefresh := strings.LastIndex(src, `refresh("status")`)
|
||||
if iInject < 0 {
|
||||
t.Fatal("injectPluginUI() is never called")
|
||||
}
|
||||
if iRefresh < 0 {
|
||||
t.Fatal("the boot sequence no longer calls refresh(\"status\")")
|
||||
}
|
||||
if iInject > iRefresh {
|
||||
t.Error("injectPluginUI() is called after the first refresh; the sidebar " +
|
||||
"and #main would be built before the plugin page exists")
|
||||
}
|
||||
// And it must be awaited, not fire-and-forget.
|
||||
window := src[iInject:]
|
||||
if !strings.Contains(window[:200], ".finally") && !strings.Contains(window[:200], "await") {
|
||||
t.Error("injectPluginUI() is not awaited before refresh; a slow response " +
|
||||
"would race the first paint")
|
||||
}
|
||||
}
|
||||
|
||||
// TestUIPluginScriptsRunAfterMarkup is the subtle one. Setting innerHTML with a
|
||||
// <script> tag does NOT execute it; appending via a template neither does. The
|
||||
// mount therefore has to be inserted first and its scripts re-created
|
||||
// afterwards, or a plugin's script runs before its own DOM exists — which is
|
||||
// exactly the "document.getElementById returns null" failure mode.
|
||||
func TestUIPluginScriptsRunAfterMarkup(t *testing.T) {
|
||||
src := uiSourceX(t)
|
||||
// A <template> is used to parse the mount without executing scripts...
|
||||
if !strings.Contains(src, "createElement(\"template\")") {
|
||||
t.Error("the mount is not parsed via <template>; scripts could execute before their DOM")
|
||||
}
|
||||
// ...and scripts are then re-created as fresh elements so they DO run.
|
||||
if !strings.Contains(src, "document.createElement(\"script\")") {
|
||||
t.Error("plugin <script> blocks are never re-created, so they never execute")
|
||||
}
|
||||
if !strings.Contains(src, "replaceWith(s)") {
|
||||
t.Error("the original inert <script> is not replaced by an executable one")
|
||||
}
|
||||
}
|
||||
|
||||
// TestUIPluginAPISurface: the documented browser API must exist with the exact
|
||||
// names docs/plugins.md promises, since plugin authors code against it.
|
||||
func TestUIPluginAPISurface(t *testing.T) {
|
||||
src := uiSourceX(t)
|
||||
for _, member := range []string{"fetchState", "postState", "onTabShown"} {
|
||||
if !strings.Contains(src, member+":") && !strings.Contains(src, member+"(") {
|
||||
t.Errorf("window.pluginAPI.%s is missing; docs/plugins.md documents it", member)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// TestUIPluginPageBecomesRealTab: a plugin page must get a pane in #main AND a
|
||||
// sidebar button wired to goTab, otherwise the page is unreachable.
|
||||
func TestUIPluginPageBecomesRealTab(t *testing.T) {
|
||||
src := uiSourceX(t)
|
||||
// pane in #main
|
||||
if !strings.Contains(src, `pane.id = "tab-" + id`) {
|
||||
t.Error("no pane is created for a plugin page")
|
||||
}
|
||||
if !strings.Contains(src, "main.appendChild(pane)") {
|
||||
t.Error("the plugin pane is not appended to #main")
|
||||
}
|
||||
// sidebar button wired to the tab router
|
||||
if !strings.Contains(src, "btn.dataset.tab = id") {
|
||||
t.Error("the sidebar button is not given a data-tab, so goTab() will not route to it")
|
||||
}
|
||||
if !strings.Contains(src, "btn.onclick = () => goTab(id)") {
|
||||
t.Error("the sidebar button is not wired to goTab()")
|
||||
}
|
||||
// and the router must know about it
|
||||
if !strings.Contains(src, "PLUGIN_PAGES.has(tab)") {
|
||||
t.Error("refresh() does not route plugin pages, so opening one renders nothing")
|
||||
}
|
||||
}
|
||||
|
||||
// TestUIPluginElementsHonorAnchor: elements declare top / bottom / before:sel /
|
||||
// after:sel. Silently ignoring the anchor would put a "top" tile at the bottom
|
||||
// of the status page, which looks like a layout bug rather than a plugin bug.
|
||||
func TestUIPluginElementsHonorAnchor(t *testing.T) {
|
||||
src := uiSourceX(t)
|
||||
for _, anchor := range []string{`anchor === "top"`, `anchor.startsWith("before:")`, `"after:"`} {
|
||||
if !strings.Contains(src, anchor) {
|
||||
t.Errorf("the anchor form %s is not handled; elements would all land at the bottom", anchor)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// TestUIPluginInjectionFailureIsNonFatal: plugins are optional, so a failed
|
||||
// /api/ui-inject must still leave a working UI (the dashboard has to render).
|
||||
// Two places have to cooperate: the function swallows the fetch error, and the
|
||||
// caller catches anything that still escapes so refresh() always runs.
|
||||
func TestUIPluginInjectionFailureIsNonFatal(t *testing.T) {
|
||||
src := uiSourceX(t)
|
||||
// inside the function: the fetch is wrapped in try/catch
|
||||
fnStart := strings.Index(src, "async function injectPluginUI()")
|
||||
if fnStart < 0 {
|
||||
t.Fatal("injectPluginUI() is not defined")
|
||||
}
|
||||
fn := src[fnStart:]
|
||||
if !strings.Contains(fn, "plugins are optional; the UI must work without them") {
|
||||
t.Error("injectPluginUI does not guard its own fetch failure")
|
||||
}
|
||||
// at the call site: the rejection cannot escape before the first render
|
||||
// LastIndex, not Index: the DEFINITION of injectPluginUI also matches, and
|
||||
// the definition has no .catch on it.
|
||||
iCall := strings.LastIndex(src, "injectPluginUI()")
|
||||
if iCall < 0 {
|
||||
t.Fatal("injectPluginUI() is never called")
|
||||
}
|
||||
// Bound the window at len(src): the call site sits near EOF and a fixed
|
||||
// slice overruns it (a panic in a test is worse than a skipped assertion).
|
||||
end := iCall + 220
|
||||
if end > len(src) {
|
||||
end = len(src)
|
||||
}
|
||||
if !strings.Contains(src[iCall:end], ".catch") {
|
||||
t.Error("a failed /api/ui-inject would reject before refresh(\"status\"), " +
|
||||
"leaving the dashboard blank")
|
||||
}
|
||||
}
|
||||
296
internal/lua/billing_test.go
Normal file
296
internal/lua/billing_test.go
Normal file
@ -0,0 +1,296 @@
|
||||
package lua
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
"testing"
|
||||
)
|
||||
|
||||
// The billing plugin ships with the gateway, so its arithmetic is a contract:
|
||||
// a wrong price silently produces wrong money. These tests drive it through the
|
||||
// real hook path and check the NUMBERS, not merely that it loads.
|
||||
|
||||
func billingVM(t *testing.T) (*Plugins, string) {
|
||||
t.Helper()
|
||||
dir := filepath.Join(t.TempDir(), "adapters")
|
||||
vm := NewVM(dir)
|
||||
if err := vm.Start(); err != nil {
|
||||
t.Fatalf("vm: %v", err)
|
||||
}
|
||||
t.Cleanup(vm.Stop)
|
||||
pdir := filepath.Join(t.TempDir(), "plugins")
|
||||
ps := NewPlugins(vm, pdir)
|
||||
if err := ps.SeedBundled(); err != nil {
|
||||
t.Fatalf("seed: %v", err)
|
||||
}
|
||||
if err := ps.LoadDir(); err != nil {
|
||||
t.Fatalf("load: %v", err)
|
||||
}
|
||||
return ps, pdir
|
||||
}
|
||||
|
||||
// stateOf reads the plugin's published state as a generic map.
|
||||
func stateOf(t *testing.T, ps *Plugins) map[string]interface{} {
|
||||
t.Helper()
|
||||
raw := ps.State("billing")
|
||||
if raw == nil {
|
||||
t.Fatal("billing published no state")
|
||||
}
|
||||
b, err := json.Marshal(raw)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
var out map[string]interface{}
|
||||
if err := json.Unmarshal(b, &out); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
func approx(t *testing.T, name string, got, want float64) {
|
||||
t.Helper()
|
||||
d := got - want
|
||||
if d < 0 {
|
||||
d = -d
|
||||
}
|
||||
if d > 1e-9 {
|
||||
t.Errorf("%s = %v, want %v (delta %v)", name, got, want, d)
|
||||
}
|
||||
}
|
||||
|
||||
// TestBillingZeroPricesIsSafe: with no configuration the plugin must still run
|
||||
// and report volume. A nil-price crash here would take out every request.
|
||||
func TestBillingZeroPricesIsSafe(t *testing.T) {
|
||||
ps, _ := billingVM(t)
|
||||
ps.Fire(StageRequestEnd, map[string]interface{}{
|
||||
"model": "m", "source": "s", "key": "***aaaaaa", "ok": true,
|
||||
"prompt_tokens": 100, "completion_tokens": 50, "time": 1750000000000,
|
||||
})
|
||||
st := stateOf(t, ps)
|
||||
total := st["total"].(map[string]interface{})
|
||||
if total["requests"].(float64) != 1 {
|
||||
t.Errorf("requests = %v, want 1", total["requests"])
|
||||
}
|
||||
approx(t, "cost with no prices", total["cost"].(float64), 0)
|
||||
}
|
||||
|
||||
// TestBillingModelTokenPricing: the core case. prompt and completion are priced
|
||||
// SEPARATELY, which is how providers publish and how the total must come out.
|
||||
func TestBillingModelTokenPricing(t *testing.T) {
|
||||
ps, _ := billingVM(t)
|
||||
// Setting prices must NOT disturb the (still empty) totals, which is the
|
||||
// whole point of the prices/state split.
|
||||
if err := ps.SetState("billing", map[string]interface{}{
|
||||
"prices": map[string]interface{}{
|
||||
"currency": "USD",
|
||||
"models": map[string]interface{}{
|
||||
"gpt-5.4": map[string]interface{}{"prompt": 1.25e-6, "completion": 1e-5},
|
||||
},
|
||||
},
|
||||
}); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
// 1000 prompt * 1.25e-6 = 0.00125 ; 500 completion * 1e-5 = 0.005
|
||||
ps.Fire(StageRequestEnd, map[string]interface{}{
|
||||
"model": "gpt-5.4", "source": "up", "key": "***aaaaaa", "ok": true,
|
||||
"prompt_tokens": 1000, "completion_tokens": 500, "time": 1750000000000,
|
||||
})
|
||||
st := stateOf(t, ps)
|
||||
approx(t, "total cost", st["total"].(map[string]interface{})["cost"].(float64), 0.00625)
|
||||
byModel := st["by_model"].(map[string]interface{})["gpt-5.4"].(map[string]interface{})
|
||||
approx(t, "model cost", byModel["cost"].(float64), 0.00625)
|
||||
if byModel["completion_tokens"].(float64) != 500 {
|
||||
t.Errorf("completion_tokens = %v, want 500", byModel["completion_tokens"])
|
||||
}
|
||||
}
|
||||
|
||||
// TestBillingPerRequestAndTokenCombine: a flat fee is ADDED to the token cost,
|
||||
// which is how an image model can be "tokens + fixed fee".
|
||||
func TestBillingPerRequestAndTokenCombine(t *testing.T) {
|
||||
ps, _ := billingVM(t)
|
||||
if err := ps.SetState("billing", map[string]interface{}{
|
||||
"prices": map[string]interface{}{
|
||||
"models": map[string]interface{}{
|
||||
"kolors": map[string]interface{}{"prompt": 1e-6, "completion": 2e-6, "per_request": 0.04},
|
||||
},
|
||||
},
|
||||
}); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
// 100*1e-6 + 50*2e-6 + 0.04 = 0.0402
|
||||
ps.Fire(StageRequestEnd, map[string]interface{}{
|
||||
"model": "kolors", "source": "sf", "key": "***bbbbbb", "ok": true,
|
||||
"prompt_tokens": 100, "completion_tokens": 50, "time": 1750000000000,
|
||||
})
|
||||
st := stateOf(t, ps)
|
||||
approx(t, "total", st["total"].(map[string]interface{})["cost"].(float64), 0.0402)
|
||||
}
|
||||
|
||||
// TestBillingPrecedence: keys > models > default for token prices.
|
||||
func TestBillingPrecedence(t *testing.T) {
|
||||
ps, _ := billingVM(t)
|
||||
_ = ps.SetState("billing", map[string]interface{}{
|
||||
"prices": map[string]interface{}{
|
||||
"default": map[string]interface{}{"prompt": 9e-6, "completion": 9e-6},
|
||||
"models": map[string]interface{}{"m": map[string]interface{}{"prompt": 2e-6, "completion": 3e-6}},
|
||||
"keys": map[string]interface{}{"***cccccc": map[string]interface{}{"prompt": 1e-6, "completion": 1.5e-6}},
|
||||
},
|
||||
})
|
||||
|
||||
// No key match -> model price.
|
||||
ps.Fire(StageRequestEnd, map[string]interface{}{
|
||||
"model": "m", "source": "s", "key": "***other", "ok": true,
|
||||
"prompt_tokens": 1000, "completion_tokens": 1000, "time": 1750000000000,
|
||||
})
|
||||
// Key match -> key price wins.
|
||||
ps.Fire(StageRequestEnd, map[string]interface{}{
|
||||
"model": "m", "source": "s", "key": "***cccccc", "ok": true,
|
||||
"prompt_tokens": 1000, "completion_tokens": 1000, "time": 1750000000000,
|
||||
})
|
||||
st := stateOf(t, ps)
|
||||
// 1000*2e-6 + 1000*3e-6 = 0.005 ; 1000*1e-6 + 1000*1.5e-6 = 0.0025
|
||||
approx(t, "total (model + key)", st["total"].(map[string]interface{})["cost"].(float64), 0.0075)
|
||||
|
||||
// An unpriced model falls back to default.
|
||||
ps2, _ := billingVM(t)
|
||||
_ = ps2.SetState("billing", map[string]interface{}{
|
||||
"prices": map[string]interface{}{
|
||||
"default": map[string]interface{}{"prompt": 9e-6, "completion": 9e-6},
|
||||
},
|
||||
})
|
||||
ps2.Fire(StageRequestEnd, map[string]interface{}{
|
||||
"model": "unknown", "source": "s", "key": "***d", "ok": true,
|
||||
"prompt_tokens": 1000, "completion_tokens": 1000, "time": 1750000000000,
|
||||
})
|
||||
approx(t, "default fallback", stateOf(t, ps2)["total"].(map[string]interface{})["cost"].(float64), 0.018)
|
||||
}
|
||||
|
||||
// TestBillingAggregatesEveryDimension: one request must land in all four
|
||||
// rollups plus the daily bucket. A missing dimension is the kind of bug a
|
||||
// dashboard hides (it just renders an empty table).
|
||||
func TestBillingAggregatesEveryDimension(t *testing.T) {
|
||||
ps, _ := billingVM(t)
|
||||
_ = ps.SetState("billing", map[string]interface{}{
|
||||
"models": map[string]interface{}{"m1": map[string]interface{}{"prompt": 1e-6, "completion": 1e-6}},
|
||||
})
|
||||
ps.Fire(StageRequestEnd, map[string]interface{}{
|
||||
"model": "m1", "source": "srcA", "key": "***key01", "ok": true,
|
||||
"prompt_tokens": 100, "completion_tokens": 100, "time": 1750000000000,
|
||||
})
|
||||
st := stateOf(t, ps)
|
||||
for _, dim := range []string{"by_source", "by_model", "by_key", "by_day"} {
|
||||
m, ok := st[dim].(map[string]interface{})
|
||||
if !ok || len(m) == 0 {
|
||||
t.Errorf("%s is empty; a dimension is missing", dim)
|
||||
}
|
||||
}
|
||||
if _, ok := st["by_source"].(map[string]interface{})["srcA"]; !ok {
|
||||
t.Error("by_source lacks srcA")
|
||||
}
|
||||
if _, ok := st["by_key"].(map[string]interface{})["***key01"]; !ok {
|
||||
t.Error("by_key lacks the gateway key")
|
||||
}
|
||||
// Milliseconds must be converted, not used as seconds: a raw 1750000000000
|
||||
// would land in a year-57000 bucket.
|
||||
days := st["by_day"].(map[string]interface{})
|
||||
found := false
|
||||
for k := range days {
|
||||
if len(k) == 10 && strings.Contains(k, "-") {
|
||||
found = true
|
||||
}
|
||||
if strings.HasPrefix(k, "5") && len(k) > 6 {
|
||||
t.Errorf("by_day key %q suggests millisecond timestamps were not converted", k)
|
||||
}
|
||||
}
|
||||
if !found {
|
||||
t.Errorf("by_day has no YYYY-MM-DD key: %v", days)
|
||||
}
|
||||
}
|
||||
|
||||
// TestBillingFailedRequestPolicy: a failed request keeps its token cost (tokens
|
||||
// really were consumed) but drops the flat per_request fee (never charged).
|
||||
func TestBillingFailedRequestPolicy(t *testing.T) {
|
||||
ps, _ := billingVM(t)
|
||||
_ = ps.SetState("billing", map[string]interface{}{
|
||||
"prices": map[string]interface{}{
|
||||
"models": map[string]interface{}{
|
||||
"m": map[string]interface{}{"prompt": 1e-6, "completion": 1e-6, "per_request": 0.5},
|
||||
},
|
||||
},
|
||||
})
|
||||
ps.Fire(StageRequestEnd, map[string]interface{}{
|
||||
"model": "m", "source": "s", "key": "***e", "ok": false, "status": 500,
|
||||
"prompt_tokens": 1000, "completion_tokens": 0, "time": 1750000000000,
|
||||
})
|
||||
st := stateOf(t, ps)
|
||||
// 1000*1e-6 = 0.001, flat dropped.
|
||||
approx(t, "failed request", st["total"].(map[string]interface{})["cost"].(float64), 0.001)
|
||||
if st["total"].(map[string]interface{})["failures"].(float64) != 1 {
|
||||
t.Error("failures not counted")
|
||||
}
|
||||
}
|
||||
|
||||
// TestBillingStateAPIReplace: the admin price update must actually change
|
||||
// subsequent pricing (not just be stored).
|
||||
func TestBillingStateAPIReplace(t *testing.T) {
|
||||
ps, _ := billingVM(t)
|
||||
_ = ps.SetState("billing", map[string]interface{}{
|
||||
"prices": map[string]interface{}{
|
||||
"models": map[string]interface{}{"m": map[string]interface{}{"prompt": 1e-6, "completion": 0}},
|
||||
},
|
||||
})
|
||||
ps.Fire(StageRequestEnd, map[string]interface{}{
|
||||
"model": "m", "source": "s", "key": "***f", "ok": true,
|
||||
"prompt_tokens": 1000, "completion_tokens": 0, "time": 1750000000000,
|
||||
})
|
||||
approx(t, "before reprice", stateOf(t, ps)["total"].(map[string]interface{})["cost"].(float64), 0.001)
|
||||
|
||||
_ = ps.SetState("billing", map[string]interface{}{
|
||||
"prices": map[string]interface{}{
|
||||
"models": map[string]interface{}{"m": map[string]interface{}{"prompt": 2e-6, "completion": 0}},
|
||||
},
|
||||
})
|
||||
ps.Fire(StageRequestEnd, map[string]interface{}{
|
||||
"model": "m", "source": "s", "key": "***f", "ok": true,
|
||||
"prompt_tokens": 1000, "completion_tokens": 0, "time": 1750000000000,
|
||||
})
|
||||
// 0.001 (old) + 0.002 (new price)
|
||||
approx(t, "after reprice", stateOf(t, ps)["total"].(map[string]interface{})["cost"].(float64), 0.003)
|
||||
}
|
||||
|
||||
// TestBillingPluginDeclaresUI: the shipped plugin must ship its dashboard, or
|
||||
// "billing is enabled" would be true while showing the user nothing.
|
||||
func TestBillingPluginDeclaresUI(t *testing.T) {
|
||||
ps, _ := billingVM(t)
|
||||
for _, row := range ps.List() {
|
||||
if row["name"] != "billing" {
|
||||
continue
|
||||
}
|
||||
ui, ok := row["ui"].(map[string]interface{})
|
||||
if !ok {
|
||||
t.Fatal("billing declares no ui")
|
||||
}
|
||||
if page, _ := ui["page"].(string); page != "billing" {
|
||||
t.Errorf("ui.page = %v, want \"billing\"", ui["page"])
|
||||
}
|
||||
if n, _ := ui["elements"].(int); n < 1 {
|
||||
t.Error("billing contributes no element to an existing page")
|
||||
}
|
||||
return
|
||||
}
|
||||
t.Fatal("billing plugin is not loaded")
|
||||
}
|
||||
|
||||
// TestBillingPluginLoadedByDefault: the shipped plugin must load with no
|
||||
// configuration, since seeding only happens on a fresh plugin dir.
|
||||
func TestBillingPluginLoadedByDefault(t *testing.T) {
|
||||
ps, pdir := billingVM(t)
|
||||
if ps.Count() != 1 {
|
||||
t.Fatalf("expected 1 bundled plugin, got %d", ps.Count())
|
||||
}
|
||||
if _, err := os.Stat(filepath.Join(pdir, "billing.lua")); err != nil {
|
||||
t.Errorf("billing.lua was not written to the plugin dir: %v", err)
|
||||
}
|
||||
}
|
||||
808
internal/lua/plugins.go
Normal file
808
internal/lua/plugins.go
Normal file
@ -0,0 +1,808 @@
|
||||
package lua
|
||||
|
||||
// Plugin runtime.
|
||||
//
|
||||
// A plugin is a single .lua file, loaded from its own directory, that extends
|
||||
// the gateway in two ways:
|
||||
//
|
||||
// 1. Hooks: it registers callbacks on the request pipeline's stages
|
||||
// (request_start, response_end, …). A hook receives a JSON table and
|
||||
// returns either nil (no opinion) or a JSON object.
|
||||
// 2. UI: at boot it returns HTML/CSS/JS fragments that the kernel injects
|
||||
// into the WebUI — either as a whole new page, or as an extra element on
|
||||
// an existing page.
|
||||
//
|
||||
// Design notes that are load-bearing (each one cost something to learn):
|
||||
//
|
||||
// - Plugins run in a SEPARATE VM from adapters, and each plugin gets its own
|
||||
// elastic pool, exactly like an adapter. Sharing one state would let a
|
||||
// plugin's globals corrupt an adapter's protocol translation (or vice
|
||||
// versa), and a plugin is third-party code while an adapter is core.
|
||||
//
|
||||
// - A plugin that errors must NEVER break request forwarding. Hook calls are
|
||||
// wrapped so a plugin error is logged and the original value is returned
|
||||
// unchanged. A broken plugin is a missing feature, not an outage — the same
|
||||
// reason transform_stream_chunk swallows errors today.
|
||||
//
|
||||
// - The billing plugin therefore cannot be trusted to be the source of truth
|
||||
// for anything the gateway must enforce. It reads usage off the hook
|
||||
// payload and accumulates in its own state; the gateway's own quota
|
||||
// accounting (internal/gateway/stats.go) stays authoritative for limits.
|
||||
// Two accounting paths that disagree is worse than one that is slightly
|
||||
// less featureful, so the split is explicit and documented.
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"sort"
|
||||
"strings"
|
||||
"sync"
|
||||
"sync/atomic"
|
||||
|
||||
golua "github.com/aarzilli/golua/lua"
|
||||
)
|
||||
|
||||
// Stage identifies a point in the request pipeline. Plugins may register a
|
||||
// function for any stage; unknown stages are ignored at call time so an old
|
||||
// plugin survives a gateway that grew new stages.
|
||||
type Stage string
|
||||
|
||||
// The pipeline stages. The order here is the order they fire in; it is the
|
||||
// contract plugins are written against.
|
||||
const (
|
||||
// StageRequestStart fires after the gateway has parsed and authorized a
|
||||
// request but BEFORE any upstream slot is chosen. payload:
|
||||
// stage, type ("chat"|"stream"|"image"), model (as requested),
|
||||
// key (masked gateway key id), role, source (empty), stream (bool),
|
||||
// messages_count, tools_count, ts (unix seconds).
|
||||
StageRequestStart Stage = "request_start"
|
||||
|
||||
// StageRouted fires once a (source, model) slot has been chosen and before
|
||||
// the upstream call. payload adds: source, model (the resolved one),
|
||||
// tier (AUTO tier, -1 for the direct path), stream.
|
||||
StageRouted Stage = "routed"
|
||||
|
||||
// StageRequestEnd fires exactly once per request, after the client response
|
||||
// has been produced (or after a failure was recorded). payload adds:
|
||||
// source, model, ok, status, latency_ms, first_byte_ms, prompt_tokens,
|
||||
// completion_tokens, cache_hit_tokens, cache_miss_tokens, image_count,
|
||||
// error ("" when ok).
|
||||
//
|
||||
// This is the stage a billing plugin should read: it carries the final
|
||||
// accounting for the request, including the upstream's own usage numbers.
|
||||
StageRequestEnd Stage = "request_end"
|
||||
)
|
||||
|
||||
// AllStages is the firing order, used by the docs and by the hook listing.
|
||||
var AllStages = []Stage{
|
||||
StageRequestStart,
|
||||
StageRouted,
|
||||
StageRequestEnd,
|
||||
}
|
||||
|
||||
// UIExtension is what a plugin contributes to the WebUI at boot.
|
||||
type UIExtension struct {
|
||||
// Page is a whole new sidebar entry + pane. Requires PageID and Title.
|
||||
// The kernel renders Page's HTML into a pane whose id is "tab-"+PageID and
|
||||
// adds a sidebar button with data-tab="<PageID>".
|
||||
Page *UIPage `json:"page,omitempty"`
|
||||
// Elements are snippets injected into EXISTING pages, keyed by target page
|
||||
// id (e.g. "status", "keys"). Order within a target is plugin load order.
|
||||
Elements []UIElement `json:"elements,omitempty"`
|
||||
}
|
||||
|
||||
// UIPage is a plugin-provided page.
|
||||
type UIPage struct {
|
||||
PageID string `json:"page_id"` // kebab-case; becomes data-tab and #tab-<id>
|
||||
Title string `json:"title"` // sidebar label
|
||||
Icon string `json:"icon"` // optional inline SVG or short glyph
|
||||
Order int `json:"order"` // sidebar sort key (default 100)
|
||||
// Mount is the page body. It may contain <script> and <style>; the kernel
|
||||
// executes scripts AFTER injecting the HTML so the DOM exists, and exposes
|
||||
// `pluginAPI` to them (see docs/plugins.md).
|
||||
Mount string `json:"mount"`
|
||||
}
|
||||
|
||||
// UIElement is a snippet injected into an existing page.
|
||||
type UIElement struct {
|
||||
// Target is the id of the pane to inject into: "status", "chat", "keys",
|
||||
// "sort", "sources" or "adapters".
|
||||
Target string `json:"target"`
|
||||
// Anchor is where in the target pane: "top", "bottom" or "before:<sel>" /
|
||||
// "after:<sel>" for a CSS selector. Empty = "bottom".
|
||||
Anchor string `json:"anchor,omitempty"`
|
||||
Order int `json:"order,omitempty"` // sort key within the target
|
||||
Mount string `json:"mount"`
|
||||
}
|
||||
|
||||
// PluginInfo is the manifest a plugin declares.
|
||||
type PluginInfo struct {
|
||||
Name string `json:"name"`
|
||||
Version string `json:"version"`
|
||||
Description string `json:"description,omitempty"`
|
||||
Author string `json:"author,omitempty"`
|
||||
}
|
||||
|
||||
// Plugin is one loaded plugin.
|
||||
type Plugin struct {
|
||||
Info PluginInfo
|
||||
Hooks map[Stage]string // stage -> exported function name
|
||||
UI *UIExtension
|
||||
// LoadError is non-empty when the plugin failed to compile or register. Such
|
||||
// a plugin is listed in the UI with its error but is never called.
|
||||
LoadError string
|
||||
dir string
|
||||
// script is the plugin's source, kept so a reload can rebuild its state.
|
||||
script string
|
||||
|
||||
// state is the plugin's SINGLE authoritative Lua state, guarded by mu.
|
||||
//
|
||||
// A plugin deliberately does NOT use the adapter's elastic pool. Adapters
|
||||
// are per-worker stateless transforms, so N independent states are correct
|
||||
// (and necessary) for them. A plugin's hook, however, typically ACCUMULATES
|
||||
// into `plugin.state` — the billing plugin's totals live there — so two
|
||||
// states would mean two divergent sets of totals, and whichever worker a
|
||||
// hook happened to get would see a different number. That was a real bug
|
||||
// found by a test: SetState wrote prices to one worker, the hook then ran on
|
||||
// another and priced everything at zero.
|
||||
//
|
||||
// A single state is sufficient because every entry point (Fire, State,
|
||||
// SetState) serializes on mu, and a plugin hook is a short synchronous call.
|
||||
// A hook that blocks for seconds would stall every other plugin's hook,
|
||||
// which is the real cost — documented as a rule in docs/plugins.md.
|
||||
mu sync.Mutex
|
||||
state *worker
|
||||
}
|
||||
|
||||
// Plugins is the loaded plugin set, owned by the VM.
|
||||
type Plugins struct {
|
||||
mu sync.RWMutex
|
||||
vm *VM
|
||||
dir string // plugin directory; "" disables plugin loading entirely
|
||||
plugins []*Plugin // load order; the index is the stable plugin id
|
||||
// ui caches the merged UI extensions so the boot payload is computed once.
|
||||
ui atomic.Pointer[UIExtension]
|
||||
// stageFuncs is the precomputed stage -> []hookCall, in plugin load order.
|
||||
stageFuncs map[Stage][]hookCall
|
||||
// hookErr records per-stage plugin failures so a silently broken plugin is
|
||||
// visible in /api/status rather than merely missing.
|
||||
hookErr *hookErrors
|
||||
}
|
||||
|
||||
type hookCall struct {
|
||||
pluginIdx int
|
||||
plugin string
|
||||
fn string
|
||||
}
|
||||
|
||||
// hookErrors counts hook failures per stage, surfaced in /api/status so a
|
||||
// silently broken plugin is visible instead of just missing.
|
||||
type hookErrors struct {
|
||||
mu sync.Mutex
|
||||
counts map[Stage]int
|
||||
last map[Stage]string
|
||||
}
|
||||
|
||||
func newHookErrors() *hookErrors {
|
||||
return &hookErrors{counts: map[Stage]int{}, last: map[Stage]string{}}
|
||||
}
|
||||
|
||||
func (h *hookErrors) note(s Stage, msg string) {
|
||||
h.mu.Lock()
|
||||
h.counts[s]++
|
||||
if len(msg) > 200 {
|
||||
msg = msg[:200]
|
||||
}
|
||||
h.last[s] = msg
|
||||
h.mu.Unlock()
|
||||
}
|
||||
|
||||
func (h *hookErrors) snapshot() map[string]map[string]interface{} {
|
||||
h.mu.Lock()
|
||||
defer h.mu.Unlock()
|
||||
out := map[string]map[string]interface{}{}
|
||||
for s, n := range h.counts {
|
||||
if n == 0 {
|
||||
continue
|
||||
}
|
||||
out[string(s)] = map[string]interface{}{"count": n, "last_error": h.last[s]}
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
// pluginGlobal is the Lua global the plugin's returned table is stored under,
|
||||
// mirroring adapterGlobal for adapters.
|
||||
const pluginGlobal = "__llmsproxy_plugin"
|
||||
|
||||
// NewPlugins creates the plugin registry for a VM. dir is the plugin directory;
|
||||
// a missing directory is not an error (plugins are optional).
|
||||
func NewPlugins(vm *VM, dir string) *Plugins {
|
||||
return &Plugins{
|
||||
vm: vm,
|
||||
stageFuncs: map[Stage][]hookCall{},
|
||||
hookErr: newHookErrors(),
|
||||
dir: dir,
|
||||
}
|
||||
}
|
||||
|
||||
// LoadDir loads every .lua file in dir as a plugin. Files are loaded in
|
||||
// lexical order so a plugin's UI order is stable across restarts.
|
||||
//
|
||||
// A plugin that fails to compile or register is NOT fatal: it is kept with its
|
||||
// LoadError so the UI can show it, and it is never called. This is the same
|
||||
// posture as adapters, except adapters are core while plugins are not.
|
||||
func (ps *Plugins) LoadDir() error {
|
||||
if ps.dir == "" {
|
||||
return nil
|
||||
}
|
||||
if _, err := os.Stat(ps.dir); os.IsNotExist(err) {
|
||||
return nil
|
||||
}
|
||||
entries, err := os.ReadDir(ps.dir)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
var names []string
|
||||
for _, e := range entries {
|
||||
if e.IsDir() || !strings.HasSuffix(e.Name(), ".lua") {
|
||||
continue
|
||||
}
|
||||
names = append(names, e.Name())
|
||||
}
|
||||
sort.Strings(names)
|
||||
for _, n := range names {
|
||||
path := filepath.Join(ps.dir, n)
|
||||
code, err := os.ReadFile(path)
|
||||
if err != nil {
|
||||
continue
|
||||
}
|
||||
if err := ps.LoadSource(strings.TrimSuffix(n, ".lua"), string(code)); err != nil {
|
||||
// LoadSource records the error on the plugin itself; keep going so
|
||||
// one bad plugin does not stop the others from loading.
|
||||
continue
|
||||
}
|
||||
}
|
||||
ps.rebuild()
|
||||
return nil
|
||||
}
|
||||
|
||||
// SeedBundled writes the plugins shipped with the gateway into dir when the
|
||||
// directory does not exist yet, mirroring the adapter seeding rule: once the
|
||||
// directory exists it is authoritative, so deleting or editing a shipped plugin
|
||||
// is a real action that survives restarts.
|
||||
func (ps *Plugins) SeedBundled() error {
|
||||
if ps.dir == "" {
|
||||
return nil
|
||||
}
|
||||
if _, err := os.Stat(ps.dir); err == nil {
|
||||
return nil
|
||||
} else if !os.IsNotExist(err) {
|
||||
return err
|
||||
}
|
||||
return writeBundledPlugins(ps.dir)
|
||||
}
|
||||
|
||||
// bootPluginState compiles one plugin into a fresh Lua state and stores its
|
||||
// returned table under pluginGlobal. It is the plugin counterpart of
|
||||
// adapterPool.boot, minus the pooling: a plugin keeps exactly one state.
|
||||
func bootPluginState(code, name string) (*worker, error) {
|
||||
L := golua.NewState()
|
||||
L.OpenLibs()
|
||||
setupGlobals(L)
|
||||
if err := L.DoString(code); err != nil {
|
||||
L.Close()
|
||||
return nil, fmt.Errorf("compile plugin %s: %w", name, err)
|
||||
}
|
||||
if L.Type(-1) != golua.LUA_TTABLE {
|
||||
L.Close()
|
||||
return nil, fmt.Errorf("plugin %s must return a table", name)
|
||||
}
|
||||
// SetGlobal POPS, so this is the only global it can set (see vm.go boot()).
|
||||
L.SetGlobal(pluginGlobal)
|
||||
L.SetTop(0)
|
||||
return &worker{L: L}, nil
|
||||
}
|
||||
|
||||
// LoadSource loads one plugin from source text. name is the plugin id (the
|
||||
// file's base name). It returns an error only for conditions the caller should
|
||||
// see; a plugin that merely registers nothing is not an error.
|
||||
func (ps *Plugins) LoadSource(name, code string) error {
|
||||
if name == "" {
|
||||
return fmt.Errorf("plugin name required")
|
||||
}
|
||||
if err := os.MkdirAll(ps.dir, 0755); err != nil && ps.dir != "" {
|
||||
return err
|
||||
}
|
||||
p := &Plugin{
|
||||
Info: PluginInfo{Name: name},
|
||||
Hooks: map[Stage]string{},
|
||||
dir: ps.dir,
|
||||
}
|
||||
|
||||
// A plugin runs in its OWN Lua state, separate from every adapter and every
|
||||
// other plugin, so an error in one cannot corrupt another. The returned
|
||||
// table is stored under pluginGlobal, exactly like adapterGlobal.
|
||||
//
|
||||
// One state, held for the plugin's lifetime (see Plugin.state): a hook that
|
||||
// accumulates into plugin.state would otherwise split its totals across
|
||||
// whichever worker it happened to run on.
|
||||
w, err := bootPluginState(code, name)
|
||||
if err != nil {
|
||||
p.LoadError = err.Error()
|
||||
ps.append(p)
|
||||
ps.rebuild()
|
||||
return err
|
||||
}
|
||||
p.script = code
|
||||
p.state = w
|
||||
|
||||
// ---- manifest ----
|
||||
w.L.GetGlobal(pluginGlobal)
|
||||
if !w.L.IsNil(-1) {
|
||||
w.L.GetField(-1, "name")
|
||||
if s := w.L.ToString(-1); s != "" {
|
||||
p.Info.Name = s
|
||||
}
|
||||
w.L.SetTop(-2)
|
||||
w.L.GetField(-1, "version")
|
||||
if s := w.L.ToString(-1); s != "" {
|
||||
p.Info.Version = s
|
||||
}
|
||||
w.L.SetTop(-2)
|
||||
w.L.GetField(-1, "description")
|
||||
if s := w.L.ToString(-1); s != "" {
|
||||
p.Info.Description = s
|
||||
}
|
||||
w.L.SetTop(-2)
|
||||
w.L.GetField(-1, "author")
|
||||
if s := w.L.ToString(-1); s != "" {
|
||||
p.Info.Author = s
|
||||
}
|
||||
w.L.SetTop(-2)
|
||||
}
|
||||
w.L.SetTop(0)
|
||||
|
||||
// ---- hooks ----
|
||||
for _, st := range AllStages {
|
||||
if fn, ok := pluginHookName(w.L, string(st)); ok {
|
||||
p.Hooks[st] = fn
|
||||
}
|
||||
}
|
||||
|
||||
// ---- UI ----
|
||||
if ui, ok := readPluginUI(w.L); ok {
|
||||
p.UI = ui
|
||||
}
|
||||
|
||||
ps.append(p)
|
||||
// Rebuild here rather than only in LoadDir: LoadSource is also the single-
|
||||
// plugin entry point (the WebUI upload path), and a caller that loads one
|
||||
// plugin and immediately fires a stage must not silently get nothing.
|
||||
ps.rebuild()
|
||||
return nil
|
||||
}
|
||||
|
||||
// pluginHookName returns the exported function name a plugin registered for a
|
||||
// stage. A plugin registers either `hooks = {request_end = "on_end"}` or a
|
||||
// direct `request_end = function(...) end` on the returned table; both forms are
|
||||
// accepted because the table form keeps the manifest tidy while the direct form
|
||||
// is shorter for a single-hook plugin.
|
||||
//
|
||||
// For the anonymous-function forms the value is re-keyed onto the plugin table
|
||||
// under a synthetic name so the hot path can fetch every hook by name.
|
||||
//
|
||||
// STACK DISCIPLINE (this binding aborts the PROCESS on a bad index — SIGABRT,
|
||||
// not a Go panic, so nothing can recover it):
|
||||
//
|
||||
// - GetField/SetField take the table by ABSOLUTE index, so the plugin table's
|
||||
// index must be re-read after every SetTop, since SetTop(0) invalidates it.
|
||||
// - Therefore each form re-pushes the plugin table and re-reads its index,
|
||||
// instead of caching one index across a reset. Getting this wrong was a
|
||||
// real crash found by running the test, not by reading the code.
|
||||
func pluginHookName(L *golua.State, stage string) (string, bool) {
|
||||
L.SetTop(0)
|
||||
L.GetGlobal(pluginGlobal)
|
||||
if L.IsNil(-1) {
|
||||
L.SetTop(0)
|
||||
return "", false
|
||||
}
|
||||
plug := L.GetTop()
|
||||
|
||||
// form 1: hooks = { request_end = "fn" } (or = function)
|
||||
L.GetField(plug, "hooks")
|
||||
if L.Type(-1) != golua.LUA_TNIL {
|
||||
hooksIdx := L.GetTop()
|
||||
L.GetField(hooksIdx, stage)
|
||||
switch L.Type(-1) {
|
||||
case golua.LUA_TSTRING:
|
||||
name := L.ToString(-1)
|
||||
L.SetTop(0)
|
||||
if name != "" {
|
||||
return name, true
|
||||
}
|
||||
return "", false
|
||||
case golua.LUA_TFUNCTION:
|
||||
// An anonymous function: key it onto the plugin table under a stable
|
||||
// per-stage name so invoke() can fetch it like any other hook.
|
||||
name := "__hook_" + stage
|
||||
L.SetField(plug, name)
|
||||
L.SetTop(0)
|
||||
return name, true
|
||||
}
|
||||
L.SetTop(0)
|
||||
}
|
||||
// form 2: request_end = function(...) end directly on the table.
|
||||
// Re-push and re-read the index: SetTop(0) above invalidated `plug`.
|
||||
L.GetGlobal(pluginGlobal)
|
||||
if L.IsNil(-1) {
|
||||
L.SetTop(0)
|
||||
return "", false
|
||||
}
|
||||
plug = L.GetTop()
|
||||
L.GetField(plug, stage)
|
||||
if L.Type(-1) == golua.LUA_TFUNCTION {
|
||||
name := "__hook_" + stage
|
||||
L.SetField(plug, name)
|
||||
L.SetTop(0)
|
||||
return name, true
|
||||
}
|
||||
L.SetTop(0)
|
||||
return "", false
|
||||
}
|
||||
|
||||
// readPluginUI reads the optional ui extension block.
|
||||
func readPluginUI(L *golua.State) (*UIExtension, bool) {
|
||||
L.GetGlobal(pluginGlobal)
|
||||
if L.IsNil(-1) {
|
||||
return nil, false
|
||||
}
|
||||
L.GetField(-1, "ui")
|
||||
if L.Type(-1) == golua.LUA_TNIL {
|
||||
L.SetTop(0)
|
||||
return nil, false
|
||||
}
|
||||
var ui UIExtension
|
||||
if err := luaToJSON(L, -1, &ui); err != nil {
|
||||
L.SetTop(0)
|
||||
return nil, false
|
||||
}
|
||||
L.SetTop(0)
|
||||
if ui.Page == nil && len(ui.Elements) == 0 {
|
||||
return nil, false
|
||||
}
|
||||
return &ui, true
|
||||
}
|
||||
|
||||
func (ps *Plugins) append(p *Plugin) {
|
||||
ps.mu.Lock()
|
||||
ps.plugins = append(ps.plugins, p)
|
||||
ps.mu.Unlock()
|
||||
}
|
||||
|
||||
// rebuild recomputes the stage dispatch table and the merged UI payload. It is
|
||||
// called after any load so the hot path (Fire) is a slice walk with no map
|
||||
// lookups or locking beyond one RLock.
|
||||
func (ps *Plugins) rebuild() {
|
||||
ps.mu.Lock()
|
||||
defer ps.mu.Unlock()
|
||||
stageFuncs := map[Stage][]hookCall{}
|
||||
for i, p := range ps.plugins {
|
||||
if p.LoadError != "" {
|
||||
continue
|
||||
}
|
||||
for _, st := range AllStages {
|
||||
if fn, ok := p.Hooks[st]; ok {
|
||||
stageFuncs[st] = append(stageFuncs[st], hookCall{pluginIdx: i, plugin: p.Info.Name, fn: fn})
|
||||
}
|
||||
}
|
||||
}
|
||||
ps.stageFuncs = stageFuncs
|
||||
|
||||
merged := &UIExtension{}
|
||||
for _, p := range ps.plugins {
|
||||
if p.LoadError != "" || p.UI == nil {
|
||||
continue
|
||||
}
|
||||
if p.UI.Page != nil {
|
||||
merged.Page = p.UI.Page
|
||||
}
|
||||
merged.Elements = append(merged.Elements, p.UI.Elements...)
|
||||
}
|
||||
ps.ui.Store(merged)
|
||||
}
|
||||
|
||||
// Count returns how many plugins loaded (including ones with LoadError).
|
||||
func (ps *Plugins) Count() int {
|
||||
ps.mu.RLock()
|
||||
defer ps.mu.RUnlock()
|
||||
return len(ps.plugins)
|
||||
}
|
||||
|
||||
// List returns a JSON-friendly view of every loaded plugin, for the status API.
|
||||
func (ps *Plugins) List() []map[string]interface{} {
|
||||
ps.mu.RLock()
|
||||
defer ps.mu.RUnlock()
|
||||
out := make([]map[string]interface{}, 0, len(ps.plugins))
|
||||
for _, p := range ps.plugins {
|
||||
stages := make([]string, 0, len(p.Hooks))
|
||||
for _, st := range AllStages {
|
||||
if _, ok := p.Hooks[st]; ok {
|
||||
stages = append(stages, string(st))
|
||||
}
|
||||
}
|
||||
row := map[string]interface{}{
|
||||
"name": p.Info.Name,
|
||||
"version": p.Info.Version,
|
||||
"description": p.Info.Description,
|
||||
"author": p.Info.Author,
|
||||
"hooks": stages,
|
||||
"loaded": p.LoadError == "",
|
||||
}
|
||||
if p.LoadError != "" {
|
||||
row["error"] = p.LoadError
|
||||
}
|
||||
if p.UI != nil {
|
||||
ui := map[string]interface{}{}
|
||||
if p.UI.Page != nil {
|
||||
ui["page"] = p.UI.Page.PageID
|
||||
}
|
||||
if len(p.UI.Elements) > 0 {
|
||||
ui["elements"] = len(p.UI.Elements)
|
||||
}
|
||||
row["ui"] = ui
|
||||
}
|
||||
out = append(out, row)
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
// HookErrors returns per-stage hook failure counts (empty when all is well).
|
||||
func (ps *Plugins) HookErrors() map[string]map[string]interface{} {
|
||||
return ps.hookErr.snapshot()
|
||||
}
|
||||
|
||||
// UI returns the merged UI extensions to inject into the WebUI.
|
||||
func (ps *Plugins) UI() *UIExtension {
|
||||
if u := ps.ui.Load(); u != nil {
|
||||
return u
|
||||
}
|
||||
return &UIExtension{}
|
||||
}
|
||||
|
||||
// State returns a plugin's own published state. A plugin publishes it by
|
||||
// setting `plugin.state = {...}` inside its hook; that is the only way a hook's
|
||||
// numbers reach the WebUI, because request_end is the LAST pipeline stage and
|
||||
// has no downstream consumer to hand a return value to.
|
||||
//
|
||||
// Returns nil when the plugin does not exist or has not published anything.
|
||||
func (ps *Plugins) State(name string) interface{} {
|
||||
ps.mu.RLock()
|
||||
p := ps.find(name)
|
||||
ps.mu.RUnlock()
|
||||
if p == nil {
|
||||
return nil
|
||||
}
|
||||
p.mu.Lock()
|
||||
w := p.state
|
||||
if w == nil {
|
||||
p.mu.Unlock()
|
||||
return nil
|
||||
}
|
||||
defer p.mu.Unlock()
|
||||
L := w.L
|
||||
L.SetTop(0)
|
||||
defer L.SetTop(0)
|
||||
L.GetGlobal(pluginGlobal)
|
||||
if L.IsNil(-1) {
|
||||
return nil
|
||||
}
|
||||
plug := L.GetTop()
|
||||
L.GetField(plug, "state")
|
||||
if L.Type(-1) == golua.LUA_TNIL {
|
||||
return nil
|
||||
}
|
||||
var out interface{}
|
||||
if err := luaToJSON(L, -1, &out); err != nil {
|
||||
return nil
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
// SetState replaces a plugin's published state (admin API). It is how a
|
||||
// configuration change (a new price for a model) reaches the plugin without
|
||||
// reloading it.
|
||||
//
|
||||
// CONTRACT: a `prices` key in the payload is applied to the plugin's SEPARATE
|
||||
// `prices` field and STRIPPED from `state`. That split is deliberate: prices are
|
||||
// configuration while state is accumulated history, and a single replaceable
|
||||
// field would make a price update wipe the totals (or make the totals carry a
|
||||
// stale price table). A plugin that keeps its config elsewhere can ignore the
|
||||
// convention and read the whole payload from `state` instead.
|
||||
func (ps *Plugins) SetState(name string, state interface{}) error {
|
||||
ps.mu.RLock()
|
||||
p := ps.find(name)
|
||||
ps.mu.RUnlock()
|
||||
if p == nil {
|
||||
return fmt.Errorf("plugin %s not loaded", name)
|
||||
}
|
||||
p.mu.Lock()
|
||||
w := p.state
|
||||
if w == nil {
|
||||
p.mu.Unlock()
|
||||
return fmt.Errorf("plugin %s has no state", name)
|
||||
}
|
||||
defer p.mu.Unlock()
|
||||
L := w.L
|
||||
L.SetTop(0)
|
||||
defer L.SetTop(0)
|
||||
L.GetGlobal(pluginGlobal)
|
||||
if L.IsNil(-1) {
|
||||
return fmt.Errorf("plugin %s has no table", name)
|
||||
}
|
||||
plug := L.GetTop()
|
||||
|
||||
// Pull `prices` out of the payload before storing the rest as state.
|
||||
body := state
|
||||
if m, ok := state.(map[string]interface{}); ok {
|
||||
if prices, has := m["prices"]; has {
|
||||
pushGoValue(L, prices)
|
||||
L.SetField(plug, "prices")
|
||||
rest := make(map[string]interface{}, len(m))
|
||||
for k, v := range m {
|
||||
if k != "prices" {
|
||||
rest[k] = v
|
||||
}
|
||||
}
|
||||
if len(rest) == 0 {
|
||||
// A prices-ONLY payload is a configuration change, not a state
|
||||
// reset. Leaving `state` untouched is what makes repricing safe:
|
||||
// replacing it with an empty table would silently erase every
|
||||
// accumulated total, so the next request would start from zero
|
||||
// and the dashboard would show a sudden drop in spend.
|
||||
L.SetTop(0)
|
||||
return nil
|
||||
}
|
||||
body = rest
|
||||
}
|
||||
}
|
||||
|
||||
pushGoValue(L, body)
|
||||
L.SetField(plug, "state")
|
||||
L.SetTop(0)
|
||||
return nil
|
||||
}
|
||||
|
||||
// Unload removes a plugin from the running set. Its states are closed so the
|
||||
// memory goes back; a subsequent LoadSource with the same name works again.
|
||||
func (ps *Plugins) Unload(name string) error {
|
||||
ps.mu.Lock()
|
||||
idx := -1
|
||||
for i, p := range ps.plugins {
|
||||
if p.Info.Name == name || strings.HasSuffix(filepath.Base(p.dir), name+".lua") {
|
||||
idx = i
|
||||
break
|
||||
}
|
||||
}
|
||||
if idx < 0 {
|
||||
ps.mu.Unlock()
|
||||
return fmt.Errorf("plugin %s not loaded", name)
|
||||
}
|
||||
p := ps.plugins[idx]
|
||||
ps.plugins = append(ps.plugins[:idx], ps.plugins[idx+1:]...)
|
||||
ps.mu.Unlock()
|
||||
p.mu.Lock()
|
||||
if p.state != nil && p.state.L != nil {
|
||||
p.state.L.Close()
|
||||
p.state = nil
|
||||
}
|
||||
p.mu.Unlock()
|
||||
ps.rebuild()
|
||||
return nil
|
||||
}
|
||||
|
||||
// find returns a loaded plugin by declared name. Caller holds ps.mu.
|
||||
func (ps *Plugins) find(name string) *Plugin {
|
||||
for _, p := range ps.plugins {
|
||||
if p.Info.Name == name {
|
||||
return p
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// Fire calls every plugin registered for a stage, in plugin load order.
|
||||
//
|
||||
// A plugin may MUTATE the payload by returning a JSON object: any keys it
|
||||
// returns are merged into the payload for the next hook and returned to the
|
||||
// caller. Returning nil or an empty table means "no opinion". This lets a
|
||||
// plugin add fields (the billing plugin adds `cost_usd`) without the gateway
|
||||
// having to know about them.
|
||||
//
|
||||
// Errors are contained: a plugin that throws is logged against the stage and
|
||||
// skipped. Forwarding never depends on plugin health.
|
||||
func (ps *Plugins) Fire(stage Stage, payload map[string]interface{}) map[string]interface{} {
|
||||
ps.mu.RLock()
|
||||
calls := ps.stageFuncs[stage]
|
||||
ps.mu.RUnlock()
|
||||
if len(calls) == 0 {
|
||||
return payload
|
||||
}
|
||||
for _, hc := range calls {
|
||||
ps.mu.RLock()
|
||||
p := ps.plugins[hc.pluginIdx]
|
||||
ps.mu.RUnlock()
|
||||
if p == nil || p.LoadError != "" {
|
||||
continue
|
||||
}
|
||||
out, err := ps.invoke(p, hc.fn, payload)
|
||||
if err != nil {
|
||||
ps.hookErr.note(stage, p.Info.Name+": "+err.Error())
|
||||
continue
|
||||
}
|
||||
if len(out) > 0 {
|
||||
for k, v := range out {
|
||||
payload[k] = v
|
||||
}
|
||||
}
|
||||
}
|
||||
return payload
|
||||
}
|
||||
|
||||
// invoke runs one plugin hook on that plugin's own state, under its pool's
|
||||
// concurrency cap. The plugin's returned table is re-fetched each call because
|
||||
// the pool is elastic: a plugin may have several states, and the hook function
|
||||
// lives in each.
|
||||
func (ps *Plugins) invoke(p *Plugin, fn string, payload map[string]interface{}) (map[string]interface{}, error) {
|
||||
p.mu.Lock()
|
||||
w := p.state
|
||||
if w == nil {
|
||||
p.mu.Unlock()
|
||||
return nil, fmt.Errorf("no state for plugin %s", p.Info.Name)
|
||||
}
|
||||
defer p.mu.Unlock()
|
||||
L := w.L
|
||||
L.SetTop(0)
|
||||
defer L.SetTop(0)
|
||||
|
||||
L.GetGlobal(pluginGlobal)
|
||||
if L.IsNil(-1) {
|
||||
return nil, fmt.Errorf("plugin table missing")
|
||||
}
|
||||
plug := L.GetTop() // absolute, so nothing below shifts
|
||||
L.GetField(plug, fn)
|
||||
if !L.IsFunction(-1) {
|
||||
L.SetTop(0)
|
||||
return nil, fmt.Errorf("hook %s missing", fn)
|
||||
}
|
||||
// The hook is called with a DECODED table, not the raw JSON string: the
|
||||
// adapter protocol passes JSON text to its transforms (they decode it
|
||||
// themselves), but a plugin hook receives a table so it can read
|
||||
// payload.model directly. Passing the string made every hook fail with
|
||||
// "attempt to index local 'payload' (a string value)".
|
||||
raw, err := json.Marshal(payload)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
var decoded interface{}
|
||||
if err := json.Unmarshal(raw, &decoded); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
pushGoValue(L, decoded)
|
||||
// Call takes NO function index: it invokes whatever sits directly below the
|
||||
// nargs values it just pushed. Passing an index here is a compile-time no-op
|
||||
// in this binding and the call lands on the argument instead
|
||||
// ("attempt to call a table value").
|
||||
if err := L.Call(1, 1); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if L.GetTop() < 1 || L.IsNil(-1) {
|
||||
return nil, nil
|
||||
}
|
||||
var out map[string]interface{}
|
||||
if err := luaToJSON(L, -1, &out); err != nil {
|
||||
return nil, nil // not a table: treat as "no opinion"
|
||||
}
|
||||
return out, nil
|
||||
}
|
||||
396
internal/lua/plugins/billing.lua
Normal file
396
internal/lua/plugins/billing.lua
Normal file
@ -0,0 +1,396 @@
|
||||
-- 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 = {
|
||||
request_end = "on_request_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
|
||||
|
||||
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 { "&": "&", "<": "<", ">": ">", '"': """ }[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],
|
||||
["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
|
||||
313
internal/lua/plugins_test.go
Normal file
313
internal/lua/plugins_test.go
Normal file
@ -0,0 +1,313 @@
|
||||
package lua
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
"testing"
|
||||
)
|
||||
|
||||
// newPluginVM builds a VM plus a plugin registry rooted at dir.
|
||||
func newPluginVM(t *testing.T) (*VM, *Plugins, string) {
|
||||
t.Helper()
|
||||
dir := filepath.Join(t.TempDir(), "adapters")
|
||||
vm := NewVM(dir)
|
||||
if err := vm.Start(); err != nil {
|
||||
t.Fatalf("vm start: %v", err)
|
||||
}
|
||||
t.Cleanup(vm.Stop)
|
||||
pdir := filepath.Join(t.TempDir(), "plugins")
|
||||
return vm, NewPlugins(vm, pdir), pdir
|
||||
}
|
||||
|
||||
// loadPlugin writes one plugin to disk and loads it.
|
||||
func loadPlugin(t *testing.T, ps *Plugins, name, code string) error {
|
||||
t.Helper()
|
||||
if err := os.MkdirAll(ps.dir, 0755); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := os.WriteFile(filepath.Join(ps.dir, name+".lua"), []byte(code), 0644); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
return ps.LoadSource(name, code)
|
||||
}
|
||||
|
||||
// TestPluginManifestAndHooks: the two hook registration forms both work and the
|
||||
// manifest is read.
|
||||
func TestPluginManifestAndHooks(t *testing.T) {
|
||||
_, ps, _ := newPluginVM(t)
|
||||
code := `
|
||||
local p = {}
|
||||
p.name = "demo"
|
||||
p.version = "1.2.3"
|
||||
p.description = "a demo plugin"
|
||||
p.author = "tester"
|
||||
p.hooks = { request_end = "on_end" }
|
||||
function p.on_end(payload)
|
||||
payload.seen = true
|
||||
payload.name_seen = "demo"
|
||||
return payload
|
||||
end
|
||||
return p
|
||||
`
|
||||
if err := loadPlugin(t, ps, "demo", code); err != nil {
|
||||
t.Fatalf("load: %v", err)
|
||||
}
|
||||
list := ps.List()
|
||||
if len(list) != 1 {
|
||||
t.Fatalf("List() = %d plugins, want 1", len(list))
|
||||
}
|
||||
if list[0]["name"] != "demo" || list[0]["version"] != "1.2.3" {
|
||||
t.Errorf("manifest not read: %+v", list[0])
|
||||
}
|
||||
hooks := list[0]["hooks"].([]string)
|
||||
if len(hooks) != 1 || hooks[0] != string(StageRequestEnd) {
|
||||
t.Errorf("hooks = %v, want [request_end]", hooks)
|
||||
}
|
||||
out := ps.Fire(StageRequestEnd, map[string]interface{}{"model": "m"})
|
||||
if out["seen"] != true || out["name_seen"] != "demo" {
|
||||
t.Errorf("hook did not mutate payload: %+v", out)
|
||||
}
|
||||
}
|
||||
|
||||
// TestPluginAnonymousHookForm: `request_end = function() end` directly on the
|
||||
// table must register too, since a single-hook plugin should not need a name.
|
||||
func TestPluginAnonymousHookForm(t *testing.T) {
|
||||
_, ps, _ := newPluginVM(t)
|
||||
code := `
|
||||
local p = { name = "anon" }
|
||||
p.request_end = function(payload)
|
||||
payload.hit = 1
|
||||
return payload
|
||||
end
|
||||
return p
|
||||
`
|
||||
if err := loadPlugin(t, ps, "anon", code); err != nil {
|
||||
t.Fatalf("load: %v", err)
|
||||
}
|
||||
out := ps.Fire(StageRequestEnd, map[string]interface{}{})
|
||||
if out["hit"] != float64(1) {
|
||||
t.Errorf("anonymous hook did not fire: %+v", out)
|
||||
}
|
||||
}
|
||||
|
||||
// TestPluginHookStagesFireInOrder: each stage reaches only its own hooks.
|
||||
func TestPluginHookStagesFireInOrder(t *testing.T) {
|
||||
_, ps, _ := newPluginVM(t)
|
||||
code := `
|
||||
local p = { name = "stages" }
|
||||
p.hooks = {
|
||||
request_start = "s1",
|
||||
routed = "s2",
|
||||
request_end = "s3",
|
||||
}
|
||||
function p.s1(x) x.order = (x.order or "") .. "1" return x end
|
||||
function p.s2(x) x.order = (x.order or "") .. "2" return x end
|
||||
function p.s3(x) x.order = (x.order or "") .. "3" return x end
|
||||
return p
|
||||
`
|
||||
if err := loadPlugin(t, ps, "stages", code); err != nil {
|
||||
t.Fatalf("load: %v", err)
|
||||
}
|
||||
payload := map[string]interface{}{}
|
||||
ps.Fire(StageRequestStart, payload)
|
||||
ps.Fire(StageRouted, payload)
|
||||
ps.Fire(StageRequestEnd, payload)
|
||||
if payload["order"] != "123" {
|
||||
t.Errorf("stage order = %v, want \"123\"", payload["order"])
|
||||
}
|
||||
}
|
||||
|
||||
// TestPluginErrorIsContained is the critical safety property: a throwing hook
|
||||
// must not propagate. Forwarding depends on it.
|
||||
func TestPluginErrorIsContained(t *testing.T) {
|
||||
_, ps, _ := newPluginVM(t)
|
||||
code := `
|
||||
local p = { name = "boom" }
|
||||
p.hooks = { request_end = "kaboom" }
|
||||
function p.kaboom(payload)
|
||||
error("intentional plugin failure")
|
||||
end
|
||||
return p
|
||||
`
|
||||
if err := loadPlugin(t, ps, "boom", code); err != nil {
|
||||
t.Fatalf("load: %v", err)
|
||||
}
|
||||
// Must not panic and must return the payload unchanged.
|
||||
out := ps.Fire(StageRequestEnd, map[string]interface{}{"model": "m"})
|
||||
if out["model"] != "m" {
|
||||
t.Errorf("payload was altered by a failing plugin: %+v", out)
|
||||
}
|
||||
// And the failure must be visible, not silent.
|
||||
errs := ps.HookErrors()
|
||||
if errs["request_end"] == nil {
|
||||
t.Error("a failing plugin left no error record; it would be silently missing")
|
||||
}
|
||||
}
|
||||
|
||||
// TestPluginFailingHookDoesNotBlockLaterPlugins: one bad plugin must not stop
|
||||
// the next one from running.
|
||||
func TestPluginFailingHookDoesNotBlockLaterPlugins(t *testing.T) {
|
||||
_, ps, _ := newPluginVM(t)
|
||||
bad := `
|
||||
local p = { name = "bad" }
|
||||
p.hooks = { request_end = "f" }
|
||||
function p.f(x) error("boom") end
|
||||
return p
|
||||
`
|
||||
good := `
|
||||
local p = { name = "good" }
|
||||
p.hooks = { request_end = "f" }
|
||||
function p.f(x) x.good = true return x end
|
||||
return p
|
||||
`
|
||||
_ = loadPlugin(t, ps, "bad", bad)
|
||||
if err := loadPlugin(t, ps, "good", good); err != nil {
|
||||
t.Fatalf("load good: %v", err)
|
||||
}
|
||||
out := ps.Fire(StageRequestEnd, map[string]interface{}{})
|
||||
if out["good"] != true {
|
||||
t.Errorf("a good plugin was blocked by a failing one: %+v", out)
|
||||
}
|
||||
}
|
||||
|
||||
// TestPluginSyntaxErrorIsIsolated: a plugin that will not compile is listed
|
||||
// with its error and is never called — it must not prevent LoadDir from loading
|
||||
// the rest.
|
||||
func TestPluginSyntaxErrorIsIsolated(t *testing.T) {
|
||||
_, ps, _ := newPluginVM(t)
|
||||
broken := "this is not lua((("
|
||||
good := `
|
||||
local p = { name = "ok" }
|
||||
p.hooks = { request_end = "f" }
|
||||
function p.f(x) x.ok = true return x end
|
||||
return p
|
||||
`
|
||||
_ = loadPlugin(t, ps, "broken", broken)
|
||||
if err := loadPlugin(t, ps, "ok", good); err != nil {
|
||||
t.Fatalf("load ok: %v", err)
|
||||
}
|
||||
if err := ps.LoadDir(); err != nil {
|
||||
t.Fatalf("LoadDir: %v", err)
|
||||
}
|
||||
// The broken plugin must not be callable and must carry an error.
|
||||
for _, row := range ps.List() {
|
||||
if row["name"] == "broken" {
|
||||
if row["loaded"] == true {
|
||||
t.Error("a plugin with a syntax error reported itself as loaded")
|
||||
}
|
||||
if row["error"] == nil || row["error"] == "" {
|
||||
t.Error("a broken plugin carries no error message")
|
||||
}
|
||||
}
|
||||
}
|
||||
// The good plugin still works.
|
||||
out := ps.Fire(StageRequestEnd, map[string]interface{}{})
|
||||
if out["ok"] != true {
|
||||
t.Errorf("good plugin stopped working: %+v", out)
|
||||
}
|
||||
}
|
||||
|
||||
// TestPluginUIExtension: a plugin can contribute a page and elements.
|
||||
func TestPluginUIExtension(t *testing.T) {
|
||||
_, ps, _ := newPluginVM(t)
|
||||
code := `
|
||||
local p = { name = "ui" }
|
||||
p.hooks = { request_end = "f" }
|
||||
function p.f(x) return x end
|
||||
p.ui = {
|
||||
page = {
|
||||
page_id = "billing",
|
||||
title = "Billing",
|
||||
icon = "💰",
|
||||
order = 50,
|
||||
mount = "<div id=billing>hi</div><script>console.log('m')</script>",
|
||||
},
|
||||
elements = {
|
||||
{ target = "status", anchor = "top", mount = "<div>cost</div>" },
|
||||
},
|
||||
}
|
||||
return p
|
||||
`
|
||||
if err := loadPlugin(t, ps, "ui", code); err != nil {
|
||||
t.Fatalf("load: %v", err)
|
||||
}
|
||||
ui := ps.UI()
|
||||
if ui.Page == nil {
|
||||
t.Fatal("no page contributed")
|
||||
}
|
||||
if ui.Page.PageID != "billing" || ui.Page.Title != "Billing" {
|
||||
t.Errorf("page = %+v", ui.Page)
|
||||
}
|
||||
if !strings.Contains(ui.Page.Mount, "console.log") {
|
||||
t.Error("mount lost its script content")
|
||||
}
|
||||
if len(ui.Elements) != 1 || ui.Elements[0].Target != "status" {
|
||||
t.Errorf("elements = %+v", ui.Elements)
|
||||
}
|
||||
}
|
||||
|
||||
// TestPluginHookReturnsNilIsNoOpinion: a hook returning nothing must leave the
|
||||
// payload untouched (plugins should not be forced to echo it back).
|
||||
func TestPluginHookReturnsNilIsNoOpinion(t *testing.T) {
|
||||
_, ps, _ := newPluginVM(t)
|
||||
code := `
|
||||
local p = { name = "silent" }
|
||||
p.hooks = { request_end = "f" }
|
||||
function p.f(payload)
|
||||
-- records nothing, returns nothing
|
||||
return nil
|
||||
end
|
||||
return p
|
||||
`
|
||||
if err := loadPlugin(t, ps, "silent", code); err != nil {
|
||||
t.Fatalf("load: %v", err)
|
||||
}
|
||||
out := ps.Fire(StageRequestEnd, map[string]interface{}{"model": "m", "ok": true})
|
||||
if out["model"] != "m" || out["ok"] != true {
|
||||
t.Errorf("a no-op hook disturbed the payload: %+v", out)
|
||||
}
|
||||
}
|
||||
|
||||
// TestPluginUIJSONShape is a wire-format guard: the kernel sends this to the
|
||||
// browser, so the shape is a contract with the WebUI.
|
||||
func TestPluginUIJSONShape(t *testing.T) {
|
||||
_, ps, _ := newPluginVM(t)
|
||||
code := `
|
||||
local p = { name = "shape" }
|
||||
p.hooks = { request_end = "f" }
|
||||
function p.f(x) return x end
|
||||
p.ui = { elements = { { target = "keys", mount = "<b>k</b>" } } }
|
||||
return p
|
||||
`
|
||||
if err := loadPlugin(t, ps, "shape", code); err != nil {
|
||||
t.Fatalf("load: %v", err)
|
||||
}
|
||||
b, err := json.Marshal(ps.UI())
|
||||
if err != nil {
|
||||
t.Fatalf("marshal UI: %v", err)
|
||||
}
|
||||
var view struct {
|
||||
Elements []struct {
|
||||
Target string `json:"target"`
|
||||
Mount string `json:"mount"`
|
||||
} `json:"elements"`
|
||||
}
|
||||
if err := json.Unmarshal(b, &view); err != nil {
|
||||
t.Fatalf("unmarshal: %v", err)
|
||||
}
|
||||
if len(view.Elements) != 1 || view.Elements[0].Target != "keys" {
|
||||
t.Errorf("UI JSON shape = %+v", view.Elements)
|
||||
}
|
||||
}
|
||||
|
||||
// TestPluginFireWithNoPluginsIsNoop: an empty registry must not allocate or fail.
|
||||
func TestPluginFireWithNoPluginsIsNoop(t *testing.T) {
|
||||
_, ps, _ := newPluginVM(t)
|
||||
in := map[string]interface{}{"a": 1}
|
||||
out := ps.Fire(StageRequestEnd, in)
|
||||
if out["a"] != 1 || ps.Count() != 0 {
|
||||
t.Errorf("empty registry misbehaved: %+v", out)
|
||||
}
|
||||
}
|
||||
@ -36,7 +36,7 @@ import (
|
||||
golua "github.com/aarzilli/golua/lua"
|
||||
)
|
||||
|
||||
//go:embed adapters/*.lua
|
||||
//go:embed adapters/*.lua plugins/*.lua
|
||||
var bundledAdapters embed.FS
|
||||
|
||||
// adapterGlobal is the reserved global holding the adapter table after the
|
||||
@ -103,6 +103,11 @@ type adapterPool struct {
|
||||
lastGrow time.Time
|
||||
idleRounds int // consecutive janitor rounds that saw reclaimable slack
|
||||
peakInUse int // high-water mark of inUse, for observability
|
||||
// pluginMode makes boot() store the returned table under pluginGlobal
|
||||
// instead of adapterGlobal. Everything else (elastic sizing, reclaim) is
|
||||
// identical, which is why plugins reuse this pool rather than getting a
|
||||
// second implementation.
|
||||
pluginMode bool
|
||||
}
|
||||
|
||||
const (
|
||||
@ -116,6 +121,13 @@ const (
|
||||
// growCooldown keeps a burst of misses from batching repeatedly while the
|
||||
// previous batch is still booting.
|
||||
growCooldown = time.Second
|
||||
// maxPluginStates caps how many concurrent VM states ONE plugin may occupy.
|
||||
// A plugin is third-party code on the request path, so its ceiling is much
|
||||
// lower than an adapter's (which is sized from the sources'
|
||||
// max_concurrent): a plugin hook is a short synchronous call, so a handful
|
||||
// of states is already far more parallelism than any real hook needs, and a
|
||||
// runaway plugin cannot balloon memory the way a per-source adapter pool can.
|
||||
maxPluginStates = 4
|
||||
// residentWorkers is how many states an adapter keeps warm once it has
|
||||
// served at least one request. Booting is milliseconds, but keeping one warm
|
||||
// removes that from the critical path of the next request. Adapters that
|
||||
@ -265,7 +277,14 @@ func (p *adapterPool) boot() (*worker, error) {
|
||||
L.Close()
|
||||
return nil, fmt.Errorf("adapter %s must return a table", p.name)
|
||||
}
|
||||
L.SetGlobal(adapterGlobal)
|
||||
// SetGlobal POPS the value off the stack, so it can only be called once per
|
||||
// boot. Plugins therefore store under pluginGlobal only, and NOT under
|
||||
// adapterGlobal: a second SetGlobal on the now-empty stack would assign nil.
|
||||
if p.pluginMode {
|
||||
L.SetGlobal(pluginGlobal)
|
||||
} else {
|
||||
L.SetGlobal(adapterGlobal)
|
||||
}
|
||||
L.SetTop(0)
|
||||
return &worker{L: L}, nil
|
||||
}
|
||||
@ -601,6 +620,55 @@ func (v *VM) Start() error {
|
||||
return nil
|
||||
}
|
||||
|
||||
// bundledPluginsDir is the directory inside the embedded FS holding the
|
||||
// plugins shipped with the gateway. It is separate from adapters/ on purpose:
|
||||
// the two are loaded by different machinery into different kinds of Lua state
|
||||
// (a protocol transform vs. request-pipeline hooks), and keeping them apart
|
||||
// makes it obvious that dropping a file in one does not affect the other.
|
||||
const bundledPluginsDir = "plugins"
|
||||
|
||||
// ReadBundledPlugin returns the source of a plugin shipped with the gateway.
|
||||
// It exists so a test (or an operator tool) can load a bundled plugin without
|
||||
// depending on whether seeding has already run for this directory.
|
||||
func ReadBundledPlugin(name string) (string, error) {
|
||||
data, err := bundledAdapters.ReadFile(bundledPluginsDir + "/" + name + ".lua")
|
||||
if err != nil {
|
||||
return "", fmt.Errorf("bundled plugin %s: %w", name, err)
|
||||
}
|
||||
return string(data), nil
|
||||
}
|
||||
|
||||
// writeBundledPlugins seeds the plugin directory with the shipped plugins.
|
||||
//
|
||||
// It runs only when the directory does not exist yet (same rule as adapters):
|
||||
// once the directory exists it is authoritative, so deleting a shipped plugin is
|
||||
// a real delete and editing one survives restarts.
|
||||
func writeBundledPlugins(dir string) error {
|
||||
if dir == "" {
|
||||
return nil
|
||||
}
|
||||
if err := os.MkdirAll(dir, 0755); err != nil {
|
||||
return fmt.Errorf("mkdir plugin dir: %w", err)
|
||||
}
|
||||
entries, err := bundledAdapters.ReadDir(bundledPluginsDir)
|
||||
if err != nil {
|
||||
return nil // nothing embedded; not an error
|
||||
}
|
||||
for _, e := range entries {
|
||||
if e.IsDir() || filepath.Ext(e.Name()) != ".lua" {
|
||||
continue
|
||||
}
|
||||
data, err := bundledAdapters.ReadFile(bundledPluginsDir + "/" + e.Name())
|
||||
if err != nil {
|
||||
continue
|
||||
}
|
||||
if err := os.WriteFile(filepath.Join(dir, e.Name()), data, 0644); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (v *VM) Stop() {
|
||||
select {
|
||||
case <-v.janitorStop:
|
||||
@ -1104,6 +1172,56 @@ func pushGoValue(L *golua.State, v interface{}) {
|
||||
}
|
||||
|
||||
func jsonEncode(v interface{}) ([]byte, error) { return json.Marshal(v) }
|
||||
|
||||
// luaToJSON converts the Lua value at idx into a Go value via json.encode, then
|
||||
// unmarshals it into out. It is the bridge used by the plugin manifest/UI
|
||||
// reader: the plugin returns a plain Lua table, and Go wants a typed struct.
|
||||
//
|
||||
// It goes through JSON rather than walking the Lua stack directly because the
|
||||
// adapter/plugin boundary already speaks JSON everywhere else (transform_request
|
||||
// gets a JSON string, hooks get a JSON string), so this keeps one representation
|
||||
// instead of two.
|
||||
func luaToJSON(L *golua.State, idx int, out interface{}) error {
|
||||
if L.GetTop() < 1 {
|
||||
return fmt.Errorf("empty stack")
|
||||
}
|
||||
abs := idx
|
||||
if abs < 0 {
|
||||
abs = L.GetTop() + 1 + abs
|
||||
}
|
||||
if abs < 1 || abs > L.GetTop() {
|
||||
return fmt.Errorf("index %d out of range (top=%d)", idx, L.GetTop())
|
||||
}
|
||||
// Absolute indices throughout: this binding aborts the process (SIGABRT)
|
||||
// on a bad index rather than panicking, so the stack is captured before any
|
||||
// push instead of being addressed relative to a shifting top.
|
||||
//
|
||||
// json.encode is pushed onto the stack and the value is pushed AFTER it, so
|
||||
// Call(1, 1) invokes it (Call takes no function index — it calls whatever
|
||||
// sits below the nargs values).
|
||||
L.GetGlobal("json")
|
||||
if L.IsNil(-1) {
|
||||
L.SetTop(0)
|
||||
return fmt.Errorf("json global missing")
|
||||
}
|
||||
L.GetField(-1, "encode")
|
||||
if L.Type(-1) != golua.LUA_TFUNCTION {
|
||||
L.SetTop(0)
|
||||
return fmt.Errorf("json.encode missing")
|
||||
}
|
||||
L.PushValue(abs)
|
||||
if err := L.Call(1, 1); err != nil {
|
||||
L.SetTop(0)
|
||||
return err
|
||||
}
|
||||
if L.GetTop() < 1 || L.Type(-1) != golua.LUA_TSTRING {
|
||||
L.SetTop(0)
|
||||
return fmt.Errorf("json.encode did not return a string")
|
||||
}
|
||||
s := L.ToString(-1)
|
||||
L.SetTop(0)
|
||||
return json.Unmarshal([]byte(s), out)
|
||||
}
|
||||
func jsonDecode(s string) (interface{}, error) {
|
||||
var v interface{}
|
||||
if err := json.Unmarshal([]byte(s), &v); err != nil {
|
||||
|
||||
Reference in New Issue
Block a user