mirror of
https://gitcode.com/JianFeeeee/ModelRouter.git
synced 2026-11-07 17:57:23 +00:00
Compare commits
149 Commits
| Author | SHA1 | Date | |
|---|---|---|---|
| c96e6051ab | |||
| 82424564b6 | |||
| d1e7fd3e77 | |||
| 8909096d80 | |||
| f6561a3809 | |||
| 2e6cc60da6 | |||
| 92aa3d6a9b | |||
| 8aa06574fc | |||
| cdd1fcb529 | |||
| 0d9fb26244 | |||
| 42eef98395 | |||
| 86ba4f57a0 | |||
| 4826837f9c | |||
| 2cf7808e63 | |||
| aaad52acc1 | |||
| 06f8c5dba4 | |||
| 4a0a1f96e4 | |||
| ac9746ec83 | |||
| 2272b4ecbd | |||
| ff8afaa43d | |||
| e5af4ec34c | |||
| 360a2a8644 | |||
| 2fd536e99b | |||
| 700f4179fa | |||
| 91e3013cb5 | |||
| 92b81ecfd3 | |||
| d119d09929 | |||
| 27b908d98c | |||
| ba10f03fc7 | |||
| cab3361073 | |||
| 35b9422232 | |||
| 2b34c7dc8a | |||
| 85c81c63cc | |||
| a2bb6de9fb | |||
| 6e63e4761d | |||
| 37545b6dbb | |||
| 314f4da8f2 | |||
| b93799a919 | |||
| c36200b78b | |||
| ba6de89d57 | |||
| 3b28fde22e | |||
| 9823773330 | |||
| 5eef5ad635 | |||
| 40636abcc2 | |||
| 22c9215b02 | |||
| e76e1623e4 | |||
| 0bf0653d00 | |||
| 6448934096 | |||
| 48e7e6b2a5 | |||
| be30ba01b3 | |||
| a21d0ca5a1 | |||
| 82d4e30b9f | |||
| cf14f66ad4 | |||
| c5079747b7 | |||
| 4e3b905b58 | |||
| 130b3f2e8d | |||
| dac0b4ce93 | |||
| 5a3b5d357c | |||
| a9d1bdeafb | |||
| d547bed457 | |||
| 97864895d4 | |||
| 31090fec81 | |||
| f9cd0770a9 | |||
| 6311913812 | |||
| aa003e84d3 | |||
| 28149bdd1a | |||
| ed9baae8e4 | |||
| 13e2397e9c | |||
| 0cd1cd51f6 | |||
| 9dc2b64013 | |||
| c66e1b5955 | |||
| 22decf2bd3 | |||
| 4b07be67b4 | |||
| ba01da937b | |||
| 913974bb6c | |||
| c8edea23f5 | |||
| ffb2b8f2f2 | |||
| 689c8cb383 | |||
| 400b6c35a5 | |||
| c19b8e6394 | |||
| 09cb215208 | |||
| d9652f479a | |||
| fe0764e375 | |||
| d0ddc9754d | |||
| c241a19b51 | |||
| d1da40e493 | |||
| aa10ee8c27 | |||
| fbdf0dea10 | |||
| 30064696b3 | |||
| ce2032435a | |||
| 8de1499c40 | |||
| 1c690611f8 | |||
| a78f7cb6c5 | |||
| d0c7465130 | |||
| cb6df0a3f0 | |||
| 42764bc99e | |||
| 8c18e0c3d7 | |||
| a51a6811a6 | |||
| a2e1adc2d8 | |||
| 7082ffb723 | |||
| 5e723b5aa5 | |||
| e10bfbb278 | |||
| ad54616bee | |||
| 18e422a943 | |||
| 980f4a0e40 | |||
| 1fe379f630 | |||
| 429afce67e | |||
| b03b12148f | |||
| 5639bb662c | |||
| a1d4365515 | |||
| bf0657bb84 | |||
| 3548745f98 | |||
| 504c5ac9a0 | |||
| 70f1c879bd | |||
| 3fe0df5bf2 | |||
| 04e544c823 | |||
| 7e33d11d15 | |||
| 8ca59b4bb2 | |||
| cd82835f25 | |||
| 26ea782350 | |||
| 980408d045 | |||
| 5c58244781 | |||
| 0121d23f91 | |||
| 4b37e1a0db | |||
| de7c372ad2 | |||
| ebe60028f5 | |||
| c51066f0b6 | |||
| 5530912d32 | |||
| cc5e725226 | |||
| 652842783f | |||
| a21ae84cbe | |||
| d072a03c9a | |||
| f814ff7468 | |||
| ef631b43dd | |||
| 9c3aabb7f9 | |||
| a7355debed | |||
| f6baa13583 | |||
| fd03ef6e4a | |||
| 97bb9c6ef4 | |||
| d81074621a | |||
| 4dd3431c26 | |||
| 3d1407e6fa | |||
| db320d2f6c | |||
| b544671f5e | |||
| 0881a49286 | |||
| 4e2b5bfda9 | |||
| 61b847da86 | |||
| 6432b61eb6 | |||
| dd708c9680 |
4
.gitignore
vendored
4
.gitignore
vendored
@ -26,3 +26,7 @@ cmd/gui/dist/
|
||||
|
||||
# local design/working notes (not part of the shipped repo)
|
||||
/plan.md
|
||||
|
||||
# scratch build/probe dirs used while debugging (GOTMPDIR, throwaway binaries)
|
||||
.build-work/
|
||||
.probe/
|
||||
|
||||
122
README.md
122
README.md
@ -24,6 +24,7 @@
|
||||
### 强大的多租户调度能力
|
||||
|
||||
- **多密钥多租户**:支持无限密钥,每个密钥独立角色、模型范围、Token 配额、重置周期
|
||||
- **按模型配额**:每把 key 的每个模型单独配 token + 请求数配额与重置周期(小时/周/月/自定义 N 小时);一个模型用满只影响该模型,同 key 其它模型照常;耗尽返 429 + `Retry-After` 并点名模型,admin key 永不受限
|
||||
- **AUTO 智能调度**:基于优先级档位的分级调度,同优先级源自动轮询负载均衡,故障自动毫秒级故障转移
|
||||
- **自愈冷却**:冷却上限 5 分钟,过半后放行 1 个探测请求,上游/额度恢复即刻回归轮询,无需等满冷却窗口
|
||||
- **Token 配额管理**:精确到模型级别的 Token 配额控制,支持小时/周/月/自定义小时周期自动重置
|
||||
@ -149,6 +150,47 @@ sources:
|
||||
timeout: 120s # 请求超时,默认 120s
|
||||
```
|
||||
|
||||
|
||||
### 自动探测模型列表
|
||||
|
||||
不想手抄模型 id 时,只填 URL 和 key,让网关去扫:
|
||||
|
||||
```bash
|
||||
# 只探测,不写任何东西
|
||||
curl -X POST http://127.0.0.1:8080/api/sources/discover \
|
||||
-H "Authorization: Bearer sk-gw-…" -H "Content-Type: application/json" \
|
||||
-d '{"base_url":"https://api.vendor.com/v1","api_key":"sk-…"}'
|
||||
# → {"ok":true,"base_url":"…","endpoint":"","models":["gpt-x",…],
|
||||
# "image_models":["gpt-image-2"],"models_path":"/models","latency_ms":1387,…}
|
||||
|
||||
# 探测后直接落库
|
||||
curl -X POST http://127.0.0.1:8080/api/sources/discover/apply \
|
||||
-H "Authorization: Bearer sk-gw-…" -H "Content-Type: application/json" \
|
||||
-d '{"base_url":"https://api.vendor.com/v1","api_key":"sk-…","name":"vendor"}'
|
||||
```
|
||||
|
||||
WebUI 源编辑框里也有「自动探测模型」按钮:填好 URL 和 key 点一下,列表自动
|
||||
填好,仍需点保存才生效。
|
||||
|
||||
几点值得知道:
|
||||
|
||||
- **URL 形态不用猜。** 厂商对「`/v1` 放 base 还是放 path」没有共识,探测会
|
||||
依次尝试若干候选并返回答对的那个,同时给出归一化后的 `base_url` +
|
||||
`endpoint`(`endpoint` 平时留空,因为 openai adapter 自带 `/chat/completions`)。
|
||||
- **`image_endpoint` 是另一套规则。** 网关的图像默认路径是**绝对**的
|
||||
`/v1/images/generations`,而 chat 默认是相对的。所以 `base_url` 带 `/v1` 时
|
||||
必须显式给 `image_endpoint: /images/generations`,否则图像请求 404。
|
||||
探测会替你填好。
|
||||
- **图像模型只看上游自己声明的 modality**,绝不从模型名猜。`gpt-image-2`
|
||||
这类 id 也出现在同时提供对话服务的目录上,猜错会让 AUTO 链选中一个只支持
|
||||
图像生成的模型,其后每一层都被它拖垮。
|
||||
- **不会逐模型探活。** 34 个模型的目录跑一轮要花钱且耗时数分钟,`/models`
|
||||
是足够便宜的信号;模型是否真能应答留给正常请求去发现。
|
||||
- **对已存在的源,必须显式 `replace_models: true`** 才会覆盖模型列表。重新
|
||||
探测顺手抹掉为 AUTO 链手工调过的 priority,是最难排查的一类数据丢失。
|
||||
- **`api_key` 留空表示沿用已存值**,不会把凭据清空。
|
||||
|
||||
|
||||
`config.yaml`、`runtime.json`、`master.key`、`adapters/` 均不提交版本库(见
|
||||
`.gitignore`),请随备份带走——`master.key` 丢失后运行时密文无法解密。
|
||||
|
||||
@ -178,6 +220,81 @@ sources:
|
||||
- 客户端用任意一个已授权的密钥明文作为 Bearer(`Authorization: Bearer <key>`)。
|
||||
- 删除密钥即从运行时存储移除,立即失效。
|
||||
|
||||
#### 密钥用量配额
|
||||
|
||||
**配额按模型单独设置**:每个密钥的 `models[]` 里,每一条模型范围各自带一份
|
||||
token 配额与请求数配额。一个模型用满只影响该模型,同一密钥的其它模型照常工作。
|
||||
|
||||
```yaml
|
||||
keys:
|
||||
- key: sk-gw-<hex>
|
||||
role: user
|
||||
name: agent-alice
|
||||
models:
|
||||
- model: deepseek-v4-flash
|
||||
token_quota: 1000000 # 本周期内该 key 用这个模型的 token 预算
|
||||
req_quota: 20000 # 本周期内的请求次数
|
||||
period: nhour # "" | hour | week | month | nhour
|
||||
hours: 6 # 仅 nhour
|
||||
- model: AUTO
|
||||
token_quota: 5000000 # AUTO 也是一条独立配额
|
||||
period: day # ← 这种写法会被拒绝(词表只有 hour/week/month/nhour)
|
||||
- model: kimi-k3 # 未列配额 = 无限
|
||||
```
|
||||
|
||||
- **没有「整钥总配额」**:这是刻意的设计。整钥总额会让「把 A 模型的额度挪给
|
||||
B 模型」变成一次全局重分配;按模型独立计费则每个模型各自可控,运维可以
|
||||
看出哪个模型吃掉了预算。
|
||||
- `period` 词表:空 = 永不过期(累计总量),`hour` / `week` / `month` = 固定窗口,
|
||||
`nhour` + `hours` = 自定义小时数。**拼错的周期在写入时就被拒**,不会静默变成
|
||||
永不过期。
|
||||
- 配额严格按密钥隔离,且**同一密钥内按模型隔离**:A 密钥用满 `m1` 不会消耗
|
||||
B 密钥的额度,同一密钥的 `m2` 也不受影响。
|
||||
- 配额统计含聊天、流式、生图,跨重启从审计日志回放(保留 40 天,覆盖最长的
|
||||
month 窗口)。
|
||||
- 配额耗尽返回 **429 + `Retry-After`**(`rate_limit_exceeded`),消息里点名是哪个
|
||||
模型用满了,客户端可等窗口重置后自动恢复;模型越权才是 403。
|
||||
**admin 密钥永不受配额限制**(其 scope 上的配额也不强制),避免把管理员
|
||||
锁在门外。
|
||||
- 窗口用量按整点小时分桶统计,实际释放比配置窗口最多晚 1 小时(配额宁可晚释放
|
||||
也不超发)。
|
||||
- `PUT /api/keys/{key}` 提交 `models` 即同时提交它们的配额(配额就是 scope 的一部分,
|
||||
不存在会与模型列表脱节的第二份预算)。显式 `0` = 解除该模型的限制。
|
||||
|
||||
##### 配额拒绝 vs 容量拒绝:两种「拒绝」含义不同
|
||||
|
||||
配额在**入口**检查,不占用任何上游槽位,因此是廉价拒绝;容量不足则要走
|
||||
`busyWait` 有界等待后才返回 503。实测(容量 8、0.6s/请求、100 并发):
|
||||
|
||||
| 情况 | 状态码 | 延迟 | 是否打到上游 |
|
||||
|---|---|---|---|
|
||||
| 配额用尽(`req_quota` 50 → 100 并发) | `429 rate_limit_exceeded` + `Retry-After` | **19–21ms** | 否 |
|
||||
| 容量打满(同一拓扑,未限配额) | `503 upstream_error` | ~2.6s(`busyWait`) | 超出部分否 |
|
||||
|
||||
客户端因此可以区分「等窗口重置」与「等上游腾出容量」:前者只需等待,后者
|
||||
通常该降并发或换源。
|
||||
|
||||
##### 配额的开销
|
||||
|
||||
配额桶按 (密钥, 模型, 整点小时) 分桶保留 40 天,实测(AMD 7840HS):
|
||||
|
||||
- 窗口查询按窗口长度而非保留总量扫描:1h 窗口 49ns、24h 窗口 55ns、
|
||||
30d 窗口 3.9μs(此前全扫保留总量,960 桶时 5.9μs)。
|
||||
- 记录一条请求:283ns、3 分配(与引入配额前相同,分配来自 ring buffer)。
|
||||
- 内存:生产形态(7 密钥 × 8 模型 × 2 源 × 满 40 天 retention)约 **3.7MB**。
|
||||
按源 pin 的 `source::model` 桶**惰性创建**——只有当某条配额真的 pin 了
|
||||
某个源时才维护,否则每条记录多写一份桶,在 20 密钥 × 8 模型 × 3 源下会
|
||||
多占 18MB。
|
||||
代价:配置 pinned 配额之前的历史用量无法事后按源拆分,所以 **pinned 配额的
|
||||
首个窗口可能少算**。
|
||||
|
||||
```bash
|
||||
# 配额耗尽时客户端看到(点名了具体模型)
|
||||
HTTP/1.1 429 Too Many Requests
|
||||
Retry-After: 2100
|
||||
{"error":{"type":"rate_limit_exceeded","message":"token quota exceeded for \"deepseek-v4-flash\" (5000000/5000000)"}}
|
||||
```
|
||||
|
||||
### 模型路由
|
||||
|
||||
`/v1/chat/completions` 的 `model` 解析顺序:
|
||||
@ -344,7 +461,10 @@ Environment=MALLOC_ARENA_MAX=2
|
||||
支持按时间范围导出 CSV;点击模型可生成 pin 到该模型的连接配置。
|
||||
- **对话页**:流式/非流式调试。
|
||||
- **密钥页**:创建/编辑网关 key,为每个 key 配模型范围(模型 + 源 + token 配额 + 周期),
|
||||
管理员管理全部 key,用户只看到自己的 key。
|
||||
管理员管理全部 key,用户只看到自己的 key。每个模型砖块显示自己的配额徽标
|
||||
(如 `1.0K 77×·1h`,未设配额显示 `∞`),点开可编辑该模型的 token 配额、
|
||||
请求数配额与重置周期 —— 配额按模型独立生效,一个用满不影响同一 key 的其它模型。
|
||||
「我的密钥」页对用户展示本 key 的模型范围与各自配额。
|
||||
- **优先级页**:拖拽积木配置 AUTO 链档位。
|
||||
- **源页**:在线增删改上游源(API key 等敏感字段加密落盘)。
|
||||
- **适配器页**:上传 / 删除 Lua 适配器脚本。
|
||||
|
||||
99
README_EN.md
99
README_EN.md
@ -38,6 +38,11 @@ Extracted and independently evolved from the multi-source LLM adapter layer of
|
||||
`reasoning_content`, `tool_calls`, `usage`).
|
||||
- **Image generation**: `POST /v1/images/generations`, routed to models with
|
||||
`kind: image`.
|
||||
- **Per-model quota**: each key gives every model its own token and request caps
|
||||
plus a reset period (hour/week/month/custom N hours). One model running out
|
||||
affects only that model — the key's other models keep working. Exhaustion
|
||||
answers 429 + `Retry-After` naming the model, so a client resumes when the
|
||||
window rolls over; admin keys are never capped.
|
||||
- **Multimodal**: `content` arrays (`image_url` etc.) pass through losslessly;
|
||||
Anthropic/Gemini/Ollama are translated automatically.
|
||||
- **LuaJIT VM**: golua-binding LuaJIT; each adapter has its own VM + worker
|
||||
@ -171,11 +176,103 @@ under the `keys` field of the runtime file (encrypted at rest):
|
||||
delete the seed key.
|
||||
- The WebUI **Keys page** creates/deletes keys. Each key has a role (`admin`
|
||||
manages everything, `user` sees only its own key) and an optional **model
|
||||
scope** (model + source + token quota + reset period).
|
||||
scope** (model + source + token quota + reset period). Each model brick shows
|
||||
its own budget badge (e.g. `1.0K 77×·1h`, `∞` when uncapped); clicking it
|
||||
edits that model's token quota, request quota and reset period. Quotas apply
|
||||
per model, so one model running out never blocks the key's others. The
|
||||
"My key" view shows a user its model scopes and their budgets.
|
||||
- Clients authenticate with any authorized key's plaintext as
|
||||
`Authorization: Bearer <key>`.
|
||||
- Deleting a key removes it from the store immediately.
|
||||
|
||||
#### Per-key usage quota
|
||||
|
||||
**Quotas are per model.** Each key's `models[]` list gives every model its own
|
||||
token budget and request budget. One model running out affects only that model
|
||||
— the key's other models keep working.
|
||||
|
||||
```yaml
|
||||
keys:
|
||||
- key: sk-gw-<hex>
|
||||
role: user
|
||||
name: agent-alice
|
||||
models:
|
||||
- model: deepseek-v4-flash
|
||||
token_quota: 1000000 # this key's token budget for this model
|
||||
req_quota: 20000 # requests within the window
|
||||
period: nhour # "" | hour | week | month | nhour
|
||||
hours: 6 # n-hour only
|
||||
- model: AUTO
|
||||
token_quota: 5000000 # AUTO is a quota entry like any other
|
||||
period: hour
|
||||
- model: kimi-k3 # no quota listed = unlimited
|
||||
```
|
||||
|
||||
- **There is deliberately no key-wide total.** A key-wide cap would make
|
||||
"move A's budget to B" a global reallocation; per-model budgets keep each
|
||||
model independently controllable, so it stays visible which model is
|
||||
actually consuming the spend.
|
||||
- `period`: empty = never resets (lifetime total); `hour` / `week` / `month` =
|
||||
fixed windows; `nhour` + `hours` = a custom hour count. **A misspelled
|
||||
period is rejected at write time** rather than silently becoming a
|
||||
never-resetting quota.
|
||||
- Quotas are isolated per key *and* per model within a key: one key exhausting
|
||||
`m1` never draws on another key's budget, and never blocks the same key's
|
||||
`m2`.
|
||||
- Usage counts chat, streaming and image requests, and survives a restart by
|
||||
replaying the audit log (40 days retained, covering the longest `month`
|
||||
window).
|
||||
- An exhausted quota returns **429 + `Retry-After`**
|
||||
(`rate_limit_exceeded`) and the message names the model that ran out, so a
|
||||
client resumes when the window rolls over; a model the key may not use stays
|
||||
403. **Admin keys are never capped** (quotas on their scopes are not
|
||||
enforced either), so a cap can never lock the operator out.
|
||||
- Buckets are whole unix hours, so a window frees up at most an hour late
|
||||
(deliberately freeing late rather than overspending).
|
||||
- `PUT /api/keys/{key}` submits quotas by submitting `models` — the caps are
|
||||
part of the scope, so there is no second budget that can drift out of sync
|
||||
with the model list. An explicit `0` lifts that model's cap.
|
||||
|
||||
##### Quota rejection vs capacity rejection
|
||||
|
||||
A quota is checked at the entry point and occupies no upstream slot, so it is a
|
||||
cheap rejection. Running out of capacity instead waits out the bounded
|
||||
`busyWait` before answering 503. Measured (capacity 8, 0.6 s/request, 100
|
||||
concurrent clients):
|
||||
|
||||
| Case | Status | Latency | Reaches upstream |
|
||||
|---|---|---|---|
|
||||
| Quota spent (`req_quota` 50, 100 concurrent) | `429 rate_limit_exceeded` + `Retry-After` | **19–21 ms** | no |
|
||||
| Capacity exhausted (same topology, no quota) | `503 upstream_error` | ~2.6 s (`busyWait`) | overflow: no |
|
||||
|
||||
So a client can tell "wait for the window to roll over" apart from "wait for
|
||||
upstream capacity" — the first only needs patience, the second usually means
|
||||
back off concurrency or switch sources.
|
||||
|
||||
##### What the quota costs
|
||||
|
||||
Buckets are kept per (key, model, whole unix hour) for 40 days. Measured on an
|
||||
AMD 7840HS:
|
||||
|
||||
- A window query scans the window, not the whole retention: 49 ns for 1 h,
|
||||
55 ns for 24 h, 3.9 µs for 30 d.
|
||||
- Recording one request: 283 ns, 3 allocations (unchanged from before this
|
||||
feature; the allocations come from the record ring buffer).
|
||||
- Memory: the production shape (7 keys x 8 models x 2 sources at full 40-day
|
||||
retention) costs about **3.7 MB**. The `source::model` bucket used by a
|
||||
source-pinned quota is created **lazily** — it is maintained only once some
|
||||
quota actually pins that source, because writing it on every record would
|
||||
cost an extra 18 MB at 20 keys x 8 models x 3 sources.
|
||||
The trade-off: usage recorded before a pinned quota was configured cannot be
|
||||
split by source afterwards, so **a pinned quota may under-count its first
|
||||
window**.
|
||||
|
||||
```
|
||||
HTTP/1.1 429 Too Many Requests
|
||||
Retry-After: 2100
|
||||
{"error":{"type":"rate_limit_exceeded","message":"token quota exceeded for \"deepseek-v4-flash\" (5000000/5000000)"}}
|
||||
```
|
||||
|
||||
### Model routing
|
||||
|
||||
`/v1/chat/completions` `model` resolution order:
|
||||
|
||||
179
cmd/gui/main.js
179
cmd/gui/main.js
@ -17,7 +17,7 @@ const fs = require("fs");
|
||||
const http = require("http");
|
||||
const os = require("os");
|
||||
const crypto = require("crypto");
|
||||
const { spawn } = require("child_process");
|
||||
const { spawn, execFileSync } = require("child_process");
|
||||
|
||||
const LOG_FILE = path.join(app.getPath("userData"), "gui.log");
|
||||
function log(msg) {
|
||||
@ -97,24 +97,99 @@ function readAdminKey() {
|
||||
// gateway_keys entries get marked seed:true by the core (-> "replace the
|
||||
// initial key" warning in the web UI), while a plain keys entry authenticates
|
||||
// without any seed warning.
|
||||
const raw = readConfigRaw();
|
||||
if (!raw) return "";
|
||||
// keys:
|
||||
// - key: sk-gw-...
|
||||
const k = /^\s*keys:\s*\n(?:\s*-\s*key:\s*"?([^\s"#]+)"?[^\n]*\n)+/m.exec(
|
||||
raw,
|
||||
);
|
||||
let v = k && k[1] ? k[1] : "";
|
||||
if (!v) {
|
||||
// legacy: gateway_keys list (pre-migration configs)
|
||||
const g = /^\s*gateway_keys:\s*\n(?:\s*-\s*"?([^\s"#]+)"?\s*\n)+/m.exec(raw);
|
||||
if (g && g[1]) v = g[1];
|
||||
if (!v) {
|
||||
const g2 = /gateway_keys:[^[\n]*\[\s*"?([^\s"\]]+)"?/m.exec(raw);
|
||||
if (g2 && g2[1]) v = g2[1];
|
||||
}
|
||||
}
|
||||
// The core seals credentials at rest (enc:v1:...), so from the second start
|
||||
// onward the file holds ciphertext. Injecting that as a cookie is a
|
||||
// guaranteed 401, which surfaces to the user as "enter your key" for a key
|
||||
// they never set. Ask the core to unseal rather than reimplementing its
|
||||
// crypto here. -show-secrets only reads, prints and exits: it starts nothing
|
||||
// and writes nothing.
|
||||
if (v && v.indexOf(SEALED_PREFIX) === 0) return unsealViaCore();
|
||||
return v;
|
||||
}
|
||||
|
||||
const SEALED_PREFIX = "enc:v1:";
|
||||
|
||||
function readConfigRaw() {
|
||||
try {
|
||||
if (!fs.existsSync(CONFIG_FILE)) return "";
|
||||
const raw = fs.readFileSync(CONFIG_FILE, "utf-8");
|
||||
// keys:
|
||||
// - key: sk-gw-...
|
||||
const k = /^\s*keys:\s*\n(?:\s*-\s*key:\s*"?([^\s"#]+)"?[^\n]*\n)+/m.exec(
|
||||
raw,
|
||||
return fs.readFileSync(CONFIG_FILE, "utf-8");
|
||||
} catch (e) {
|
||||
return "";
|
||||
}
|
||||
}
|
||||
|
||||
// Returns the admin key in the clear, or "" if the core cannot be asked.
|
||||
//
|
||||
// cwd must be the directory that holds master.key. The core derives that path
|
||||
// from the config's runtime_file, so a call made from any other directory has
|
||||
// the core generate a *second* master key in the CWD and then fail to decrypt
|
||||
// ("master key changed?"). Electron's CWD is not the profile dir, so this is
|
||||
// not optional: without it the desktop build cannot read its own key back.
|
||||
//
|
||||
// readAdminKey runs on every outbound request, and this spawns a process, so
|
||||
// the result is cached against the config's mtime. Editing the config (or the
|
||||
// port rewrite, or the first write) changes the mtime and invalidates it, which
|
||||
// keeps the "read live so ordering never matters" property the auth rule
|
||||
// depends on.
|
||||
function unsealViaCore() {
|
||||
const st = safeStat(CONFIG_FILE);
|
||||
const stamp = st ? st.mtimeMs : 0;
|
||||
if (unsealCache && unsealCache.stamp === stamp) return unsealCache.key;
|
||||
let key = "";
|
||||
try {
|
||||
const out = execFileSync(
|
||||
CORE_EXE,
|
||||
["-config", CONFIG_FILE, "-show-secrets"],
|
||||
{
|
||||
encoding: "utf-8",
|
||||
timeout: 10000,
|
||||
windowsHide: true,
|
||||
cwd: PROFILE_DIR,
|
||||
},
|
||||
);
|
||||
if (k && k[1]) return k[1];
|
||||
// legacy: gateway_keys list (pre-migration configs)
|
||||
const g = /^\s*gateway_keys:\s*\n(?:\s*-\s*"?([^\s"#]+)"?\s*\n)+/m.exec(
|
||||
raw,
|
||||
);
|
||||
if (g && g[1]) return g[1];
|
||||
const g2 = /gateway_keys:[^[\n]*\[\s*"?([^\s"\]]+)"?/m.exec(raw);
|
||||
if (g2 && g2[1]) return g2[1];
|
||||
} catch (e) {}
|
||||
return "";
|
||||
// -show-secrets prints one line per credential:
|
||||
// key <name> role=<role> <value>
|
||||
const admin = /^key\s+\S+\s+role=admin\s+(\S+)\s*$/m.exec(out);
|
||||
if (admin && admin[1]) key = admin[1];
|
||||
// Tolerate a field-order or spacing change rather than locking the user out.
|
||||
if (!key) {
|
||||
const loose = /role=admin\s+(\S+)/.exec(out);
|
||||
if (loose && loose[1]) key = loose[1];
|
||||
}
|
||||
} catch (e) {
|
||||
console.error("unseal via core failed:", e.message);
|
||||
}
|
||||
// Only cache a success. A transient failure must not pin an empty key until
|
||||
// the config next changes, or a momentary spawn error locks the user out.
|
||||
if (key) unsealCache = { stamp, key };
|
||||
return key;
|
||||
}
|
||||
|
||||
let unsealCache = null;
|
||||
|
||||
function safeStat(p) {
|
||||
try {
|
||||
return fs.statSync(p);
|
||||
} catch (e) {
|
||||
return null;
|
||||
}
|
||||
}
|
||||
|
||||
const embeddedBaseUrl = () =>
|
||||
@ -596,6 +671,78 @@ ipcMain.handle(
|
||||
(e) => !!BrowserWindow.fromWebContents(e.sender)?.isMaximized(),
|
||||
);
|
||||
|
||||
// ---- plugin management over IPC -------------------------------------------
|
||||
//
|
||||
// The renderer cannot call the embedded core directly: it has no key and no
|
||||
// network identity, and the core binds a loopback port that only the main
|
||||
// process knows about. So every plugin action is proxied through the main
|
||||
// process, which already knows how to obtain the admin key (unsealViaCore).
|
||||
//
|
||||
// The proxy is deliberately a raw (method, path, body) pass-through rather than
|
||||
// a fixed set of commands. A fixed set would have to be extended for every new
|
||||
// plugin endpoint, and the one thing worse than "no button for this" is "a
|
||||
// button that silently does nothing" — with a pass-through the renderer can talk
|
||||
// to any /api/plugins route the core grows, and the path is validated here so
|
||||
// this channel cannot be used to reach arbitrary endpoints.
|
||||
function pluginProxy(req) {
|
||||
const { method, path, body } = req || {};
|
||||
const M = ["GET", "POST", "PUT", "DELETE"];
|
||||
if (!M.includes(method)) throw new Error("bad method: " + method);
|
||||
// The path must stay inside the plugin namespace. A prefix check alone would
|
||||
// still allow /api/plugins/../keys, so reject any traversal outright.
|
||||
if (typeof path !== "string" || !path.startsWith("/api/plugins")) {
|
||||
throw new Error("path must start with /api/plugins");
|
||||
}
|
||||
if (path.includes("..") || path.includes("\\")) {
|
||||
throw new Error("path traversal rejected");
|
||||
}
|
||||
const key = unsealViaCore();
|
||||
if (!key) throw new Error("no admin key available yet");
|
||||
return new Promise((resolve, reject) => {
|
||||
const u = new URL(embeddedBaseUrl() + path);
|
||||
const data = body == null ? null : JSON.stringify(body);
|
||||
const headers = { Authorization: "Bearer " + key };
|
||||
if (data) {
|
||||
headers["Content-Type"] = "application/json";
|
||||
headers["Content-Length"] = Buffer.byteLength(data);
|
||||
}
|
||||
const r = http.request(
|
||||
{
|
||||
hostname: u.hostname,
|
||||
port: u.port,
|
||||
path: u.pathname + u.search,
|
||||
method,
|
||||
headers,
|
||||
},
|
||||
(res) => {
|
||||
let raw = "";
|
||||
res.setEncoding("utf8");
|
||||
res.on("data", (c) => (raw += c));
|
||||
res.on("end", () => {
|
||||
let parsed = null;
|
||||
try {
|
||||
parsed = raw ? JSON.parse(raw) : null;
|
||||
} catch (e) {
|
||||
parsed = { raw };
|
||||
}
|
||||
if (res.statusCode >= 400) {
|
||||
const msg =
|
||||
(parsed && parsed.error && parsed.error.message) ||
|
||||
"HTTP " + res.statusCode;
|
||||
reject(new Error(msg));
|
||||
return;
|
||||
}
|
||||
resolve(parsed);
|
||||
});
|
||||
},
|
||||
);
|
||||
r.on("error", reject);
|
||||
if (data) r.write(data);
|
||||
r.end();
|
||||
});
|
||||
}
|
||||
|
||||
ipcMain.handle("plugins:proxy", (_e, req) => pluginProxy(req));
|
||||
ipcMain.handle("core:state", () => ({
|
||||
running: coreStarted() && coreReady,
|
||||
ready: coreReady,
|
||||
|
||||
8
cmd/gui/package-lock.json
generated
8
cmd/gui/package-lock.json
generated
@ -1,12 +1,12 @@
|
||||
{
|
||||
"name": "modelrouter-gui",
|
||||
"version": "1.4.2",
|
||||
"version": "1.9.3",
|
||||
"lockfileVersion": 3,
|
||||
"requires": true,
|
||||
"packages": {
|
||||
"": {
|
||||
"name": "modelrouter-gui",
|
||||
"version": "1.4.2",
|
||||
"version": "1.9.3",
|
||||
"license": "Proprietary",
|
||||
"devDependencies": {
|
||||
"electron": "^33.0.0",
|
||||
@ -582,7 +582,7 @@
|
||||
}
|
||||
},
|
||||
"node_modules/@peculiar/webcrypto": {
|
||||
"version": "1.7.1",
|
||||
"version": "1.7.5",
|
||||
"resolved": "https://registry.npmjs.org/@peculiar/webcrypto/-/webcrypto-1.7.1.tgz",
|
||||
"integrity": "sha512-ODOov0sGMJMf3jPonOkgGqPknTsu+DdQ7kD++gz8aI+aFMOMHFbWAA2taqXXVTdP+OTOQR/znGvSpmkeI0WTYQ==",
|
||||
"dev": true,
|
||||
@ -3335,7 +3335,7 @@
|
||||
}
|
||||
},
|
||||
"node_modules/resedit": {
|
||||
"version": "1.7.2",
|
||||
"version": "1.7.5",
|
||||
"resolved": "https://registry.npmjs.org/resedit/-/resedit-1.7.2.tgz",
|
||||
"integrity": "sha512-vHjcY2MlAITJhC0eRD/Vv8Vlgmu9Sd3LX9zZvtGzU5ZImdTN3+d6e/4mnTyV8vEbyf1sgNIrWxhWlrys52OkEA==",
|
||||
"dev": true,
|
||||
|
||||
@ -1,6 +1,6 @@
|
||||
{
|
||||
"name": "modelrouter-gui",
|
||||
"version": "1.4.2",
|
||||
"version": "1.9.7",
|
||||
"description": "ModelRouter Desktop — embedded ModelRouter core with tray",
|
||||
"author": "ModelRouter",
|
||||
"main": "main.js",
|
||||
|
||||
@ -16,6 +16,13 @@ contextBridge.exposeInMainWorld("modelrouter", {
|
||||
key: () => ipcRenderer.invoke("core:key"),
|
||||
onState: (cb) => ipcRenderer.on("core:state", (_e, d) => cb(d)),
|
||||
},
|
||||
plugins: {
|
||||
// Raw pass-through to the embedded core's /api/plugins surface. The main
|
||||
// process validates the path and attaches the admin key; the renderer never
|
||||
// sees either.
|
||||
request: (method, path, body) =>
|
||||
ipcRenderer.invoke("plugins:proxy", { method, path, body }),
|
||||
},
|
||||
settings: {
|
||||
get: () => ipcRenderer.invoke("settings:get"),
|
||||
set: (patch) => ipcRenderer.invoke("settings:set", patch),
|
||||
|
||||
@ -155,6 +155,7 @@ async function openSettings() {
|
||||
$("#set-tray").checked = !!state.settings.minimizeToTray;
|
||||
$("#settings-overlay").style.display = "flex";
|
||||
renderRail();
|
||||
loadPlugins();
|
||||
}
|
||||
function closeSettings() {
|
||||
$("#settings-overlay").style.display = "none";
|
||||
@ -182,6 +183,296 @@ async function saveSettings() {
|
||||
}
|
||||
}
|
||||
|
||||
// esc / escAttr escape text for innerHTML. The WebUI has its own copies; the
|
||||
// shell needs its own because renderer/app.js is a separate document that
|
||||
// never loads index.html's script.
|
||||
function esc(s) {
|
||||
return String(s == null ? "" : s).replace(
|
||||
/[&<>"']/g,
|
||||
(c) => ({ "&": "&", "<": "<", ">": ">", '"': """, "'": "'" })[c],
|
||||
);
|
||||
}
|
||||
function escAttr(s) {
|
||||
return esc(s).replace(/`/g, "`");
|
||||
}
|
||||
|
||||
// ===== plugin management =====
|
||||
//
|
||||
// The desktop shell manages plugins through the embedded core's /api/plugins
|
||||
// surface, proxied over IPC (see plugins:proxy in the main process). The
|
||||
// renderer never holds the admin key.
|
||||
//
|
||||
// Scope note: the desktop build has no plugin_dir configured by default, so this
|
||||
// panel normally reports "plugins disabled" with the one-line fix. That is a
|
||||
// deliberate state, not an error — the packaged profile is a per-user directory
|
||||
// and seeding a plugin tree into someone's home without asking would be rude.
|
||||
|
||||
async function loadPlugins() {
|
||||
const list = $("#pl-list");
|
||||
const hint = $("#set-plugins-hint");
|
||||
if (!list || !hint) return;
|
||||
let j;
|
||||
try {
|
||||
j = await window.modelrouter.plugins.request("GET", "/api/plugins");
|
||||
} catch (e) {
|
||||
hint.textContent = "内核未就绪:" + e.message;
|
||||
list.innerHTML = "";
|
||||
return;
|
||||
}
|
||||
if (!j.plugin_dir) {
|
||||
hint.innerHTML =
|
||||
'未配置 <code>plugin_dir</code>,插件功能未启用。在 config.yaml 加一行后重启内核即可。';
|
||||
list.innerHTML = "";
|
||||
return;
|
||||
}
|
||||
const rows = j.on_disk || [];
|
||||
window.__plRows = rows; // read by editPlugin/deletePlugin for the builtin flag
|
||||
const active = rows.filter((p) => p.loaded && !p.disabled).length;
|
||||
const broken = rows.filter((p) => !p.loaded).length;
|
||||
hint.textContent =
|
||||
`${rows.length} 个插件 · ${active} 个启用中` +
|
||||
(broken ? ` · ${broken} 个加载失败` : "");
|
||||
list.innerHTML = rows.length
|
||||
? rows
|
||||
.map((p) => {
|
||||
const cls = !p.loaded ? "pl-broken" : p.disabled ? "pl-off" : "pl-on";
|
||||
const label = !p.loaded
|
||||
? "加载失败"
|
||||
: p.disabled
|
||||
? "已禁用"
|
||||
: "启用中";
|
||||
const btn = p.loaded
|
||||
? `<button class="ghost" data-act="toggle" data-name="${escAttr(
|
||||
p.name,
|
||||
)}" data-en="${p.disabled ? "1" : "0"}">${
|
||||
p.disabled ? "启用" : "禁用"
|
||||
}</button>`
|
||||
: "";
|
||||
const builtin = p.builtin
|
||||
? '<span class="pl-builtin">内置</span>'
|
||||
: "";
|
||||
// Edit / view source / delete are the operations an operator
|
||||
// actually needs after enable/disable: a plugin that cannot be
|
||||
// read back cannot be fixed, and one that cannot be removed cannot
|
||||
// be rolled back.
|
||||
const acts = [
|
||||
btn,
|
||||
`<button class="ghost" data-act="code" data-name="${escAttr(p.name)}">${
|
||||
p.builtin ? "查看源码" : "编辑"
|
||||
}</button>`,
|
||||
builtin
|
||||
? ""
|
||||
: `<button class="ghost danger" data-act="del" data-name="${escAttr(
|
||||
p.name,
|
||||
)}">删除</button>`,
|
||||
]
|
||||
.filter(Boolean)
|
||||
.join("");
|
||||
const pages = Array.isArray(p.pages) ? p.pages : [];
|
||||
return `<div class="pl-item ${cls}">
|
||||
<div class="pl-head"><b>${esc(p.name)}</b>${builtin}<span class="pl-state">${label}</span></div>
|
||||
${p.description ? `<div class="pl-desc">${esc(p.description)}</div>` : ""}
|
||||
${pages.length ? `<div class="pl-desc">页面:${esc(pages.join(", "))}</div>` : ""}
|
||||
${p.error ? `<div class="pl-err">${esc(String(p.error).slice(0, 160))}</div>` : ""}
|
||||
<div class="pl-acts">${acts}</div>
|
||||
</div>`;
|
||||
})
|
||||
.join("")
|
||||
: '<div class="pl-empty">插件目录为空</div>';
|
||||
list.querySelectorAll('button[data-act="toggle"]').forEach((b) => {
|
||||
b.onclick = () => togglePlugin(b.dataset.name, b.dataset.en === "1");
|
||||
});
|
||||
list.querySelectorAll('button[data-act="code"]').forEach((b) => {
|
||||
b.onclick = () => editPlugin(b.dataset.name);
|
||||
});
|
||||
list.querySelectorAll('button[data-act="del"]').forEach((b) => {
|
||||
b.onclick = () => deletePlugin(b.dataset.name);
|
||||
});
|
||||
}
|
||||
|
||||
// editPlugin shows the plugin source in a modal. Built-in plugins are
|
||||
// read-only: their file ships inside the binary, so a save would either be
|
||||
// ignored or silently do nothing — showing an editable text area that
|
||||
// discards on save is worse than saying so up front.
|
||||
function editPlugin(name) {
|
||||
// The plugin list is the only place `builtin` is reported; keep the last
|
||||
// fetch so the modal does not need a second round trip (and cannot guess).
|
||||
const row = (window.__plRows || []).find((p) => p.name === name);
|
||||
const isBuiltin = !!(row && row.builtin);
|
||||
const wrap = document.createElement("div");
|
||||
wrap.className = "overlay";
|
||||
wrap.innerHTML = `<div class="overlay-card pl-wide">
|
||||
<div class="overlay-head"><h2>${esc(name)}${isBuiltin ? "(内置,只读)" : ""}</h2></div>
|
||||
${isBuiltin
|
||||
? `<p class="hint">内置插件随内核发布,源码在二进制内,不可编辑。</p>`
|
||||
: `<textarea id="pl-code" spellcheck="false"></textarea>
|
||||
<p class="hint" id="pl-code-err"></p>`}
|
||||
<div class="row actions">
|
||||
${isBuiltin ? "" : `<button id="pl-code-save">保存并热重载</button>`}
|
||||
<button class="ghost" id="pl-code-close">关闭</button>
|
||||
</div>
|
||||
</div>`;
|
||||
document.body.appendChild(wrap);
|
||||
wrap.addEventListener("click", (e) => {
|
||||
if (e.target === wrap) wrap.remove();
|
||||
});
|
||||
$("#pl-code-close").onclick = () => wrap.remove();
|
||||
const save = $("#pl-code-save");
|
||||
if (save) {
|
||||
save.onclick = async () => {
|
||||
save.disabled = true;
|
||||
try {
|
||||
// Saving writes through POST /api/plugins — the same install path the
|
||||
// installer uses, which validates the name, compiles the source and
|
||||
// hot-reloads. PUT /{name} only toggles enable/disable, so pointing the
|
||||
// save there would silently disable nothing and store nothing.
|
||||
await window.modelrouter.plugins.request("POST", "/api/plugins", {
|
||||
name,
|
||||
code: $("#pl-code").value,
|
||||
});
|
||||
toast(`${name} 已保存`);
|
||||
wrap.remove();
|
||||
await loadPlugins();
|
||||
} catch (e) {
|
||||
$("#pl-code-err").textContent = e.message;
|
||||
save.disabled = false;
|
||||
}
|
||||
};
|
||||
}
|
||||
if (isBuiltin) return;
|
||||
(async () => {
|
||||
const ta = $("#pl-code");
|
||||
try {
|
||||
const j = await window.modelrouter.plugins.request(
|
||||
"GET",
|
||||
`/api/plugins/${encodeURIComponent(name)}`,
|
||||
);
|
||||
ta.value = j.code || "";
|
||||
const err = $("#pl-code-err");
|
||||
if (err && j.parse_error) err.textContent = String(j.parse_error);
|
||||
} catch (e) {
|
||||
$("#pl-code-err").textContent = e.message;
|
||||
}
|
||||
})();
|
||||
}
|
||||
|
||||
function deletePlugin(name) {
|
||||
const wrap = document.createElement("div");
|
||||
wrap.className = "overlay";
|
||||
wrap.innerHTML = `<div class="overlay-card">
|
||||
<div class="overlay-head"><h2>删除插件</h2></div>
|
||||
<p>确定删除 <b>${esc(name)}</b>?插件目录中的文件会被移除,核心将不再加载它。</p>
|
||||
<div class="row actions">
|
||||
<button class="danger" id="pl-del-ok">删除</button>
|
||||
<button class="ghost" id="pl-del-no">取消</button>
|
||||
</div>
|
||||
</div>`;
|
||||
document.body.appendChild(wrap);
|
||||
wrap.addEventListener("click", (e) => {
|
||||
if (e.target === wrap) wrap.remove();
|
||||
});
|
||||
$("#pl-del-no").onclick = () => wrap.remove();
|
||||
$("#pl-del-ok").onclick = async () => {
|
||||
try {
|
||||
await window.modelrouter.plugins.request(
|
||||
"DELETE",
|
||||
`/api/plugins/${encodeURIComponent(name)}`,
|
||||
);
|
||||
toast(`已删除 ${name}`);
|
||||
wrap.remove();
|
||||
await loadPlugins();
|
||||
} catch (e) {
|
||||
toast(e.message, true);
|
||||
}
|
||||
};
|
||||
}
|
||||
|
||||
// installPlugin creates a new plugin from source. The name becomes the file
|
||||
// name, so it is validated client-side to save a round trip on the obvious
|
||||
// mistakes (empty, path separators, non-portable characters).
|
||||
function installPlugin() {
|
||||
const wrap = document.createElement("div");
|
||||
wrap.className = "overlay";
|
||||
wrap.innerHTML = `<div class="overlay-card pl-wide">
|
||||
<div class="overlay-head"><h2>安装插件</h2></div>
|
||||
<div class="row"><label>名称</label><input id="pl-new-name" placeholder="my_plugin" /></div>
|
||||
<p class="hint" id="pl-new-err"></p>
|
||||
<textarea id="pl-new-code" spellcheck="false" placeholder="-- 单文件 Lua 插件 plugin = {} ..."></textarea>
|
||||
<div class="row actions">
|
||||
<button id="pl-new-ok">安装</button>
|
||||
<button class="ghost" id="pl-new-no">取消</button>
|
||||
</div>
|
||||
</div>`;
|
||||
document.body.appendChild(wrap);
|
||||
wrap.addEventListener("click", (e) => {
|
||||
if (e.target === wrap) wrap.remove();
|
||||
});
|
||||
$("#pl-new-no").onclick = () => wrap.remove();
|
||||
$("#pl-new-ok").onclick = async () => {
|
||||
const name = $("#pl-new-name").value.trim();
|
||||
const code = $("#pl-new-code").value;
|
||||
const err = $("#pl-new-err");
|
||||
if (!/^[A-Za-z0-9_-]+$/.test(name)) {
|
||||
err.textContent = "名称只能包含字母、数字、下划线和连字符";
|
||||
return;
|
||||
}
|
||||
if (!code.trim()) {
|
||||
err.textContent = "源码不能为空";
|
||||
return;
|
||||
}
|
||||
$("#pl-new-ok").disabled = true;
|
||||
try {
|
||||
await window.modelrouter.plugins.request("POST", "/api/plugins", {
|
||||
// The backend appends .lua itself (installPlugin writes name+".lua"),
|
||||
// so passing it here would create name.lua.lua.
|
||||
name,
|
||||
code,
|
||||
});
|
||||
toast(`已安装 ${name}`);
|
||||
wrap.remove();
|
||||
await loadPlugins();
|
||||
} catch (e) {
|
||||
err.textContent = e.message;
|
||||
$("#pl-new-ok").disabled = false;
|
||||
}
|
||||
};
|
||||
}
|
||||
|
||||
async function togglePlugin(name, disabled) {
|
||||
try {
|
||||
await window.modelrouter.plugins.request("PUT", `/api/plugins/${encodeURIComponent(name)}`, {
|
||||
enabled: disabled,
|
||||
});
|
||||
toast(disabled ? `已禁用 ${name}` : `已启用 ${name}`);
|
||||
await loadPlugins();
|
||||
} catch (e) {
|
||||
toast(e.message, true);
|
||||
}
|
||||
}
|
||||
|
||||
async function enableAllPlugins() {
|
||||
let j;
|
||||
try {
|
||||
j = await window.modelrouter.plugins.request("GET", "/api/plugins");
|
||||
} catch (e) {
|
||||
return toast(e.message, true);
|
||||
}
|
||||
const off = (j.on_disk || []).filter((p) => p.loaded && p.disabled);
|
||||
for (const p of off) {
|
||||
try {
|
||||
await window.modelrouter.plugins.request(
|
||||
"PUT",
|
||||
`/api/plugins/${encodeURIComponent(p.name)}`,
|
||||
{ enabled: true },
|
||||
);
|
||||
} catch (e) {
|
||||
toast(`${p.name}: ${e.message}`, true);
|
||||
}
|
||||
}
|
||||
toast(off.length ? `已启用 ${off.length} 个插件` : "没有处于禁用状态的插件");
|
||||
await loadPlugins();
|
||||
}
|
||||
|
||||
// ===== theme =====
|
||||
function applyTheme() {
|
||||
document.documentElement.dataset.theme = state.theme;
|
||||
@ -197,6 +488,12 @@ function init() {
|
||||
$("#tb-close").onclick = () => window.modelrouter.win.close();
|
||||
$("#tb-settings").onclick = openSettings;
|
||||
$("#rail-settings").onclick = openSettings;
|
||||
const plReload = document.getElementById("pl-reload");
|
||||
if (plReload) plReload.onclick = loadPlugins;
|
||||
const plAll = document.getElementById("pl-toggle-all");
|
||||
if (plAll) plAll.onclick = enableAllPlugins;
|
||||
const plInstall = document.getElementById("pl-install");
|
||||
if (plInstall) plInstall.onclick = installPlugin;
|
||||
$("#rail-autostart").onclick = toggleAutoStart;
|
||||
$("#rail-silent").onclick = toggleSilent;
|
||||
$("#rail-theme").onclick = () => {
|
||||
|
||||
@ -207,6 +207,16 @@
|
||||
><input type="checkbox" id="set-tray" /> 关闭时最小化到托盘</label
|
||||
><span class="hint">点关闭按钮隐藏到系统托盘</span>
|
||||
</div>
|
||||
<div class="row">
|
||||
<label>插件</label>
|
||||
<span class="hint" id="set-plugins-hint">加载中…</span>
|
||||
</div>
|
||||
<div id="pl-list" class="pl-list"></div>
|
||||
<div class="row actions">
|
||||
<button class="ghost" id="pl-toggle-all">全部启用</button>
|
||||
<button class="ghost" id="pl-install">安装插件</button>
|
||||
<button class="ghost" id="pl-reload">刷新</button>
|
||||
</div>
|
||||
<div class="row actions">
|
||||
<button class="ghost" id="set-dir">打开数据目录</button>
|
||||
<button class="ghost" id="set-log">查看日志</button>
|
||||
|
||||
@ -596,3 +596,129 @@ html[data-theme="dark"] .overlay {
|
||||
#toast.err {
|
||||
border-color: var(--danger);
|
||||
}
|
||||
|
||||
/* ===== plugin management panel =========================================
|
||||
* The existing .ghost/.primary rules are scoped to `.form .actions`, so a
|
||||
* button outside that selector gets browser defaults. The plugin rows live in
|
||||
* their own list, hence their own rules — reusing a scoped class here would have
|
||||
* produced unstyled buttons that still worked, which is the kind of thing that
|
||||
* looks fine until someone themes the shell.
|
||||
*/
|
||||
.pl-list {
|
||||
display: flex;
|
||||
flex-direction: column;
|
||||
gap: 8px;
|
||||
margin: 8px 0 4px;
|
||||
}
|
||||
.pl-item {
|
||||
border: 1px solid var(--line);
|
||||
border-radius: 9px;
|
||||
padding: 10px 12px;
|
||||
}
|
||||
.pl-item.pl-broken {
|
||||
border-color: var(--danger, #d1435b);
|
||||
}
|
||||
.pl-head {
|
||||
display: flex;
|
||||
align-items: center;
|
||||
gap: 8px;
|
||||
font-size: 13px;
|
||||
}
|
||||
.pl-builtin {
|
||||
font-size: 10px;
|
||||
padding: 1px 6px;
|
||||
border-radius: 999px;
|
||||
background: var(--primary-50);
|
||||
color: var(--primary-h);
|
||||
}
|
||||
.pl-state {
|
||||
margin-left: auto;
|
||||
font-size: 11px;
|
||||
color: var(--muted);
|
||||
}
|
||||
.pl-item.pl-broken .pl-state {
|
||||
color: var(--danger, #d1435b);
|
||||
}
|
||||
.pl-desc {
|
||||
font-size: 12px;
|
||||
color: var(--muted);
|
||||
margin-top: 3px;
|
||||
}
|
||||
.pl-err {
|
||||
font-size: 11px;
|
||||
color: var(--danger, #d1435b);
|
||||
margin-top: 4px;
|
||||
word-break: break-word;
|
||||
}
|
||||
.pl-acts {
|
||||
margin-top: 8px;
|
||||
display: flex;
|
||||
gap: 8px;
|
||||
}
|
||||
.pl-acts button {
|
||||
padding: 5px 12px;
|
||||
font-size: 12px;
|
||||
border-radius: 7px;
|
||||
border: 1px solid var(--line);
|
||||
background: var(--bg-s2, #fff);
|
||||
color: var(--fg, inherit);
|
||||
cursor: pointer;
|
||||
transition: all 0.15s;
|
||||
}
|
||||
.pl-acts button:hover {
|
||||
border-color: var(--primary);
|
||||
color: var(--primary-h);
|
||||
}
|
||||
.pl-empty {
|
||||
font-size: 12px;
|
||||
color: var(--muted);
|
||||
padding: 10px 0;
|
||||
}
|
||||
|
||||
/* ---- plugin editor / installer --------------------------------------------
|
||||
The plugin source editor and the install dialog reuse the settings overlay
|
||||
(overlay / overlay-card) rather than a second popup system. They need three
|
||||
things the settings card does not: a wider column for code, a monospace
|
||||
textarea, and a destructive button for delete. */
|
||||
.overlay-card.pl-wide {
|
||||
width: 760px;
|
||||
}
|
||||
.pl-wide textarea {
|
||||
width: 100%;
|
||||
min-height: 380px;
|
||||
resize: vertical;
|
||||
padding: 10px 12px;
|
||||
border-radius: 10px;
|
||||
border: 1px solid var(--line-strong);
|
||||
background: var(--bg-sunken, rgba(127, 127, 127, 0.06));
|
||||
color: inherit;
|
||||
font-family: ui-monospace, SFMono-Regular, Menlo, Consolas, monospace;
|
||||
font-size: 12.5px;
|
||||
line-height: 1.55;
|
||||
tab-size: 2;
|
||||
}
|
||||
.pl-wide textarea:focus {
|
||||
outline: none;
|
||||
border-color: var(--primary);
|
||||
}
|
||||
.pl-wide .row input {
|
||||
flex: 1;
|
||||
padding: 7px 10px;
|
||||
border-radius: 9px;
|
||||
border: 1px solid var(--line-strong);
|
||||
background: var(--bg-sunken, rgba(127, 127, 127, 0.06));
|
||||
color: inherit;
|
||||
}
|
||||
.pl-editor .hint,
|
||||
.pl-wide .hint {
|
||||
min-height: 16px;
|
||||
color: var(--muted, #888);
|
||||
}
|
||||
button.danger {
|
||||
color: #e5484d;
|
||||
border-color: rgba(229, 72, 77, 0.4);
|
||||
}
|
||||
button.danger:hover {
|
||||
border-color: #e5484d;
|
||||
color: #e5484d;
|
||||
}
|
||||
|
||||
@ -24,6 +24,7 @@ import (
|
||||
func main() {
|
||||
cfgPath := flag.String("config", "config.yaml", "path to gateway config file")
|
||||
checkOnly := flag.Bool("check", false, "validate the config file and exit (0 = valid, 1 = invalid); nothing is started and no file is written")
|
||||
showSecrets := flag.Bool("show-secrets", false, "print the config's credentials in the clear and exit; nothing is started and no file is written")
|
||||
flag.Parse()
|
||||
|
||||
// -check is the deploy-time preflight: parse and validate the config
|
||||
@ -41,6 +42,20 @@ func main() {
|
||||
return
|
||||
}
|
||||
|
||||
// -show-secrets is the operator escape hatch for a config whose
|
||||
// credentials are sealed at rest: it prints them to stdout and exits
|
||||
// without starting anything or writing anything. Nothing else in the
|
||||
// program prints a credential, and this path never logs to a file.
|
||||
if *showSecrets {
|
||||
if _, err := os.Stat(*cfgPath); err != nil {
|
||||
log.Fatalf("[llmsproxy] show-secrets: %v", err)
|
||||
}
|
||||
if err := config.PrintSecrets(*cfgPath); err != nil {
|
||||
log.Fatalf("[llmsproxy] show-secrets: %v", err)
|
||||
}
|
||||
return
|
||||
}
|
||||
|
||||
created, err := config.EnsureDefault(*cfgPath)
|
||||
if err != nil {
|
||||
log.Fatalf("[llmsproxy] config: %v", err)
|
||||
@ -58,19 +73,26 @@ func main() {
|
||||
}
|
||||
}
|
||||
|
||||
gw, err := gateway.New(c, c.GatewayKeys())
|
||||
gw, err := gateway.New(c)
|
||||
if err != nil {
|
||||
log.Fatalf("[llmsproxy] gateway: %v", err)
|
||||
}
|
||||
|
||||
// Ops hygiene: surface the two most common footguns instead of silently
|
||||
// running with them.
|
||||
if keys := c.GatewayKeys(); len(keys) == 0 {
|
||||
log.Printf("[llmsproxy] WARNING: gateway_keys is EMPTY — without a key every request is rejected")
|
||||
//
|
||||
// Read the AUTHORITATIVE source. Auth uses core.ListKeys() (cfg.Keys), not
|
||||
// the legacy cfg.GatewayKeys list: seedKeys copies gateway_keys into keys[]
|
||||
// on first start and the YAML list is ignored for auth afterwards. Checking
|
||||
// GatewayKeys() here made the warning lie — once the legacy list was empty
|
||||
// (e.g. after rotating the starter key away) it reported "every request is
|
||||
// rejected" while seven working keys, one of them admin, were in service.
|
||||
if keys := c.ListKeys(); len(keys) == 0 {
|
||||
log.Printf("[llmsproxy] WARNING: no gateway key configured — every request will be rejected. Add one in the WebUI or set gateway_keys")
|
||||
} else {
|
||||
for _, k := range keys {
|
||||
if k == "sk-gw-local-0001" || k == "sk-local-0001" {
|
||||
log.Printf("[llmsproxy] WARNING: gateway key %q looks like the starter/example key — rotate it before exposing the gateway", k)
|
||||
if k.Key == "sk-gw-local-0001" || k.Key == "sk-local-0001" {
|
||||
log.Printf("[llmsproxy] WARNING: gateway key %q looks like the starter/example key — rotate it before exposing the gateway", k.Key)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
116
deploy.sh
116
deploy.sh
@ -244,6 +244,117 @@ verify_config() {
|
||||
return 1
|
||||
}
|
||||
|
||||
# ---------- 密钥可解封校验 ----------
|
||||
# config.yaml 里的凭据以 enc:v1: 密文落盘,解密钥匙是 master.key(或
|
||||
# LLMS_PROXY_MASTER_KEY)。一旦钥匙丢失/错位,配置在**启动时**才会报错,
|
||||
# 那时服务已经处于「二进制已替换 + 配置已替换」的半死状态,只能靠回滚。
|
||||
# 所以在替换任何东西之前先验:能不能用现有的钥匙解开现有的密文。
|
||||
#
|
||||
# 两种情形要分开:
|
||||
# 1. 配置里**有**密文 ⇒ 钥匙必须能解开(否则拒绝部署)
|
||||
# 2. 配置里**没有**密文 ⇒ 全新安装/首次加密,无需钥匙(但会顺手记一条,
|
||||
# 因为它意味着这次部署会把明文写成密文,而钥匙必须同时就位)
|
||||
verify_master_key() {
|
||||
log "密钥可解封校验"
|
||||
[[ -f "$TARGET_CONFIG" ]] || return 0 # 配置都不在,交给 verify_config 报错
|
||||
|
||||
local sealed
|
||||
sealed=$(grep -c 'enc:v1:' "$TARGET_CONFIG" 2>/dev/null || echo 0)
|
||||
if [[ "$sealed" -eq 0 ]]; then
|
||||
warn "配置内无密文(尚未加密或全为明文)——首次加密需要 master.key 可写"
|
||||
return 0
|
||||
fi
|
||||
|
||||
# 钥匙来源与 config 包一致:env 优先,否则 runtime_file 同目录的 master.key
|
||||
local runtime_file key_file
|
||||
runtime_file=$(awk '/^runtime_file:/{print $2; exit}' "$TARGET_CONFIG")
|
||||
[[ -n "$runtime_file" ]] || runtime_file="$(dirname "$TARGET_CONFIG")/runtime.json"
|
||||
key_file="$(dirname "$runtime_file")/master.key"
|
||||
|
||||
if [[ -n "${LLMS_PROXY_MASTER_KEY:-}" ]]; then
|
||||
log " 钥匙来源: LLMS_PROXY_MASTER_KEY(环境变量)"
|
||||
# 走真实解密路径验证,而不是只检查格式
|
||||
if out=$("$TARGET_BIN" -show-secrets -config "$TARGET_CONFIG" 2>&1); then
|
||||
log "✓ 凭据可解封($(echo "$out" | grep -cE '^(source|key) ') 项)"
|
||||
return 0
|
||||
fi
|
||||
err "凭据无法用 LLMS_PROXY_MASTER_KEY 解封,未替换任何文件:"
|
||||
printf '%s\n' "$out" >&2
|
||||
return 1
|
||||
fi
|
||||
|
||||
if [[ ! -f "$key_file" ]]; then
|
||||
err "配置含 $sealed 处密文,但钥匙文件不存在: $key_file"
|
||||
err " 没有它,部署后服务将无法启动。恢复 master.key 或设置 LLMS_PROXY_MASTER_KEY 后重试。"
|
||||
return 1
|
||||
fi
|
||||
log " 钥匙文件: $key_file ($(stat -c '%A' "$key_file" 2>/dev/null || echo '?'),$(wc -c < "$key_file") 字节)"
|
||||
|
||||
if out=$("$TARGET_BIN" -show-secrets -config "$TARGET_CONFIG" 2>&1); then
|
||||
log "✓ 凭据可解封($(echo "$out" | grep -cE '^(source|key) ') 项)"
|
||||
return 0
|
||||
fi
|
||||
err "凭据无法解密(钥匙与配置不匹配),未替换任何文件:"
|
||||
printf '%s\n' "$out" | head -5 >&2
|
||||
err " 密钥来源: $key_file"
|
||||
err " 常见原因:master.key 被重新生成、恢复了不匹配的备份,或换了机器。"
|
||||
return 1
|
||||
}
|
||||
|
||||
# ---------- 适配器备份保留策略 ----------
|
||||
# 每次部署都新建一个 adapters.bak.<时间戳> 目录,而回滚只用到最近一次
|
||||
# ($BACKUP_BIN / $BACKUP_CONFIG 都是单文件覆盖)。多出来的目录从没人清理,
|
||||
# 实测线上已积累 71 个,/etc/llmsproxy 因此涨到 122M。
|
||||
#
|
||||
# 保留最近 KEEP_ADAPTER_BACKUPS 份足够回滚,同时给目录数设上限——否则
|
||||
# 一次误配置(比如 adapter_dir 指错)就可能在几秒内造出成百上千个目录。
|
||||
# 只删名字严格匹配 adapters.bak.<14位时间戳> 的目录,避免误伤人工放的目录。
|
||||
KEEP_ADAPTER_BACKUPS=5
|
||||
|
||||
prune_adapter_backups() {
|
||||
local base="/etc/llmsproxy"
|
||||
[[ -d "$base" ]] || return 0
|
||||
|
||||
# 先按数量上限硬裁:即使时间戳排序失效也不会无上限增长。
|
||||
local all
|
||||
mapfile -t all < <(find "$base" -maxdepth 1 -type d -name 'adapters.bak.*' -printf '%f\n' | sort)
|
||||
local cap=$((KEEP_ADAPTER_BACKUPS * 4))
|
||||
if (( ${#all[@]} > cap )); then
|
||||
warn "适配器备份目录有 ${#all[@]} 个(异常),裁到 $cap"
|
||||
local i=0
|
||||
for d in "${all[@]}"; do
|
||||
i=$((i + 1))
|
||||
# 从最旧的开始删(sort 后升序)。名字不规范的跳过不删。
|
||||
if (( i <= ${#all[@]} - cap )); then
|
||||
if [[ "$d" =~ ^adapters\.bak\.[0-9]{14}$ ]]; then
|
||||
rm -rf "${base:?}/$d"
|
||||
else
|
||||
warn "跳过名字不规范的备份目录(不删): $d"
|
||||
fi
|
||||
fi
|
||||
done
|
||||
mapfile -t all < <(find "$base" -maxdepth 1 -type d -name 'adapters.bak.*' -printf '%f\n' | sort)
|
||||
fi
|
||||
|
||||
# 再按时间保留最近 KEEP_ADAPTER_BACKUPS 份(sort 后最新在末尾)。
|
||||
local total=${#all[@]} i=0
|
||||
for d in "${all[@]}"; do
|
||||
i=$((i + 1))
|
||||
# i <= total-KEEP 的是较旧的,要删。
|
||||
if (( i <= total - KEEP_ADAPTER_BACKUPS )); then
|
||||
if [[ "$d" =~ ^adapters\.bak\.[0-9]{14}$ ]]; then
|
||||
rm -rf "${base:?}/$d"
|
||||
else
|
||||
warn "跳过名字不规范的备份目录(不删): $d"
|
||||
fi
|
||||
fi
|
||||
done
|
||||
|
||||
local left
|
||||
left=$(find "$base" -maxdepth 1 -type d -name 'adapters.bak.*' | wc -l)
|
||||
log " 适配器备份保留最近 $KEEP_ADAPTER_BACKUPS 份(当前剩 $left 个)"
|
||||
}
|
||||
|
||||
# ---------- 同步适配器 ----------
|
||||
sync_adapters() {
|
||||
log "同步适配器"
|
||||
@ -261,6 +372,7 @@ sync_adapters() {
|
||||
cp -rf "$TARGET_ADAPTERS/"*.lua "$BACKUP_DIR/" 2>/dev/null || true
|
||||
log " 旧适配器已备份到 $BACKUP_DIR"
|
||||
fi
|
||||
prune_adapter_backups
|
||||
|
||||
cp -f "$SRC_ADAPTERS/"*.lua "$TARGET_ADAPTERS/"
|
||||
chmod 0644 "$TARGET_ADAPTERS/"*.lua
|
||||
@ -393,6 +505,10 @@ main() {
|
||||
parse_args "$@"
|
||||
log "═══ llmsproxy 原子部署开始 ═══"
|
||||
precheck
|
||||
# 密钥校验放在 build/替换之前:它的失败后果最重(配置里的密文永久解不开),
|
||||
# 必须在动任何文件之前就挡住。此阶段用的是**当前线上二进制**,它已经
|
||||
# 支持 -show-secrets(首次引入该子命令的版本除外,此时跳过并提示)。
|
||||
verify_master_key
|
||||
build
|
||||
sync_adapters
|
||||
backup_old_state
|
||||
|
||||
@ -15,7 +15,7 @@ deployable.
|
||||
|---|---|---|---|
|
||||
| Main | `main` | permanent | Only long-lived branch. Always deployable. Accumulates the next version. |
|
||||
| Feature | `feature/<desc>` | short (dev → merge → delete) | New features / ordinary fixes. Born from `main`, merged back into `main`. |
|
||||
| Release | `release/vX.Y.Z` | one version cycle | Cut from `main`, tagged for release. Version-specific hotfixes land here. |
|
||||
| Release | `release/vX.Y.x` | one version cycle | Cut from `main`, tagged for release. Version-specific hotfixes land here. |
|
||||
|
||||
## Change flow (important)
|
||||
|
||||
@ -26,15 +26,21 @@ deployable.
|
||||
│ │
|
||||
│ cut │ cut
|
||||
▼ ▼
|
||||
release/v1.4.2 release/v1.4.3
|
||||
release/v1.6.x release/v1.7.x
|
||||
│ │
|
||||
tag: v1.4.2 tag: v1.4.3
|
||||
tag: v1.6.0 tag: v1.7.6
|
||||
│ │
|
||||
hotfix ◄─────┘ hotfix ◄─────┘
|
||||
│ │
|
||||
└── cherry-pick back ──────────┘
|
||||
```
|
||||
|
||||
The patch position in the branch name is a literal `x`, while the tag carries
|
||||
the concrete version: one `release/v1.7.x` can hold tags v1.7.0 … v1.7.6.
|
||||
Spanning several patches on a single minor branch is deliberate — patches are
|
||||
revisions of the same feature batch, hotfixes land on one branch, and back-port
|
||||
to main never has to resolve dependencies between several release branches.
|
||||
|
||||
### Key rules
|
||||
|
||||
1. **main is always deployable**: never leave half-done work on `main`.
|
||||
@ -43,9 +49,9 @@ deployable.
|
||||
then `git merge --no-ff feature/xxx` (or squash) when done.
|
||||
3. **Release = cut a release branch from main + tag**:
|
||||
```bash
|
||||
git checkout -b release/v1.4.2 main
|
||||
git tag -a v1.4.2 -m "ModelRouter v1.4.2"
|
||||
git push origin release/v1.4.2 v1.4.2
|
||||
git checkout -b release/v1.7.x main
|
||||
git tag -a v1.7.0 -m "ModelRouter v1.7.0"
|
||||
git push origin release/v1.7.x v1.7.0
|
||||
```
|
||||
Build installers and upload the GitCode Release from this tag so the
|
||||
published state is exactly reproducible.
|
||||
@ -54,7 +60,7 @@ deployable.
|
||||
an already-released branch (unless you deliberately ship a minor revision).
|
||||
5. **Hotfixes MUST flow back to main**:
|
||||
```bash
|
||||
git checkout release/v1.4.2 # fix in the release branch
|
||||
git checkout release/v1.7.x # fix in the release branch
|
||||
git commit -m "fix: ..."
|
||||
git checkout main
|
||||
git cherry-pick <hotfix-commit> # and into main
|
||||
@ -67,12 +73,23 @@ deployable.
|
||||
When the next version ships, the previous release branch retires:
|
||||
|
||||
- **Default: delete the remote release branch**
|
||||
(`git push origin :release/v1.4.2`). All hotfixes were already
|
||||
(`git push origin :release/v1.7.x`). All hotfixes were already
|
||||
cherry-picked into main, so main contains everything; no merge needed.
|
||||
- **Long-term maintenance** (e.g. an enterprise client pinned to an old
|
||||
version): keep the branch, accept only security fixes, keep the
|
||||
commit-then-cherry-pick loop.
|
||||
|
||||
> **Where practice diverged from this section (checked 2026-10-01)**:
|
||||
> `release/v1.4.x` and `release/v1.5.x` still exist locally and on the remote,
|
||||
> so "retire the previous branch when the next version ships" was never
|
||||
> carried out. Keeping them is harmless (hotfixes were back-ported), but it
|
||||
> contradicts the rule above and makes a reader wonder whether they should be
|
||||
> there at all. **Feature branches, by contrast, are cleaned up**:
|
||||
> `feature/key-quota-control`, `feature/toolcall-id-sanitize`,
|
||||
> `feature/anthropic-usage-cache` and `feature/agentrouter-id-sanitize` were
|
||||
> deleted on 2026-10-01 after confirming with a per-commit `git patch-id`
|
||||
> comparison that their work had already landed in main.
|
||||
|
||||
## Explicit non-goals
|
||||
|
||||
- **Never rebase main**: main's history stays append-only; anyone pulling gets
|
||||
@ -134,8 +151,8 @@ pain points:
|
||||
2. No feature branches meant two independent efforts could not proceed in
|
||||
parallel without colliding.
|
||||
|
||||
With release branches: the published state = `release/vX.Y.Z` branch +
|
||||
`vX.Y.Z` tag, exactly reproducible; hotfixes have a clear landing spot; main
|
||||
stays "latest + all fixes + deployable".
|
||||
With release branches: the published state = `release/vX.Y.x` branch +
|
||||
the concrete `vX.Y.Z` tag, exactly reproducible; hotfixes have a clear landing
|
||||
spot; main stays "latest + all fixes + deployable".
|
||||
|
||||
> 中文版见 [docs/git-workflow.md](git-workflow.md)。
|
||||
@ -13,7 +13,7 @@ ModelRouter 采用 **GitHub Flow + 发布分支** 模型:`main` 是唯一长
|
||||
|---|---|---|---|
|
||||
| 主分支 | `main` | 永久 | 唯一长命分支。永远可部署。积攒下一个版本的功能。 |
|
||||
| 特性分支 | `feature/<描述>` | 短命(开发→合并即删) | 新特性 / 一般 bug 修复。从 `main` 开出,完成后合回 `main`。 |
|
||||
| 发布分支 | `release/vX.Y.Z` | 一个版本周期 | 从 `main` 分出,打 tag 发布。该版本生命周期内的 hotfix 都提交在此分支。 |
|
||||
| 发布分支 | `release/vX.Y.x` | 一个版本周期 | 从 `main` 分出,打 tag 发布。该版本生命周期内的 hotfix 都提交在此分支。 |
|
||||
|
||||
## 变更流向(重要)
|
||||
|
||||
@ -24,15 +24,20 @@ ModelRouter 采用 **GitHub Flow + 发布分支** 模型:`main` 是唯一长
|
||||
│ │
|
||||
│ 切出 │ 切出
|
||||
▼ ▼
|
||||
release/v1.4.2 release/v1.4.3
|
||||
release/v1.6.x release/v1.7.x
|
||||
│ │
|
||||
tag: v1.4.2 tag: v1.4.3
|
||||
tag: v1.6.0 tag: v1.7.6
|
||||
│ │
|
||||
hotfix ◄─────┘ hotfix ◄─────┘
|
||||
│ │
|
||||
└── cherry-pick 回 main ───────┘
|
||||
```
|
||||
|
||||
分支名里 patch 位是**字面的 x**,而 tag 打具体版本号:`release/v1.7.x` 这一条
|
||||
发布分支上的 tag 可以有 v1.7.0 … v1.7.6 多个。一条 minor 分支跨多个 patch 是
|
||||
刻意的:patch 是同一批功能的不同修订,hotfix 落在同一条分支上,回流 main 时
|
||||
也不必处理多条 release 分支之间的依赖。
|
||||
|
||||
### 关键规则
|
||||
|
||||
1. **main 永远可部署**:不在 main 上留半成品。任何未完成的工作必须在特性分支上。
|
||||
@ -40,16 +45,16 @@ ModelRouter 采用 **GitHub Flow + 发布分支** 模型:`main` 是唯一长
|
||||
开发完 `git merge --no-ff feature/xxx` 或 squash 合回。
|
||||
3. **发布 = 从 main 切 release 分支 + 打 tag**:
|
||||
```bash
|
||||
git checkout -b release/v1.4.2 main
|
||||
git tag -a v1.4.2 -m "ModelRouter v1.4.2"
|
||||
git push origin release/v1.4.2 v1.4.2
|
||||
git checkout -b release/v1.7.x main
|
||||
git tag -a v1.7.0 -m "ModelRouter v1.7.0"
|
||||
git push origin release/v1.7.x v1.7.0
|
||||
```
|
||||
构建安装包、上传 GitCode Release 都基于这个 tag,保证可精确回溯发布态。
|
||||
4. **版本生命周期内只收该版本的 hotfix**:新特性一律并入 `main` 等下一个版本,
|
||||
绝不塞进已发布的 release 分支(除非主动选择在该版本内发次要版)。
|
||||
5. **hotfix 必须回流 main**:
|
||||
```bash
|
||||
git checkout release/v1.4.2 # 在发布分支提交修复
|
||||
git checkout release/v1.7.x # 在发布分支提交修复
|
||||
git commit -m "fix: ..."
|
||||
git checkout main
|
||||
git cherry-pick <hotfix-commit> # 回主分支
|
||||
@ -61,11 +66,20 @@ ModelRouter 采用 **GitHub Flow + 发布分支** 模型:`main` 是唯一长
|
||||
|
||||
下一个版本发布时,上一个 release 分支退役:
|
||||
|
||||
- **默认:直接删除远端 release 分支**(`git push origin :release/v1.4.2`)。
|
||||
- **默认:直接删除远端 release 分支**(`git push origin :release/v1.7.x`)。
|
||||
因为 hotfix 都已逐个 cherry-pick 回 main,main 已包含全部修复,无需再合并。
|
||||
- **如需要长期维护旧版**(例如企业大客户卡在旧版本):保留分支,仅 stopship 接受
|
||||
该版本的安全修复,继续走「提交 + cherry-pick 回 main」循环。
|
||||
|
||||
> **实践与本节的历史出入(2026-10-01 核对)**:`release/v1.4.x` 与 `release/v1.5.x`
|
||||
> 至今仍在本地与远端,说明"下一个版本发布就删上一个分支"实际没有执行。
|
||||
> 保留无害(hotfix 已回流),但它与上面写的规则不一致,读文档的人会以为
|
||||
> 这些分支不该存在。**特性分支则确实在清理**:`feature/key-quota-control`、
|
||||
> `feature/toolcall-id-sanitize`、`feature/anthropic-usage-cache`、
|
||||
> `feature/agentrouter-id-sanitize` 四个分支在 2026-10-01 删除——它们的工作
|
||||
> 早已全部进入 main(逐提交用 `git patch-id` 比对确认),留着只是给下个版本
|
||||
> 制造 cherry-pick/merge 陷阱。
|
||||
|
||||
## 明确不做的事
|
||||
|
||||
- **不 rebase main**:`main` 的历史保持追加式,任何人拉取后 `git pull` 都得到直接可用的历史。
|
||||
|
||||
507
docs/plugins.md
Normal file
507
docs/plugins.md
Normal file
@ -0,0 +1,507 @@
|
||||
# 插件系统(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` / `fireImageStart` | 每个请求一次(chat 与生图各一条),在配额闸门**之前** |
|
||||
| `chain_step` | `gateway/chat.go` `chainTraceSink` | **仅 AUTO 路径**,每步一次 |
|
||||
| `routed` | `singleChat` / `streamChat` / `singleChatAuto` / `streamChatAuto` | 成功选定源之后,各一次 |
|
||||
| `request_end` | `gateway/chat.go` `writeRec` | **所有出口的唯一汇合点**,每个请求一次 |
|
||||
|
||||
### 3.1b 为什么单独有 `chain_step`
|
||||
|
||||
`routed` 只在**遍历结束后**触发一次,只带最终胜出的槽位。所以
|
||||
"tier 1 冷却所以降级到 tier 3"和"tier 1 正常接单"在它眼里**完全一样**——
|
||||
而这恰恰是优先级链存在的全部理由。
|
||||
|
||||
`chain_step` 补上这条信息,四种 `kind`:
|
||||
|
||||
| kind | 含义 | 何时产生 |
|
||||
|---|---|---|
|
||||
| `tier_skip` | 整档被跳过 | 该档所有槽位冷却中/配额用尽 |
|
||||
| `slot_fail` | 某个槽位硬失败 | 上游报错 / 适配器输出不可用 |
|
||||
| `tier_busy` | 整档全忙且有界等待超时 | 2s 内没等到空位 |
|
||||
| `selected` | 这个槽位接了单 | 每次成功遍历**恰好一次**,且是最后一步 |
|
||||
|
||||
顺序保证:所有 `chain_step` 都在 `routed` 之前,`selected` 是最后一步。
|
||||
所以只订阅 `request_end` 的插件也能拿到轨迹摘要(见下)。
|
||||
|
||||
### 3.1c `request_end` 里的轨迹摘要
|
||||
|
||||
除了逐个 `chain_step`,`request_end` 还带三个便于做报表的字段:
|
||||
|
||||
| 字段 | 含义 |
|
||||
|---|---|
|
||||
| `chain_walk` | 整个遍历的步骤数组(上限 12 步,超出截断) |
|
||||
| `degraded` | 布尔。`true` = 有过跳过/失败,即**发生了降级** |
|
||||
| `tier_served` | 实际服务的那一档;直连或全失败时为 `-1` |
|
||||
|
||||
> **计费口径**:按**实际服务的模型**计费。降级到 tier 3 仍按 tier 3 的价算,
|
||||
> `chain_step` / `degraded` / `tier_served` 只作**观测**,不参与计价。
|
||||
> 理由见 §7.5。
|
||||
|
||||
`request_end` 放在 `writeRec` 是因为四条入口路径(直连/AUTO × 流式/非流式)都
|
||||
经过它,既不会漏(流式的 token 数只有流结束才知道),也不会重复。
|
||||
|
||||
**hot path 注意事项**:没有插件注册某 stage 时,`Fire` 立刻返回(一次 `RLock`
|
||||
加一次 map 查找)。装了插件之后,每个请求会在该 stage 上多一次 Lua 调用——
|
||||
这是同步的,在关键路径上。计费插件那种"每请求一次"是正常的;把重活放进钩子是
|
||||
反模式。
|
||||
|
||||
### 3.2 钩子的返回值
|
||||
|
||||
- 返回 `nil` 或不返回 = **没有意见**,payload 原样传给下一个插件
|
||||
- 返回 table = 其中的键会**合并进 payload**,并作为 `Fire` 的返回值
|
||||
|
||||
**同一 stage 的多个插件并行执行**,但返回值按**插件加载顺序**合并,所以结果是
|
||||
确定的(不依赖 goroutine 调度)。代价是一个插件看不到另一个插件刚加的字段:
|
||||
每个钩子拿到的是**同一份 payload 快照**。
|
||||
|
||||
这与早期版本不同 —— 早期是顺序执行,后一个插件能看到前一个的返回值。它从未被
|
||||
实际依赖(随核心发布的 billing 在每个 stage 都 `return nil`,注释里写着
|
||||
"nobody downstream would read a return value"),但这是一处**契约变化**:如果你的
|
||||
插件依赖「读到前一个插件写的字段」,并行的两个插件之间必须改用外部通信
|
||||
(例如各自写 `plugin.state`,由 `/api/plugins/<name>/state` 读取)。
|
||||
|
||||
单个插件时不启 goroutine,直接调用。
|
||||
|
||||
前三个 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 又收固定费。
|
||||
**`cache_discount`** 见 §7.7。
|
||||
|
||||
### 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`。
|
||||
|
||||
另有两个**降级观测**维度(来自 `chain_step`):
|
||||
|
||||
| 字段 | 含义 |
|
||||
|---|---|
|
||||
| `degraded_reqs` | 发生过降级的请求数 |
|
||||
| `by_tier_served` | 各档实际接单数(`{"1": 812, "2": 37}`) |
|
||||
| `skip_reasons` | 跳过原因计数,等待时长已归一(`no free slot within <wait>`) |
|
||||
|
||||
这三项是"网关是不是在悄悄降级"的核心指标:一个持续降级的网关,账单结构和健康
|
||||
网关看起来一模一样——**除非**单独统计降级次数。
|
||||
|
||||
### 7.4 计费策略:失败的请求怎么算
|
||||
|
||||
**保留 token 费用,丢弃固定费用。** 理由:上游在生成后才 500,token 确实被消耗
|
||||
了;但那个从未真正发生的固定费不该收。
|
||||
|
||||
这是本插件里**最可争议的一条**。想改成"失败也收固定费":
|
||||
|
||||
```json
|
||||
{ "prices": { ... }, "count_failures": true }
|
||||
```
|
||||
|
||||
### 7.5 ⚠️ 计费插件只报表,不执法
|
||||
|
||||
**网关自己的配额会计(`internal/gateway/stats.go`,入口处强制)才是限额权威。**
|
||||
本插件不参与任何路由或配额决策。
|
||||
|
||||
理由:两套独立的会计路径如果对不上,比一套功能略少的更糟。计费是**观察**,
|
||||
配额是**控制**,二者分开。
|
||||
|
||||
### 7.6 未定价流量(重要)
|
||||
|
||||
**任何维度都没配价的请求,成本记 0。** 这是最危险的失败模式:账单照样能加总,
|
||||
只是**悄悄少报**,而且没有任何报错。
|
||||
|
||||
所以插件单独统计它们:
|
||||
|
||||
| 字段 | 含义 |
|
||||
|---|---|
|
||||
| `unpriced_reqs` | 没有任何价目覆盖的请求数 |
|
||||
| `unpriced_models` | 按模型点名(`{"MYSTERY-MODEL": 12}`)——直接告诉你价目表缺了哪一行 |
|
||||
|
||||
仪表盘上有 "Unpriced" 卡片。**这个数应该是 0**;不是 0 就去补价目。
|
||||
|
||||
注意"未定价"不等于"免费":这些请求的 `requests` / token 数**照常计入**
|
||||
`total` 与各维度,只有金额是 0。
|
||||
|
||||
### 7.7 提示缓存计价
|
||||
|
||||
**缓存命中的 prompt token 不按全价算。** 绝大多数 provider 对缓存读给很深的折扣
|
||||
(常见是 1/10),而 agent 流量会反复重放长前缀——正是缓存要让它便宜的那类流量。
|
||||
|
||||
```
|
||||
fresh = prompt_tokens - cache_hit_tokens → 全价
|
||||
cached = cache_hit_tokens → 全价 × cache_discount
|
||||
```
|
||||
|
||||
`cache_discount` 默认 **0.1**(10 倍折扣),因为 DeepSeek / Qwen / Kimi 等都是这个
|
||||
量级。它是**每个 provider 的事实、不是自然常数**,所以可以按条目覆盖:
|
||||
|
||||
```json
|
||||
"models": { "gpt-5.4": { "prompt": 1.25e-6, "completion": 1e-5, "cache_discount": 0.25 } }
|
||||
```
|
||||
|
||||
设成 `1` 恢复成旧的"prompt 一律全价"行为,设成 `0` 表示该 provider 不打折。
|
||||
|
||||
优先级与 token 价一致(`keys` > `models` > `default`)。
|
||||
|
||||
> **这一条改过行为。** 修复前缓存命中按全价算:100 万 prompt token 里 90 万是
|
||||
> 缓存命中,会算出 10 USD 而不是 ~1.9——**高估约 10 倍**,而且恰好发生在缓存
|
||||
> 最有价值的高频流量上。
|
||||
|
||||
### 7.8 数据从哪来
|
||||
|
||||
`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)。
|
||||
141
e2e/e2e_test.go
141
e2e/e2e_test.go
@ -326,3 +326,144 @@ func statusOf(resp *http.Response) int {
|
||||
}
|
||||
return resp.StatusCode
|
||||
}
|
||||
|
||||
// The startup warning about key configuration must describe the key set that
|
||||
// AUTH actually uses (cfg.Keys, via core.ListKeys), not the legacy
|
||||
// gateway_keys YAML list.
|
||||
//
|
||||
// It used to check GatewayKeys() only. seedKeys copies gateway_keys into keys[]
|
||||
// on first start and the YAML list stops being consulted, so once that list was
|
||||
// emptied — e.g. after rotating the starter key away — the process logged
|
||||
// "every request will be rejected" while seven working keys, one of them admin,
|
||||
// were in service. Observed in production: the warning appeared and that same
|
||||
// key returned HTTP 200 on /v1/models.
|
||||
//
|
||||
// The fixture reproduces exactly that shape: gateway_keys EMPTY, with the
|
||||
// working key present only in keys[].
|
||||
func TestStartupWarningReflectsRealKeysNotLegacyList(t *testing.T) {
|
||||
up := newMockUpstream(t)
|
||||
dir := t.TempDir()
|
||||
|
||||
// Seed a real install once so keys[] gets populated from gateway_keys, then
|
||||
// rewrite the file to look like a rotated install: legacy list cleared, the
|
||||
// seeded key still present in keys[].
|
||||
path := writeConfig(t, dir, "127.0.0.1:0", map[string]*mockUpstream{"good": up})
|
||||
bin := buildBinary(t)
|
||||
|
||||
// First run creates the sealed config with keys[] populated.
|
||||
seedPort := freePort(t)
|
||||
cfg1 := writeConfig(t, dir, "127.0.0.1:"+seedPort, map[string]*mockUpstream{"good": up})
|
||||
g1 := startGateway(t, bin, "127.0.0.1:"+seedPort, cfg1)
|
||||
g1.stop()
|
||||
|
||||
data, err := os.ReadFile(cfg1)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if !strings.Contains(string(data), "keys:") {
|
||||
t.Fatalf("first run did not populate keys[]:\n%s", data)
|
||||
}
|
||||
|
||||
// Now clear the legacy gateway_keys list, keeping keys[] intact. The
|
||||
// rewritten file may use different indentation, so match loosely.
|
||||
cleared, n := clearGatewayKeys(string(data))
|
||||
if n == 0 {
|
||||
t.Fatalf("could not find a gateway_keys entry to clear in:\n%s", data)
|
||||
}
|
||||
if err := os.WriteFile(cfg1, []byte(cleared), 0o644); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
// Start again and inspect the log. The config pins its listen address, so
|
||||
// rewrite that field to the port we will actually poll.
|
||||
port2 := freePort(t)
|
||||
rehosted, ok := replaceListen(string(cleared), "127.0.0.1:"+port2)
|
||||
if !ok {
|
||||
t.Fatalf("could not rewrite listen in:\n%s", cleared)
|
||||
}
|
||||
cfg2 := filepath.Join(dir, "config2.yaml")
|
||||
if err := os.WriteFile(cfg2, []byte(rehosted), 0o644); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
g2 := startGateway(t, bin, "127.0.0.1:"+port2, cfg2)
|
||||
defer g2.stop()
|
||||
|
||||
// Sanity: the key really does authenticate, so the warning would be a lie.
|
||||
resp, _ := g2.do("GET", "/v1/models", "", true)
|
||||
if resp == nil || resp.StatusCode != http.StatusOK {
|
||||
t.Fatalf("key from keys[] did not authenticate (status %v); test fixture is wrong",
|
||||
resp)
|
||||
}
|
||||
|
||||
logs := g2.logString()
|
||||
if strings.Contains(logs, "gateway_keys is EMPTY") {
|
||||
t.Errorf("startup logged 'gateway_keys is EMPTY ... every request is rejected' while "+
|
||||
"a working admin key was in service — the warning must read the authoritative "+
|
||||
"key set (core.ListKeys), not the legacy YAML list.\nlog:\n%s", logs)
|
||||
}
|
||||
_ = path
|
||||
_ = g1
|
||||
}
|
||||
|
||||
func freePort(t *testing.T) string {
|
||||
t.Helper()
|
||||
l, err := net.Listen("tcp", "127.0.0.1:0")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer l.Close()
|
||||
_, port, _ := net.SplitHostPort(l.Addr().String())
|
||||
return port
|
||||
}
|
||||
|
||||
// stop terminates the gateway and waits, so its log buffer is complete.
|
||||
func (g *gatewayUnderTest) stop() {
|
||||
if g == nil || g.cmd == nil || g.cmd.Process == nil {
|
||||
return
|
||||
}
|
||||
_ = g.cmd.Process.Kill()
|
||||
_, _ = g.cmd.Process.Wait()
|
||||
}
|
||||
|
||||
// logString returns everything the gateway has written so far.
|
||||
func (g *gatewayUnderTest) logString() string { return g.log.String() }
|
||||
|
||||
// clearGatewayKeys empties the gateway_keys list in a config, tolerating any
|
||||
// indentation the YAML writer chose. Returns the new text and how many entries
|
||||
// it removed.
|
||||
func clearGatewayKeys(src string) (string, int) {
|
||||
lines := strings.Split(src, "\n")
|
||||
var out []string
|
||||
removed := 0
|
||||
inList := false
|
||||
for _, l := range lines {
|
||||
trimmed := strings.TrimSpace(l)
|
||||
if strings.HasPrefix(l, "gateway_keys:") {
|
||||
out = append(out, "gateway_keys: []")
|
||||
inList = true
|
||||
removed++
|
||||
continue
|
||||
}
|
||||
if inList {
|
||||
if strings.HasPrefix(trimmed, "- ") {
|
||||
removed++
|
||||
continue
|
||||
}
|
||||
inList = false
|
||||
}
|
||||
out = append(out, l)
|
||||
}
|
||||
return strings.Join(out, "\n"), removed
|
||||
}
|
||||
|
||||
// replaceListen rewrites the top-level listen: value.
|
||||
func replaceListen(src, addr string) (string, bool) {
|
||||
lines := strings.Split(src, "\n")
|
||||
for i, l := range lines {
|
||||
if strings.HasPrefix(l, "listen:") {
|
||||
lines[i] = "listen: " + addr
|
||||
return strings.Join(lines, "\n"), true
|
||||
}
|
||||
}
|
||||
return src, false
|
||||
}
|
||||
|
||||
323
internal/billing/compile.go
Normal file
323
internal/billing/compile.go
Normal file
@ -0,0 +1,323 @@
|
||||
package billing
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"strings"
|
||||
|
||||
"llmsproxy/internal/config"
|
||||
)
|
||||
|
||||
// Compile turns one billing profile into the prices table the billing plugin
|
||||
// expects, resolving URL rules against the gateway's actual sources.
|
||||
//
|
||||
// WHY URL RULES NEED RESOLVING AT ALL: the operator declares pricing by URL
|
||||
// because that is what a provider's price list is keyed on, and because several
|
||||
// sources can point at the same URL. The plugin, however, looks up by
|
||||
// `payload.source` (a source NAME) and by model — it has no idea what URL a
|
||||
// request went to. So the URL match happens here, at compile time, where the
|
||||
// config's name->base_url mapping is known, and the result is expressed in the
|
||||
// dimensions the plugin already supports.
|
||||
//
|
||||
// Profiles are the "let the user choose" axis: the same URL can appear in
|
||||
// several profiles and switching recomputes this table, so a gateway can be
|
||||
// repriced without editing the plugin.
|
||||
func Compile(profile *config.BillingProfile, sources []config.Source) (map[string]interface{}, error) {
|
||||
return CompileOpts(profile, sources, false)
|
||||
}
|
||||
|
||||
// CompileOpts is Compile with an explicit policy for rules that match no
|
||||
// source.
|
||||
//
|
||||
// The default is STRICT and that is correct for startup: a profile with a
|
||||
// typo'd URL prices nothing, every request on it is recorded as unpriced, and
|
||||
// the bill silently comes out wrong. Failing the load is the right response.
|
||||
//
|
||||
// The rules editor needs the opposite. An operator editing a half-finished
|
||||
// profile — adding a rule before the source that will use it exists, or
|
||||
// renaming a URL while typing — must be able to save and then see the warning,
|
||||
// not be blocked by an error they can only resolve by guessing. Lenient mode
|
||||
// keeps the offending rule in the config and drops it from the compiled table,
|
||||
// so it is preserved as text and clearly reported, without pretending it
|
||||
// prices anything.
|
||||
func CompileOpts(profile *config.BillingProfile, sources []config.Source, lenient bool) (map[string]interface{}, error) {
|
||||
if profile == nil {
|
||||
return nil, fmt.Errorf("no billing profile")
|
||||
}
|
||||
prices := map[string]interface{}{
|
||||
"currency": "USD",
|
||||
"default": map[string]interface{}{"prompt": 0.0, "completion": 0.0, "per_request": 0.0},
|
||||
"sources": map[string]interface{}{},
|
||||
"models": map[string]interface{}{},
|
||||
"keys": map[string]interface{}{},
|
||||
}
|
||||
if profile.Default != "" {
|
||||
prices["default_mode"] = profile.Default
|
||||
}
|
||||
|
||||
// Currency is a profile-level statement; a rule may override it.
|
||||
for i := range profile.Rules {
|
||||
if c := profile.Rules[i].Currency; c != "" {
|
||||
prices["currency"] = c
|
||||
break
|
||||
}
|
||||
}
|
||||
|
||||
matched := map[string]bool{}
|
||||
for i := range profile.Rules {
|
||||
rule := &profile.Rules[i]
|
||||
targets := matchSources(rule.URL, sources)
|
||||
if len(targets) == 0 && rule.URL != "*" {
|
||||
// A rule for a URL no source uses is almost always a typo (or a
|
||||
// source that was removed). Failing loudly beats a profile that
|
||||
// silently prices nothing — the whole reason this is a config file
|
||||
// instead of a hand-written JSON blob.
|
||||
if lenient {
|
||||
continue // kept in the config text, reported as a warning
|
||||
}
|
||||
return nil, fmt.Errorf("rule url %q matches no configured source base_url", rule.URL)
|
||||
}
|
||||
for _, srcName := range targets {
|
||||
if matched[srcName] {
|
||||
// First rule wins. Two rules matching one source is ambiguous,
|
||||
// and silently letting the later one win makes the file's
|
||||
// meaning depend on ordering the operator cannot see.
|
||||
continue
|
||||
}
|
||||
matched[srcName] = true
|
||||
if err := applyRule(prices, rule, srcName, sources); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Sources no rule matched fall through to the profile default, which the
|
||||
// plugin applies via prices.default. Recording them explicitly means the
|
||||
// UI can say "this source is unpriced" instead of leaving the operator to
|
||||
// infer it from a zero.
|
||||
var unmatched []string
|
||||
for _, src := range sources {
|
||||
if !matched[src.Name] {
|
||||
unmatched = append(unmatched, src.Name)
|
||||
}
|
||||
}
|
||||
if len(unmatched) > 0 {
|
||||
prices["unmatched_sources"] = unmatched
|
||||
}
|
||||
return prices, nil
|
||||
}
|
||||
|
||||
// matchSources returns the source names whose base_url matches pattern.
|
||||
// "*" matches every source (used as a catch-all default rule).
|
||||
func matchSources(pattern string, sources []config.Source) []string {
|
||||
var out []string
|
||||
if pattern == "*" {
|
||||
for _, s := range sources {
|
||||
out = append(out, s.Name)
|
||||
}
|
||||
return out
|
||||
}
|
||||
want := normalizeURL(pattern)
|
||||
for _, s := range sources {
|
||||
if normalizeURL(s.BaseURL) == want {
|
||||
out = append(out, s.Name)
|
||||
}
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
// normalizeURL compares URLs the way an operator expects: trailing slashes and
|
||||
// case in the host are differences the provider's price list does not care
|
||||
// about, and requiring an exact byte match would make the config brittle.
|
||||
func normalizeURL(u string) string {
|
||||
u = strings.TrimSpace(u)
|
||||
u = strings.TrimRight(u, "/")
|
||||
return strings.ToLower(u)
|
||||
}
|
||||
|
||||
// applyRule writes one rule's pricing for one source into the prices table.
|
||||
func applyRule(prices map[string]interface{}, rule *config.BillingRule, srcName string, sources []config.Source) error {
|
||||
switch rule.Mode {
|
||||
case "free":
|
||||
// Explicitly priced at zero. This is NOT the same as unpriced: a source
|
||||
// the operator says is free must not appear in the unpriced warnings,
|
||||
// or those warnings become noise and stop being read.
|
||||
prices["sources"].(map[string]interface{})[srcName] = map[string]interface{}{
|
||||
"prompt": 0.0, "completion": 0.0, "per_request": 0.0,
|
||||
}
|
||||
return nil
|
||||
|
||||
case "subscription":
|
||||
// A fixed monthly commitment: the per-request MARGINAL cost is zero, and
|
||||
// the flat fee is reported separately. Spreading a monthly fee across
|
||||
// requests would invent a per-request number the provider never charges,
|
||||
// and it would change every time traffic did.
|
||||
prices["sources"].(map[string]interface{})[srcName] = map[string]interface{}{
|
||||
"prompt": 0.0, "completion": 0.0, "per_request": 0.0,
|
||||
}
|
||||
fixed, _ := prices["subscriptions"].(map[string]interface{})
|
||||
if fixed == nil {
|
||||
fixed = map[string]interface{}{}
|
||||
prices["subscriptions"] = fixed
|
||||
}
|
||||
cur := "USD"
|
||||
if rule.Currency != "" {
|
||||
cur = rule.Currency
|
||||
}
|
||||
fixed[srcName] = map[string]interface{}{
|
||||
"monthly": rule.Subscription, "currency": cur,
|
||||
}
|
||||
return nil
|
||||
|
||||
case "unpriced":
|
||||
// Deliberately left out of `sources` so the plugin's unpriced_models /
|
||||
// unpriced_reqs counters catch it. That is the point: a subscription
|
||||
// plan whose credits cannot be converted to tokens must be VISIBLE as
|
||||
// unbilled, not quietly estimated.
|
||||
return nil
|
||||
|
||||
case "token":
|
||||
if len(rule.Models) == 0 {
|
||||
return fmt.Errorf("rule for url %q: mode token requires at least one model", rule.URL)
|
||||
}
|
||||
models := prices["models"].(map[string]interface{})
|
||||
for modelID, t := range rule.Models {
|
||||
entry := map[string]interface{}{}
|
||||
prompt, err := usdPerM(t.Prompt)
|
||||
if err != nil {
|
||||
return fmt.Errorf("model %q prompt: %w", modelID, err)
|
||||
}
|
||||
completion, err := usdPerM(t.Completion)
|
||||
if err != nil {
|
||||
return fmt.Errorf("model %q completion: %w", modelID, err)
|
||||
}
|
||||
entry["prompt"] = prompt
|
||||
entry["completion"] = completion
|
||||
if t.CacheDiscount != nil {
|
||||
entry["cache_discount"] = *t.CacheDiscount
|
||||
}
|
||||
// reasoning is optional and stays ABSENT (not nil, not 0) when the
|
||||
// rule does not set it, which is what makes the billing plugin skip
|
||||
// the repricing branch entirely. Writing an explicit 0 here would
|
||||
// mean "thinking is free" and silently under-bill every request
|
||||
// from a provider that bills it like ordinary output.
|
||||
if strings.TrimSpace(t.Reasoning) != "" {
|
||||
rp, err := usdPerM(t.Reasoning)
|
||||
if err != nil {
|
||||
return fmt.Errorf("model %q reasoning: %w", modelID, err)
|
||||
}
|
||||
entry["reasoning_price"] = rp
|
||||
}
|
||||
if lc := t.LongContext; lc != nil {
|
||||
if lc.Above <= 0 {
|
||||
return fmt.Errorf("model %q long_context: above must be > 0", modelID)
|
||||
}
|
||||
above, err := longContextRate(lc, prompt)
|
||||
if err != nil {
|
||||
return fmt.Errorf("model %q long_context: %w", modelID, err)
|
||||
}
|
||||
entry["long_context"] = map[string]interface{}{
|
||||
"above": lc.Above,
|
||||
// Marginal, not blended. The plugin charges the tokens past the
|
||||
// threshold at this rate and the rest at the base rate, so the
|
||||
// threshold composes with a cache discount (which applies to the
|
||||
// cached portion) instead of forcing the plugin to guess how
|
||||
// the two interact.
|
||||
"prompt": above,
|
||||
}
|
||||
}
|
||||
if rule.Peak != nil {
|
||||
entry["peak"] = compilePeak(rule.Peak)
|
||||
}
|
||||
models[modelID] = entry
|
||||
}
|
||||
// A specific-URL token rule needs a source entry so a model served
|
||||
// from THIS url but absent from `rule.Models` is still recognised as
|
||||
// belonging to a priced provider.
|
||||
//
|
||||
// A WILDCARD rule ("*") must NOT get one. "*" is a model catalogue — it
|
||||
// says "these model ids have these prices, wherever they come from" —
|
||||
// and marking every source priced at {0,0} would make priceFor() set
|
||||
// priced=true for every model the catalogue does NOT list. Paid
|
||||
// sources (commandcode, alittokenplan, …) would then record $0 instead
|
||||
// of showing up in unpriced_reqs, which is precisely the silent
|
||||
// zero-billing failure this whole plugin is built to make visible.
|
||||
// The model lookup p.models[payload.model] already marks a listed
|
||||
// model priced on its own, so the wildcard loses nothing.
|
||||
if rule.URL != "*" {
|
||||
if _, ok := prices["sources"].(map[string]interface{})[srcName]; !ok {
|
||||
prices["sources"].(map[string]interface{})[srcName] = map[string]interface{}{
|
||||
"prompt": 0.0, "completion": 0.0, "per_request": 0.0,
|
||||
}
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
return fmt.Errorf("unknown mode %q", rule.Mode)
|
||||
}
|
||||
|
||||
// usdPerM converts a USD-per-million string into the per-token rate the plugin
|
||||
// expects. Delegates to config.ParseFloatUSDPerM so the DSL validator and the
|
||||
// compiler agree byte-for-byte on what a valid price is — a value that passes
|
||||
// Validate() but fails here (or vice versa) would be the worst kind of drift.
|
||||
func usdPerM(s string) (float64, error) {
|
||||
return config.ParseFloatUSDPerM(s)
|
||||
}
|
||||
|
||||
// compilePeak renders the peak window in the shape the plugin reads:
|
||||
// { multiplier, windows = { { days = {...}, hours = { {lo,hi}, ... } } } }.
|
||||
//
|
||||
// The plugin reads `days` / `hours` pairs. The first version of the deployed
|
||||
// price table used { start, end, weekdays } — a shape nothing reads — so peak
|
||||
// traffic was billed at off-peak rates with no error anywhere. Compiling from
|
||||
// typed config fields removes the chance of writing the wrong key names by hand.
|
||||
func compilePeak(p *config.BillingPeak) map[string]interface{} {
|
||||
win := map[string]interface{}{}
|
||||
if len(p.Weekdays) > 0 {
|
||||
days := make([]interface{}, 0, len(p.Weekdays))
|
||||
for _, d := range p.Weekdays {
|
||||
days = append(days, d)
|
||||
}
|
||||
win["days"] = days
|
||||
}
|
||||
if len(p.Hours) > 0 {
|
||||
hours := make([]interface{}, 0, len(p.Hours))
|
||||
for _, h := range p.Hours {
|
||||
hours = append(hours, []interface{}{h[0], h[1]})
|
||||
}
|
||||
win["hours"] = hours
|
||||
}
|
||||
out := map[string]interface{}{"multiplier": p.Multiplier}
|
||||
if len(win) > 0 {
|
||||
out["windows"] = []interface{}{win}
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
// longContextRate resolves the marginal per-token rate for the tokens ABOVE a
|
||||
// threshold, given the model's base per-million prompt price.
|
||||
//
|
||||
// Two ways to express it are supported because providers publish both: an
|
||||
// absolute rate (OpenAI quotes \$3.00/M past 200k) and a multiple (DeepSeek
|
||||
// quotes exactly 2x past 128k). Neither may be missing: a step with no rate is a
|
||||
// silent no-op, and an operator who wrote one meant something.
|
||||
func longContextRate(lc *config.BillingLongContext, basePerToken float64) (float64, error) {
|
||||
hasPrompt := strings.TrimSpace(lc.Prompt) != ""
|
||||
hasMult := lc.Multiplier > 0
|
||||
switch {
|
||||
case hasPrompt && hasMult:
|
||||
p, err := usdPerM(lc.Prompt)
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
// Both given: the absolute rate wins, because it is the one the provider
|
||||
// actually quotes. Silently multiplying instead would produce a number
|
||||
// the operator never wrote.
|
||||
return p, nil
|
||||
case hasPrompt:
|
||||
return usdPerM(lc.Prompt)
|
||||
case hasMult:
|
||||
return basePerToken * lc.Multiplier, nil
|
||||
default:
|
||||
return 0, fmt.Errorf("needs either prompt or multiplier")
|
||||
}
|
||||
}
|
||||
214
internal/billing/wildcard_test.go
Normal file
214
internal/billing/wildcard_test.go
Normal file
@ -0,0 +1,214 @@
|
||||
package billing
|
||||
|
||||
import (
|
||||
"testing"
|
||||
|
||||
"llmsproxy/internal/config"
|
||||
)
|
||||
|
||||
// The wildcard token rule must not blanket-mark every source as priced.
|
||||
//
|
||||
// Why this is a hazard and not a style question: priceFor() sets priced=true as
|
||||
// soon as the request's SOURCE has an entry in the price table, and that flag
|
||||
// is what keeps a request OUT of unpriced_reqs. A "*" rule is a model
|
||||
// catalogue — it prices specific model ids wherever they come from — but
|
||||
// compiling it into a {0,0,0} entry on all 21 sources made every model the
|
||||
// catalogue does not list look "priced at zero". Paid providers (commandcode,
|
||||
// alittokenplan, …) then stopped appearing in the unpriced warnings, which is
|
||||
// exactly the silent zero-billing failure the plugin exists to surface.
|
||||
func TestWildcardTokenRuleDoesNotMarkEverySourcePriced(t *testing.T) {
|
||||
sources := []config.Source{
|
||||
{Name: "free1", BaseURL: "https://free1.example.com/v1"},
|
||||
{Name: "paid1", BaseURL: "https://paid1.example.com/v1"},
|
||||
{Name: "paid2", BaseURL: "https://paid2.example.com/v1"},
|
||||
}
|
||||
profile := &config.BillingProfile{
|
||||
ID: "p",
|
||||
Rules: []config.BillingRule{
|
||||
{URL: "https://free1.example.com/v1", Mode: "free"},
|
||||
{
|
||||
URL: "*", Mode: "token",
|
||||
Models: map[string]config.BillingToken{
|
||||
"only-catalogued-model": {Prompt: "1.00", Completion: "2.00"},
|
||||
},
|
||||
},
|
||||
},
|
||||
}
|
||||
prices, err := Compile(profile, sources)
|
||||
if err != nil {
|
||||
t.Fatalf("compile: %v", err)
|
||||
}
|
||||
srcs := prices["sources"].(map[string]interface{})
|
||||
|
||||
// The explicitly free source keeps its entry (free is a deliberate 0).
|
||||
if _, ok := srcs["free1"]; !ok {
|
||||
t.Error("the explicit free rule lost its source entry")
|
||||
}
|
||||
// Paid sources must NOT be marked, or their unlisted models bill $0.
|
||||
for _, name := range []string{"paid1", "paid2"} {
|
||||
if _, ok := srcs[name]; ok {
|
||||
t.Errorf("★ wildcard token rule marked %q as priced at 0 — its unlisted "+
|
||||
"models will bill $0 instead of appearing as unpriced", name)
|
||||
}
|
||||
}
|
||||
// The catalogue itself must still be present, or the rule prices nothing.
|
||||
models := prices["models"].(map[string]interface{})
|
||||
if _, ok := models["only-catalogued-model"]; !ok {
|
||||
t.Error("the wildcard rule did not publish its model price")
|
||||
}
|
||||
}
|
||||
|
||||
// The inverse guard: a token rule for a SPECIFIC url still marks that source,
|
||||
// because it genuinely declares pricing for the models served from it.
|
||||
func TestSpecificURLTokenRuleStillMarksItsSource(t *testing.T) {
|
||||
sources := []config.Source{
|
||||
{Name: "mine", BaseURL: "https://mine.example.com/v1"},
|
||||
{Name: "other", BaseURL: "https://other.example.com/v1"},
|
||||
}
|
||||
profile := &config.BillingProfile{
|
||||
ID: "p",
|
||||
Rules: []config.BillingRule{{
|
||||
URL: "https://mine.example.com/v1", Mode: "token",
|
||||
Models: map[string]config.BillingToken{
|
||||
"m": {Prompt: "1.00", Completion: "2.00"},
|
||||
},
|
||||
}},
|
||||
}
|
||||
prices, err := Compile(profile, sources)
|
||||
if err != nil {
|
||||
t.Fatalf("compile: %v", err)
|
||||
}
|
||||
srcs := prices["sources"].(map[string]interface{})
|
||||
if _, ok := srcs["mine"]; !ok {
|
||||
t.Error("a specific-url token rule must still mark its own source priced")
|
||||
}
|
||||
if _, ok := srcs["other"]; ok {
|
||||
t.Error("an unrelated source was marked priced by a specific-url rule")
|
||||
}
|
||||
}
|
||||
|
||||
// TestReasoningPriceIsOmittedUnlessSet pins the absence convention for the
|
||||
// thinking-token rate.
|
||||
//
|
||||
// A rule that says nothing about reasoning must produce a price entry with NO
|
||||
// reasoning_price key, so the billing plugin skips the repricing branch and
|
||||
// bills thinking at the ordinary output rate — which is what every provider
|
||||
// actually charges. Writing an explicit 0 instead would read as "thinking is
|
||||
// free" and silently under-bill every request from that model.
|
||||
func TestReasoningPriceIsOmittedUnlessSet(t *testing.T) {
|
||||
sources := []config.Source{{Name: "up", BaseURL: "https://up.example/v1"}}
|
||||
|
||||
// Case 1: no reasoning declared → the key must not appear.
|
||||
p1 := &config.BillingProfile{ID: "p", Rules: []config.BillingRule{{
|
||||
URL: "https://up.example/v1", Mode: "token",
|
||||
Models: map[string]config.BillingToken{
|
||||
"m": {Prompt: "1.00", Completion: "2.00"},
|
||||
},
|
||||
}}}
|
||||
prices, err := Compile(p1, sources)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
entry := prices["models"].(map[string]interface{})["m"].(map[string]interface{})
|
||||
if _, present := entry["reasoning_price"]; present {
|
||||
t.Errorf("reasoning_price present (%v) without a rule; the plugin would treat 0 as free",
|
||||
entry["reasoning_price"])
|
||||
}
|
||||
|
||||
// Case 2: declared → present, and converted to per-token like every other rate.
|
||||
p2 := &config.BillingProfile{ID: "p", Rules: []config.BillingRule{{
|
||||
URL: "https://up.example/v1", Mode: "token",
|
||||
Models: map[string]config.BillingToken{
|
||||
"m": {Prompt: "1.00", Completion: "2.00", Reasoning: "6.00"},
|
||||
},
|
||||
}}}
|
||||
prices2, err := Compile(p2, sources)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
e2 := prices2["models"].(map[string]interface{})["m"].(map[string]interface{})
|
||||
rp, ok := e2["reasoning_price"].(float64)
|
||||
if !ok {
|
||||
t.Fatalf("reasoning_price missing or wrong type: %#v", e2["reasoning_price"])
|
||||
}
|
||||
if rp != 6e-6 {
|
||||
t.Errorf("reasoning_price = %v, want 6e-6 (USD per TOKEN)", rp)
|
||||
}
|
||||
// A malformed rate must be rejected, not silently ignored: an operator who
|
||||
// typed "6 USD" expects either that value or an error, never a rule that
|
||||
// looks applied and bills nothing.
|
||||
p3 := &config.BillingProfile{ID: "p", Rules: []config.BillingRule{{
|
||||
URL: "https://up.example/v1", Mode: "token",
|
||||
Models: map[string]config.BillingToken{
|
||||
"m": {Prompt: "1.00", Completion: "2.00", Reasoning: "six dollars"},
|
||||
},
|
||||
}}}
|
||||
if _, err := Compile(p3, sources); err == nil {
|
||||
t.Error("a non-numeric reasoning rate must fail the compile")
|
||||
}
|
||||
}
|
||||
|
||||
// TestLongContextStepCompiles pins the long-context step's two accepted forms
|
||||
// and its rejection of a step that would silently do nothing.
|
||||
//
|
||||
// Providers publish both shapes — DeepSeek says "2x past 128k", OpenAI quotes
|
||||
// an absolute "$3.00/M past 200k" — so both must compile. Neither may be
|
||||
// missing: a step with no rate is a rule that LOOKS like it handles long
|
||||
// context and does not, which under-bills without any visible sign.
|
||||
func TestLongContextStepCompiles(t *testing.T) {
|
||||
sources := []config.Source{{Name: "up", BaseURL: "https://up.example/v1"}}
|
||||
|
||||
compile := func(lc config.BillingLongContext) (map[string]interface{}, error) {
|
||||
p := &config.BillingProfile{ID: "p", Rules: []config.BillingRule{{
|
||||
URL: "https://up.example/v1", Mode: "token",
|
||||
Models: map[string]config.BillingToken{
|
||||
"m": {Prompt: "1.00", Completion: "2.00", LongContext: &lc},
|
||||
},
|
||||
}}}
|
||||
prices, err := Compile(p, sources)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
e := prices["models"].(map[string]interface{})["m"].(map[string]interface{})
|
||||
return e["long_context"].(map[string]interface{}), nil
|
||||
}
|
||||
|
||||
// Absolute rate (OpenAI's published shape).
|
||||
lc1, err := compile(config.BillingLongContext{Above: 200000, Prompt: "3.00"})
|
||||
if err != nil {
|
||||
t.Fatalf("absolute long_context rate rejected: %v", err)
|
||||
}
|
||||
if lc1["above"] != 200000 || lc1["prompt"] != 3e-6 {
|
||||
t.Errorf("absolute step = %#v, want above=200000 prompt=3e-6", lc1)
|
||||
}
|
||||
|
||||
// Multiplier (DeepSeek's shape): 2x of the 1.00/M base.
|
||||
lc2, err := compile(config.BillingLongContext{Above: 128000, Multiplier: 2})
|
||||
if err != nil {
|
||||
t.Fatalf("multiplier long_context rejected: %v", err)
|
||||
}
|
||||
if lc2["prompt"] != 2e-6 {
|
||||
t.Errorf("multiplier step = %#v, want prompt=2e-6", lc2)
|
||||
}
|
||||
|
||||
// Both given: the absolute rate wins, because that is the number the
|
||||
// provider quotes. Multiplying instead would produce a figure the operator
|
||||
// never wrote.
|
||||
lc3, err := compile(config.BillingLongContext{Above: 128000, Prompt: "5.00", Multiplier: 2})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if lc3["prompt"] != 5e-6 {
|
||||
t.Errorf("both given = %#v, want the quoted absolute rate 5e-6", lc3)
|
||||
}
|
||||
|
||||
// Neither: a no-op step must fail, not compile into a rule that silently
|
||||
// bills long context at the base rate.
|
||||
if _, err := compile(config.BillingLongContext{Above: 128000}); err == nil {
|
||||
t.Error("a step with neither prompt nor multiplier must be rejected")
|
||||
}
|
||||
// A zero threshold would make every request "long".
|
||||
if _, err := compile(config.BillingLongContext{Above: 0, Multiplier: 2}); err == nil {
|
||||
t.Error("above=0 must be rejected")
|
||||
}
|
||||
}
|
||||
@ -7,6 +7,8 @@ import (
|
||||
"fmt"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strconv"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"gopkg.in/yaml.v3"
|
||||
@ -19,6 +21,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
|
||||
@ -27,6 +30,29 @@ type Config struct {
|
||||
Auto []ModelScope `yaml:"auto,omitempty"` // AUTO 调度链规则(WebUI 优先级页编辑,chat)
|
||||
AutoImage []ModelScope `yaml:"auto_image,omitempty"` // AUTO 生图调度链规则(WebUI 优先级页·生图)
|
||||
Keys []GWKey `yaml:"keys,omitempty"` // 网关密钥(WebUI 密钥页管理)
|
||||
// BillingDSL declares per-URL pricing profiles for the billing plugin.
|
||||
//
|
||||
// Declarative, because hand-writing the prices JSON was error-prone in ways
|
||||
// that looked successful: the first PUT went out without the required
|
||||
// `prices` wrapper and silently REPLACED the accumulated totals with the
|
||||
// price table; the second used a peak-window shape the plugin does not read,
|
||||
// so peak traffic was billed at off-peak rates — no error, the bill just
|
||||
// quietly halved. A schema-checked config file catches both at load.
|
||||
BillingDSL *BillingDSL `yaml:"billing,omitempty"`
|
||||
// PluginVisibility overrides the audience a plugin declared for itself,
|
||||
// keyed by plugin name: "admin" hides a plugin (and its state) from user
|
||||
// keys, "all" widens one that shipped admin-only, "user,admin" names roles
|
||||
// explicitly. An absent name leaves the plugin's own declaration in force.
|
||||
//
|
||||
// This exists because a plugin author cannot know the operator's policy, and
|
||||
// the operator should not have to edit a bundled .lua to change one — the
|
||||
// next upgrade overwrites it.
|
||||
PluginVisibility map[string]string `yaml:"plugin_visibility,omitempty" json:"plugin_visibility,omitempty"`
|
||||
// box seals credentials (sources' api_key/headers, keys' key) at rest.
|
||||
// In-memory values are always plaintext; only the bytes on disk are sealed.
|
||||
// Wired by AttachSecretBox — Load leaves it nil so `-check` and tests stay
|
||||
// filesystem-free.
|
||||
box *SecretBox
|
||||
}
|
||||
|
||||
// Defaults applied to any source (YAML or runtime) that leaves a field unset.
|
||||
@ -127,6 +153,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)
|
||||
@ -192,10 +219,15 @@ func RemoveSourceFromYAML(path, name string) error {
|
||||
|
||||
// UpsertSourceInYAML adds or updates a source entry in the YAML config file.
|
||||
// Uses yaml.Node to preserve the rest of the file's comments and formatting.
|
||||
func UpsertSourceInYAML(path, name string, src Source) error {
|
||||
// When box is non-nil the source's credentials are sealed before writing, so a
|
||||
// newly added source never lands in the file as plaintext.
|
||||
func UpsertSourceInYAML(path, name string, src Source, box *SecretBox) error {
|
||||
if path == "" {
|
||||
return fmt.Errorf("config path is empty")
|
||||
}
|
||||
if err := sealSource(&src, box); err != nil {
|
||||
return err
|
||||
}
|
||||
data, err := os.ReadFile(path)
|
||||
if err != nil {
|
||||
return err
|
||||
@ -253,11 +285,17 @@ func UpsertSourceInYAML(path, name string, src Source) error {
|
||||
|
||||
// Save writes the current config back to the YAML file (preserving comments
|
||||
// via yaml.Node round-trip when possible, or full marshaling as fallback).
|
||||
// Credentials are sealed on the way out and the in-memory copy is restored to
|
||||
// plaintext afterwards, so callers keep working with usable values.
|
||||
func (c *Config) Save() error {
|
||||
if c.Path == "" {
|
||||
return fmt.Errorf("config path is empty")
|
||||
}
|
||||
if err := c.sealInPlace(c.box); err != nil {
|
||||
return err
|
||||
}
|
||||
out, err := yaml.Marshal(c)
|
||||
c.unsealAfterWrite(c.box)
|
||||
if err != nil {
|
||||
return fmt.Errorf("marshal config: %w", err)
|
||||
}
|
||||
@ -272,6 +310,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"
|
||||
}
|
||||
@ -352,24 +393,271 @@ type SourceTemplate struct {
|
||||
// (full management) or "user" (sees only its own key); Models is the allowed
|
||||
// model scope with per-model token quota (0 = unlimited).
|
||||
type GWKey struct {
|
||||
Key string `yaml:"key" json:"key"`
|
||||
Role string `yaml:"role" json:"role"`
|
||||
Name string `yaml:"name,omitempty" json:"name,omitempty"`
|
||||
Models []ModelScope `yaml:"models,omitempty" json:"models,omitempty"`
|
||||
Note string `yaml:"note,omitempty" json:"note,omitempty"`
|
||||
CreatedAt int64 `yaml:"created_at,omitempty" json:"created_at,omitempty"`
|
||||
Seed bool `yaml:"seed,omitempty" json:"seed,omitempty"` // true if migrated from config gateway_keys
|
||||
Key string `yaml:"key" json:"key"`
|
||||
Role string `yaml:"role" json:"role"`
|
||||
Name string `yaml:"name,omitempty" json:"name,omitempty"`
|
||||
// Auto is this key's own AUTO scheduling chain. It overrides the global
|
||||
// cfg.Auto for requests authenticated with this key; nil (or an empty
|
||||
// list) means "inherit the global chain", so keys created before per-key
|
||||
// chains existed keep working untouched.
|
||||
Auto []ModelScope `yaml:"auto,omitempty" json:"auto,omitempty"`
|
||||
Models []ModelScope `yaml:"models,omitempty" json:"models,omitempty"`
|
||||
Note string `yaml:"note,omitempty" json:"note,omitempty"`
|
||||
CreatedAt int64 `yaml:"created_at,omitempty" json:"created_at,omitempty"`
|
||||
Seed bool `yaml:"seed,omitempty" json:"seed,omitempty"` // true if migrated from config gateway_keys
|
||||
}
|
||||
|
||||
// ModelScope is one allowed model for a key, or one AUTO scheduling slot,
|
||||
// with an optional token quota and reset period. TokenQuota 0 = unlimited;
|
||||
// Period "" = never resets; "hour"/"week"/"month" are fixed windows; "nhour"
|
||||
// uses Hours as the window length in hours.
|
||||
// HasOwnAuto reports whether the key declares its own AUTO chain rather than
|
||||
// inheriting the global one. An explicitly empty list still counts as "no own
|
||||
// chain" so a user can clear an override back to the global default.
|
||||
func (k GWKey) HasOwnAuto() bool { return len(k.Auto) > 0 }
|
||||
|
||||
// BillingDSL holds declarative per-URL pricing profiles for the billing plugin.
|
||||
//
|
||||
// Profiles are the "let the user pick" axis: the same upstream URL can be
|
||||
// declared under multiple profiles, and the operator chooses which is active.
|
||||
// Switching recomputes and re-injects the billing plugin's prices table.
|
||||
type BillingDSL struct {
|
||||
// Active is the id of the profile in effect. Empty => the first profile.
|
||||
Active string `yaml:"active,omitempty" json:"active,omitempty"`
|
||||
Profiles []BillingProfile `yaml:"profiles,omitempty" json:"profiles,omitempty"`
|
||||
}
|
||||
|
||||
// BillingProfile is one named set of per-URL pricing rules.
|
||||
type BillingProfile struct {
|
||||
ID string `yaml:"id" json:"id"`
|
||||
Label string `yaml:"label,omitempty" json:"label,omitempty"`
|
||||
// Default applies when no rule matches: "free" or "unpriced".
|
||||
Default string `yaml:"default,omitempty" json:"default,omitempty"`
|
||||
Rules []BillingRule `yaml:"rules,omitempty" json:"rules,omitempty"`
|
||||
}
|
||||
|
||||
// BillingRule declares the pricing for everything under one URL.
|
||||
type BillingRule struct {
|
||||
// URL matches a source's base_url. "*" = any unmatched URL.
|
||||
URL string `yaml:"url" json:"url"`
|
||||
// Mode is one of: free, token, subscription, unpriced.
|
||||
Mode string `yaml:"mode" json:"mode"`
|
||||
Currency string `yaml:"currency,omitempty" json:"currency,omitempty"`
|
||||
// Subscription is the flat monthly cost (for mode=subscription).
|
||||
Subscription float64 `yaml:"subscription,omitempty" json:"subscription,omitempty"`
|
||||
// Models declares per-token prices when mode=token.
|
||||
Models map[string]BillingToken `yaml:"models,omitempty" json:"models,omitempty"`
|
||||
// Free is true when mode=free (kept for explicitness).
|
||||
// Peak applies to ALL token-priced requests under this URL.
|
||||
Peak *BillingPeak `yaml:"peak,omitempty" json:"peak,omitempty"`
|
||||
}
|
||||
|
||||
// BillingToken is the per-model price.
|
||||
type BillingToken struct {
|
||||
Prompt string `yaml:"prompt" json:"prompt"` // USD per million tokens, e.g. "0.15"
|
||||
Completion string `yaml:"completion" json:"completion"` // USD per million tokens
|
||||
CacheDiscount *float64 `yaml:"cache_discount,omitempty" json:"cache_discount,omitempty"`
|
||||
// Reasoning is the USD-per-million rate for thinking tokens, for the rare
|
||||
// provider that charges them differently from ordinary output. Omit it and
|
||||
// reasoning is billed at Completion — which is what OpenAI, DeepSeek and
|
||||
// Anthropic all do, and what every source on this deployment reports.
|
||||
//
|
||||
// It is a rate for a SUBSET of completion tokens, never an extra charge: the
|
||||
// billing plugin applies only the difference between this and Completion, so
|
||||
// setting it to a multiple of Completion reprices the thinking share rather
|
||||
// than billing it twice.
|
||||
Reasoning string `yaml:"reasoning,omitempty" json:"reasoning,omitempty"`
|
||||
// LongContext is the prompt-price step above a token threshold. Nil = one
|
||||
// flat rate, which is what a rule means when it says nothing.
|
||||
LongContext *BillingLongContext `yaml:"long_context,omitempty" json:"long_context,omitempty"`
|
||||
}
|
||||
|
||||
// BillingPeak declares peak-hour pricing.
|
||||
type BillingPeak struct {
|
||||
Multiplier float64 `yaml:"multiplier" json:"multiplier"`
|
||||
// Weekdays: 1=Mon..5=Fri (Lua os.date !%w, 0=Sun)
|
||||
Weekdays []int `yaml:"weekdays,omitempty" json:"weekdays,omitempty"`
|
||||
// Hours: each pair is [lo,hi] inclusive
|
||||
Hours [][2]int `yaml:"hours,omitempty" json:"hours,omitempty"`
|
||||
}
|
||||
|
||||
// BillingLongContext is a prompt-price step for long inputs.
|
||||
//
|
||||
// Every major provider charges the input tokens ABOVE a threshold at a higher
|
||||
// rate — DeepSeek doubles past 128k, OpenAI and Anthropic step up as well. A
|
||||
// flat prompt price therefore under-bills exactly the traffic an agent workload
|
||||
// generates, and it under-bills it silently: the number is plausible, the bill
|
||||
// adds up, and nothing says the rate was wrong.
|
||||
//
|
||||
// The step applies to the PROMPT only. Output and cache reads are billed at
|
||||
// their own rates regardless of input length, which is how providers publish
|
||||
// it — the higher rate is for the extra context you made the model read, not for
|
||||
// what it wrote back.
|
||||
type BillingLongContext struct {
|
||||
// Above is the token count the step starts at. A request with exactly this
|
||||
// many prompt tokens is still at the base rate; the step covers the tokens
|
||||
// BEYOND it. Inclusive-below matches how the providers document it.
|
||||
Above int `yaml:"above" json:"above"`
|
||||
// Prompt is the USD-per-million rate for the tokens above Above.
|
||||
Prompt string `yaml:"prompt" json:"prompt"`
|
||||
// Multiplier is an alternative to Prompt: a factor applied to the base rate
|
||||
// (DeepSeek's published step is exactly 2x). When both are set Prompt wins,
|
||||
// and setting neither is a validation error rather than a silent no-op.
|
||||
Multiplier float64 `yaml:"multiplier,omitempty" json:"multiplier,omitempty"`
|
||||
}
|
||||
|
||||
// Validate rejects a BillingDSL that could not work as written.
|
||||
func (b *BillingDSL) Validate() error {
|
||||
if b == nil {
|
||||
return nil
|
||||
}
|
||||
ids := map[string]bool{}
|
||||
for i := range b.Profiles {
|
||||
p := &b.Profiles[i]
|
||||
if p.ID == "" {
|
||||
return fmt.Errorf("billing profile [%d] has no id", i)
|
||||
}
|
||||
if ids[p.ID] {
|
||||
return fmt.Errorf("billing profile id %q appears more than once", p.ID)
|
||||
}
|
||||
ids[p.ID] = true
|
||||
if p.Default != "" && p.Default != "free" && p.Default != "unpriced" {
|
||||
return fmt.Errorf("profile %q: default must be free or unpriced, got %q", p.ID, p.Default)
|
||||
}
|
||||
for j := range p.Rules {
|
||||
r := &p.Rules[j]
|
||||
if r.URL == "" {
|
||||
return fmt.Errorf("profile %q rule [%d]: url is required", p.ID, j)
|
||||
}
|
||||
switch r.Mode {
|
||||
case "free", "token", "subscription", "unpriced":
|
||||
default:
|
||||
return fmt.Errorf("profile %q rule [%d]: mode must be free/token/subscription/unpriced, got %q", p.ID, j, r.Mode)
|
||||
}
|
||||
if r.Mode == "token" {
|
||||
for m, t := range r.Models {
|
||||
if t.Prompt == "" || t.Completion == "" {
|
||||
return fmt.Errorf("profile %q model %q: prompt and completion are required for token mode", p.ID, m)
|
||||
}
|
||||
if _, err := ParseFloatUSDPerM(t.Prompt); err != nil {
|
||||
return fmt.Errorf("profile %q model %q prompt: %w", p.ID, m, err)
|
||||
}
|
||||
if _, err := ParseFloatUSDPerM(t.Completion); err != nil {
|
||||
return fmt.Errorf("profile %q model %q completion: %w", p.ID, m, err)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
// Active must resolve to a profile id.
|
||||
if b.Active != "" && !ids[b.Active] {
|
||||
return fmt.Errorf("billing.active = %q but no profile has that id", b.Active)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// Resolve returns the profile with the given id; empty id (or an id that does
|
||||
// not exist) falls back to the FIRST profile. A config with profiles but no
|
||||
// active marker should mean "use the one I wrote first" rather than "price
|
||||
// nothing", which is what a nil here would silently do.
|
||||
func (b *BillingDSL) Resolve(id string) *BillingProfile {
|
||||
if b == nil || len(b.Profiles) == 0 {
|
||||
return nil
|
||||
}
|
||||
if id != "" {
|
||||
for i := range b.Profiles {
|
||||
if b.Profiles[i].ID == id {
|
||||
return &b.Profiles[i]
|
||||
}
|
||||
}
|
||||
}
|
||||
return &b.Profiles[0]
|
||||
}
|
||||
|
||||
// parseFloatUSDPerM parses a USD-per-million price string. Returns the
|
||||
// per-single-token rate the plugin expects (price / 1e6).
|
||||
//
|
||||
// Kept as a string in the DSL rather than a float: "0.15" survives round-trips
|
||||
// through YAML editors, and a malformed value must fail VALIDATION (with the
|
||||
// file and field named) rather than parse as 0 and under-bill silently — which
|
||||
// is exactly the failure mode that motivated the DSL.
|
||||
func ParseFloatUSDPerM(s string) (float64, error) {
|
||||
f, err := strconv.ParseFloat(strings.TrimSpace(s), 64)
|
||||
if err != nil {
|
||||
return 0, fmt.Errorf("not a number: %q", s)
|
||||
}
|
||||
if f < 0 {
|
||||
return 0, fmt.Errorf("negative price: %q", s)
|
||||
}
|
||||
return f / 1e6, nil
|
||||
}
|
||||
|
||||
// KeyQuota is retained only to carry a scope entry's caps through the admin
|
||||
// API. Quotas are per model, never per key: there is deliberately no key-wide
|
||||
// total, so exhausting one model's budget never blocks the others.
|
||||
type KeyQuota struct {
|
||||
TokenQuota int64 `json:"token_quota"`
|
||||
ReqQuota int64 `json:"req_quota"`
|
||||
Period string `json:"period"`
|
||||
Hours int64 `json:"hours"`
|
||||
}
|
||||
|
||||
// NormalizeRole defaults an empty role to "user", so a key can never end up in
|
||||
// a state where no role means "neither admin nor user".
|
||||
func NormalizeRole(role string) string {
|
||||
if role == "admin" {
|
||||
return "admin"
|
||||
}
|
||||
return "user"
|
||||
}
|
||||
|
||||
// Validate rejects a quota configuration that could not work as written. A
|
||||
// period is only meaningful when at least one cap is set, and a cap of zero
|
||||
// means "unlimited" rather than "deny everything", so those are the only two
|
||||
// things worth rejecting.
|
||||
func (q KeyQuota) Validate() error {
|
||||
if q.TokenQuota < 0 {
|
||||
return fmt.Errorf("token_quota must be >= 0 (0 = unlimited)")
|
||||
}
|
||||
if q.ReqQuota < 0 {
|
||||
return fmt.Errorf("req_quota must be >= 0 (0 = unlimited)")
|
||||
}
|
||||
if q.Hours < 0 {
|
||||
return fmt.Errorf("hours must be >= 0")
|
||||
}
|
||||
if q.TokenQuota > 0 || q.ReqQuota > 0 {
|
||||
if err := ValidatePeriod(q.Period, q.Hours); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// ValidatePeriod accepts the quota period vocabulary: "" (never resets),
|
||||
// "hour", "week", "month", or "nhour" with hours >= 1. An unknown period is
|
||||
// rejected rather than silently treated as "never resets", which would turn a
|
||||
// typo into an all-time quota — the opposite of what the operator typed.
|
||||
func ValidatePeriod(period string, hours int64) error {
|
||||
switch period {
|
||||
case "", "hour", "week", "month":
|
||||
return nil
|
||||
case "nhour":
|
||||
if hours < 1 {
|
||||
return fmt.Errorf("period %q needs hours >= 1", period)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
return fmt.Errorf("period must be one of \"\", hour, week, month, nhour (got %q)", period)
|
||||
}
|
||||
|
||||
// ModelScope is one allowed model for a key, or one AUTO scheduling slot.
|
||||
// Its TokenQuota and ReqQuota cap THAT entry only, independently of every
|
||||
// other entry on the same key: a model that runs out of budget stops being
|
||||
// served while the key's other models keep working. TokenQuota 0 / ReqQuota 0
|
||||
// = unlimited. Period "" = never resets; "hour"/"week"/"month" are fixed
|
||||
// windows; "nhour" uses Hours.
|
||||
type ModelScope struct {
|
||||
Model string `yaml:"model" json:"model"`
|
||||
Source string `yaml:"source,omitempty" json:"source,omitempty"` // optional: pin to one upstream source; "" = any source
|
||||
Tier int `yaml:"tier,omitempty" json:"tier,omitempty"`
|
||||
TokenQuota int64 `yaml:"token_quota" json:"token_quota"`
|
||||
ReqQuota int64 `yaml:"req_quota,omitempty" json:"req_quota,omitempty"`
|
||||
Period string `yaml:"period,omitempty" json:"period,omitempty"`
|
||||
Hours int64 `yaml:"hours,omitempty" json:"hours,omitempty"`
|
||||
}
|
||||
@ -388,6 +676,7 @@ func (m *ModelScope) UnmarshalJSON(b []byte) error {
|
||||
Source string `json:"source"`
|
||||
Tier int `json:"tier"`
|
||||
TokenQuota int64 `json:"token_quota"`
|
||||
ReqQuota int64 `json:"req_quota"`
|
||||
Period string `json:"period"`
|
||||
Hours int64 `json:"hours"`
|
||||
}
|
||||
@ -398,6 +687,7 @@ func (m *ModelScope) UnmarshalJSON(b []byte) error {
|
||||
m.Source = o.Source
|
||||
m.Tier = o.Tier
|
||||
m.TokenQuota = o.TokenQuota
|
||||
m.ReqQuota = o.ReqQuota
|
||||
m.Period = o.Period
|
||||
m.Hours = o.Hours
|
||||
return nil
|
||||
|
||||
317
internal/config/secret_config.go
Normal file
317
internal/config/secret_config.go
Normal file
@ -0,0 +1,317 @@
|
||||
package config
|
||||
|
||||
// Secret handling for config.yaml.
|
||||
//
|
||||
// config.yaml is 0644 world-readable by design (ops need to inspect it), so any
|
||||
// credential in it must not sit there in plaintext. Sources' api_key / headers
|
||||
// and gateway keys are therefore sealed at rest with the same SecretBox used by
|
||||
// the runtime store, and unsealed in memory at load time.
|
||||
//
|
||||
// The invariant that makes this safe: **in-memory values are always plaintext**,
|
||||
// and the enc:v1: prefix is what marks a file value as sealed. Read paths that
|
||||
// predate this (core.resolveSourceKey) already unseal, so only the write side and
|
||||
// the load-time normalize step are new.
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"log"
|
||||
"os"
|
||||
"strings"
|
||||
)
|
||||
|
||||
// sealSource seals one source's credentials in place (used by the YAML upsert
|
||||
// path, which is a package function and therefore has no Config to borrow a box
|
||||
// from).
|
||||
func sealSource(s *Source, box *SecretBox) error {
|
||||
if box == nil {
|
||||
return nil
|
||||
}
|
||||
if s.APIKey != "" && !strings.HasPrefix(s.APIKey, encPrefix) {
|
||||
v, err := box.Encrypt(s.APIKey)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
s.APIKey = v
|
||||
}
|
||||
for k, v := range s.Headers {
|
||||
if v == "" || strings.HasPrefix(v, encPrefix) {
|
||||
continue
|
||||
}
|
||||
e, err := box.Encrypt(v)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
s.Headers[k] = e
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// normalizeSecrets unseals every credential in the freshly parsed config so the
|
||||
// rest of the program only ever sees plaintext. A value without the enc:v1:
|
||||
// prefix is left untouched, which keeps hand-written plaintext configs working
|
||||
// (and is what a pre-encryption config file looks like).
|
||||
//
|
||||
// A value that carries the prefix but fails to decrypt is a hard error, not
|
||||
// something to paper over: returning the ciphertext (MustDecrypt's behavior)
|
||||
// would let the next Save re-seal it and turn one bad value into permanent,
|
||||
// compounding corruption. Losing the master key must be loud.
|
||||
func (c *Config) normalizeSecrets(box *SecretBox) error {
|
||||
if box == nil {
|
||||
return nil
|
||||
}
|
||||
for i := range c.Sources {
|
||||
s := &c.Sources[i]
|
||||
if strings.HasPrefix(s.APIKey, encPrefix) {
|
||||
v, err := box.Decrypt(s.APIKey)
|
||||
if err != nil {
|
||||
return fmt.Errorf("source %q api_key: %w", s.Name, err)
|
||||
}
|
||||
s.APIKey = v
|
||||
}
|
||||
for k, v := range s.Headers {
|
||||
if strings.HasPrefix(v, encPrefix) {
|
||||
d, err := box.Decrypt(v)
|
||||
if err != nil {
|
||||
return fmt.Errorf("source %q header %q: %w", s.Name, k, err)
|
||||
}
|
||||
s.Headers[k] = d
|
||||
}
|
||||
}
|
||||
}
|
||||
for i := range c.Keys {
|
||||
if strings.HasPrefix(c.Keys[i].Key, encPrefix) {
|
||||
v, err := box.Decrypt(c.Keys[i].Key)
|
||||
if err != nil {
|
||||
return fmt.Errorf("gateway key %q: %w", c.Keys[i].Name, err)
|
||||
}
|
||||
c.Keys[i].Key = v
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// sealInPlace replaces plaintext credentials with ciphertext for writing. It is
|
||||
// deliberately a separate step from Marshal: callers that need the plaintext
|
||||
// (auth comparisons, log output, returning a key to the operator who just
|
||||
// created it) must not be handed a sealed config by accident.
|
||||
func (c *Config) sealInPlace(box *SecretBox) error {
|
||||
if box == nil {
|
||||
return nil
|
||||
}
|
||||
for i := range c.Sources {
|
||||
s := &c.Sources[i]
|
||||
if s.APIKey != "" && !strings.HasPrefix(s.APIKey, encPrefix) {
|
||||
v, err := box.Encrypt(s.APIKey)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
s.APIKey = v
|
||||
}
|
||||
if len(s.Headers) > 0 {
|
||||
sealed := make(map[string]string, len(s.Headers))
|
||||
for k, v := range s.Headers {
|
||||
if v == "" || strings.HasPrefix(v, encPrefix) {
|
||||
sealed[k] = v
|
||||
continue
|
||||
}
|
||||
e, err := box.Encrypt(v)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
sealed[k] = e
|
||||
}
|
||||
s.Headers = sealed
|
||||
}
|
||||
}
|
||||
for i := range c.Keys {
|
||||
k := &c.Keys[i]
|
||||
if k.Key != "" && !strings.HasPrefix(k.Key, encPrefix) {
|
||||
v, err := box.Encrypt(k.Key)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
k.Key = v
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// unsealAfterWrite restores plaintext after a sealed marshal so the live process
|
||||
// keeps working on plaintext values (mirrors Store.persistLocked's dance).
|
||||
func (c *Config) unsealAfterWrite(box *SecretBox) {
|
||||
// Save just encrypted every value it can see, so a failure here is
|
||||
// impossible; ignore the error rather than panic in a write path.
|
||||
_ = c.normalizeSecrets(box)
|
||||
}
|
||||
|
||||
// hasPlaintextSecrets reports whether any credential in the config is still in
|
||||
// the clear. Used to decide whether a startup migration write is needed, and to
|
||||
// warn (without leaking values) when no master key is available.
|
||||
func (c *Config) hasPlaintextSecrets() bool {
|
||||
for _, s := range c.Sources {
|
||||
if s.APIKey != "" && !strings.HasPrefix(s.APIKey, encPrefix) {
|
||||
return true
|
||||
}
|
||||
for _, v := range s.Headers {
|
||||
if v != "" && !strings.HasPrefix(v, encPrefix) {
|
||||
return true
|
||||
}
|
||||
}
|
||||
}
|
||||
for _, k := range c.Keys {
|
||||
if k.Key != "" && !strings.HasPrefix(k.Key, encPrefix) {
|
||||
return true
|
||||
}
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
// countPlaintextSecrets returns how many credentials are still in the clear, for
|
||||
// an operator-facing migration log line that must not print the values.
|
||||
func (c *Config) countPlaintextSecrets() int {
|
||||
n := 0
|
||||
for _, s := range c.Sources {
|
||||
if s.APIKey != "" && !strings.HasPrefix(s.APIKey, encPrefix) {
|
||||
n++
|
||||
}
|
||||
for _, v := range s.Headers {
|
||||
if v != "" && !strings.HasPrefix(v, encPrefix) {
|
||||
n++
|
||||
}
|
||||
}
|
||||
}
|
||||
for _, k := range c.Keys {
|
||||
if k.Key != "" && !strings.HasPrefix(k.Key, encPrefix) {
|
||||
n++
|
||||
}
|
||||
}
|
||||
return n
|
||||
}
|
||||
|
||||
// NormalizeSecretsForRun unseals the loaded config and then seals it back on
|
||||
// disk if anything was still in the clear. Order matters: Load() read the file
|
||||
// with ciphertext still in place, so the unseal has to happen before the
|
||||
// registry (and any Save the startup path performs) sees the values.
|
||||
func (c *Config) NormalizeSecretsForRun(box *SecretBox) error {
|
||||
hadPlaintext, err := c.UnsealSecrets(box)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
return c.SealIfNeeded(hadPlaintext)
|
||||
}
|
||||
|
||||
// UnsealSecrets decrypts every sealed credential in memory and reports whether
|
||||
// the config ON DISK still held plaintext (i.e. whether a sealing write is
|
||||
// needed). It never writes; a failed decrypt (wrong master key) is returned so
|
||||
// the process refuses to start instead of running with unusable credentials.
|
||||
//
|
||||
// The return value must be computed BEFORE unsealing and from the disk state,
|
||||
// not from memory: after a Save the in-memory values are always plaintext, so a
|
||||
// "is anything plaintext?" test run afterwards is unconditionally true and a
|
||||
// caller would rewrite the file on every start. That was the actual behaviour -
|
||||
// migratePlaintextSecrets() claimed to be idempotent in a comment but rewrote
|
||||
// config.yaml on every boot.
|
||||
//
|
||||
// Callers that consume credentials (the provider registry, key seeding) must
|
||||
// unseal FIRST. Seeding compares cfg.Keys[i].Key against the plaintext
|
||||
// gateway_keys entries; running it while keys are still ciphertext made the
|
||||
// dedupe never match, so every restart appended another copy of the same admin
|
||||
// key (production accumulated four).
|
||||
func (c *Config) UnsealSecrets(box *SecretBox) (bool, error) {
|
||||
if box == nil {
|
||||
return false, nil
|
||||
}
|
||||
c.AttachSecretBox(box)
|
||||
hadPlaintext := c.hasPlaintextSecrets()
|
||||
if err := c.normalizeSecrets(box); err != nil {
|
||||
return hadPlaintext, err
|
||||
}
|
||||
return hadPlaintext, nil
|
||||
}
|
||||
|
||||
// SealIfNeeded writes the config back once if hadPlaintext reported that the
|
||||
// file still held clear-text credentials. When it is false the file is left
|
||||
// untouched, which is what makes startup a no-op for an already-sealed config.
|
||||
func (c *Config) SealIfNeeded(hadPlaintext bool) error {
|
||||
if !hadPlaintext || c.box == nil || c.Path == "" {
|
||||
return nil
|
||||
}
|
||||
n := c.countPlaintextSecrets()
|
||||
if err := c.Save(); err != nil {
|
||||
return err
|
||||
}
|
||||
log.Printf("[config] sealed %d plaintext credential(s) in %s", n, c.Path)
|
||||
return nil
|
||||
}
|
||||
|
||||
// AttachSecretBox wires the encryption box into the config so Save can seal
|
||||
// credentials. Kept as an explicit call (rather than a constructor argument) so
|
||||
// config.Load stays usable in contexts that have no filesystem secrets (tests,
|
||||
// `-check`).
|
||||
func (c *Config) AttachSecretBox(box *SecretBox) { c.box = box }
|
||||
|
||||
// SecretBox returns the wired encryption box, or nil when none is attached.
|
||||
func (c *Config) SecretBox() *SecretBox { return c.box }
|
||||
|
||||
// migratePlaintextSecrets seals any credential still in the clear and writes the
|
||||
// file once. Kept for callers that attach the box themselves; it decides from
|
||||
// the in-memory state, which is why UnsealSecrets + SealIfNeeded (which decide
|
||||
// from the on-disk state) are preferred on the startup path.
|
||||
func (c *Config) migratePlaintextSecrets() error {
|
||||
if c.box == nil || c.Path == "" {
|
||||
return nil
|
||||
}
|
||||
if !c.hasPlaintextSecrets() {
|
||||
return nil
|
||||
}
|
||||
n := c.countPlaintextSecrets()
|
||||
if err := c.Save(); err != nil {
|
||||
return err
|
||||
}
|
||||
log.Printf("[config] sealed %d plaintext credential(s) in %s", n, c.Path)
|
||||
return nil
|
||||
}
|
||||
|
||||
// PrintSecrets writes the config's credentials to stdout in the clear and
|
||||
// returns an error only when the file cannot be read or the master key does not
|
||||
// match. It is the operator counterpart to sealing-at-rest: with keys stored as
|
||||
// ciphertext, "which key is this source using" must still be answerable.
|
||||
//
|
||||
// It deliberately takes a path rather than a *Config so the caller cannot
|
||||
// accidentally hand it a config that has already been unsealed in memory, and it
|
||||
// never writes anything.
|
||||
func PrintSecrets(path string) error {
|
||||
cfg, err := Load(path)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
box, err := NewSecretBox(cfg.RuntimeFile)
|
||||
if err != nil {
|
||||
return fmt.Errorf("%w (is master.key present and intact?)", err)
|
||||
}
|
||||
if err := cfg.normalizeSecrets(box); err != nil {
|
||||
return err
|
||||
}
|
||||
w := os.Stdout
|
||||
fmt.Fprintf(w, "# %s — %d source(s), %d gateway key(s)\n", path, len(cfg.Sources), len(cfg.Keys))
|
||||
for _, s := range cfg.Sources {
|
||||
fmt.Fprintf(w, "source %-16s api_key=%s\n", s.Name, orNone(s.APIKey))
|
||||
for k, v := range s.Headers {
|
||||
if v != "" {
|
||||
fmt.Fprintf(w, "source %-16s header[%s]=%s\n", s.Name, k, v)
|
||||
}
|
||||
}
|
||||
}
|
||||
for _, k := range cfg.Keys {
|
||||
fmt.Fprintf(w, "key %-16s role=%-5s %s\n", k.Name, k.Role, k.Key)
|
||||
}
|
||||
fmt.Fprintf(w, "gateway_keys (legacy): %s\n", strings.Join(cfg.GatewayKeys, " "))
|
||||
return nil
|
||||
}
|
||||
|
||||
func orNone(s string) string {
|
||||
if s == "" {
|
||||
return "(unset)"
|
||||
}
|
||||
return s
|
||||
}
|
||||
323
internal/config/secret_config_test.go
Normal file
323
internal/config/secret_config_test.go
Normal file
@ -0,0 +1,323 @@
|
||||
package config
|
||||
|
||||
import (
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
"testing"
|
||||
)
|
||||
|
||||
// writeMasterKey plants a deterministic master.key next to the given runtime
|
||||
// file so tests exercise the real box instead of the env-var shortcut.
|
||||
func writeMasterKey(t *testing.T, runtimeFile string) *SecretBox {
|
||||
t.Helper()
|
||||
if err := os.WriteFile(filepath.Join(filepath.Dir(runtimeFile), "master.key"),
|
||||
[]byte(strings.Repeat("ab", 32)), 0600); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
box, err := NewSecretBox(runtimeFile)
|
||||
if err != nil {
|
||||
t.Fatalf("NewSecretBox: %v", err)
|
||||
}
|
||||
return box
|
||||
}
|
||||
|
||||
func TestSaveSealsCredentialsAndLoadUnseals(t *testing.T) {
|
||||
dir := t.TempDir()
|
||||
path := filepath.Join(dir, "config.yaml")
|
||||
runtime := filepath.Join(dir, "runtime.json")
|
||||
box := writeMasterKey(t, runtime)
|
||||
|
||||
cfg := &Config{
|
||||
Path: path,
|
||||
Listen: "127.0.0.1:0",
|
||||
Sources: []Source{{
|
||||
Name: "up",
|
||||
BaseURL: "http://up/v1",
|
||||
APIKey: "sk-plaintext-secret",
|
||||
Adapter: "openai",
|
||||
Headers: map[string]string{"X-Extra": "header-secret"},
|
||||
Models: []Model{{ID: "m", Kind: "chat"}},
|
||||
}},
|
||||
Keys: []GWKey{{Key: "sk-gw-plain", Role: "admin", Name: "admin"}},
|
||||
}
|
||||
cfg.AttachSecretBox(box)
|
||||
if err := cfg.Save(); err != nil {
|
||||
t.Fatalf("Save: %v", err)
|
||||
}
|
||||
|
||||
raw, err := os.ReadFile(path)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
text := string(raw)
|
||||
if strings.Contains(text, "sk-plaintext-secret") {
|
||||
t.Error("source api_key written in plaintext")
|
||||
}
|
||||
if strings.Contains(text, "header-secret") {
|
||||
t.Error("source header written in plaintext")
|
||||
}
|
||||
if strings.Contains(text, "sk-gw-plain") {
|
||||
t.Error("gateway key written in plaintext")
|
||||
}
|
||||
if n := strings.Count(text, encPrefix); n != 3 {
|
||||
t.Errorf("expected 3 sealed values, found %d", n)
|
||||
}
|
||||
|
||||
// The live in-memory config must still hold plaintext after Save.
|
||||
if cfg.Sources[0].APIKey != "sk-plaintext-secret" {
|
||||
t.Errorf("in-memory api_key mutated by Save: %q", cfg.Sources[0].APIKey)
|
||||
}
|
||||
if cfg.Keys[0].Key != "sk-gw-plain" {
|
||||
t.Errorf("in-memory gateway key mutated by Save: %q", cfg.Keys[0].Key)
|
||||
}
|
||||
|
||||
// A fresh load must hand back plaintext again.
|
||||
loaded, err := Load(path)
|
||||
if err != nil {
|
||||
t.Fatalf("Load: %v", err)
|
||||
}
|
||||
loaded.AttachSecretBox(box)
|
||||
loaded.normalizeSecrets(box)
|
||||
if loaded.Sources[0].APIKey != "sk-plaintext-secret" {
|
||||
t.Errorf("loaded api_key = %q, want plaintext", loaded.Sources[0].APIKey)
|
||||
}
|
||||
if loaded.Sources[0].Headers["X-Extra"] != "header-secret" {
|
||||
t.Errorf("loaded header = %q, want plaintext", loaded.Sources[0].Headers["X-Extra"])
|
||||
}
|
||||
if loaded.Keys[0].Key != "sk-gw-plain" {
|
||||
t.Errorf("loaded gateway key = %q, want plaintext", loaded.Keys[0].Key)
|
||||
}
|
||||
}
|
||||
|
||||
func TestPlaintextConfigStillLoads(t *testing.T) {
|
||||
dir := t.TempDir()
|
||||
path := filepath.Join(dir, "config.yaml")
|
||||
// A hand-written, pre-encryption config: no enc:v1: anywhere.
|
||||
body := `listen: 127.0.0.1:0
|
||||
sources:
|
||||
- name: legacy
|
||||
base_url: http://legacy/v1
|
||||
api_key: sk-written-by-hand
|
||||
adapter: openai
|
||||
models:
|
||||
- id: m
|
||||
kind: chat
|
||||
`
|
||||
if err := os.WriteFile(path, []byte(body), 0644); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
cfg, err := Load(path)
|
||||
if err != nil {
|
||||
t.Fatalf("Load: %v", err)
|
||||
}
|
||||
box := writeMasterKey(t, filepath.Join(dir, "runtime.json"))
|
||||
cfg.AttachSecretBox(box)
|
||||
cfg.normalizeSecrets(box)
|
||||
if cfg.Sources[0].APIKey != "sk-written-by-hand" {
|
||||
t.Errorf("plaintext config value changed: %q", cfg.Sources[0].APIKey)
|
||||
}
|
||||
}
|
||||
|
||||
// Sealing a config is a one-time migration: the FIRST run writes, every later
|
||||
// start must leave the file alone.
|
||||
//
|
||||
// This test previously called migratePlaintextSecrets() twice and compared
|
||||
// ModTime. That assertion was both weak and wrong:
|
||||
//
|
||||
// - migratePlaintextSecrets decides from IN-MEMORY state, and Save() ends by
|
||||
// unsealing memory so the running process keeps working. So every call sees
|
||||
// "plaintext present" and rewrites the file. Comparing ModTime hid this
|
||||
// because two writes inside one filesystem timestamp tick look identical —
|
||||
// the test flaked (~1 in 6) instead of failing.
|
||||
// - it also exercised a function production no longer calls on the startup
|
||||
// path, which now uses UnsealSecrets + SealIfNeeded (they decide from the
|
||||
// ON-DISK state).
|
||||
//
|
||||
// So the test targets that real path and compares file CONTENT, which cannot be
|
||||
// fooled by timestamp granularity.
|
||||
func TestSealingIsIdempotentAcrossStarts(t *testing.T) {
|
||||
dir := t.TempDir()
|
||||
path := filepath.Join(dir, "config.yaml")
|
||||
runtime := filepath.Join(dir, "runtime.json")
|
||||
box := writeMasterKey(t, runtime)
|
||||
|
||||
cfg := &Config{
|
||||
Path: path,
|
||||
Listen: "127.0.0.1:0",
|
||||
Sources: []Source{{Name: "up", BaseURL: "http://up/v1", APIKey: "sk-clear", Adapter: "openai", Models: []Model{{ID: "m"}}}},
|
||||
}
|
||||
cfg.AttachSecretBox(box)
|
||||
|
||||
hadPlaintext, err := cfg.UnsealSecrets(box)
|
||||
if err != nil {
|
||||
t.Fatalf("unseal: %v", err)
|
||||
}
|
||||
if !hadPlaintext {
|
||||
t.Fatal("a config written with a clear-text credential must report plaintext")
|
||||
}
|
||||
if err := cfg.SealIfNeeded(hadPlaintext); err != nil {
|
||||
t.Fatalf("seal: %v", err)
|
||||
}
|
||||
first, err := os.ReadFile(path)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if !strings.Contains(string(first), encPrefix) {
|
||||
t.Fatal("sealing did not encrypt the value on disk")
|
||||
}
|
||||
// In-memory must be plaintext so the running process keeps working.
|
||||
if cfg.Sources[0].APIKey != "sk-clear" {
|
||||
t.Errorf("in-memory api_key = %q, want plaintext", cfg.Sources[0].APIKey)
|
||||
}
|
||||
|
||||
// Simulate the next start: load from disk, unseal, seal-if-needed. The file
|
||||
// was already sealed, so nothing may be written.
|
||||
cfg2, err := Load(path)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
cfg2.AttachSecretBox(box)
|
||||
hadPlaintext2, err := cfg2.UnsealSecrets(box)
|
||||
if err != nil {
|
||||
t.Fatalf("second unseal: %v", err)
|
||||
}
|
||||
if hadPlaintext2 {
|
||||
t.Error("a sealed config reported plaintext — SealIfNeeded would rewrite " +
|
||||
"the file on every start")
|
||||
}
|
||||
if err := cfg2.SealIfNeeded(hadPlaintext2); err != nil {
|
||||
t.Fatalf("second seal: %v", err)
|
||||
}
|
||||
second, err := os.ReadFile(path)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if string(first) != string(second) {
|
||||
t.Error("the second start rewrote an already-sealed config; startup must be a no-op")
|
||||
}
|
||||
}
|
||||
|
||||
func TestSealedValueDoesNotDoubleEncrypt(t *testing.T) {
|
||||
dir := t.TempDir()
|
||||
box := writeMasterKey(t, filepath.Join(dir, "runtime.json"))
|
||||
cfg := &Config{Path: filepath.Join(dir, "config.yaml"), Sources: []Source{{Name: "a", BaseURL: "http://a"}}}
|
||||
cfg.AttachSecretBox(box)
|
||||
|
||||
once, err := box.Encrypt("sk-x")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
cfg.Sources[0].APIKey = once
|
||||
if err := cfg.Save(); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
got, err := box.Decrypt(cfg.Sources[0].APIKey)
|
||||
if err != nil {
|
||||
t.Fatalf("value became undecryptable: %v", err)
|
||||
}
|
||||
if got != "sk-x" {
|
||||
t.Errorf("round trip = %q, want sk-x", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestUpsertSourceInYAMLSealsCredentials(t *testing.T) {
|
||||
dir := t.TempDir()
|
||||
path := filepath.Join(dir, "config.yaml")
|
||||
box := writeMasterKey(t, filepath.Join(dir, "runtime.json"))
|
||||
|
||||
if err := os.WriteFile(path, []byte("listen: 127.0.0.1:0\nsources: []\n"), 0644); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
src := Source{Name: "new", BaseURL: "http://new/v1", APIKey: "sk-fresh", Adapter: "openai", Models: []Model{{ID: "m"}}}
|
||||
if err := UpsertSourceInYAML(path, src.Name, src, box); err != nil {
|
||||
t.Fatalf("UpsertSourceInYAML: %v", err)
|
||||
}
|
||||
raw, _ := os.ReadFile(path)
|
||||
if strings.Contains(string(raw), "sk-fresh") {
|
||||
t.Error("newly added source stored its api_key in plaintext")
|
||||
}
|
||||
if !strings.Contains(string(raw), encPrefix) {
|
||||
t.Error("newly added source was not sealed")
|
||||
}
|
||||
}
|
||||
|
||||
func TestCountPlaintextSecrets(t *testing.T) {
|
||||
cfg := &Config{
|
||||
Sources: []Source{
|
||||
{Name: "a", APIKey: encPrefix + "abc", Headers: map[string]string{"H": encPrefix + "x"}},
|
||||
{Name: "b", APIKey: "sk-in-clear", Headers: map[string]string{"H2": "clear-too"}},
|
||||
},
|
||||
Keys: []GWKey{{Key: "sk-gw-clear"}},
|
||||
}
|
||||
if got := cfg.countPlaintextSecrets(); got != 3 {
|
||||
t.Errorf("countPlaintextSecrets = %d, want 3", got)
|
||||
}
|
||||
if !cfg.hasPlaintextSecrets() {
|
||||
t.Error("hasPlaintextSecrets = false, want true")
|
||||
}
|
||||
cfg.Sources[1].APIKey = encPrefix + "y"
|
||||
cfg.Sources[1].Headers["H2"] = encPrefix + "z"
|
||||
cfg.Keys[0].Key = encPrefix + "k"
|
||||
if cfg.hasPlaintextSecrets() {
|
||||
t.Error("hasPlaintextSecrets = true after sealing everything")
|
||||
}
|
||||
}
|
||||
|
||||
// TestNormalizeSecretsFailsLoudlyOnWrongMasterKey is the negative case that
|
||||
// guards the migration: with a wrong key, a sealed value must NOT be handed
|
||||
// back as ciphertext for a later re-seal (that compounds corruption and turns
|
||||
// one lost key file into permanently unusable config).
|
||||
func TestNormalizeSecretsFailsLoudlyOnWrongMasterKey(t *testing.T) {
|
||||
dir := t.TempDir()
|
||||
path := filepath.Join(dir, "config.yaml")
|
||||
runtime := filepath.Join(dir, "runtime.json")
|
||||
good := writeMasterKey(t, runtime)
|
||||
|
||||
cfg := &Config{
|
||||
Path: path,
|
||||
Listen: "127.0.0.1:0",
|
||||
Sources: []Source{{Name: "up", BaseURL: "http://up/v1", APIKey: "sk-secret", Adapter: "openai", Models: []Model{{ID: "m"}}}},
|
||||
}
|
||||
cfg.AttachSecretBox(good)
|
||||
if err := cfg.Save(); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
sealedOnce, err := os.ReadFile(path)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
// A different master key must be rejected, loudly. NewSecretBox derives the
|
||||
// key path from the RUNTIME file's directory, so the wrong key has to live
|
||||
// in a different directory — otherwise it would read the good master.key.
|
||||
otherDir := t.TempDir()
|
||||
if err := os.WriteFile(filepath.Join(otherDir, "master.key"),
|
||||
[]byte(strings.Repeat("cd", 32)), 0600); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
badBox, err := NewSecretBox(filepath.Join(otherDir, "runtime.json"))
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
loaded, err := Load(path)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
loaded.AttachSecretBox(badBox)
|
||||
if err := loaded.NormalizeSecretsForRun(badBox); err == nil {
|
||||
t.Fatal("expected an error with the wrong master key, got nil")
|
||||
}
|
||||
|
||||
// Crucially: the file must be untouched — no second seal on top.
|
||||
after, err := os.ReadFile(path)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if string(after) != string(sealedOnce) {
|
||||
t.Error("config file was rewritten despite the decrypt failure")
|
||||
}
|
||||
if n := strings.Count(string(after), encPrefix); n != strings.Count(string(sealedOnce), encPrefix) {
|
||||
t.Errorf("sealed count changed %d -> %d (double encryption)", strings.Count(string(sealedOnce), encPrefix), n)
|
||||
}
|
||||
}
|
||||
@ -7,14 +7,17 @@ import (
|
||||
"crypto/rand"
|
||||
"encoding/hex"
|
||||
"fmt"
|
||||
"log"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"sort"
|
||||
"net/url"
|
||||
"strings"
|
||||
"sync"
|
||||
"sync/atomic"
|
||||
"time"
|
||||
|
||||
"llmsproxy/internal/billing"
|
||||
"llmsproxy/internal/config"
|
||||
"llmsproxy/internal/lua"
|
||||
"llmsproxy/internal/provider"
|
||||
@ -29,11 +32,18 @@ type Core struct {
|
||||
mu sync.Mutex
|
||||
cfg *config.Config
|
||||
vm *lua.VM
|
||||
plugins *lua.Plugins
|
||||
store *config.Store
|
||||
scheduler *scheduler.Scheduler
|
||||
registry *provider.Registry
|
||||
autoChain atomic.Pointer[scheduler.Chain] // chat AUTO chain
|
||||
autoImageChain atomic.Pointer[scheduler.Chain] // image-generation AUTO chain
|
||||
|
||||
// keyAutoChains caches per-key AUTO chains, keyed by the gateway key.
|
||||
// Entries are rebuilt by rebuildKeyAutoChains whenever sources or key
|
||||
// config change; AutoChainFor reads them without taking c.mu so request
|
||||
// handling stays lock-free like the global chain lookup.
|
||||
keyAutoChains atomic.Pointer[map[string]*scheduler.Chain]
|
||||
}
|
||||
|
||||
// New builds the core from a config file plus runtime overlay.
|
||||
@ -52,7 +62,74 @@ 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)
|
||||
// Plugin outbound HTTP is scoped to the hosts of configured sources. Set
|
||||
// before LoadDir so a plugin that requests on its first hook already sees
|
||||
// the real allowlist rather than the deny-all placeholder.
|
||||
c.plugins.SetOutboundHosts(sourceHosts(cfg))
|
||||
// Let the plugin layer validate access specs against the keys that exist.
|
||||
// A stored allowlist naming a key that does not exist refuses every caller,
|
||||
// which looks exactly like a broken endpoint.
|
||||
lua.SetKeyNameLookup(func() []string {
|
||||
// ListKeys is the structured view (bare `gateway_keys:` entries are
|
||||
// expanded into records with names); cfg.GatewayKeys is still []string.
|
||||
keys := c.ListKeys()
|
||||
out := make([]string, 0, len(keys))
|
||||
for _, k := range keys {
|
||||
if k.Name != "" {
|
||||
out = append(out, k.Name)
|
||||
}
|
||||
}
|
||||
return out
|
||||
})
|
||||
// Operator overrides win over the plugin's own declared audience.
|
||||
for name, aud := range cfg.PluginVisibility {
|
||||
c.plugins.SetPluginVisibility(name, aud)
|
||||
}
|
||||
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.applyBillingDSL()
|
||||
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
|
||||
// ciphertext (if it was sealed); unseal it NOW, before anything reads a
|
||||
// credential.
|
||||
//
|
||||
// ORDER IS LOAD-BEARING. These steps each consume secrets and must run
|
||||
// after the unseal:
|
||||
// - seedKeys compares cfg.Keys[i].Key against the plaintext
|
||||
// gateway_keys entries; with ciphertext keys the comparison never
|
||||
// matched and every restart appended another duplicate admin key
|
||||
// (production had four copies of the same admin key).
|
||||
// - rebuildRegistry hands cfg.Sources[i].APIKey to the providers; with
|
||||
// ciphertext it built every upstream client with "enc:v1:..." as its
|
||||
// bearer token.
|
||||
// Previously the unseal happened at the END of this function, and things
|
||||
// only appeared to work because seedKeys' Save() unsealed memory as a side
|
||||
// effect. Removing the redundant saves exposed the real ordering bug.
|
||||
cfg.AttachSecretBox(c.store.SecretBox())
|
||||
hadPlaintextSecrets, err := cfg.UnsealSecrets(c.store.SecretBox())
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("unseal secrets: %w", err)
|
||||
}
|
||||
if err := c.store.Load(); err != nil {
|
||||
return nil, fmt.Errorf("runtime store: %w", err)
|
||||
}
|
||||
@ -74,6 +151,12 @@ func NewFromConfig(cfg *config.Config) (*Core, error) {
|
||||
if err := c.seedPresetTemplates(); err != nil {
|
||||
return nil, fmt.Errorf("seed preset templates: %w", err)
|
||||
}
|
||||
// Seal any credential still in the clear in config.yaml. Startup writes
|
||||
// nothing when the file was already sealed (hadPlaintextSecrets is decided
|
||||
// from the on-disk state, before the unseal above).
|
||||
if err := cfg.SealIfNeeded(hadPlaintextSecrets); err != nil {
|
||||
return nil, fmt.Errorf("normalize secrets: %w", err)
|
||||
}
|
||||
return c, nil
|
||||
}
|
||||
|
||||
@ -192,6 +275,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 }
|
||||
@ -199,6 +287,68 @@ func (c *Core) Registry() *provider.Registry { return c.registry }
|
||||
// AutoChain returns the current AUTO scheduling chain.
|
||||
func (c *Core) AutoChain() *scheduler.Chain { return c.autoChain.Load() }
|
||||
|
||||
// AutoChainFor returns the AUTO chain to use for a request authenticated with
|
||||
// key. A key that declares its own chain gets that chain; every other key
|
||||
// inherits the global chain, so keys created before per-key chains existed
|
||||
// keep their current behaviour with no configuration change.
|
||||
//
|
||||
// It returns (chain, ok). ok is false only when the key has its own chain but
|
||||
// it compiled to nothing usable (e.g. every slot names a model that no longer
|
||||
// exists) — the caller must surface that rather than silently downgrading the
|
||||
// user to the global chain, which would be a confusing, invisible switch.
|
||||
func (c *Core) AutoChainFor(key string) (*scheduler.Chain, bool) {
|
||||
if m := c.keyAutoChains.Load(); m != nil {
|
||||
if ch, ok := (*m)[key]; ok {
|
||||
return ch, true
|
||||
}
|
||||
}
|
||||
return c.autoChain.Load(), true
|
||||
}
|
||||
|
||||
// rebuildKeyAutoChains recompiles every per-key AUTO chain from the current
|
||||
// config and provider registry, then swaps the cache in one shot so requests
|
||||
// never observe a half-built map. Keys without their own chain are absent
|
||||
// from the map and fall back to the global chain.
|
||||
func (c *Core) rebuildKeyAutoChains() {
|
||||
m := make(map[string]*scheduler.Chain)
|
||||
for _, k := range c.cfg.Keys {
|
||||
if !k.HasOwnAuto() {
|
||||
continue
|
||||
}
|
||||
m[k.Key] = c.chainForRules(k.Auto, false)
|
||||
}
|
||||
c.keyAutoChains.Store(&m)
|
||||
}
|
||||
|
||||
// SaveKeyAuto persists one key's own AUTO chain and recompiles the per-key
|
||||
// chain cache. Passing an empty list clears the override so the key inherits
|
||||
// the global chain again.
|
||||
func (c *Core) SaveKeyAuto(key string, entries []config.ModelScope) error {
|
||||
c.mu.Lock()
|
||||
defer c.mu.Unlock()
|
||||
idx := -1
|
||||
for i, k := range c.cfg.Keys {
|
||||
if k.Key == key {
|
||||
idx = i
|
||||
break
|
||||
}
|
||||
}
|
||||
if idx < 0 {
|
||||
return fmt.Errorf("key not found")
|
||||
}
|
||||
if err := ValidateScopeQuotas(entries); err != nil {
|
||||
return err
|
||||
}
|
||||
c.cfg.Keys[idx].Auto = cleanScopes(entries)
|
||||
if err := c.saveConfig(); err != nil {
|
||||
return err
|
||||
}
|
||||
// Rebuild the whole cache so a cleared override stops shadowing the
|
||||
// global chain for this key.
|
||||
c.rebuildKeyAutoChains()
|
||||
return nil
|
||||
}
|
||||
|
||||
// AutoImageChain returns the persisted image-generation AUTO chain snapshot.
|
||||
func (c *Core) AutoImageChain() *scheduler.Chain { return c.autoImageChain.Load() }
|
||||
|
||||
@ -238,10 +388,15 @@ func (c *Core) FindKey(key string) (config.GWKey, bool) {
|
||||
}
|
||||
|
||||
// CreateKey builds a new random gateway key and persists it to config.yaml.
|
||||
// Quotas live on the model scope entries, so a new key's budget is whatever
|
||||
// its scopes carry.
|
||||
func (c *Core) CreateKey(name, role string, models []config.ModelScope, note string) (config.GWKey, error) {
|
||||
c.mu.Lock()
|
||||
defer c.mu.Unlock()
|
||||
models = cleanScopes(models)
|
||||
if err := ValidateScopeQuotas(models); err != nil {
|
||||
return config.GWKey{}, err
|
||||
}
|
||||
key := make([]byte, 16)
|
||||
if _, err := rand.Read(key); err != nil {
|
||||
return config.GWKey{}, err
|
||||
@ -254,9 +409,7 @@ func (c *Core) CreateKey(name, role string, models []config.ModelScope, note str
|
||||
Note: note,
|
||||
CreatedAt: time.Now().Unix(),
|
||||
}
|
||||
if rec.Role == "" {
|
||||
rec.Role = "user"
|
||||
}
|
||||
rec.Role = config.NormalizeRole(rec.Role)
|
||||
c.cfg.Keys = append(c.cfg.Keys, rec)
|
||||
if err := c.saveConfig(); err != nil {
|
||||
return config.GWKey{}, err
|
||||
@ -264,10 +417,16 @@ func (c *Core) CreateKey(name, role string, models []config.ModelScope, note str
|
||||
return rec, nil
|
||||
}
|
||||
|
||||
// UpdateKey mutates a key's name/role/model scope and persists it.
|
||||
// UpdateKey mutates a key's name/role/model scope and persists it. The scope
|
||||
// entries carry their own quotas, so replacing the scope replaces the budgets.
|
||||
func (c *Core) UpdateKey(key, name, role string, models []config.ModelScope, note string) (config.GWKey, error) {
|
||||
c.mu.Lock()
|
||||
defer c.mu.Unlock()
|
||||
if models != nil {
|
||||
if err := ValidateScopeQuotas(models); err != nil {
|
||||
return config.GWKey{}, err
|
||||
}
|
||||
}
|
||||
for i, k := range c.cfg.Keys {
|
||||
if k.Key == key {
|
||||
if name != "" {
|
||||
@ -439,6 +598,11 @@ func (c *Core) mergedSources() []config.Source {
|
||||
}
|
||||
|
||||
// resolveSourceKey applies api_key_env and decrypts enc:v1: ciphertext.
|
||||
// Config values are already unsealed at startup (Config.NormalizeSecretsForRun),
|
||||
// so the decrypt branch is the belt-and-braces path for a source that arrived
|
||||
// already sealed through another route. A failure there leaves the ciphertext in
|
||||
// place, which makes the upstream reject the key loudly rather than sending an
|
||||
// empty Authorization header that might look like a config-less source.
|
||||
func (c *Core) resolveSourceKey(s config.Source) config.Source {
|
||||
if s.APIKeyEnv != "" {
|
||||
if v := os.Getenv(s.APIKeyEnv); v != "" {
|
||||
@ -447,7 +611,11 @@ func (c *Core) resolveSourceKey(s config.Source) config.Source {
|
||||
return s
|
||||
}
|
||||
if box := c.store.SecretBox(); box != nil && strings.HasPrefix(s.APIKey, "enc:v1:") {
|
||||
s.APIKey = box.MustDecrypt(s.APIKey)
|
||||
if v, err := box.Decrypt(s.APIKey); err == nil {
|
||||
s.APIKey = v
|
||||
} else {
|
||||
log.Printf("[core] source %s: api_key decrypt failed: %v", s.Name, err)
|
||||
}
|
||||
}
|
||||
return s
|
||||
}
|
||||
@ -470,26 +638,64 @@ func (c *Core) rebuildRegistry() error {
|
||||
}
|
||||
c.buildAutoChain()
|
||||
c.buildAutoImageChain()
|
||||
c.rebuildKeyAutoChains()
|
||||
// A source edit can add or remove an upstream host, which changes what a
|
||||
// plugin is allowed to reach. Keeping this beside the chain rebuilds means
|
||||
// one call site owns "config changed, refresh derived state".
|
||||
c.plugins.SetOutboundHosts(sourceHosts(c.cfg))
|
||||
return nil
|
||||
}
|
||||
|
||||
// sourceHosts extracts the host of every configured source's base_url — the
|
||||
// allowlist for plugin outbound HTTP. Unparseable URLs are skipped rather than
|
||||
// failing startup: a malformed base_url is already rejected when the provider is
|
||||
// built, and it must not take plugins down with it.
|
||||
func sourceHosts(cfg *config.Config) []string {
|
||||
if cfg == nil {
|
||||
return nil
|
||||
}
|
||||
out := make([]string, 0, len(cfg.Sources))
|
||||
for _, s := range cfg.Sources {
|
||||
if s.BaseURL == "" {
|
||||
continue
|
||||
}
|
||||
u, err := url.Parse(s.BaseURL)
|
||||
if err != nil || u.Hostname() == "" {
|
||||
continue
|
||||
}
|
||||
out = append(out, u.Hostname())
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
// buildAutoChain rebuilds the AUTO chain snapshot from config.yaml rules
|
||||
// against the current providers. Slot model ids are normalized to the exact
|
||||
// configured spelling (case-insensitive match), otherwise ModelFor would fall
|
||||
// back to the source's best chat model and cooldown/quota bookkeeping would
|
||||
// key on a name that never matches.
|
||||
func (c *Core) buildAutoChain() {
|
||||
c.autoChain.Store(c.chainForRules(c.cfg.Auto, false))
|
||||
}
|
||||
|
||||
// chainForRules compiles a chat AUTO chain from a rule list. kindFilter is
|
||||
// false for the chat chain (image models are skipped) and true for the image
|
||||
// chain (only image models participate). Per-key chains always compile as
|
||||
// chat chains: a user's key configures the models *they* talk to.
|
||||
func (c *Core) chainForRules(rules []config.ModelScope, imageOnly bool) *scheduler.Chain {
|
||||
prov := func(model, source string) scheduler.Provider {
|
||||
p := c.registry.ProviderForSlot(model, source)
|
||||
if p == nil {
|
||||
return nil
|
||||
}
|
||||
if m := p.ModelByID(model); m != nil && m.Kind == "image" {
|
||||
m := p.ModelByID(model)
|
||||
if imageOnly && (m == nil || m.Kind != "image") {
|
||||
return nil
|
||||
}
|
||||
if !imageOnly && m != nil && m.Kind == "image" {
|
||||
return nil
|
||||
}
|
||||
return p
|
||||
}
|
||||
rules := c.cfg.Auto
|
||||
sr := make([]scheduler.Rule, 0, len(rules))
|
||||
for _, e := range rules {
|
||||
model, source := e.Model, e.Source
|
||||
@ -498,7 +704,11 @@ func (c *Core) buildAutoChain() {
|
||||
model = exact
|
||||
}
|
||||
if m := p.ModelByID(model); m != nil && m.Kind == "image" {
|
||||
continue // image-kind models never join the chat AUTO chain
|
||||
if !imageOnly {
|
||||
continue // image-kind models never join the chat AUTO chain
|
||||
}
|
||||
} else if imageOnly {
|
||||
continue
|
||||
}
|
||||
if source == "" {
|
||||
source = p.Name()
|
||||
@ -513,7 +723,7 @@ func (c *Core) buildAutoChain() {
|
||||
Hours: e.Hours,
|
||||
})
|
||||
}
|
||||
c.autoChain.Store(scheduler.BuildChain(sr, prov))
|
||||
return scheduler.BuildChain(sr, prov)
|
||||
}
|
||||
|
||||
// buildAutoImageChain rebuilds the image-generation AUTO chain snapshot.
|
||||
@ -655,7 +865,7 @@ func (c *Core) AddSource(src config.Source) error {
|
||||
if err := normalizeSource(&src); err != nil {
|
||||
return err
|
||||
}
|
||||
if err := config.UpsertSourceInYAML(c.cfg.Path, src.Name, src); err != nil {
|
||||
if err := config.UpsertSourceInYAML(c.cfg.Path, src.Name, src, c.cfg.SecretBox()); err != nil {
|
||||
return err
|
||||
}
|
||||
// Update in-memory Sources so mergedSources() finds the entry.
|
||||
@ -752,7 +962,88 @@ func normalizeSource(s *config.Source) error {
|
||||
|
||||
// Close releases resources.
|
||||
func (c *Core) Close() {
|
||||
// Plugin state must be flushed BEFORE the VM stops. The saver's final write
|
||||
// reads each plugin's Lua tables; once vm.Stop() has closed those states the
|
||||
// read finds nothing and the last interval of accumulation is lost — which
|
||||
// is the exact failure this persistence was added to prevent.
|
||||
if c.plugins != nil {
|
||||
c.plugins.Close()
|
||||
}
|
||||
if c.vm != nil {
|
||||
c.vm.Stop()
|
||||
}
|
||||
}
|
||||
|
||||
// applyBillingDSL compiles the configured billing profile and injects it into
|
||||
// the billing plugin.
|
||||
//
|
||||
// Configured pricing replaces the hand-written JSON blob that used to be PUT
|
||||
// through the state API. That path was used twice in production and both times
|
||||
// it "succeeded" while being wrong: once without the required `prices` wrapper
|
||||
// (silently wiping the accumulated totals), once with a peak-window shape the
|
||||
// plugin does not read (peak traffic billed at off-peak rates, no error). A
|
||||
// typed config with load-time validation turns both mistakes into startup
|
||||
// errors naming the file and field.
|
||||
//
|
||||
// A DSL with no profile, or one that fails to compile, leaves the plugin's
|
||||
// prices untouched: billing keeps running on whatever it had, and the problem
|
||||
// is logged rather than becoming a startup failure — the gateway must forward
|
||||
// even when its cost accounting is misconfigured.
|
||||
func (c *Core) applyBillingDSL() {
|
||||
ps := c.plugins
|
||||
if ps == nil || ps.Count() == 0 {
|
||||
return
|
||||
}
|
||||
// The billing plugin may simply not be installed; a DSL without it is a
|
||||
// no-op, not an error.
|
||||
var have bool
|
||||
for _, row := range ps.List() {
|
||||
if row["name"] == "billing" && row["loaded"] == true {
|
||||
have = true
|
||||
break
|
||||
}
|
||||
}
|
||||
if !have {
|
||||
return
|
||||
}
|
||||
dsl := c.cfg.BillingDSL
|
||||
if dsl == nil || len(dsl.Profiles) == 0 {
|
||||
return
|
||||
}
|
||||
if err := dsl.Validate(); err != nil {
|
||||
log.Printf("[billing] config invalid, profile not applied: %v", err)
|
||||
return
|
||||
}
|
||||
profile := dsl.Resolve(dsl.Active)
|
||||
if profile == nil {
|
||||
log.Printf("[billing] no billing profile resolved (active=%q)", dsl.Active)
|
||||
return
|
||||
}
|
||||
prices, err := billing.Compile(profile, c.cfg.Sources)
|
||||
if err != nil {
|
||||
log.Printf("[billing] profile %q not applied: %v", profile.ID, err)
|
||||
return
|
||||
}
|
||||
if err := ps.SetState("billing", map[string]interface{}{"prices": prices}); err != nil {
|
||||
log.Printf("[billing] applying profile %q: %v", profile.ID, err)
|
||||
return
|
||||
}
|
||||
log.Printf("[billing] profile %q applied (%d rules, %d sources)", profile.ID, len(profile.Rules), len(c.cfg.Sources))
|
||||
}
|
||||
|
||||
// ValidateScopeQuotas checks every scope entry's caps before they are stored.
|
||||
// A typo in a period must be rejected at write time rather than silently
|
||||
// becoming a never-resetting budget — the opposite of what was typed.
|
||||
func ValidateScopeQuotas(entries []config.ModelScope) error {
|
||||
for _, e := range entries {
|
||||
if err := (config.KeyQuota{
|
||||
TokenQuota: e.TokenQuota,
|
||||
ReqQuota: e.ReqQuota,
|
||||
Period: e.Period,
|
||||
Hours: e.Hours,
|
||||
}).Validate(); err != nil {
|
||||
return fmt.Errorf("model %q: %w", e.Model, err)
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
189
internal/core/key_auto_concurrency_test.go
Normal file
189
internal/core/key_auto_concurrency_test.go
Normal file
@ -0,0 +1,189 @@
|
||||
package core
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"sync"
|
||||
"sync/atomic"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"llmsproxy/internal/config"
|
||||
)
|
||||
|
||||
// AutoChainFor sits on the request hot path and reads keyAutoChains with no
|
||||
// lock, while SaveKeyAuto / rebuildRegistry replace that map from the admin
|
||||
// path. A "call it twice" test never interleaves those, so it would pass even
|
||||
// if the map were mutated in place. These tests run both sides at once under
|
||||
// -race and assert what readers actually observe.
|
||||
//
|
||||
// The property: a reader sees a chain that is either the key's own or the
|
||||
// global one — never a torn or empty one. If rebuildKeyAutoChains built the
|
||||
// map in place instead of swapping a fresh pointer, readers would sometimes
|
||||
// see a partially-filled map and silently fall back to the global chain, i.e.
|
||||
// a user quietly losing their configured chain under load.
|
||||
|
||||
// concurrentChainConfig builds N keys, half with their own chain pointing at a
|
||||
// distinct model so a fallback to the global chain is distinguishable.
|
||||
func concurrentChainConfig(t *testing.T, n int) *config.Config {
|
||||
t.Helper()
|
||||
cfg := perKeyTestConfig(t)
|
||||
cfg.Keys = nil
|
||||
cfg.GatewayKeys = []string{"sk-gw-admin"}
|
||||
cfg.Auto = []config.ModelScope{{Model: "gpt-4o", Source: "s1", Tier: 1}}
|
||||
for i := 0; i < n; i++ {
|
||||
k := config.GWKey{
|
||||
Key: fmt.Sprintf("sk-gw-k%03d", i),
|
||||
Role: "user",
|
||||
Name: fmt.Sprintf("user%d", i),
|
||||
}
|
||||
if i%2 == 0 {
|
||||
k.Auto = []config.ModelScope{{Model: "gpt-4o-mini", Source: "s1", Tier: 1}}
|
||||
}
|
||||
cfg.Keys = append(cfg.Keys, k)
|
||||
}
|
||||
return cfg
|
||||
}
|
||||
|
||||
func TestAutoChainForConcurrentWithSaveKeyAuto(t *testing.T) {
|
||||
cfg := concurrentChainConfig(t, 40)
|
||||
c := newTestCore(t, cfg)
|
||||
|
||||
var stop atomic.Bool
|
||||
var wg sync.WaitGroup
|
||||
var bad atomic.Int64
|
||||
|
||||
for r := 0; r < 8; r++ {
|
||||
wg.Add(1)
|
||||
go func() {
|
||||
defer wg.Done()
|
||||
for !stop.Load() {
|
||||
for i := 0; i < 40; i++ {
|
||||
key := fmt.Sprintf("sk-gw-k%03d", i)
|
||||
ch, ok := c.AutoChainFor(key)
|
||||
if !ok || ch == nil || len(ch.Tiers) == 0 {
|
||||
bad.Add(1)
|
||||
continue
|
||||
}
|
||||
// Legal values are this key's own chain model (either
|
||||
// side of the writer's flip) or the global chain's. The
|
||||
// assertion is about the legal set, not a fixed value:
|
||||
// the writers legitimately change the chain underneath.
|
||||
got := ch.Tiers[0].Slots[0].Model
|
||||
if got != "gpt-4o" && got != "gpt-4o-mini" {
|
||||
bad.Add(1)
|
||||
}
|
||||
}
|
||||
}
|
||||
}()
|
||||
}
|
||||
|
||||
for w := 0; w < 2; w++ {
|
||||
wg.Add(1)
|
||||
go func() {
|
||||
defer wg.Done()
|
||||
model := "gpt-4o-mini"
|
||||
for i := 0; i < 25 && !stop.Load(); i++ {
|
||||
for j := 0; j < 40; j += 2 {
|
||||
if err := c.SaveKeyAuto(fmt.Sprintf("sk-gw-k%03d", j), []config.ModelScope{
|
||||
{Model: model, Source: "s1", Tier: 1},
|
||||
}); err != nil {
|
||||
bad.Add(1)
|
||||
return
|
||||
}
|
||||
}
|
||||
if model == "gpt-4o-mini" {
|
||||
model = "gpt-4o"
|
||||
} else {
|
||||
model = "gpt-4o-mini"
|
||||
}
|
||||
}
|
||||
}()
|
||||
}
|
||||
|
||||
time.Sleep(400 * time.Millisecond)
|
||||
stop.Store(true)
|
||||
wg.Wait()
|
||||
|
||||
if n := bad.Load(); n > 0 {
|
||||
t.Fatalf("%d reads saw a chain outside {own, global} — the cache is "+
|
||||
"being swapped in place instead of atomically", n)
|
||||
}
|
||||
}
|
||||
|
||||
// The registry rebuild path (AddSource / RemoveSource) also replaces the
|
||||
// cache, so readers must keep a usable chain throughout.
|
||||
//
|
||||
// Churned sources use a "churn" prefix on purpose: deleting a source a chain
|
||||
// references legitimately empties that chain (BuildChain drops unresolvable
|
||||
// slots), and asserting otherwise would be testing a fiction rather than the
|
||||
// code.
|
||||
func TestAutoChainForConcurrentWithSourceChanges(t *testing.T) {
|
||||
cfg := concurrentChainConfig(t, 20)
|
||||
c := newTestCore(t, cfg)
|
||||
|
||||
var stop atomic.Bool
|
||||
var wg sync.WaitGroup
|
||||
var bad atomic.Int64
|
||||
|
||||
for r := 0; r < 6; r++ {
|
||||
wg.Add(1)
|
||||
go func() {
|
||||
defer wg.Done()
|
||||
for !stop.Load() {
|
||||
for i := 0; i < 20; i++ {
|
||||
ch, ok := c.AutoChainFor(fmt.Sprintf("sk-gw-k%03d", i))
|
||||
if !ok || ch == nil || len(ch.Tiers) == 0 {
|
||||
bad.Add(1)
|
||||
continue
|
||||
}
|
||||
if got := ch.Tiers[0].Slots[0].Model; got != "gpt-4o" && got != "gpt-4o-mini" {
|
||||
bad.Add(1)
|
||||
}
|
||||
}
|
||||
}
|
||||
}()
|
||||
}
|
||||
|
||||
wg.Add(1)
|
||||
go func() {
|
||||
defer wg.Done()
|
||||
for n := 0; n < 30 && !stop.Load(); n++ {
|
||||
name := fmt.Sprintf("churn%d", n)
|
||||
_ = c.AddSource(config.Source{
|
||||
Name: name, BaseURL: "http://127.0.0.1:9/v1", Adapter: "openai",
|
||||
Models: []config.Model{{ID: "m-" + name, Priority: 10}}},
|
||||
)
|
||||
_ = c.RemoveSource(name)
|
||||
}
|
||||
}()
|
||||
|
||||
time.Sleep(500 * time.Millisecond)
|
||||
stop.Store(true)
|
||||
wg.Wait()
|
||||
|
||||
if n := bad.Load(); n > 0 {
|
||||
t.Fatalf("%d reads lost the chain while unrelated sources were churned", n)
|
||||
}
|
||||
}
|
||||
|
||||
// Many distinct keys: the lookup must stay a map hit and must not degrade into
|
||||
// returning the global chain as the key count grows.
|
||||
func TestAutoChainForManyKeys(t *testing.T) {
|
||||
cfg := concurrentChainConfig(t, 500)
|
||||
c := newTestCore(t, cfg)
|
||||
|
||||
for i := 0; i < 500; i++ {
|
||||
key := fmt.Sprintf("sk-gw-k%03d", i)
|
||||
ch, _ := c.AutoChainFor(key)
|
||||
if ch == nil || len(ch.Tiers) == 0 {
|
||||
t.Fatalf("key %s lost its chain with 500 keys configured", key)
|
||||
}
|
||||
want := "gpt-4o"
|
||||
if i%2 == 0 {
|
||||
want = "gpt-4o-mini"
|
||||
}
|
||||
if got := ch.Tiers[0].Slots[0].Model; got != want {
|
||||
t.Fatalf("key %s got %q, want %q", key, got, want)
|
||||
}
|
||||
}
|
||||
}
|
||||
221
internal/core/key_auto_test.go
Normal file
221
internal/core/key_auto_test.go
Normal file
@ -0,0 +1,221 @@
|
||||
package core
|
||||
|
||||
import (
|
||||
"path/filepath"
|
||||
"testing"
|
||||
|
||||
"llmsproxy/internal/config"
|
||||
"llmsproxy/internal/scheduler"
|
||||
)
|
||||
|
||||
// Per-key AUTO chains let an admin give one user their own scheduling chain
|
||||
// while every other key keeps the global one. These tests pin the three
|
||||
// properties that make that safe:
|
||||
//
|
||||
// 1. a key WITH its own chain gets it;
|
||||
// 2. a key WITHOUT one inherits the global chain (backward compatible);
|
||||
// 3. the override is actually persisted, so it survives a restart.
|
||||
//
|
||||
// A test that only checked (1) would pass even if AutoChainFor ignored the key
|
||||
// and always returned the global chain whenever the two happened to be equal,
|
||||
// so (2) uses deliberately *different* chains.
|
||||
|
||||
func perKeyTestConfig(t *testing.T) *config.Config {
|
||||
t.Helper()
|
||||
dir := t.TempDir()
|
||||
return &config.Config{
|
||||
Path: filepath.Join(dir, "config.yaml"),
|
||||
AdapterDir: filepath.Join(dir, "adapters"),
|
||||
RuntimeFile: filepath.Join(dir, "runtime.json"),
|
||||
Listen: "127.0.0.1:0",
|
||||
DefaultModel: "AUTO",
|
||||
GatewayKeys: []string{"sk-gw-admin"},
|
||||
Sources: []config.Source{
|
||||
{Name: "s1", BaseURL: "http://127.0.0.1:1/v1", Adapter: "openai",
|
||||
Models: []config.Model{{ID: "gpt-4o", Priority: 10}, {ID: "gpt-4o-mini", Priority: 5}}},
|
||||
},
|
||||
// Global chain points at gpt-4o.
|
||||
Auto: []config.ModelScope{{Model: "gpt-4o", Source: "s1", Tier: 1}},
|
||||
Keys: []config.GWKey{
|
||||
{Key: "sk-gw-admin", Role: "admin"},
|
||||
{Key: "sk-gw-inherits", Role: "user", Name: "inherits"},
|
||||
// This key's own chain points at a DIFFERENT model, so returning
|
||||
// the global chain by mistake is detectable.
|
||||
{Key: "sk-gw-own", Role: "user", Name: "own",
|
||||
Auto: []config.ModelScope{{Model: "gpt-4o-mini", Source: "s1", Tier: 1}}},
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
func firstSlotModel(t *testing.T, c *Core, key string) string {
|
||||
t.Helper()
|
||||
ch, _ := c.AutoChainFor(key)
|
||||
if ch == nil || len(ch.Tiers) == 0 {
|
||||
t.Fatalf("key %s: no chain tiers", key)
|
||||
}
|
||||
slots := ch.Tiers[0].Slots
|
||||
if len(slots) == 0 {
|
||||
t.Fatalf("key %s: tier 1 has no slots", key)
|
||||
}
|
||||
return slots[0].Model
|
||||
}
|
||||
|
||||
// A key that declares its own chain must be scheduled by it, not the global.
|
||||
func TestAutoChainForUsesKeyOwnChain(t *testing.T) {
|
||||
c := newTestCore(t, perKeyTestConfig(t))
|
||||
|
||||
if got := firstSlotModel(t, c, "sk-gw-own"); got != "gpt-4o-mini" {
|
||||
t.Fatalf("key with own chain must use it, got %q (want gpt-4o-mini)", got)
|
||||
}
|
||||
// Sanity: the global chain really is different, so the assertion above
|
||||
// cannot pass by accident.
|
||||
if got := firstSlotModel(t, c, "sk-gw-admin"); got != "gpt-4o" {
|
||||
t.Fatalf("admin must use global chain, got %q (want gpt-4o)", got)
|
||||
}
|
||||
}
|
||||
|
||||
// A key with no own chain inherits the global one. This is what keeps every
|
||||
// pre-existing deployment working with no config change.
|
||||
func TestAutoChainForInheritsGlobalWhenUnset(t *testing.T) {
|
||||
c := newTestCore(t, perKeyTestConfig(t))
|
||||
|
||||
if got := firstSlotModel(t, c, "sk-gw-inherits"); got != "gpt-4o" {
|
||||
t.Fatalf("key without own chain must inherit global, got %q (want gpt-4o)", got)
|
||||
}
|
||||
}
|
||||
|
||||
// A key that is not in the key table at all (e.g. a seed key) must also
|
||||
// inherit the global chain rather than erroring or returning nil.
|
||||
func TestAutoChainForUnknownKeyInheritsGlobal(t *testing.T) {
|
||||
c := newTestCore(t, perKeyTestConfig(t))
|
||||
|
||||
ch, ok := c.AutoChainFor("sk-gw-never-configured")
|
||||
if !ok || ch == nil || len(ch.Tiers) == 0 {
|
||||
t.Fatalf("unknown key must fall back to global chain, got ok=%v chain=%v", ok, ch)
|
||||
}
|
||||
}
|
||||
|
||||
// The per-key chain must survive a save/reload round trip: an admin configuring
|
||||
// a chain in the WebUI and restarting the gateway must not silently lose it.
|
||||
func TestSaveKeyAutoPersistsAndApplies(t *testing.T) {
|
||||
cfg := perKeyTestConfig(t)
|
||||
c := newTestCore(t, cfg)
|
||||
|
||||
if err := c.SaveKeyAuto("sk-gw-inherits", []config.ModelScope{
|
||||
{Model: "gpt-4o-mini", Source: "s1", Tier: 1},
|
||||
}); err != nil {
|
||||
t.Fatalf("SaveKeyAuto: %v", err)
|
||||
}
|
||||
if got := firstSlotModel(t, c, "sk-gw-inherits"); got != "gpt-4o-mini" {
|
||||
t.Fatalf("saved chain must apply immediately, got %q", got)
|
||||
}
|
||||
|
||||
// Reload from disk exactly like a restart does.
|
||||
reloaded, err := config.Load(cfg.Path)
|
||||
if err != nil {
|
||||
t.Fatalf("reload config: %v", err)
|
||||
}
|
||||
c2 := newTestCore(t, reloaded)
|
||||
if got := firstSlotModel(t, c2, "sk-gw-inherits"); got != "gpt-4o-mini" {
|
||||
t.Fatalf("per-key chain must survive restart, got %q", got)
|
||||
}
|
||||
// The other key must be unaffected by its neighbour's override.
|
||||
if got := firstSlotModel(t, c2, "sk-gw-admin"); got != "gpt-4o" {
|
||||
t.Fatalf("override must not leak to other keys, got %q", got)
|
||||
}
|
||||
}
|
||||
|
||||
// Clearing the override must return the key to the global chain, both in
|
||||
// memory and after a restart. This is what the WebUI's "use global" toggle
|
||||
// does, so a stale shadowing entry would silently pin the user to their old
|
||||
// chain forever.
|
||||
func TestSaveKeyAutoClearRestoresGlobal(t *testing.T) {
|
||||
cfg := perKeyTestConfig(t)
|
||||
c := newTestCore(t, cfg)
|
||||
|
||||
if err := c.SaveKeyAuto("sk-gw-own", nil); err != nil {
|
||||
t.Fatalf("clear: %v", err)
|
||||
}
|
||||
if got := firstSlotModel(t, c, "sk-gw-own"); got != "gpt-4o" {
|
||||
t.Fatalf("cleared key must inherit global again, got %q", got)
|
||||
}
|
||||
|
||||
reloaded, err := config.Load(cfg.Path)
|
||||
if err != nil {
|
||||
t.Fatalf("reload: %v", err)
|
||||
}
|
||||
c2 := newTestCore(t, reloaded)
|
||||
if got := firstSlotModel(t, c2, "sk-gw-own"); got != "gpt-4o" {
|
||||
t.Fatalf("clear must survive restart, got %q", got)
|
||||
}
|
||||
}
|
||||
|
||||
// Saving an unknown key must fail loudly rather than creating a shadow entry
|
||||
// that no request can ever match.
|
||||
func TestSaveKeyAutoUnknownKeyFails(t *testing.T) {
|
||||
c := newTestCore(t, perKeyTestConfig(t))
|
||||
|
||||
if err := c.SaveKeyAuto("sk-gw-nope", []config.ModelScope{{Model: "gpt-4o", Source: "s1"}}); err == nil {
|
||||
t.Fatal("SaveKeyAuto on unknown key must return an error")
|
||||
}
|
||||
}
|
||||
|
||||
// Per-key quota validation must run, exactly like the model scope path: an
|
||||
// override is another place an operator can write a bad budget.
|
||||
func TestSaveKeyAutoRejectsBadQuota(t *testing.T) {
|
||||
c := newTestCore(t, perKeyTestConfig(t))
|
||||
|
||||
err := c.SaveKeyAuto("sk-gw-inherits", []config.ModelScope{
|
||||
{Model: "gpt-4o", Source: "s1", TokenQuota: -5},
|
||||
})
|
||||
if err == nil {
|
||||
t.Fatal("negative token quota in a per-key auto chain must be rejected")
|
||||
}
|
||||
}
|
||||
|
||||
// Source changes must recompile per-key chains. A key whose chain names a
|
||||
// model that does not exist yet compiles to an empty chain (BuildChain drops
|
||||
// slots it cannot resolve — sending the request into a black hole would be
|
||||
// worse), and must start working as soon as the source appears.
|
||||
func TestKeyAutoChainsRebuildOnSourceChange(t *testing.T) {
|
||||
cfg := perKeyTestConfig(t)
|
||||
cfg.Keys = append(cfg.Keys, config.GWKey{
|
||||
Key: "sk-gw-late", Role: "user", Name: "late",
|
||||
Auto: []config.ModelScope{{Model: "new-model", Source: "s2", Tier: 1}},
|
||||
})
|
||||
c := newTestCore(t, cfg)
|
||||
|
||||
// s2 does not exist yet, so the chain has no usable slot.
|
||||
if ch, _ := c.AutoChainFor("sk-gw-late"); chainSlotsOf(ch) != 0 {
|
||||
t.Fatalf("chain naming an unknown model must compile empty, got %d slots",
|
||||
chainSlotsOf(ch))
|
||||
}
|
||||
// The other keys are unaffected by the broken one.
|
||||
if got := firstSlotModel(t, c, "sk-gw-admin"); got != "gpt-4o" {
|
||||
t.Fatalf("broken per-key chain must not affect others, got %q", got)
|
||||
}
|
||||
|
||||
// Add the source; the per-key chain must pick it up without a restart.
|
||||
cfg.Sources = append(cfg.Sources, config.Source{
|
||||
Name: "s2", BaseURL: "http://127.0.0.1:2/v1", Adapter: "openai",
|
||||
Models: []config.Model{{ID: "new-model", Priority: 10}},
|
||||
})
|
||||
if err := c.Reload(); err != nil {
|
||||
t.Fatalf("Reload: %v", err)
|
||||
}
|
||||
if got := firstSlotModel(t, c, "sk-gw-late"); got != "new-model" {
|
||||
t.Fatalf("per-key chain must resolve after the source is added, got %q", got)
|
||||
}
|
||||
}
|
||||
|
||||
// chainSlotsOf counts usable slots, tolerating a nil chain.
|
||||
func chainSlotsOf(ch *scheduler.Chain) int {
|
||||
if ch == nil {
|
||||
return 0
|
||||
}
|
||||
n := 0
|
||||
for _, tn := range ch.Tiers {
|
||||
n += len(tn.Slots)
|
||||
}
|
||||
return n
|
||||
}
|
||||
178
internal/core/startup_secrets_test.go
Normal file
178
internal/core/startup_secrets_test.go
Normal file
@ -0,0 +1,178 @@
|
||||
package core
|
||||
|
||||
import (
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"llmsproxy/internal/config"
|
||||
)
|
||||
|
||||
// Startup must be idempotent, and every credential consumer must see PLAINTEXT.
|
||||
//
|
||||
// Two bugs shared one root cause: NewFromConfig unsealed the config at the END,
|
||||
// after the steps that read credentials had already run.
|
||||
//
|
||||
// 1. seedKeys() deduped cfg.Keys[i].Key against the plaintext gateway_keys
|
||||
// entries. With ciphertext keys the comparison never matched, so every
|
||||
// restart appended another copy of the same admin key — production reached
|
||||
// four identical admin keys.
|
||||
// 2. rebuildRegistry() handed cfg.Sources[i].APIKey to provider.New, so with
|
||||
// ciphertext it built every upstream client with "enc:v1:..." as its bearer
|
||||
// token and every upstream call would 401.
|
||||
//
|
||||
// They masked each other: seedKeys' Save() unsealed memory as a side effect, so
|
||||
// bug 2 was invisible until the redundant saves were removed. These tests drive
|
||||
// the real entry point (core.New -> config.Load off disk), because building a
|
||||
// fresh in-memory config per attempt hides both bugs — that is exactly how the
|
||||
// first draft of this test failed to reproduce anything.
|
||||
|
||||
// writeSealedInstall creates a config on disk the way a real install looks
|
||||
// after its first run: credentials already sealed, gateway_keys still plaintext.
|
||||
func writeSealedInstall(t *testing.T, dir string) string {
|
||||
t.Helper()
|
||||
path := filepath.Join(dir, "config.yaml")
|
||||
cfg := &config.Config{
|
||||
Path: path,
|
||||
AdapterDir: filepath.Join(dir, "adapters"),
|
||||
RuntimeFile: filepath.Join(dir, "runtime.json"),
|
||||
Listen: "127.0.0.1:0",
|
||||
DefaultModel: "AUTO",
|
||||
GatewayKeys: []string{"sk-gw-seed-me"},
|
||||
Sources: []config.Source{{
|
||||
Name: "up", BaseURL: "http://127.0.0.1:1/v1", APIKey: "sk-upstream-secret",
|
||||
Adapter: "openai", Models: []config.Model{{ID: "gpt-4o", Priority: 10}},
|
||||
}},
|
||||
}
|
||||
c, err := NewFromConfig(cfg)
|
||||
if err != nil {
|
||||
t.Fatalf("initial install: %v", err)
|
||||
}
|
||||
c.Close()
|
||||
|
||||
raw, err := os.ReadFile(path)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if !strings.Contains(string(raw), "enc:v1:") {
|
||||
t.Fatalf("fixture is wrong: on-disk config holds no sealed credential, "+
|
||||
"so this test cannot detect ciphertext leaking into consumers:\n%s", raw)
|
||||
}
|
||||
return path
|
||||
}
|
||||
|
||||
// writeSealedInstallNoSeed is writeSealedInstall without gateway_keys, so no
|
||||
// key is ever seeded and no startup Save() can unseal the config as a side
|
||||
// effect. That isolation is what makes the ciphertext leak observable.
|
||||
func writeSealedInstallNoSeed(t *testing.T, dir string) string {
|
||||
t.Helper()
|
||||
path := filepath.Join(dir, "config.yaml")
|
||||
cfg := &config.Config{
|
||||
Path: path,
|
||||
AdapterDir: filepath.Join(dir, "adapters"),
|
||||
RuntimeFile: filepath.Join(dir, "runtime.json"),
|
||||
Listen: "127.0.0.1:0",
|
||||
DefaultModel: "AUTO",
|
||||
Sources: []config.Source{{
|
||||
Name: "up", BaseURL: "http://127.0.0.1:1/v1", APIKey: "sk-upstream-secret",
|
||||
Adapter: "openai", Models: []config.Model{{ID: "gpt-4o", Priority: 10}},
|
||||
}},
|
||||
}
|
||||
c, err := NewFromConfig(cfg)
|
||||
if err != nil {
|
||||
t.Fatalf("initial install: %v", err)
|
||||
}
|
||||
c.Close()
|
||||
raw, err := os.ReadFile(path)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if !strings.Contains(string(raw), "enc:v1:") {
|
||||
t.Fatalf("fixture is wrong: no sealed credential on disk:\n%s", raw)
|
||||
}
|
||||
return path
|
||||
}
|
||||
|
||||
// countKeysWithSecret loads the config, unseals it, and counts keys whose
|
||||
// plaintext equals one of the gateway_keys entries.
|
||||
func countKeysWithSecret(t *testing.T, path string) (total, matching int) {
|
||||
t.Helper()
|
||||
cfg, err := config.Load(path)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
box, err := config.NewSecretBox(cfg.RuntimeFile)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if _, err := cfg.UnsealSecrets(box); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
want := map[string]bool{}
|
||||
for _, g := range cfg.GatewayKeys {
|
||||
want[g] = true
|
||||
}
|
||||
for _, k := range cfg.Keys {
|
||||
total++
|
||||
if want[k.Key] {
|
||||
matching++
|
||||
}
|
||||
}
|
||||
return total, matching
|
||||
}
|
||||
|
||||
// Restarting must not grow the key list. Before the fix this went 2 -> 3 -> 4.
|
||||
func TestRestartDoesNotDuplicateSeededKeys(t *testing.T) {
|
||||
dir := t.TempDir()
|
||||
path := writeSealedInstall(t, dir)
|
||||
|
||||
if _, n := countKeysWithSecret(t, path); n != 1 {
|
||||
t.Fatalf("after install: %d keys carry the seeded secret, want 1", n)
|
||||
}
|
||||
|
||||
for i := 2; i <= 4; i++ {
|
||||
c, err := New(path) // production path: reads the file
|
||||
if err != nil {
|
||||
t.Fatalf("restart %d: %v", i, err)
|
||||
}
|
||||
c.Close()
|
||||
total, n := countKeysWithSecret(t, path)
|
||||
if n != 1 {
|
||||
t.Fatalf("restart %d: %d keys carry the seeded secret (total %d), want "+
|
||||
"exactly 1 — seedKeys is comparing ciphertext against plaintext and "+
|
||||
"appending a duplicate on every start", i, n, total)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Every provider must hold the real upstream credential, not its ciphertext.
|
||||
//
|
||||
// This deliberately runs with gateway_keys EMPTY. The buggy order was masked
|
||||
// whenever seedKeys had work to do, because its Save() unsealed memory as a
|
||||
// side effect — providers then happened to see plaintext. With no key to seed,
|
||||
// no Save runs, and the ciphertext reaches the registry directly. An earlier
|
||||
// version of this test kept gateway_keys populated and was insensitive: it
|
||||
// passed even with the bug present.
|
||||
func TestProvidersGetPlaintextCredentials(t *testing.T) {
|
||||
dir := t.TempDir()
|
||||
path := writeSealedInstallNoSeed(t, dir)
|
||||
|
||||
c, err := New(path)
|
||||
if err != nil {
|
||||
t.Fatalf("restart: %v", err)
|
||||
}
|
||||
defer c.Close()
|
||||
|
||||
for _, p := range c.registry.Providers() {
|
||||
key := p.Config().APIKey
|
||||
if strings.HasPrefix(key, "enc:v1:") {
|
||||
t.Errorf("provider %q holds CIPHERTEXT api_key %q — every upstream call "+
|
||||
"would 401", p.Name(), key)
|
||||
}
|
||||
if key != "sk-upstream-secret" {
|
||||
t.Errorf("provider %q api_key = %q, want the plaintext upstream secret",
|
||||
p.Name(), key)
|
||||
}
|
||||
}
|
||||
}
|
||||
@ -3,6 +3,7 @@ package gateway
|
||||
import (
|
||||
"encoding/csv"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"io"
|
||||
"log"
|
||||
"net/http"
|
||||
@ -78,6 +79,177 @@ type sourcePayload struct {
|
||||
MaxTokens int `json:"max_tokens"`
|
||||
MaxConcurrent int `json:"max_concurrent"`
|
||||
RPM int `json:"rpm"` // optional requests-per-minute cap, 0 = unlimited
|
||||
// ModelIDs is an additive-only shorthand: when set, the named models are
|
||||
// merged into an existing source instead of Models replacing the list.
|
||||
// It exists because "add one model" is the most common scripted edit and a
|
||||
// full Models list cannot be written without reading the source first.
|
||||
ModelIDs []string `json:"model_ids,omitempty"`
|
||||
|
||||
// The four fields below are POINTERS, and that is the whole point.
|
||||
//
|
||||
// POST/PUT is an upsert that rewrites the whole source, so any field the
|
||||
// payload cannot express is silently reset to its zero value. That already
|
||||
// destroyed api_key once (fixed with resolveAPIKey) and would equally
|
||||
// destroy proxy_url, api_key_env, timeout and queue_timeout.
|
||||
//
|
||||
// A plain string/duration cannot distinguish "the caller did not mention
|
||||
// this field" from "the caller asked for the empty value", and only the
|
||||
// former may inherit. So:
|
||||
//
|
||||
// nil → field absent from the request: keep the current value
|
||||
// &"x" → present: store exactly "x" (including "" to clear it)
|
||||
//
|
||||
// api_key deliberately keeps its older "empty means inherit" rule rather
|
||||
// than gaining a pointer: that rule is already published (v1.7.6) and
|
||||
// scripts depend on it. Changing it now would let a script that echoes an
|
||||
// empty api_key erase a live credential, which is the exact failure this
|
||||
// whole area exists to prevent.
|
||||
ProxyURL *string `json:"proxy_url,omitempty"`
|
||||
APIKeyEnv *string `json:"api_key_env,omitempty"`
|
||||
Timeout *string `json:"timeout,omitempty"` // duration string, e.g. "300s"
|
||||
QueueTimeout *string `json:"queue_timeout,omitempty"` // duration string, e.g. "90s"
|
||||
}
|
||||
|
||||
// optionalSourceFields carries the Source fields a partial payload must not
|
||||
// reset, each as a pointer so "absent" and "explicitly empty" stay distinct.
|
||||
//
|
||||
// nil → absent from the request: keep the current value
|
||||
// &"" → present: clear it
|
||||
// &"x" → present: store "x"
|
||||
//
|
||||
// The alternative — a plain value with "empty means inherit" — cannot express
|
||||
// "clear this field", so emptying the proxy input in the UI would silently keep
|
||||
// the old proxy. It is exactly why api_key is NOT modelled this way: for a
|
||||
// credential, failing to keep the old value is worse than failing to clear it,
|
||||
// and that rule is already published (v1.7.6). Two fields with opposite failure
|
||||
// modes therefore get opposite rules, and both are spelled out here.
|
||||
type optionalSourceFields struct {
|
||||
ProxyURL *string
|
||||
APIKeyEnv *string
|
||||
Timeout *string
|
||||
QueueTimeout *string
|
||||
}
|
||||
|
||||
// resolve overlays the four optional fields onto payload.
|
||||
//
|
||||
// A nil pointer means the request did not mention the field, so it keeps the
|
||||
// value already stored in cur (nil cur = a source being created, where the
|
||||
// payload's own zero value is correct). A non-nil pointer always wins, including
|
||||
// when it points at the empty string, which is how a UI form clears a field.
|
||||
//
|
||||
// Only these four fields are overlaid. Everything else in payload is taken as
|
||||
// sent: an upsert that inherited the whole record would make it impossible to
|
||||
// change anything.
|
||||
func (o optionalSourceFields) resolve(payload config.Source, cur *config.Source) (config.Source, error) {
|
||||
// cur is dereferenced exactly once here so the create path (nil) cannot
|
||||
// panic on the per-field lookups below.
|
||||
var storedProxyURL, storedKeyEnv string
|
||||
var storedTimeout, storedQueueTimeout time.Duration
|
||||
if cur != nil {
|
||||
storedProxyURL, storedKeyEnv = cur.ProxyURL, cur.APIKeyEnv
|
||||
storedTimeout, storedQueueTimeout = cur.Timeout, cur.QueueTimeout
|
||||
}
|
||||
pick := func(p *string, stored string) string {
|
||||
if p != nil {
|
||||
return *p
|
||||
}
|
||||
return stored
|
||||
}
|
||||
payload.ProxyURL = pick(o.ProxyURL, storedProxyURL)
|
||||
payload.APIKeyEnv = pick(o.APIKeyEnv, storedKeyEnv)
|
||||
if err := overlayDuration(&payload.Timeout, o.Timeout, "timeout", storedTimeout); err != nil {
|
||||
return payload, err
|
||||
}
|
||||
if err := overlayDuration(&payload.QueueTimeout, o.QueueTimeout, "queue_timeout", storedQueueTimeout); err != nil {
|
||||
return payload, err
|
||||
}
|
||||
return payload, nil
|
||||
}
|
||||
|
||||
// overlayDuration applies the absent/present rule to one duration field: a nil
|
||||
// pointer keeps stored, a set pointer replaces it (with "" / "0s" clearing).
|
||||
func overlayDuration(dst *time.Duration, p *string, field string, stored time.Duration) error {
|
||||
if p == nil {
|
||||
*dst = stored
|
||||
return nil
|
||||
}
|
||||
d, err := parseOptionalDuration(*p, field)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
*dst = d
|
||||
return nil
|
||||
}
|
||||
|
||||
// parseOptionalDuration parses a duration written as a Go duration string.
|
||||
// An empty string clears the field back to the config default (0), which is
|
||||
// what a UI form submitting an empty timeout box should mean. Parsing happens
|
||||
// only for fields the request actually mentions, so a typo can never surface
|
||||
// as a silent reset of something else.
|
||||
func parseOptionalDuration(v, field string) (time.Duration, error) {
|
||||
v = strings.TrimSpace(v)
|
||||
if v == "" {
|
||||
return 0, nil
|
||||
}
|
||||
d, err := time.ParseDuration(v)
|
||||
if err != nil {
|
||||
return 0, fmt.Errorf("%s %q is not a duration (use e.g. 120s, 5m, 1h)", field, v)
|
||||
}
|
||||
if d < 0 {
|
||||
return 0, fmt.Errorf("%s must be >= 0 (got %s)", field, v)
|
||||
}
|
||||
return d, nil
|
||||
}
|
||||
|
||||
// keepExistingAPIKey is the mask a client sends when it means "keep the
|
||||
// current credential". POST /api/sources is an upsert, so a script that edits
|
||||
// one field and echoes a placeholder back would otherwise persist that
|
||||
// placeholder as the live api_key — which fails closed as 401 on the next
|
||||
// request, long after the script exited successfully. See TestSourceUpsertKeepsAPIKey.
|
||||
const keepExistingAPIKey = "__KEEP__"
|
||||
|
||||
// resolveAPIKey returns the credential to store, inheriting the current one
|
||||
// when the payload does not supply a new one.
|
||||
//
|
||||
// Two distinct "no new credential" signals exist and they must not collapse:
|
||||
// - the explicit mask keepExistingAPIKey — always inherit.
|
||||
// - the empty string — inherit only for a source that already has a
|
||||
// credential. Empty is a legitimate value for a *new* credential-less
|
||||
// source (dshcodebuddy uses "dsh-local-noauth", local servers use ""),
|
||||
// so a fresh source with an empty key must stay empty rather than inherit
|
||||
// nothing and fail differently.
|
||||
func (g *Gateway) resolveAPIKey(name, want string, exists bool) string {
|
||||
if want == keepExistingAPIKey {
|
||||
return g.currentAPIKey(name)
|
||||
}
|
||||
if exists && want == "" {
|
||||
return g.currentAPIKey(name)
|
||||
}
|
||||
return want
|
||||
}
|
||||
|
||||
func (g *Gateway) currentAPIKey(name string) string {
|
||||
if s := g.sourceByName(name); s != nil {
|
||||
return s.APIKey
|
||||
}
|
||||
return ""
|
||||
}
|
||||
|
||||
// sourceByName returns a live view of one source, or nil when it does not
|
||||
// exist. It goes through core.Sources() so the key is already unsealed and
|
||||
// defaults applied, matching what the routes actually use.
|
||||
func (g *Gateway) sourceByName(name string) *config.Source {
|
||||
for _, s := range g.core.Sources() {
|
||||
if s.Name == name {
|
||||
ss := s
|
||||
return &ss
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (g *Gateway) sourceExists(name string) bool {
|
||||
return g.sourceByName(name) != nil
|
||||
}
|
||||
|
||||
func (g *Gateway) handleSourcesAPI(w http.ResponseWriter, r *http.Request) {
|
||||
@ -91,17 +263,34 @@ func (g *Gateway) handleSourcesAPI(w http.ResponseWriter, r *http.Request) {
|
||||
switch r.Method {
|
||||
case http.MethodGet:
|
||||
writeJSON(w, http.StatusOK, map[string]interface{}{"sources": g.core.Sources()})
|
||||
case http.MethodPost:
|
||||
case http.MethodPost, http.MethodPut:
|
||||
var p sourcePayload
|
||||
body, _ := io.ReadAll(r.Body)
|
||||
if err := json.Unmarshal(body, &p); err != nil {
|
||||
writeError(w, http.StatusBadRequest, "invalid_request", "invalid json: "+err.Error())
|
||||
return
|
||||
}
|
||||
// PUT addresses the source by path; the body name is only a fallback so
|
||||
// a single-field edit does not have to echo the name back. A body that
|
||||
// disagrees with the path is a caller bug, not something to silently
|
||||
// resolve — it usually means the wrong source is about to be written.
|
||||
if r.Method == http.MethodPut {
|
||||
if path == "" {
|
||||
writeError(w, http.StatusBadRequest, "invalid_request", "source name required in path")
|
||||
return
|
||||
}
|
||||
if p.Name != "" && p.Name != path {
|
||||
writeError(w, http.StatusBadRequest, "invalid_request",
|
||||
"source name in body ("+p.Name+") does not match path ("+path+")")
|
||||
return
|
||||
}
|
||||
p.Name = path
|
||||
}
|
||||
exists := g.sourceExists(p.Name)
|
||||
src := config.Source{
|
||||
Name: p.Name,
|
||||
BaseURL: p.BaseURL,
|
||||
APIKey: p.APIKey,
|
||||
APIKey: g.resolveAPIKey(p.Name, p.APIKey, exists),
|
||||
Adapter: p.Adapter,
|
||||
Endpoint: p.Endpoint,
|
||||
ImageEndpoint: p.ImageEndpoint,
|
||||
@ -113,6 +302,44 @@ func (g *Gateway) handleSourcesAPI(w http.ResponseWriter, r *http.Request) {
|
||||
MaxConcurrent: p.MaxConcurrent,
|
||||
RPM: p.RPM,
|
||||
}
|
||||
// The four upsert-fragile fields are applied last: a nil pointer means
|
||||
// "not in the request", so they are inherited from the stored source
|
||||
// instead of being reset to the payload's zero value. This runs whether
|
||||
// or not the request mentioned them, because resolve() itself decides
|
||||
// per field — passing them unconditionally keeps the rule in one place.
|
||||
src, err := optionalSourceFields{
|
||||
ProxyURL: p.ProxyURL,
|
||||
APIKeyEnv: p.APIKeyEnv,
|
||||
Timeout: p.Timeout,
|
||||
QueueTimeout: p.QueueTimeout,
|
||||
}.resolve(src, g.sourceByName(p.Name))
|
||||
if err != nil {
|
||||
writeError(w, http.StatusBadRequest, "source_error", err.Error())
|
||||
return
|
||||
}
|
||||
// model_ids is additive: "add these models" is the common scripted edit
|
||||
// and it must not require reading (and echoing) the whole list back.
|
||||
// A request that omits models entirely is therefore a pure add, not a
|
||||
// request to empty the list -- that distinction is the whole point.
|
||||
if len(p.ModelIDs) > 0 {
|
||||
if len(src.Models) == 0 {
|
||||
if existing := g.sourceByName(p.Name); existing != nil {
|
||||
src.Models = append(src.Models, existing.Models...)
|
||||
}
|
||||
}
|
||||
have := map[string]bool{}
|
||||
for _, m := range src.Models {
|
||||
have[m.ID] = true
|
||||
}
|
||||
for _, id := range p.ModelIDs {
|
||||
id = strings.TrimSpace(id)
|
||||
if id == "" || have[id] {
|
||||
continue
|
||||
}
|
||||
have[id] = true
|
||||
src.Models = append(src.Models, config.Model{ID: id, Kind: "chat"})
|
||||
}
|
||||
}
|
||||
if err := g.core.AddSource(src); err != nil {
|
||||
writeError(w, http.StatusBadRequest, "source_error", err.Error())
|
||||
return
|
||||
@ -206,6 +433,55 @@ func (g *Gateway) handleStatsAPI(w http.ResponseWriter, r *http.Request) {
|
||||
}
|
||||
}
|
||||
key := exportKey(r)
|
||||
// ?period=day|week|month|all switches the whole payload to a calendar
|
||||
// window (UTC) aggregated off the audit files, instead of the
|
||||
// since-process-start totals. The CSV exports below are unaffected: they
|
||||
// take an explicit from/to range and stream, so a period selector there
|
||||
// would only be a second way to spell the same bounds.
|
||||
if p := periodFromQuery(r.URL.Query()); p != PeriodAll {
|
||||
if !ValidPeriod(p) {
|
||||
writeError(w, http.StatusBadRequest, "bad_period",
|
||||
"period must be one of day, week, month, all")
|
||||
return
|
||||
}
|
||||
out := g.stats.PeriodSnapshot(p, key, time.Now())
|
||||
// The dashboard paints its records table from this payload's `records`
|
||||
// field (paintRecords(st.records)), so a period view that omits the key
|
||||
// renders an empty table no matter how much traffic the window had. The
|
||||
// aggregate above is period-scoped; the record list is the same
|
||||
// newest-first page the scroll handler pulls from /api/stats/records,
|
||||
// which is deliberately not period-filtered (the cursor, not a date
|
||||
// range, is its state) — that asymmetry is what the paging code
|
||||
// already assumes.
|
||||
page := g.stats.AuditPage("", limit, key)
|
||||
writeJSON(w, http.StatusOK, map[string]interface{}{
|
||||
"period": out.Period,
|
||||
"from": out.From,
|
||||
"total": out.Total,
|
||||
"by_key": out.ByKey,
|
||||
"by_model": out.Models,
|
||||
"by_source": out.Srcs,
|
||||
"by_status": out.Status,
|
||||
"buckets": out.Bucket,
|
||||
"truncated": out.Truncated,
|
||||
// "Active requests" is a live gauge (requests in flight right now),
|
||||
// not a period aggregate. Omitting it from the period response made
|
||||
// the KPI read `st.active || 0` → 0 on every day/week/month view while
|
||||
// the all-time view showed the truth, which read as "the gateway is
|
||||
// idle" precisely when it was busiest. A calendar window cannot
|
||||
// filter it, so it travels with every response unfiltered.
|
||||
"active": g.stats.Active(),
|
||||
"records": page.Records,
|
||||
// The paging cursor has to travel with the first page, otherwise
|
||||
// the scroll handler has nothing to page from and the table is
|
||||
// stuck at one screen forever.
|
||||
"next_cursor": page.Next,
|
||||
"has_more": page.HasMore,
|
||||
"rotated": page.Rotated,
|
||||
"key_names": g.keyNamesFor(),
|
||||
})
|
||||
return
|
||||
}
|
||||
if r.URL.Query().Get("export") == "csv" {
|
||||
from, _ := strconv.ParseInt(r.URL.Query().Get("from"), 10, 64)
|
||||
to, _ := strconv.ParseInt(r.URL.Query().Get("to"), 10, 64)
|
||||
@ -313,14 +589,39 @@ func (g *Gateway) handleStatsAPI(w http.ResponseWriter, r *http.Request) {
|
||||
return
|
||||
}
|
||||
snap := g.stats.Snapshot(limit, key)
|
||||
keyNames := map[string]string{}
|
||||
for _, k := range g.core.ListKeys() {
|
||||
keyNames[keyID(k.Key)] = k.Name
|
||||
}
|
||||
snap["key_names"] = keyNames
|
||||
snap["key_names"] = g.keyNamesFor()
|
||||
// Records come from the audit files via the same AuditPage the scroll
|
||||
// handler uses, not from Snapshot's in-memory ring. Two reasons:
|
||||
//
|
||||
// - Ordering. The ring is oldest-first, AuditPage is newest-first. The
|
||||
// first screen and the paged continuation have to agree, or scrolling
|
||||
// splices new rows in above old ones.
|
||||
// - Depth. The ring holds maxRecs entries (a few hundred); the audit
|
||||
// files hold everything, so "all" actually means all and the cursor
|
||||
// can keep paging.
|
||||
//
|
||||
// It also carries next_cursor, without which the table can never load a
|
||||
// second page.
|
||||
page := g.stats.AuditPage("", limit, key)
|
||||
snap["records"] = page.Records
|
||||
snap["next_cursor"] = page.Next
|
||||
snap["has_more"] = page.HasMore
|
||||
snap["rotated"] = page.Rotated
|
||||
writeJSON(w, http.StatusOK, snap)
|
||||
}
|
||||
|
||||
// keyNamesFor is the masked-id -> display-name map every stats payload needs.
|
||||
// It is keyed by keyID (the mask), not the raw key, because that is what the
|
||||
// aggregate rows carry — building it in one place stops the period branch and
|
||||
// the lifetime branch from drifting apart.
|
||||
func (g *Gateway) keyNamesFor() map[string]string {
|
||||
names := map[string]string{}
|
||||
for _, k := range g.core.ListKeys() {
|
||||
names[keyID(k.Key)] = k.Name
|
||||
}
|
||||
return names
|
||||
}
|
||||
|
||||
// handleStatsRecordsAPI pages the request records straight off the audit files.
|
||||
// The dashboard loads only its first screen and asks for the next page as the
|
||||
// user scrolls, so neither side holds the full history: the server keeps no
|
||||
|
||||
474
internal/gateway/apiv1.go
Normal file
474
internal/gateway/apiv1.go
Normal file
@ -0,0 +1,474 @@
|
||||
package gateway
|
||||
|
||||
// Agent-facing management API.
|
||||
//
|
||||
// The Web UI has always driven the gateway through /api/*, so the management
|
||||
// surface already exists. What it lacked was anything an agent could *rely* on:
|
||||
// no way to discover the contract, no single call that answers "what is the
|
||||
// current state", and per-endpoint response shapes that make scripting brittle.
|
||||
//
|
||||
// This file adds a versioned, self-describing facade under /api/v1 without
|
||||
// touching the existing endpoints, so the Web UI keeps working unchanged while
|
||||
// agents get a stable contract:
|
||||
//
|
||||
// GET /api/v1 — machine-readable index of every endpoint
|
||||
// GET /api/v1/overview — one call: sources + auto chain + keys + health
|
||||
// GET /api/v1/sources — sources, credentials masked
|
||||
// GET /api/v1/sources/{name} — one source
|
||||
// GET /api/v1/auto — scheduling chain + live slot states
|
||||
// GET /api/v1/keys — key metadata (never the secret)
|
||||
// GET /api/v1/models — every model id the gateway can route
|
||||
// GET /api/v1/health — per-source health, no auth needed beyond the key
|
||||
//
|
||||
// Auth is the same gateway key as everywhere else (Authorization: Bearer, or
|
||||
// ?api_key=, or the gw_key cookie). Read endpoints accept any role; writes
|
||||
// still require admin, enforced by the same middleware the UI goes through.
|
||||
|
||||
import (
|
||||
"errors"
|
||||
"net/http"
|
||||
"net/url"
|
||||
"sort"
|
||||
"strings"
|
||||
|
||||
"llmsproxy/internal/config"
|
||||
)
|
||||
|
||||
// apiV1Routes dispatches /api/v1/*. Kept separate from routes() so the original
|
||||
// switch stays readable and the versioned surface can grow on its own.
|
||||
func (g *Gateway) apiV1Routes(w http.ResponseWriter, r *http.Request) {
|
||||
path := strings.TrimPrefix(r.URL.Path, "/api/v1")
|
||||
path = strings.Trim(path, "/")
|
||||
|
||||
switch {
|
||||
case path == "":
|
||||
g.apiV1Index(w, r)
|
||||
case path == "overview":
|
||||
g.apiV1Overview(w, r)
|
||||
case path == "health":
|
||||
g.apiV1Health(w, r)
|
||||
case path == "models":
|
||||
g.apiV1Models(w, r)
|
||||
case path == "sources":
|
||||
if r.Method != http.MethodGet {
|
||||
writeError(w, http.StatusMethodNotAllowed, "method_not_allowed", "use GET")
|
||||
return
|
||||
}
|
||||
writeJSON(w, http.StatusOK, map[string]interface{}{"sources": maskSources(g.core.Sources())})
|
||||
case strings.HasPrefix(path, "sources/"):
|
||||
name, err := decodePathSegment(strings.TrimPrefix(path, "sources/"))
|
||||
if err != nil {
|
||||
writeError(w, http.StatusBadRequest, "invalid_request", err.Error())
|
||||
return
|
||||
}
|
||||
if r.Method != http.MethodGet {
|
||||
writeError(w, http.StatusMethodNotAllowed, "method_not_allowed",
|
||||
"use PUT/DELETE on /api/sources/{name} to modify a source")
|
||||
return
|
||||
}
|
||||
for _, s := range g.core.Sources() {
|
||||
if s.Name == name {
|
||||
// ?reveal=credentials is admin-only and is what the Web UI's
|
||||
// edit dialog uses: a form that round-trips a source must be
|
||||
// able to show the current key, otherwise saving an unrelated
|
||||
// field would blank it. Everything else stays masked.
|
||||
if r.URL.Query().Get("reveal") == "credentials" {
|
||||
if reqRole(r.Context()) != "admin" {
|
||||
writeError(w, http.StatusForbidden, "forbidden",
|
||||
"admin role required to reveal credentials")
|
||||
return
|
||||
}
|
||||
// The edit form round-trips this response, so it needs the credential in the
|
||||
// clear AND the durations — see maskSourceFields.
|
||||
writeJSON(w, http.StatusOK, map[string]interface{}{"source": maskSourceFields(s, true, true)})
|
||||
return
|
||||
}
|
||||
writeJSON(w, http.StatusOK, map[string]interface{}{"source": maskSource(s)})
|
||||
return
|
||||
}
|
||||
}
|
||||
writeError(w, http.StatusNotFound, "not_found", "no such source: "+name)
|
||||
case path == "auto":
|
||||
if r.Method != http.MethodGet {
|
||||
writeError(w, http.StatusMethodNotAllowed, "method_not_allowed", "use GET")
|
||||
return
|
||||
}
|
||||
writeJSON(w, http.StatusOK, map[string]interface{}{
|
||||
"rules": g.core.AutoRules(),
|
||||
"image_rules": g.core.AutoImageRules(),
|
||||
"states": g.core.AutoSlotStates(),
|
||||
})
|
||||
case path == "keys":
|
||||
if reqRole(r.Context()) != "admin" {
|
||||
writeError(w, http.StatusForbidden, "forbidden", "admin role required")
|
||||
return
|
||||
}
|
||||
if r.Method != http.MethodGet {
|
||||
writeError(w, http.StatusMethodNotAllowed, "method_not_allowed",
|
||||
"use POST/PUT/DELETE on /api/keys to manage keys")
|
||||
return
|
||||
}
|
||||
keys := g.core.ListKeys()
|
||||
out := make([]map[string]interface{}, 0, len(keys))
|
||||
for _, k := range keys {
|
||||
out = append(out, map[string]interface{}{
|
||||
"name": k.Name,
|
||||
"role": k.Role,
|
||||
"models": k.Models,
|
||||
"note": k.Note,
|
||||
"created_at": k.CreatedAt,
|
||||
"seed": k.Seed,
|
||||
// Quotas live on the scope entries (k.Models), echoed above;
|
||||
// there is deliberately no key-wide total.
|
||||
// The secret itself is never echoed. An operator that needs it
|
||||
// already has it from creation time or from config.yaml.
|
||||
"key_prefix": maskKey(k.Key),
|
||||
})
|
||||
}
|
||||
writeJSON(w, http.StatusOK, map[string]interface{}{"keys": out})
|
||||
default:
|
||||
g.apiV1Index(w, r)
|
||||
}
|
||||
}
|
||||
|
||||
// apiV1Index is the discovery document: it names every endpoint with its
|
||||
// method, auth requirement and purpose, so an agent does not have to scrape the
|
||||
// HTML to find out what it can call.
|
||||
func (g *Gateway) apiV1Index(w http.ResponseWriter, r *http.Request) {
|
||||
type ep struct {
|
||||
Method string `json:"method"`
|
||||
Path string `json:"path"`
|
||||
Auth string `json:"auth"`
|
||||
Summary string `json:"summary"`
|
||||
WriteEffect string `json:"writes,omitempty"`
|
||||
}
|
||||
index := map[string]interface{}{
|
||||
"api_version": "v1",
|
||||
"description": "llmsproxy management API. Authenticate with a gateway key: " +
|
||||
"'Authorization: Bearer <key>', '?api_key=<key>', or the gw_key cookie. " +
|
||||
"Read endpoints accept any role; mutations require an admin key.",
|
||||
"discovery": []ep{
|
||||
{Method: "GET", Path: "/api/v1", Auth: "any", Summary: "this document"},
|
||||
{Method: "GET", Path: "/api/v1/overview", Auth: "any",
|
||||
Summary: "current state in one call: sources, auto chain, key count, source health"},
|
||||
{Method: "GET", Path: "/api/v1/health", Auth: "any",
|
||||
Summary: "per-source health snapshot only"},
|
||||
{Method: "GET", Path: "/api/v1/models", Auth: "any",
|
||||
Summary: "every model id the gateway can route, grouped by owning source"},
|
||||
{Method: "GET", Path: "/api/v1/sources", Auth: "any",
|
||||
Summary: "all sources with credentials masked"},
|
||||
{Method: "GET", Path: "/api/v1/sources/{name}", Auth: "any", Summary: "one source"},
|
||||
{Method: "GET", Path: "/api/v1/auto", Auth: "any",
|
||||
Summary: "AUTO scheduling chain and live per-slot state"},
|
||||
{Method: "GET", Path: "/api/v1/keys", Auth: "admin",
|
||||
Summary: "gateway key metadata; secrets are never returned"},
|
||||
|
||||
{Method: "GET", Path: "/api/sources", Auth: "admin", Summary: "raw source list (includes api_key)"},
|
||||
{Method: "POST", Path: "/api/sources", Auth: "admin", Summary: "add or replace a source",
|
||||
WriteEffect: "writes config.yaml (api_key sealed at rest)"},
|
||||
{Method: "PUT", Path: "/api/sources/{name}", Auth: "admin", Summary: "update one source",
|
||||
WriteEffect: "writes config.yaml (api_key sealed at rest)"},
|
||||
{Method: "DELETE", Path: "/api/sources/{name}", Auth: "admin", Summary: "delete a source",
|
||||
WriteEffect: "writes config.yaml"},
|
||||
|
||||
{Method: "GET", Path: "/api/auto", Auth: "any", Summary: "scheduling chain (readable by any role)"},
|
||||
{Method: "PUT", Path: "/api/auto", Auth: "admin", Summary: "replace the scheduling chain",
|
||||
WriteEffect: "writes config.yaml (the AUTO chain lives here)"},
|
||||
|
||||
{Method: "GET", Path: "/api/keys", Auth: "admin", Summary: "gateway keys"},
|
||||
{Method: "POST", Path: "/api/keys", Auth: "admin", Summary: "create a gateway key",
|
||||
WriteEffect: "writes config.yaml"},
|
||||
{Method: "PUT", Path: "/api/keys/{key}", Auth: "admin",
|
||||
Summary: "update a gateway key (name, role, model scopes and their per-model quotas)",
|
||||
WriteEffect: "writes config.yaml"},
|
||||
{Method: "DELETE", Path: "/api/keys/{key}", Auth: "admin",
|
||||
Summary: "delete a gateway key — the path segment is the KEY itself, not its name",
|
||||
WriteEffect: "writes config.yaml"},
|
||||
|
||||
{Method: "GET", Path: "/api/status", Auth: "any", Summary: "per-source health detail"},
|
||||
{Method: "GET", Path: "/api/stats", Auth: "any", Summary: "usage aggregates"},
|
||||
{Method: "GET", Path: "/api/stats/records", Auth: "any", Summary: "paged request records"},
|
||||
{Method: "GET", Path: "/api/adapters", Auth: "admin", Summary: "installed Lua adapters"},
|
||||
{Method: "POST", Path: "/api/sources/discover", Auth: "admin",
|
||||
Summary: "probe a base_url+api_key for its model catalog",
|
||||
// Probing is read-only with respect to configuration: it writes
|
||||
// nothing, which is why it is safe to call repeatedly while the
|
||||
// operator is still typing a key.
|
||||
WriteEffect: "none",
|
||||
},
|
||||
{Method: "POST", Path: "/api/sources/discover/apply", Auth: "admin",
|
||||
Summary: "probe and write the source in one step",
|
||||
// Requires replace_models=true when the source already exists,
|
||||
// so re-running discovery cannot silently discard a hand-tuned
|
||||
// model list. An empty api_key keeps the stored credential.
|
||||
WriteEffect: "creates or replaces a source and its model list",
|
||||
},
|
||||
{Method: "GET", Path: "/api/source_templates", Auth: "admin", Summary: "source templates"},
|
||||
{Method: "POST", Path: "/v1/chat/completions", Auth: "any", Summary: "OpenAI-compatible inference"},
|
||||
},
|
||||
"conventions": map[string]interface{}{
|
||||
"errors": "{ \"error\": { \"type\": <code>, \"message\": <text> } }",
|
||||
"path_escape": "URL-encode source and key names; {name} is a single path segment",
|
||||
"idempotency": "POST /api/sources and PUT /api/sources/{name} both upsert by name; " +
|
||||
"PUT takes the name from the path and rejects a body name that disagrees",
|
||||
"partial_update": "api_key may be omitted or sent as the literal \"__KEEP__\" to inherit the " +
|
||||
"current credential; model_ids adds models to the existing list instead of replacing it, " +
|
||||
"so a one-field edit never needs to read the source first",
|
||||
"optional_fields": "proxy_url, api_key_env, timeout and queue_timeout use presence semantics: " +
|
||||
"OMITTED from the request keeps the stored value, present (even as \"\") overwrites it. " +
|
||||
"The upsert rewrites the whole source, so without this a one-field edit would silently reset " +
|
||||
"them — and api_key_env in particular is only visible as a credential failure on the NEXT " +
|
||||
"upstream call. api_key deliberately keeps the older \"empty means inherit\" rule instead, " +
|
||||
"because losing a credential breaks the source while losing a proxy only changes its route.",
|
||||
"durations": "timeout and queue_timeout are Go duration strings (\"300s\", \"5m\", \"1h\"); " +
|
||||
"the empty string clears them back to the config defaults",
|
||||
"config_truth": "all configuration lives in config.yaml; API writes are persisted immediately",
|
||||
"credentials": "credentials are masked by default. GET /api/v1/sources/{name}?reveal=credentials " +
|
||||
"returns them in the clear and is admin-only — the Web UI edit dialog uses it, because a form " +
|
||||
"that round-trips a source must show the current key or saving another field would blank it.",
|
||||
},
|
||||
}
|
||||
writeJSON(w, http.StatusOK, index)
|
||||
}
|
||||
|
||||
// apiV1Overview is the "what is the current state" call an agent makes first.
|
||||
// One round trip instead of five.
|
||||
func (g *Gateway) apiV1Overview(w http.ResponseWriter, r *http.Request) {
|
||||
if r.Method != http.MethodGet {
|
||||
writeError(w, http.StatusMethodNotAllowed, "method_not_allowed", "use GET")
|
||||
return
|
||||
}
|
||||
sources := g.core.Sources()
|
||||
names := make([]string, 0, len(sources))
|
||||
modelCount := 0
|
||||
for _, s := range sources {
|
||||
names = append(names, s.Name)
|
||||
modelCount += len(s.Models)
|
||||
}
|
||||
keys := g.core.ListKeys()
|
||||
adminCount, userCount := 0, 0
|
||||
for _, k := range keys {
|
||||
if k.Role == "admin" {
|
||||
adminCount++
|
||||
} else {
|
||||
userCount++
|
||||
}
|
||||
}
|
||||
writeJSON(w, http.StatusOK, map[string]interface{}{
|
||||
"api_version": "v1",
|
||||
"summary": map[string]interface{}{
|
||||
"source_count": len(sources),
|
||||
"model_count": modelCount,
|
||||
"gateway_key_count": len(keys),
|
||||
"admin_key_count": adminCount,
|
||||
"user_key_count": userCount,
|
||||
"auto_slots": len(g.core.AutoRules()),
|
||||
},
|
||||
"sources": maskSources(sources),
|
||||
"auto": g.core.AutoRules(),
|
||||
"auto_image": g.core.AutoImageRules(),
|
||||
"health": g.sourceHealthBrief(),
|
||||
"caller": map[string]interface{}{
|
||||
"role": reqRole(r.Context()),
|
||||
"key": maskKey(reqKey(r.Context())),
|
||||
},
|
||||
})
|
||||
}
|
||||
|
||||
func (g *Gateway) apiV1Health(w http.ResponseWriter, r *http.Request) {
|
||||
if r.Method != http.MethodGet {
|
||||
writeError(w, http.StatusMethodNotAllowed, "method_not_allowed", "use GET")
|
||||
return
|
||||
}
|
||||
writeJSON(w, http.StatusOK, map[string]interface{}{"sources": g.sourceHealthBrief()})
|
||||
}
|
||||
|
||||
// apiV1Models lists routable model ids per source, which is what an agent needs
|
||||
// to build a valid request — /v1/models flattens them and hides the owner.
|
||||
func (g *Gateway) apiV1Models(w http.ResponseWriter, r *http.Request) {
|
||||
if r.Method != http.MethodGet {
|
||||
writeError(w, http.StatusMethodNotAllowed, "method_not_allowed", "use GET")
|
||||
return
|
||||
}
|
||||
bySource := map[string][]string{}
|
||||
var all []string
|
||||
for _, s := range g.core.Sources() {
|
||||
for _, m := range s.Models {
|
||||
bySource[s.Name] = append(bySource[s.Name], m.ID)
|
||||
all = append(all, m.ID)
|
||||
}
|
||||
}
|
||||
sort.Strings(all)
|
||||
all = dedupeStrings(all)
|
||||
writeJSON(w, http.StatusOK, map[string]interface{}{
|
||||
"count": len(all),
|
||||
"models": all,
|
||||
"by_source": bySource,
|
||||
"note": "request a model as \"<source>:<model>\" to pin one source, or the bare id to let the gateway choose",
|
||||
})
|
||||
}
|
||||
|
||||
// ---- helpers ----
|
||||
|
||||
// maskSources returns sources with credentials replaced by a presence marker.
|
||||
// The UI still uses /api/sources (which returns real values, because its edit
|
||||
// form round-trips them); this is the safe view for programmatic callers.
|
||||
func maskSources(srcs []config.Source) []map[string]interface{} {
|
||||
out := make([]map[string]interface{}, 0, len(srcs))
|
||||
for _, s := range srcs {
|
||||
out = append(out, maskSource(s))
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
func maskSource(s config.Source) map[string]interface{} {
|
||||
return maskSourceFields(s, false, false)
|
||||
}
|
||||
|
||||
// maskSourceFields builds the API view of a source.
|
||||
//
|
||||
// revealCredentials swaps the masked api_key for the real one — admin-only, and
|
||||
// used by the WebUI edit dialog, whose form has to round-trip the whole source
|
||||
// or saving an unrelated field would blank the key.
|
||||
//
|
||||
// exposeDurations adds the two durations as strings, matching what the write
|
||||
// path accepts. It rides along with the credential reveal because both exist for
|
||||
// the same reason: config.Source tags Timeout/QueueTimeout `json:"-"`, so a
|
||||
// plain marshal of the struct omits them. A form that cannot SEE the stored
|
||||
// timeout would clear it on every save, since the timeout box is always sent.
|
||||
func maskSourceFields(s config.Source, revealCredentials, exposeDurations bool) map[string]interface{} {
|
||||
key := maskKey(s.APIKey)
|
||||
if revealCredentials {
|
||||
key = s.APIKey
|
||||
}
|
||||
m := map[string]interface{}{
|
||||
"name": s.Name,
|
||||
"base_url": s.BaseURL,
|
||||
"adapter": s.Adapter,
|
||||
"endpoint": s.Endpoint,
|
||||
"image_endpoint": s.ImageEndpoint,
|
||||
"api_key": key,
|
||||
"api_key_set": s.APIKey != "",
|
||||
"models": s.Models,
|
||||
"headers": maskHeaders(s.Headers),
|
||||
"proxy_url": s.ProxyURL,
|
||||
"api_key_env": s.APIKeyEnv,
|
||||
"meta": s.Meta,
|
||||
"temperature": s.Temperature,
|
||||
"max_tokens": s.MaxTokens,
|
||||
"max_concurrent": s.MaxConcurrent,
|
||||
"rpm": s.RPM,
|
||||
}
|
||||
if exposeDurations {
|
||||
m["timeout"] = s.Timeout.String()
|
||||
m["queue_timeout"] = s.QueueTimeout.String()
|
||||
}
|
||||
return m
|
||||
}
|
||||
|
||||
func maskKey(k string) string {
|
||||
if k == "" {
|
||||
return ""
|
||||
}
|
||||
if strings.HasPrefix(k, "enc:v1:") {
|
||||
return "(sealed)"
|
||||
}
|
||||
if len(k) <= 10 {
|
||||
return k[:2] + "…"
|
||||
}
|
||||
return k[:6] + "…" + k[len(k)-4:]
|
||||
}
|
||||
|
||||
func maskHeaders(h map[string]string) map[string]string {
|
||||
if h == nil {
|
||||
return nil
|
||||
}
|
||||
out := make(map[string]string, len(h))
|
||||
for k, v := range h {
|
||||
out[k] = maskKey(v)
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
// sourceHealthBrief reuses the same registry status the UI's /api/status shows,
|
||||
// so the agent view and the UI view cannot drift apart. Recent-traffic counters
|
||||
// come from the same stats window, because a source that is actually serving
|
||||
// traffic must never look down just because a probe was rate-limited.
|
||||
func (g *Gateway) sourceHealthBrief() []map[string]interface{} {
|
||||
sts := g.core.Registry().Status()
|
||||
recent := g.stats.SourceRecent(300)
|
||||
avgs := g.stats.SourceAverages(300)
|
||||
for i := range sts {
|
||||
if v, ok := recent[sts[i].Name]; ok {
|
||||
sts[i].RecentOK = v[0]
|
||||
sts[i].RecentErr = v[1]
|
||||
}
|
||||
if a, ok := avgs[sts[i].Name]; ok {
|
||||
sts[i].AvgFirstByteMs = a.AvgFirstByteMs
|
||||
sts[i].AvgTokPerS = a.AvgTokPerS
|
||||
}
|
||||
}
|
||||
out := make([]map[string]interface{}, 0, len(sts))
|
||||
for _, s := range sts {
|
||||
row := map[string]interface{}{
|
||||
"name": s.Name,
|
||||
"adapter": s.Adapter,
|
||||
"healthy": s.Healthy,
|
||||
"available": s.Available,
|
||||
"live_available": s.LiveAvailable,
|
||||
"model_count": len(s.Models),
|
||||
}
|
||||
if s.LastError != "" {
|
||||
row["last_error"] = s.LastError
|
||||
}
|
||||
if s.LastChecked > 0 {
|
||||
row["last_checked"] = s.LastChecked
|
||||
}
|
||||
if s.RecentOK > 0 || s.RecentErr > 0 {
|
||||
row["recent_ok"] = s.RecentOK
|
||||
row["recent_err"] = s.RecentErr
|
||||
}
|
||||
if s.AvgFirstByteMs > 0 {
|
||||
row["avg_first_byte_ms"] = s.AvgFirstByteMs
|
||||
}
|
||||
if s.AvgTokPerS > 0 {
|
||||
row["avg_tok_per_s"] = s.AvgTokPerS
|
||||
}
|
||||
if s.FailCount > 0 {
|
||||
row["fail_count"] = s.FailCount
|
||||
}
|
||||
if s.Permanent {
|
||||
row["permanent"] = true
|
||||
}
|
||||
out = append(out, row)
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
// decodePathSegment URL-decodes one path segment and rejects an empty result,
|
||||
// so a name containing a slash cannot be silently mis-resolved.
|
||||
func decodePathSegment(seg string) (string, error) {
|
||||
s, err := url.PathUnescape(seg)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
if s == "" {
|
||||
return "", errEmptySegment
|
||||
}
|
||||
return s, nil
|
||||
}
|
||||
|
||||
func dedupeStrings(in []string) []string {
|
||||
out := in[:0]
|
||||
var last string
|
||||
for i, s := range in {
|
||||
if i == 0 || s != last {
|
||||
out = append(out, s)
|
||||
}
|
||||
last = s
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
// errEmptySegment marks a path like /api/v1/sources/ with no name after it.
|
||||
var errEmptySegment = errors.New("empty path segment")
|
||||
299
internal/gateway/apiv1_test.go
Normal file
299
internal/gateway/apiv1_test.go
Normal file
@ -0,0 +1,299 @@
|
||||
package gateway
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"llmsproxy/internal/config"
|
||||
)
|
||||
|
||||
// v1Gateway builds a gateway with one source and two keys, and returns the
|
||||
// secrets so a test can assert they never appear in a response.
|
||||
func v1Gateway(t *testing.T) (*Gateway, string, string) {
|
||||
t.Helper()
|
||||
g := newTestGateway(t, config.Source{
|
||||
Name: "up",
|
||||
BaseURL: "http://up.example/v1",
|
||||
APIKey: "sk-up-secret",
|
||||
Adapter: "openai",
|
||||
Models: []config.Model{{ID: "m1", Kind: "chat"}},
|
||||
})
|
||||
// The shared helper only seeds one admin key; add a user key through the
|
||||
// core so role-gated endpoints have something to reject.
|
||||
if _, err := g.core.CreateKey("agent", "user", nil, "test agent key"); err != nil {
|
||||
t.Fatalf("CreateKey: %v", err)
|
||||
}
|
||||
rec := doReq(t, g, http.MethodGet, "/api/keys", "")
|
||||
var doc struct {
|
||||
Keys []config.GWKey `json:"keys"`
|
||||
}
|
||||
if err := json.Unmarshal(rec.Body.Bytes(), &doc); err != nil {
|
||||
t.Fatalf("decode keys: %v", err)
|
||||
}
|
||||
userKey := ""
|
||||
for _, k := range doc.Keys {
|
||||
if k.Role == "user" {
|
||||
userKey = k.Key
|
||||
}
|
||||
}
|
||||
if userKey == "" {
|
||||
t.Fatal("no user key was created")
|
||||
}
|
||||
return g, "sk-test", userKey
|
||||
}
|
||||
|
||||
func decodeJSON(t *testing.T, body string, v interface{}) {
|
||||
t.Helper()
|
||||
if err := json.Unmarshal([]byte(body), v); err != nil {
|
||||
t.Fatalf("decode %s: %v", body, err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestAPIV1IndexIsDiscoverable(t *testing.T) {
|
||||
g, _, _ := v1Gateway(t)
|
||||
rec := doReq(t, g, http.MethodGet, "/api/v1", "")
|
||||
if rec.Code != http.StatusOK {
|
||||
t.Fatalf("status = %d, want 200: %s", rec.Code, rec.Body.String())
|
||||
}
|
||||
var doc struct {
|
||||
APIVersion string `json:"api_version"`
|
||||
Description string `json:"description"`
|
||||
Discovery []struct {
|
||||
Method string `json:"method"`
|
||||
Path string `json:"path"`
|
||||
Auth string `json:"auth"`
|
||||
Summary string `json:"summary"`
|
||||
} `json:"discovery"`
|
||||
Conventions map[string]interface{} `json:"conventions"`
|
||||
}
|
||||
decodeJSON(t, rec.Body.String(), &doc)
|
||||
if doc.APIVersion != "v1" {
|
||||
t.Errorf("api_version = %q", doc.APIVersion)
|
||||
}
|
||||
if doc.Description == "" {
|
||||
t.Error("index should carry a usage description")
|
||||
}
|
||||
if len(doc.Discovery) == 0 {
|
||||
t.Fatal("discovery list is empty")
|
||||
}
|
||||
want := map[string]bool{
|
||||
"/api/v1/overview": false, "/api/v1/sources": false,
|
||||
"/api/v1/auto": false, "/api/v1/models": false,
|
||||
"/api/v1/health": false, "/api/v1/keys": false,
|
||||
"/api/sources": false, "/api/auto": false, "/v1/chat/completions": false,
|
||||
}
|
||||
for _, e := range doc.Discovery {
|
||||
if _, ok := want[e.Path]; ok {
|
||||
want[e.Path] = true
|
||||
}
|
||||
if e.Method == "" || e.Path == "" || e.Auth == "" || e.Summary == "" {
|
||||
t.Errorf("incomplete discovery entry: %+v", e)
|
||||
}
|
||||
}
|
||||
for p, found := range want {
|
||||
if !found {
|
||||
t.Errorf("discovery is missing %s", p)
|
||||
}
|
||||
}
|
||||
if doc.Conventions["errors"] == nil {
|
||||
t.Error("conventions should document the error shape")
|
||||
}
|
||||
}
|
||||
|
||||
func TestAPIV1OverviewSummarisesState(t *testing.T) {
|
||||
g, _, _ := v1Gateway(t)
|
||||
rec := doReq(t, g, http.MethodGet, "/api/v1/overview", "")
|
||||
if rec.Code != http.StatusOK {
|
||||
t.Fatalf("status = %d: %s", rec.Code, rec.Body.String())
|
||||
}
|
||||
var doc struct {
|
||||
Summary struct {
|
||||
SourceCount int `json:"source_count"`
|
||||
ModelCount int `json:"model_count"`
|
||||
GatewayKeyCount int `json:"gateway_key_count"`
|
||||
} `json:"summary"`
|
||||
Sources []map[string]interface{} `json:"sources"`
|
||||
Auto []map[string]interface{} `json:"auto"`
|
||||
Health []map[string]interface{} `json:"health"`
|
||||
Caller struct {
|
||||
Role string `json:"role"`
|
||||
} `json:"caller"`
|
||||
}
|
||||
decodeJSON(t, rec.Body.String(), &doc)
|
||||
if doc.Summary.SourceCount != 1 {
|
||||
t.Errorf("source_count = %d, want 1", doc.Summary.SourceCount)
|
||||
}
|
||||
if doc.Summary.ModelCount != 1 {
|
||||
t.Errorf("model_count = %d, want 1", doc.Summary.ModelCount)
|
||||
}
|
||||
if doc.Summary.GatewayKeyCount < 2 {
|
||||
t.Errorf("gateway_key_count = %d, want >= 2", doc.Summary.GatewayKeyCount)
|
||||
}
|
||||
if len(doc.Sources) != 1 || doc.Sources[0]["name"] != "up" {
|
||||
t.Errorf("sources = %+v", doc.Sources)
|
||||
}
|
||||
if doc.Caller.Role != "admin" {
|
||||
t.Errorf("caller.role = %q, want admin", doc.Caller.Role)
|
||||
}
|
||||
}
|
||||
|
||||
func TestAPIV1NeverEchoesSecrets(t *testing.T) {
|
||||
g, admin, user := v1Gateway(t)
|
||||
for _, path := range []string{"/api/v1", "/api/v1/overview", "/api/v1/sources", "/api/v1/sources/up", "/api/v1/keys", "/api/v1/models", "/api/v1/health", "/api/v1/auto"} {
|
||||
rec := doReq(t, g, http.MethodGet, path, "")
|
||||
if rec.Code != http.StatusOK {
|
||||
t.Fatalf("%s status = %d: %s", path, rec.Code, rec.Body.String())
|
||||
}
|
||||
if strings.Contains(rec.Body.String(), "sk-up-secret") {
|
||||
t.Errorf("%s leaked the source api_key", path)
|
||||
}
|
||||
if admin != "" && strings.Contains(rec.Body.String(), admin) {
|
||||
t.Errorf("%s leaked the admin gateway key", path)
|
||||
}
|
||||
if strings.Contains(rec.Body.String(), user) {
|
||||
t.Errorf("%s leaked the user gateway key", path)
|
||||
}
|
||||
}
|
||||
// Masking must still be useful: it says a key is configured.
|
||||
rec := doReq(t, g, http.MethodGet, "/api/v1/sources/up", "")
|
||||
if !strings.Contains(rec.Body.String(), `"api_key_set":true`) {
|
||||
t.Error("masked source should report that a key is configured")
|
||||
}
|
||||
}
|
||||
|
||||
func TestAPIV1KeysRequiresAdminAndMasks(t *testing.T) {
|
||||
g, _, user := v1Gateway(t)
|
||||
// Authenticate as the user key by swapping the shared helper's header.
|
||||
req := newAuthedRequest(t, http.MethodGet, "/api/v1/keys", user)
|
||||
rec := serveViaHandler(t, g, req)
|
||||
if rec.Code != http.StatusForbidden {
|
||||
t.Errorf("user role on /api/v1/keys = %d, want 403: %s", rec.Code, rec.Body.String())
|
||||
}
|
||||
}
|
||||
|
||||
func TestAPIV1SourcesNotFound(t *testing.T) {
|
||||
g, _, _ := v1Gateway(t)
|
||||
rec := doReq(t, g, http.MethodGet, "/api/v1/sources/nope", "")
|
||||
if rec.Code != http.StatusNotFound {
|
||||
t.Errorf("status = %d, want 404: %s", rec.Code, rec.Body.String())
|
||||
}
|
||||
}
|
||||
|
||||
func TestAPIV1RejectsWrites(t *testing.T) {
|
||||
g, _, _ := v1Gateway(t)
|
||||
for _, m := range []string{http.MethodPost, http.MethodPut, http.MethodDelete, http.MethodPatch} {
|
||||
rec := doReq(t, g, m, "/api/v1/sources", `{}`)
|
||||
if rec.Code == http.StatusOK {
|
||||
t.Errorf("%s /api/v1/sources returned 200; the v1 read facade must not mutate", m)
|
||||
}
|
||||
}
|
||||
// And the source must be untouched.
|
||||
rec := doReq(t, g, http.MethodGet, "/api/v1/sources", "")
|
||||
if !strings.Contains(rec.Body.String(), `"name":"up"`) {
|
||||
t.Error("the source disappeared after a rejected write")
|
||||
}
|
||||
}
|
||||
|
||||
func TestAPIV1ModelsGroupsBySource(t *testing.T) {
|
||||
g, _, _ := v1Gateway(t)
|
||||
rec := doReq(t, g, http.MethodGet, "/api/v1/models", "")
|
||||
if rec.Code != http.StatusOK {
|
||||
t.Fatalf("status = %d", rec.Code)
|
||||
}
|
||||
var doc struct {
|
||||
Count int `json:"count"`
|
||||
Models []string `json:"models"`
|
||||
BySource map[string][]string `json:"by_source"`
|
||||
}
|
||||
decodeJSON(t, rec.Body.String(), &doc)
|
||||
if doc.Count != 1 || doc.Models[0] != "m1" {
|
||||
t.Errorf("models = %+v", doc.Models)
|
||||
}
|
||||
if len(doc.BySource["up"]) != 1 || doc.BySource["up"][0] != "m1" {
|
||||
t.Errorf("by_source = %+v", doc.BySource)
|
||||
}
|
||||
}
|
||||
|
||||
func TestAPIV1RequiresAuthentication(t *testing.T) {
|
||||
g, _, _ := v1Gateway(t)
|
||||
for _, path := range []string{"/api/v1", "/api/v1/overview", "/api/v1/sources", "/api/v1/health", "/api/v1/models"} {
|
||||
req := newRequest(http.MethodGet, path, "")
|
||||
rec := serveViaHandler(t, g, req)
|
||||
if rec.Code != http.StatusUnauthorized {
|
||||
t.Errorf("%s without a key = %d, want 401", path, rec.Code)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestMaskKeyHidesTheSecret(t *testing.T) {
|
||||
cases := map[string]string{
|
||||
"": "",
|
||||
"ab": "ab…",
|
||||
"sk-gw-abcdefghijkl": "sk-gw-…ijkl",
|
||||
"enc:v1:AAAA": "(sealed)",
|
||||
}
|
||||
for in, want := range cases {
|
||||
if got := maskKey(in); got != want {
|
||||
t.Errorf("maskKey(%q) = %q, want %q", in, got, want)
|
||||
}
|
||||
}
|
||||
// The masked form must not reveal the middle of the secret.
|
||||
full := maskKey("sk-gw-1234567890abcdef")
|
||||
if strings.Contains(full, "4567890abcdef") {
|
||||
t.Errorf("maskKey leaked the middle: %q", full)
|
||||
}
|
||||
}
|
||||
|
||||
// newRequest builds an unauthenticated request.
|
||||
func newRequest(method, path, body string) *http.Request {
|
||||
req, _ := http.NewRequest(method, path, strings.NewReader(body))
|
||||
if body != "" {
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
}
|
||||
return req
|
||||
}
|
||||
|
||||
// newAuthedRequest builds a request carrying the given bearer key.
|
||||
func newAuthedRequest(t *testing.T, method, path, key string) *http.Request {
|
||||
t.Helper()
|
||||
req := newRequest(method, path, "")
|
||||
req.Header.Set("Authorization", "Bearer "+key)
|
||||
return req
|
||||
}
|
||||
|
||||
// serveViaHandler pushes a request through the full handler chain (auth included).
|
||||
func serveViaHandler(t *testing.T, g *Gateway, req *http.Request) *httptest.ResponseRecorder {
|
||||
t.Helper()
|
||||
rec := httptest.NewRecorder()
|
||||
g.Handler().ServeHTTP(rec, req)
|
||||
return rec
|
||||
}
|
||||
|
||||
// TestAPIV1RevealRequiresAdmin: the only way to read a key in the clear is the
|
||||
// explicit admin-only reveal. Without the role check this endpoint would hand
|
||||
// every source credential to any valid (even user-scoped) key.
|
||||
func TestAPIV1RevealRequiresAdmin(t *testing.T) {
|
||||
g, _, user := v1Gateway(t)
|
||||
rec := serveViaHandler(t, g, newAuthedRequest(t, http.MethodGet, "/api/v1/sources/up?reveal=credentials", user))
|
||||
if rec.Code != http.StatusForbidden {
|
||||
t.Errorf("user reveal = %d, want 403: %s", rec.Code, rec.Body.String())
|
||||
}
|
||||
if strings.Contains(rec.Body.String(), "sk-up-secret") {
|
||||
t.Error("a forbidden reveal still leaked the key")
|
||||
}
|
||||
rec = doReq(t, g, http.MethodGet, "/api/v1/sources/up?reveal=credentials", "")
|
||||
if rec.Code != http.StatusOK {
|
||||
t.Fatalf("admin reveal = %d: %s", rec.Code, rec.Body.String())
|
||||
}
|
||||
if !strings.Contains(rec.Body.String(), "sk-up-secret") {
|
||||
t.Error("admin reveal should return the key in the clear (the edit form needs it)")
|
||||
}
|
||||
// The masked default view must still hide it on the same path.
|
||||
rec = doReq(t, g, http.MethodGet, "/api/v1/sources/up", "")
|
||||
if strings.Contains(rec.Body.String(), "sk-up-secret") {
|
||||
t.Error("the default view leaked the key")
|
||||
}
|
||||
}
|
||||
282
internal/gateway/billing_rules_api.go
Normal file
282
internal/gateway/billing_rules_api.go
Normal file
@ -0,0 +1,282 @@
|
||||
package gateway
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"net/http"
|
||||
|
||||
"llmsproxy/internal/billing"
|
||||
"llmsproxy/internal/config"
|
||||
)
|
||||
|
||||
// Billing rules API.
|
||||
//
|
||||
// GET /api/plugins/billing/rules the active profile's rules + all profiles
|
||||
// PUT /api/plugins/billing/rules replace the whole DSL and re-inject prices
|
||||
// POST /api/plugins/billing/rules { profile, rule } append one rule
|
||||
// DELETE /api/plugins/billing/rules { profile, url } remove one rule
|
||||
//
|
||||
// Why a separate endpoint instead of the generic plugin state: prices are
|
||||
// plugin state, but they are also a durable, reviewable CONFIGURATION. Routing
|
||||
// them through /state would let an admin key PUT arbitrary plugin state and
|
||||
// silently reset the accumulated accounting. Here the server owns the shape:
|
||||
// it validates the DSL, writes it back to config.yaml, recompiles, and hands
|
||||
// the plugin its prices table. The operator edits rules; the plugin keeps
|
||||
// numbers. The two are never mixed in one payload.
|
||||
//
|
||||
// Editing means "what the operator typed is what config.yaml holds". A
|
||||
// billing mistake found next week must be traceable to a reviewable file, not
|
||||
// to a blob in the plugin's state sidecar.
|
||||
func (g *Gateway) handleBillingRules(w http.ResponseWriter, r *http.Request) {
|
||||
if reqRole(r.Context()) != "admin" {
|
||||
writeError(w, http.StatusForbidden, "forbidden", "admin role required")
|
||||
return
|
||||
}
|
||||
switch r.Method {
|
||||
case http.MethodGet:
|
||||
g.getBillingRules(w)
|
||||
case http.MethodPut:
|
||||
g.putBillingRules(w, r)
|
||||
case http.MethodPost:
|
||||
g.addBillingRule(w, r)
|
||||
case http.MethodDelete:
|
||||
g.delBillingRule(w, r)
|
||||
default:
|
||||
writeError(w, http.StatusMethodNotAllowed, "method_not_allowed", "use GET/PUT/POST/DELETE")
|
||||
}
|
||||
}
|
||||
|
||||
// billingRulesPayload is the UI-facing view: the profiles as declared, plus
|
||||
// which one is active and the source URLs they can match — the editor needs
|
||||
// the URL list to offer suggestions, and an operator typing a URL that matches
|
||||
// no source is exactly the silent-no-pricing failure this feature exists to
|
||||
// remove.
|
||||
type billingRulesPayload struct {
|
||||
DSL *config.BillingDSL `json:"billing"`
|
||||
Active string `json:"active"`
|
||||
URLs []string `json:"urls"`
|
||||
Warnings []string `json:"warnings"`
|
||||
}
|
||||
|
||||
func (g *Gateway) billingRulesPayload() billingRulesPayload {
|
||||
out := billingRulesPayload{}
|
||||
cfg := g.core.Config()
|
||||
if cfg != nil && cfg.BillingDSL != nil {
|
||||
// Copy so the response cannot be mutated back through the pointer.
|
||||
cp := *cfg.BillingDSL
|
||||
out.DSL = &cp
|
||||
if p := cp.Resolve(""); p != nil {
|
||||
out.Active = p.ID
|
||||
}
|
||||
}
|
||||
for _, s := range g.sourceURLs() {
|
||||
out.URLs = append(out.URLs, s)
|
||||
}
|
||||
out.Warnings = g.billingRuleWarnings()
|
||||
return out
|
||||
}
|
||||
|
||||
func (g *Gateway) sourceURLs() []string {
|
||||
cfg := g.core.Config()
|
||||
if cfg == nil {
|
||||
return nil
|
||||
}
|
||||
seen := map[string]bool{}
|
||||
var out []string
|
||||
for _, s := range cfg.Sources {
|
||||
if s.BaseURL == "" || seen[s.BaseURL] {
|
||||
continue
|
||||
}
|
||||
seen[s.BaseURL] = true
|
||||
out = append(out, s.BaseURL)
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
// billingRuleWarnings reports rules that match no configured source. This is
|
||||
// the silent killer of hand-written price tables: the rule parses, validates,
|
||||
// and prices nothing at all, so every request on that URL lands in unpriced.
|
||||
// Saying so out loud is the difference between a five-second fix and an
|
||||
// afternoon wondering why the bill is zero.
|
||||
func (g *Gateway) billingRuleWarnings() []string {
|
||||
cfg := g.core.Config()
|
||||
if cfg == nil || cfg.BillingDSL == nil {
|
||||
return nil
|
||||
}
|
||||
urls := map[string]bool{}
|
||||
for _, u := range g.sourceURLs() {
|
||||
urls[u] = true
|
||||
}
|
||||
var warns []string
|
||||
for _, p := range cfg.BillingDSL.Profiles {
|
||||
for _, r := range p.Rules {
|
||||
if r.URL == "*" || urls[r.URL] {
|
||||
continue
|
||||
}
|
||||
warns = append(warns, "profile "+p.ID+": rule url "+r.URL+" matches no configured source")
|
||||
}
|
||||
}
|
||||
return warns
|
||||
}
|
||||
|
||||
func (g *Gateway) getBillingRules(w http.ResponseWriter) {
|
||||
writeJSON(w, http.StatusOK, g.billingRulesPayload())
|
||||
}
|
||||
|
||||
func (g *Gateway) putBillingRules(w http.ResponseWriter, r *http.Request) {
|
||||
var body struct {
|
||||
Billing *config.BillingDSL `json:"billing"`
|
||||
}
|
||||
if err := json.NewDecoder(r.Body).Decode(&body); err != nil {
|
||||
writeError(w, http.StatusBadRequest, "invalid_request", "invalid json: "+err.Error())
|
||||
return
|
||||
}
|
||||
if body.Billing == nil {
|
||||
writeError(w, http.StatusBadRequest, "invalid_request", "billing is required")
|
||||
return
|
||||
}
|
||||
g.applyBillingDSLUpdate(w, body.Billing)
|
||||
}
|
||||
|
||||
func (g *Gateway) addBillingRule(w http.ResponseWriter, r *http.Request) {
|
||||
var body struct {
|
||||
Profile string `json:"profile"`
|
||||
Rule config.BillingRule `json:"rule"`
|
||||
}
|
||||
if err := json.NewDecoder(r.Body).Decode(&body); err != nil {
|
||||
writeError(w, http.StatusBadRequest, "invalid_request", "invalid json: "+err.Error())
|
||||
return
|
||||
}
|
||||
cfg := g.core.Config()
|
||||
if cfg == nil || cfg.BillingDSL == nil || len(cfg.BillingDSL.Profiles) == 0 {
|
||||
writeError(w, http.StatusBadRequest, "no_profile", "no billing profiles configured; create one first")
|
||||
return
|
||||
}
|
||||
next := cloneBillingDSL(cfg.BillingDSL)
|
||||
idx := profileIndex(next, body.Profile)
|
||||
if idx < 0 {
|
||||
writeError(w, http.StatusBadRequest, "unknown_profile", "no such profile: "+body.Profile)
|
||||
return
|
||||
}
|
||||
p := &next.Profiles[idx]
|
||||
// Replace rather than append when the URL is already declared: two rules
|
||||
// for one URL would silently make the first unreachable (first match wins),
|
||||
// which is the exact ambiguity an editor should not be able to create.
|
||||
replaced := false
|
||||
for i := range p.Rules {
|
||||
if p.Rules[i].URL == body.Rule.URL {
|
||||
p.Rules[i] = body.Rule
|
||||
replaced = true
|
||||
break
|
||||
}
|
||||
}
|
||||
if !replaced {
|
||||
p.Rules = append(p.Rules, body.Rule)
|
||||
}
|
||||
g.applyBillingDSLUpdate(w, next)
|
||||
}
|
||||
|
||||
func (g *Gateway) delBillingRule(w http.ResponseWriter, r *http.Request) {
|
||||
var body struct {
|
||||
Profile string `json:"profile"`
|
||||
URL string `json:"url"`
|
||||
}
|
||||
if err := json.NewDecoder(r.Body).Decode(&body); err != nil {
|
||||
writeError(w, http.StatusBadRequest, "invalid_request", "invalid json: "+err.Error())
|
||||
return
|
||||
}
|
||||
cfg := g.core.Config()
|
||||
if cfg == nil || cfg.BillingDSL == nil {
|
||||
writeError(w, http.StatusBadRequest, "no_profile", "no billing profiles configured")
|
||||
return
|
||||
}
|
||||
next := cloneBillingDSL(cfg.BillingDSL)
|
||||
idx := profileIndex(next, body.Profile)
|
||||
if idx < 0 {
|
||||
writeError(w, http.StatusBadRequest, "unknown_profile", "no such profile: "+body.Profile)
|
||||
return
|
||||
}
|
||||
p := &next.Profiles[idx]
|
||||
out := p.Rules[:0]
|
||||
found := false
|
||||
for _, rule := range p.Rules {
|
||||
if rule.URL == body.URL {
|
||||
found = true
|
||||
continue
|
||||
}
|
||||
out = append(out, rule)
|
||||
}
|
||||
if !found {
|
||||
writeError(w, http.StatusNotFound, "not_found", "no rule for url "+body.URL)
|
||||
return
|
||||
}
|
||||
p.Rules = out
|
||||
g.applyBillingDSLUpdate(w, next)
|
||||
}
|
||||
|
||||
// applyBillingDSLUpdate is the single write path: validate, persist to
|
||||
// config.yaml, recompile, inject, and only then report success. If the config
|
||||
// cannot be written the change is NOT applied in memory either — a rules editor
|
||||
// that says "saved" and loses the edit on the next restart is worse than one
|
||||
// that refuses.
|
||||
func (g *Gateway) applyBillingDSLUpdate(w http.ResponseWriter, next *config.BillingDSL) {
|
||||
if err := next.Validate(); err != nil {
|
||||
writeError(w, http.StatusBadRequest, "invalid_rules", err.Error())
|
||||
return
|
||||
}
|
||||
cfg := g.core.Config()
|
||||
if cfg == nil {
|
||||
writeError(w, http.StatusInternalServerError, "no_config", "no config loaded")
|
||||
return
|
||||
}
|
||||
// Assign BEFORE saving. The previous order saved first and assigned after,
|
||||
// so every rule edit reported success while writing a config.yaml with no
|
||||
// billing section at all — the operator's edit vanished on restart and
|
||||
// nothing in the API response said so.
|
||||
prev := cfg.BillingDSL
|
||||
cfg.BillingDSL = next
|
||||
|
||||
// Compile before persisting: a profile that cannot be turned into prices
|
||||
// must not reach the config file, or the next restart fails to load the
|
||||
// very rules the editor just accepted.
|
||||
prices, err := billing.CompileOpts(next.Resolve(next.Active), cfg.Sources, true)
|
||||
if err != nil {
|
||||
cfg.BillingDSL = prev // no half-applied state
|
||||
writeError(w, http.StatusBadRequest, "compile_failed", err.Error())
|
||||
return
|
||||
}
|
||||
if err := cfg.Save(); err != nil {
|
||||
cfg.BillingDSL = prev
|
||||
writeError(w, http.StatusInternalServerError, "persist_failed", err.Error())
|
||||
return
|
||||
}
|
||||
ps := g.core.Plugins()
|
||||
if ps != nil {
|
||||
if err := ps.SetState("billing", map[string]interface{}{"prices": prices}); err != nil {
|
||||
writeError(w, http.StatusBadRequest, "plugin_error", err.Error())
|
||||
return
|
||||
}
|
||||
}
|
||||
p := g.billingRulesPayload()
|
||||
writeJSON(w, http.StatusOK, map[string]interface{}{
|
||||
"ok": true, "billing": p.DSL, "active": p.Active, "warnings": p.Warnings,
|
||||
})
|
||||
}
|
||||
|
||||
func cloneBillingDSL(d *config.BillingDSL) *config.BillingDSL {
|
||||
out := &config.BillingDSL{Active: d.Active}
|
||||
out.Profiles = make([]config.BillingProfile, len(d.Profiles))
|
||||
copy(out.Profiles, d.Profiles)
|
||||
return out
|
||||
}
|
||||
|
||||
func profileIndex(d *config.BillingDSL, id string) int {
|
||||
if id == "" {
|
||||
return -1
|
||||
}
|
||||
for i, p := range d.Profiles {
|
||||
if p.ID == id {
|
||||
return i
|
||||
}
|
||||
}
|
||||
return -1
|
||||
}
|
||||
281
internal/gateway/billing_rules_api_test.go
Normal file
281
internal/gateway/billing_rules_api_test.go
Normal file
@ -0,0 +1,281 @@
|
||||
package gateway
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"os"
|
||||
"sort"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"llmsproxy/internal/config"
|
||||
)
|
||||
|
||||
// rulesGateway needs a real billing plugin (the API injects prices into it)
|
||||
// AND two sources whose base_urls the rules will match, so it cannot reuse
|
||||
// either existing helper: gatewayWithBilling has no sources, and
|
||||
// newTestGateway loads no plugins.
|
||||
func rulesGateway(t *testing.T) *Gateway {
|
||||
t.Helper()
|
||||
g := gatewayWithBilling(t)
|
||||
cfg := g.core.Config()
|
||||
cfg.Sources = []config.Source{
|
||||
{Name: "localzen", BaseURL: "https://free.example.com/v1"},
|
||||
{Name: "commandcode", BaseURL: "https://api.commandcode.ai/v1"},
|
||||
}
|
||||
return g
|
||||
}
|
||||
|
||||
func rulesDo(t *testing.T, g *Gateway, method, body string) *httptest.ResponseRecorder {
|
||||
t.Helper()
|
||||
return doReq(t, g, method, "/api/plugins/billing/rules", body)
|
||||
}
|
||||
|
||||
// TestBillingRulesCRUD covers the whole editor loop against the real HTTP
|
||||
// surface: read, add, replace-by-URL, delete. The replace-by-URL case is the
|
||||
// one worth pinning — two rules for one URL would make the first unreachable
|
||||
// (first match wins) without any error, which is exactly the kind of silent
|
||||
// ambiguity an editor must not be able to create.
|
||||
func TestBillingRulesCRUD(t *testing.T) {
|
||||
g := rulesGateway(t)
|
||||
seed := `{"billing":{"active":"p1","profiles":[{"id":"p1","rules":[
|
||||
{"url":"https://free.example.com/v1","mode":"free"}]}]}}`
|
||||
if rr := rulesDo(t, g, http.MethodPut, seed); rr.Code != http.StatusOK {
|
||||
t.Fatalf("seed rules = %d: %s", rr.Code, rr.Body.String())
|
||||
}
|
||||
|
||||
// Add a token-priced rule for the second URL.
|
||||
add := `{"profile":"p1","rule":{"url":"https://api.commandcode.ai/v1","mode":"token",
|
||||
"models":{"deepseek-v4.1-flash":{"prompt":"0.15","completion":"0.60"}}}}`
|
||||
if rr := rulesDo(t, g, http.MethodPost, add); rr.Code != http.StatusOK {
|
||||
t.Fatalf("add rule = %d: %s", rr.Code, rr.Body.String())
|
||||
}
|
||||
var got billingRulesPayload
|
||||
if err := json.Unmarshal(rulesDo(t, g, http.MethodGet, "").Body.Bytes(), &got); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(got.DSL.Profiles[0].Rules) != 2 {
|
||||
t.Fatalf("after add there are %d rules, want 2", len(got.DSL.Profiles[0].Rules))
|
||||
}
|
||||
|
||||
// Adding the SAME url again must replace, not append.
|
||||
if rr := rulesDo(t, g, http.MethodPost, `{"profile":"p1","rule":{"url":"https://api.commandcode.ai/v1","mode":"free"}}`); rr.Code != http.StatusOK {
|
||||
t.Fatalf("re-add same url = %d: %s", rr.Code, rr.Body.String())
|
||||
}
|
||||
if err := json.Unmarshal(rulesDo(t, g, http.MethodGet, "").Body.Bytes(), &got); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if n := len(got.DSL.Profiles[0].Rules); n != 2 {
|
||||
t.Fatalf("★ re-adding an existing url appended instead of replacing: %d rules", n)
|
||||
}
|
||||
var mode string
|
||||
for _, r := range got.DSL.Profiles[0].Rules {
|
||||
if r.URL == "https://api.commandcode.ai/v1" {
|
||||
mode = r.Mode
|
||||
}
|
||||
}
|
||||
if mode != "free" {
|
||||
t.Errorf("re-added rule mode = %q, want free (the replacement must win)", mode)
|
||||
}
|
||||
|
||||
// Delete by URL.
|
||||
if rr := rulesDo(t, g, http.MethodDelete, `{"profile":"p1","url":"https://api.commandcode.ai/v1"}`); rr.Code != http.StatusOK {
|
||||
t.Fatalf("delete rule = %d: %s", rr.Code, rr.Body.String())
|
||||
}
|
||||
if err := json.Unmarshal(rulesDo(t, g, http.MethodGet, "").Body.Bytes(), &got); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(got.DSL.Profiles[0].Rules) != 1 {
|
||||
t.Errorf("after delete there are %d rules, want 1", len(got.DSL.Profiles[0].Rules))
|
||||
}
|
||||
}
|
||||
|
||||
// TestBillingRulesRejectUnknownURLWithWarning is the silent-failure guard: a
|
||||
// rule whose URL matches no source prices nothing, so every request on it
|
||||
// falls into unpriced. The API must say so instead of accepting it silently.
|
||||
func TestBillingRulesRejectUnknownURLWithWarning(t *testing.T) {
|
||||
g := rulesGateway(t)
|
||||
if rr := rulesDo(t, g, http.MethodPut, `{"billing":{"active":"p","profiles":[{"id":"p",
|
||||
"rules":[{"url":"https://typo.example.com/v1","mode":"free"}]}]}}`); rr.Code != http.StatusOK {
|
||||
t.Fatalf("seed = %d: %s", rr.Code, rr.Body.String())
|
||||
}
|
||||
var got billingRulesPayload
|
||||
if err := json.Unmarshal(rulesDo(t, g, http.MethodGet, "").Body.Bytes(), &got); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(got.Warnings) == 0 {
|
||||
t.Fatal("★ a rule matching no configured source was accepted with no warning — " +
|
||||
"it will price nothing and every request will be recorded as unpriced")
|
||||
}
|
||||
if !strings.Contains(got.Warnings[0], "typo.example.com") {
|
||||
t.Errorf("warning does not name the offending url: %v", got.Warnings)
|
||||
}
|
||||
// The editor needs the real URL list to offer suggestions.
|
||||
if len(got.URLs) != 2 {
|
||||
t.Errorf("urls offered to the editor = %v, want the 2 configured base_urls", got.URLs)
|
||||
}
|
||||
}
|
||||
|
||||
// TestBillingRulesInvalidIsRejectedNotPersisted checks the write path refuses
|
||||
// bad input and leaves the previous rules in place — an editor that clears the
|
||||
// table on a typo loses the price table.
|
||||
func TestBillingRulesInvalidIsRejectedNotPersisted(t *testing.T) {
|
||||
g := rulesGateway(t)
|
||||
seed := `{"billing":{"active":"p","profiles":[{"id":"p","rules":[
|
||||
{"url":"https://free.example.com/v1","mode":"free"}]}]}}`
|
||||
if rr := rulesDo(t, g, http.MethodPut, seed); rr.Code != http.StatusOK {
|
||||
t.Fatalf("seed = %d: %s", rr.Code, rr.Body.String())
|
||||
}
|
||||
// An unknown mode must be refused.
|
||||
if rr := rulesDo(t, g, http.MethodPost, `{"profile":"p","rule":{"url":"https://x/v1","mode":"banana"}}`); rr.Code != http.StatusBadRequest {
|
||||
t.Errorf("invalid mode accepted with %d, want 400", rr.Code)
|
||||
}
|
||||
var got billingRulesPayload
|
||||
if err := json.Unmarshal(rulesDo(t, g, http.MethodGet, "").Body.Bytes(), &got); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(got.DSL.Profiles[0].Rules) != 1 {
|
||||
t.Errorf("★ a rejected rule was persisted anyway: %d rules", len(got.DSL.Profiles[0].Rules))
|
||||
}
|
||||
// An unparseable price must be refused too, not silently priced at zero.
|
||||
if rr := rulesDo(t, g, http.MethodPost, `{"profile":"p","rule":{"url":"https://y/v1","mode":"token",
|
||||
"models":{"m":{"prompt":"free","completion":"0.6"}}}}`); rr.Code != http.StatusBadRequest {
|
||||
t.Errorf("non-numeric price accepted with %d, want 400", rr.Code)
|
||||
}
|
||||
}
|
||||
|
||||
// TestBillingRulesPersistAndReachThePlugin is the point of the whole
|
||||
// endpoint, and mutation verification showed the earlier CRUD tests missed it
|
||||
// twice: dropping cfg.Save() and dropping the SetState injection both left
|
||||
// every test green. Both failures are invisible in a GET — the response echoes
|
||||
// the in-memory rules either way — so this test checks the two consequences
|
||||
// that actually matter:
|
||||
//
|
||||
// 1. persistence: the rule must be in the config file on disk, because an
|
||||
// edit that evaporates on restart is a lie the UI told the operator;
|
||||
// 2. effect: the billing plugin must now PRICE the source (its published
|
||||
// prices table has a non-zero entry), because a rule that is stored but
|
||||
// never injected prices nothing at all.
|
||||
func TestBillingRulesPersistAndReachThePlugin(t *testing.T) {
|
||||
g := rulesGateway(t)
|
||||
cfg := g.core.Config()
|
||||
|
||||
seed := `{"billing":{"active":"p","profiles":[{"id":"p","rules":[
|
||||
{"url":"https://free.example.com/v1","mode":"free"},
|
||||
{"url":"https://api.commandcode.ai/v1","mode":"token",
|
||||
"models":{"deepseek-v4.1-flash":{"prompt":"0.15","completion":"0.60"}}}]}]}}`
|
||||
if rr := rulesDo(t, g, http.MethodPut, seed); rr.Code != http.StatusOK {
|
||||
t.Fatalf("seed = %d: %s", rr.Code, rr.Body.String())
|
||||
}
|
||||
|
||||
// 1. persisted to the file the process loaded from.
|
||||
raw, err := os.ReadFile(cfg.Path)
|
||||
if err != nil {
|
||||
t.Fatalf("read config: %v", err)
|
||||
}
|
||||
if !strings.Contains(string(raw), "commandcode.ai") {
|
||||
t.Errorf("★ the rule is not in config.yaml on disk — an edit that does not "+
|
||||
"survive a restart:\n%s", raw)
|
||||
}
|
||||
reloaded, err := config.Load(cfg.Path)
|
||||
if err != nil {
|
||||
t.Fatalf("reload saved config: %v", err)
|
||||
}
|
||||
if reloaded.BillingDSL == nil || len(reloaded.BillingDSL.Profiles) == 0 {
|
||||
t.Fatalf("saved config has no billing DSL: %s", raw)
|
||||
}
|
||||
// A round-trip through YAML must not lose the price strings.
|
||||
found := false
|
||||
for _, r := range reloaded.BillingDSL.Profiles[0].Rules {
|
||||
if strings.Contains(r.URL, "commandcode") {
|
||||
found = true
|
||||
if m, ok := r.Models["deepseek-v4.1-flash"]; !ok || m.Prompt != "0.15" {
|
||||
t.Errorf("price did not survive the YAML round-trip: %+v", m)
|
||||
}
|
||||
}
|
||||
}
|
||||
if !found {
|
||||
t.Errorf("the token rule is missing after reload:\n%s", raw)
|
||||
}
|
||||
|
||||
// 2. reached the plugin. Prices live in a separate Lua field from `state`,
|
||||
// and the GET /state endpoint returns the published state together with
|
||||
// the prices sidecar — which is exactly what the plugin's own UI reads.
|
||||
state := g.core.Plugins().State("billing")
|
||||
if state == nil {
|
||||
t.Fatalf("billing plugin published no state")
|
||||
}
|
||||
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())
|
||||
}
|
||||
var envelope struct {
|
||||
Prices map[string]interface{} `json:"prices"`
|
||||
}
|
||||
if err := json.Unmarshal(rr.Body.Bytes(), &envelope); err != nil {
|
||||
t.Fatalf("decode state envelope: %v", err)
|
||||
}
|
||||
prices := envelope.Prices
|
||||
if prices == nil {
|
||||
t.Fatalf("state payload carries no prices")
|
||||
}
|
||||
sources, ok := prices["sources"].(map[string]interface{})
|
||||
if !ok {
|
||||
t.Fatalf("prices has no sources map: %#v", prices)
|
||||
}
|
||||
cc, ok := sources["commandcode"]
|
||||
if !ok {
|
||||
t.Fatalf("★ the rule was saved but never injected: prices has %v, "+
|
||||
"no commandcode entry. Every request on it stays unpriced.", keysOf(sources))
|
||||
}
|
||||
if cc == nil {
|
||||
t.Error("commandcode price entry is nil")
|
||||
}
|
||||
// The free rule must price localzen, and the two must not be conflated.
|
||||
if _, ok := sources["localzen"]; !ok {
|
||||
t.Errorf("the free rule did not reach the plugin: %v", keysOf(sources))
|
||||
}
|
||||
// The injected price must be per-token, i.e. 0.15 USD/M => 1.5e-7.
|
||||
// Asserting the exact magnitude catches a unit error (per-million vs
|
||||
// per-token), which would be off by a factor of a million and still look
|
||||
// like "a number".
|
||||
models, ok := cc.(map[string]interface{})["models"].(map[string]interface{})
|
||||
if ok {
|
||||
if m, ok := models["deepseek-v4.1-flash"].(map[string]interface{}); ok {
|
||||
p, _ := m["prompt"].(float64)
|
||||
if p <= 0 || p > 1e-6 {
|
||||
t.Errorf("injected prompt price = %v, want per-token (~1.5e-7); "+
|
||||
"a per-million value here would over-bill by 1e6x", p)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func keysOf(m map[string]interface{}) []string {
|
||||
out := make([]string, 0, len(m))
|
||||
for k := range m {
|
||||
out = append(out, k)
|
||||
}
|
||||
sort.Strings(out)
|
||||
return out
|
||||
}
|
||||
|
||||
// TestBillingRulesRequireAdmin keeps price rewriting behind the admin role,
|
||||
// like every other plugin write: prices are configuration, not user data.
|
||||
func TestBillingRulesRequireAdmin(t *testing.T) {
|
||||
g := rulesGateway(t)
|
||||
rec, err := g.core.CreateKey("viewer", "user", nil, "")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
userKey := rec.Key
|
||||
for _, m := range []string{http.MethodPut, http.MethodPost, http.MethodDelete} {
|
||||
req, _ := http.NewRequest(m, "/api/plugins/billing/rules", strings.NewReader(`{}`))
|
||||
req.Header.Set("Authorization", "Bearer "+userKey)
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
if rr := newRecorderFor(t, g, req); rr.Code != http.StatusForbidden {
|
||||
t.Errorf("user %s rules = %d, want 403", m, rr.Code)
|
||||
}
|
||||
}
|
||||
}
|
||||
@ -7,11 +7,14 @@ import (
|
||||
"fmt"
|
||||
"log"
|
||||
"net/http"
|
||||
"regexp"
|
||||
"strconv"
|
||||
"strings"
|
||||
"sync/atomic"
|
||||
"time"
|
||||
|
||||
"llmsproxy/internal/config"
|
||||
"llmsproxy/internal/lua"
|
||||
"llmsproxy/internal/provider"
|
||||
"llmsproxy/internal/scheduler"
|
||||
"llmsproxy/internal/types"
|
||||
@ -28,6 +31,9 @@ type chatRequest struct {
|
||||
ToolChoice interface{} `json:"tool_choice,omitempty"`
|
||||
DisableThinking bool `json:"disable_thinking"`
|
||||
ExtraBody map[string]interface{} `json:"extra_body,omitempty"`
|
||||
|
||||
// PromptCacheKey 是 OpenAI 原生的缓存键;部分客户端用它携带会话标识。
|
||||
PromptCacheKey string `json:"prompt_cache_key,omitempty"`
|
||||
}
|
||||
|
||||
// ChatCompletion is the non-streaming OpenAI response object.
|
||||
@ -38,6 +44,10 @@ type ChatCompletion struct {
|
||||
Model string `json:"model"`
|
||||
Choices []ChatChoice `json:"choices"`
|
||||
Usage *types.TokenUsage `json:"usage,omitempty"`
|
||||
// Cost is the upstream-reported charge for this request, passed through
|
||||
// verbatim (OpenCode reports it as a decimal string, always "0" on the
|
||||
// flat-rate Go subscription). Absent when the upstream reports nothing.
|
||||
Cost string `json:"cost,omitempty"`
|
||||
}
|
||||
|
||||
type ChatChoice struct {
|
||||
@ -62,6 +72,10 @@ type ChatChunk struct {
|
||||
// Usage is sent in the final chunk of a stream (empty choices) so
|
||||
// OpenAI-compatible clients can read token usage.
|
||||
Usage *types.TokenUsage `json:"usage,omitempty"`
|
||||
// Cost is the upstream-reported charge for this request, passed through
|
||||
// verbatim (OpenCode reports it as a decimal string). Emitted on the
|
||||
// terminal chunk, mirroring OpenCode's own {"choices":[],"cost":"0"}.
|
||||
Cost string `json:"cost,omitempty"`
|
||||
}
|
||||
|
||||
type ChunkChoice struct {
|
||||
@ -115,13 +129,15 @@ func (g *Gateway) resolveCands(ctx context.Context, req *chatRequest) ([]*provid
|
||||
}
|
||||
|
||||
// filterCandsByModels keeps only providers exposing at least one model of the
|
||||
// scope (used for user keys with a restricted model scope). An "AUTO" scope
|
||||
// entry means the key is allowed to use any model. Scope entries with a
|
||||
// Source pinned to a specific upstream narrow the candidates to that source
|
||||
// scope (used for user keys with a restricted model scope). Scope entries with
|
||||
// a Source pinned to a specific upstream narrow the candidates to that source
|
||||
// for the matching model.
|
||||
//
|
||||
// An "AUTO" scope entry only allows the AUTO routing mode; it does NOT grant
|
||||
// access to specific models.
|
||||
func filterCandsByModels(cands []*provider.Provider, allow []config.ModelScope) []*provider.Provider {
|
||||
for _, m := range allow {
|
||||
if m.Model == "" || strings.EqualFold(m.Model, "AUTO") {
|
||||
if m.Model == "" {
|
||||
return cands
|
||||
}
|
||||
}
|
||||
@ -152,20 +168,18 @@ func filterCandsByModels(cands []*provider.Provider, allow []config.ModelScope)
|
||||
return out
|
||||
}
|
||||
|
||||
// intersectModels restricts a model list to the scope (preserving order). An
|
||||
// "AUTO" scope entry grants every model.
|
||||
// intersectModels restricts a model list to the scope (preserving order).
|
||||
// An "AUTO" scope entry only allows the AUTO routing mode; it does NOT grant
|
||||
// every model.
|
||||
func intersectModels(models []string, allow []config.ModelScope) []string {
|
||||
allowed := make(map[string]bool, len(allow))
|
||||
any := false
|
||||
for _, m := range allow {
|
||||
if m.Model == "" || strings.EqualFold(m.Model, "AUTO") {
|
||||
any = true
|
||||
break
|
||||
if m.Model == "" {
|
||||
return models
|
||||
}
|
||||
if !strings.EqualFold(m.Model, "AUTO") {
|
||||
allowed[m.Model] = true
|
||||
}
|
||||
allowed[m.Model] = true
|
||||
}
|
||||
if any {
|
||||
return models
|
||||
}
|
||||
out := make([]string, 0, len(models))
|
||||
seen := map[string]bool{}
|
||||
@ -179,48 +193,110 @@ func intersectModels(models []string, allow []config.ModelScope) []string {
|
||||
}
|
||||
|
||||
// checkModelScope validates the effective model against the key's model scope
|
||||
// and token quota. Returns an error message when rejected. A scope entry with
|
||||
// model "AUTO" grants all models; its quota caps the key's total tokens.
|
||||
// and token quota. Returns an error message when rejected.
|
||||
//
|
||||
// A scope entry with model "AUTO" only allows requests where the effective
|
||||
// model is AUTO (the routing mode). It does NOT grant access to specific model
|
||||
// ids — that requires an explicit scope entry for the model.
|
||||
// checkModelScope keeps its string-returning signature for callers that only
|
||||
// need to know whether the request may proceed.
|
||||
func (g *Gateway) checkModelScope(ctx context.Context, model string) string {
|
||||
if q := g.checkQuota(ctx, model); q != nil {
|
||||
return q.msg
|
||||
}
|
||||
return ""
|
||||
}
|
||||
|
||||
// quotaRejection is a quota verdict: the message plus whether the client
|
||||
// should retry. A spent quota is a rate limit (429 + Retry-After), not a
|
||||
// permission failure (403): a client that sees 403 gives up on the key, while
|
||||
// one that sees 429 with a retry hint waits and resumes when the window rolls
|
||||
// over.
|
||||
type quotaRejection struct {
|
||||
msg string
|
||||
retry int64 // seconds until the window resets; 0 = unknown
|
||||
}
|
||||
|
||||
func (q *quotaRejection) Error() string { return q.msg }
|
||||
|
||||
// checkQuota validates the effective model against the key's model scope and
|
||||
// that entry's quota. Returns nil when the request may proceed.
|
||||
//
|
||||
// Quotas are per scope entry, never key-wide: a model whose budget is spent
|
||||
// is refused on its own while the key's other models keep working. The verdict
|
||||
// carries the remaining seconds of the reset window so a spent budget answers
|
||||
// 429 + Retry-After (come back when it rolls over) instead of 403 (which reads
|
||||
// as "this key may never use this model" and makes clients give up).
|
||||
func (g *Gateway) checkQuota(ctx context.Context, model string) *quotaRejection {
|
||||
allow := g.allowedModels(ctx)
|
||||
if allow == nil {
|
||||
return ""
|
||||
}
|
||||
for _, sc := range allow {
|
||||
if sc.Model != "" && strings.EqualFold(sc.Model, "AUTO") {
|
||||
if sc.TokenQuota > 0 {
|
||||
used := g.scopeTokens(ctx, sc)
|
||||
if used >= sc.TokenQuota {
|
||||
return fmt.Sprintf("token quota exceeded (%d/%d)", used, sc.TokenQuota)
|
||||
}
|
||||
}
|
||||
return ""
|
||||
}
|
||||
return nil
|
||||
}
|
||||
for _, sc := range allow {
|
||||
if sc.Model != model {
|
||||
continue
|
||||
}
|
||||
win := AutoPeriodSeconds(sc.Period, sc.Hours)
|
||||
k := keyID(reqKey(ctx))
|
||||
if sc.TokenQuota > 0 {
|
||||
used := g.scopeTokens(ctx, sc)
|
||||
if used >= sc.TokenQuota {
|
||||
return fmt.Sprintf("token quota exceeded for %q (%d/%d)", model, used, sc.TokenQuota)
|
||||
if used := g.scopeTokens(ctx, sc); used >= sc.TokenQuota {
|
||||
return "aRejection{
|
||||
msg: fmt.Sprintf("token quota exceeded for %q (%d/%d)", model, used, sc.TokenQuota),
|
||||
retry: AutoSecondsToReset(sc.Period, sc.Hours),
|
||||
}
|
||||
}
|
||||
}
|
||||
return ""
|
||||
if sc.ReqQuota > 0 {
|
||||
if used := g.stats.KeyWindowReqs(k, win); used >= sc.ReqQuota {
|
||||
return "aRejection{
|
||||
msg: fmt.Sprintf("request quota exceeded for %q (%d/%d)", model, used, sc.ReqQuota),
|
||||
retry: AutoSecondsToReset(sc.Period, sc.Hours),
|
||||
}
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
return fmt.Sprintf("model %q is not allowed for this key", model)
|
||||
// A scope that simply does not list the model is a different problem from
|
||||
// a quota being exhausted, and the message has to say which: telling an
|
||||
// operator "not allowed for this key" when they configured a quota sends
|
||||
// them to the wrong setting. AUTO is checked here before the chain is even
|
||||
// consulted, so this is the message that surfaces for AUTO.
|
||||
if isAuto(model) {
|
||||
return "aRejection{msg: fmt.Sprintf(
|
||||
"model %q is not in this key's model scope, so it cannot use AUTO; "+
|
||||
"add %q to the key's models or remove the scope restriction", model, model)}
|
||||
}
|
||||
return "aRejection{msg: fmt.Sprintf("model %q is not allowed for this key", model)}
|
||||
}
|
||||
|
||||
// scopeTokens returns the tokens a scope entry has consumed within its reset
|
||||
// window (total for AUTO / per model otherwise).
|
||||
// quotaWindowSuffix describes a quota's reset window for an error message, so
|
||||
// a rejected caller can tell a permanent block from one that clears in an hour.
|
||||
func quotaWindowSuffix(period string, hours int64) string {
|
||||
switch {
|
||||
case period == "hour":
|
||||
return ", resets hourly"
|
||||
case period == "week":
|
||||
return ", resets weekly"
|
||||
case period == "month":
|
||||
return ", resets monthly"
|
||||
case period == "nhour" && hours > 1:
|
||||
return fmt.Sprintf(", resets every %dh", hours)
|
||||
}
|
||||
return ""
|
||||
} // scopeTokens returns the tokens this key consumed on the scope entry's model
|
||||
// within its reset window, isolated per key. For an AUTO entry the cap covers
|
||||
// everything the key routed through AUTO; for a model entry it covers that
|
||||
// model only.
|
||||
//
|
||||
// It reads the per-key hourly buckets rather than the key-blind model
|
||||
// buckets, so one key's usage can never exhaust another's quota.
|
||||
func (g *Gateway) scopeTokens(ctx context.Context, sc config.ModelScope) int64 {
|
||||
k := keyID(reqKey(ctx))
|
||||
if sc.Model != "" && strings.EqualFold(sc.Model, "AUTO") {
|
||||
return g.stats.KeyTokens(k)
|
||||
}
|
||||
win := AutoPeriodSeconds(sc.Period, sc.Hours)
|
||||
return g.stats.WindowTokens(sc.Model, sc.Source, win)
|
||||
if sc.Model != "" && strings.EqualFold(sc.Model, "AUTO") {
|
||||
return g.stats.KeyWindowTokens(k, win)
|
||||
}
|
||||
return g.stats.KeyWindowModelTokens(k, sc.Model, sc.Source, win)
|
||||
}
|
||||
|
||||
// hasScopeModel reports whether a model (possibly with a "source-model" /
|
||||
@ -229,18 +305,41 @@ func (g *Gateway) scopeTokens(ctx context.Context, sc config.ModelScope) int64 {
|
||||
// source that actually serves the bare model (via Registry.EffectiveModel), so
|
||||
// model ids that themselves contain separators (e.g. "deepseek-v4-flash-free")
|
||||
// are never corrupted (P10-2).
|
||||
//
|
||||
// A scope entry with model "AUTO" only matches the literal AUTO routing mode;
|
||||
// it does NOT grant access to specific model ids.
|
||||
func (g *Gateway) hasScopeModel(list []config.ModelScope, s string) bool {
|
||||
if r := g.core.Registry(); r != nil {
|
||||
s = r.EffectiveModel(s)
|
||||
}
|
||||
for _, x := range list {
|
||||
if x.Model == s || (x.Model != "" && strings.EqualFold(x.Model, "AUTO")) {
|
||||
if x.Model == s {
|
||||
return true
|
||||
}
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
// writeScopeReject answers a model-scope or quota rejection. A spent quota is
|
||||
// 429 (rate_limit_exceeded) with Retry-After, so a client waits and resumes
|
||||
// after the reset; a model the key may not use stays 403 (model_not_allowed),
|
||||
// because retrying cannot help.
|
||||
func (g *Gateway) writeReject(w http.ResponseWriter, q *quotaRejection) {
|
||||
if q.retry > 0 {
|
||||
w.Header().Set("Retry-After", strconv.FormatInt(q.retry, 10))
|
||||
writeError(w, http.StatusTooManyRequests, "rate_limit_exceeded", q.msg)
|
||||
return
|
||||
}
|
||||
// no window to wait for: the cap is either permanent or key-wide with no
|
||||
// period. "rate_limit_exceeded" still says "come back after the operator
|
||||
// raises the cap", which 403 would not.
|
||||
if strings.Contains(q.msg, "quota exceeded") {
|
||||
writeError(w, http.StatusTooManyRequests, "rate_limit_exceeded", q.msg)
|
||||
return
|
||||
}
|
||||
writeError(w, http.StatusForbidden, "model_not_allowed", q.msg)
|
||||
}
|
||||
|
||||
func (g *Gateway) resolveByModel(model string) ([]*provider.Provider, string) {
|
||||
if isAuto(model) {
|
||||
return g.core.Registry().Resolve("AUTO"), ""
|
||||
@ -287,6 +386,8 @@ func (g *Gateway) handleChat(w http.ResponseWriter, r *http.Request) {
|
||||
writeError(w, http.StatusBadRequest, "invalid_request", "invalid json: "+err.Error())
|
||||
return
|
||||
}
|
||||
// 客户端自带的会话标识(若有)比推导出来的更准,见 clientSessionFromRequest。
|
||||
clientSession := clientSessionFromRequest(r, req.PromptCacheKey)
|
||||
if len(req.Messages) == 0 {
|
||||
writeError(w, http.StatusBadRequest, "invalid_request", "messages is required")
|
||||
return
|
||||
@ -295,14 +396,28 @@ func (g *Gateway) handleChat(w http.ResponseWriter, r *http.Request) {
|
||||
if model == "" {
|
||||
model = g.core.DefaultModel()
|
||||
}
|
||||
// request_start fires for EVERY chat request, on both the AUTO and the
|
||||
// direct path, and it fires BEFORE the quota / model-scope gates on
|
||||
// purpose: a plugin that counts volume or audits traffic must also see the
|
||||
// requests the gateway rejected, otherwise "requests accepted" would be all
|
||||
// it could ever report. It sits after authentication (so the key and role in
|
||||
// the payload are real) and after the messages check (a body with no
|
||||
// messages is not a chat request at all).
|
||||
//
|
||||
// Calling it here rather than inside each branch is what keeps the two paths
|
||||
// honest: an earlier version called it only from the AUTO branch, so every
|
||||
// direct (model-pinned) request silently skipped it. That was caught by
|
||||
// TestHooksFireOnRealDirectChat, not by reading the code.
|
||||
g.fireStart(r.Context(), &req, "chat", model, len(req.Messages), len(req.Tools))
|
||||
if isAuto(model) {
|
||||
chain := g.core.AutoChain()
|
||||
// The key's own chain wins; keys without one inherit the global chain.
|
||||
chain, _ := g.core.AutoChainFor(reqKey(r.Context()))
|
||||
if chain == nil || len(chain.Tiers) == 0 {
|
||||
writeError(w, http.StatusServiceUnavailable, "no_provider", "no auto slot configured")
|
||||
return
|
||||
}
|
||||
if msg := g.checkModelScope(r.Context(), "AUTO"); msg != "" {
|
||||
writeError(w, http.StatusForbidden, "model_not_allowed", msg)
|
||||
if q := g.checkQuota(r.Context(), "AUTO"); q != nil {
|
||||
g.writeReject(w, q)
|
||||
return
|
||||
}
|
||||
ctx := r.Context()
|
||||
@ -317,6 +432,7 @@ func (g *Gateway) handleChat(w http.ResponseWriter, r *http.Request) {
|
||||
ToolChoice: req.ToolChoice,
|
||||
DisableThinking: req.DisableThinking,
|
||||
ExtraBody: req.ExtraBody,
|
||||
ClientSession: clientSession,
|
||||
}
|
||||
rec := &Req{
|
||||
Key: keyID(reqKey(ctx)),
|
||||
@ -346,6 +462,11 @@ func (g *Gateway) handleChat(w http.ResponseWriter, r *http.Request) {
|
||||
return
|
||||
}
|
||||
}
|
||||
// Note: a key whose scope excludes AUTO is rejected earlier, by
|
||||
// checkQuota(ctx, "AUTO") above, which produces the same message. That
|
||||
// ordering matters — checkQuota runs before the chain is consulted, so it
|
||||
// is the check an AUTO request actually hits. Duplicating it here would be
|
||||
// dead code.
|
||||
cands, effective := g.resolveCands(r.Context(), &req)
|
||||
if len(cands) == 0 {
|
||||
writeError(w, http.StatusNotFound, "model_not_found", fmt.Sprintf("model %q is not configured", model))
|
||||
@ -354,8 +475,8 @@ func (g *Gateway) handleChat(w http.ResponseWriter, r *http.Request) {
|
||||
if effective == "" {
|
||||
effective = firstModel(cands[0])
|
||||
}
|
||||
if msg := g.checkModelScope(r.Context(), effective); msg != "" {
|
||||
writeError(w, http.StatusForbidden, "model_not_allowed", msg)
|
||||
if q := g.checkQuota(r.Context(), effective); q != nil {
|
||||
g.writeReject(w, q)
|
||||
return
|
||||
}
|
||||
ctx := r.Context()
|
||||
@ -372,6 +493,7 @@ func (g *Gateway) handleChat(w http.ResponseWriter, r *http.Request) {
|
||||
ToolChoice: req.ToolChoice,
|
||||
DisableThinking: req.DisableThinking,
|
||||
ExtraBody: req.ExtraBody,
|
||||
ClientSession: clientSession,
|
||||
}
|
||||
rec := &Req{
|
||||
Key: keyID(reqKey(ctx)),
|
||||
@ -517,21 +639,115 @@ func upstreamErrStatus(err error) int {
|
||||
return http.StatusBadGateway
|
||||
}
|
||||
|
||||
// overflowMarkers 匹配上游「上下文超窗」类措辞。
|
||||
//
|
||||
// 上游写法五花八门,且**不在 pi 客户端的识别列表里**。pi 靠
|
||||
// @earendil-works/pi-ai 的 OVERFLOW_PATTERNS 判断超窗并据此触发压缩重试,
|
||||
// 而 justworker 返回的是「请精简对话历史…(Context window is full…)」——
|
||||
// 与那 25 条正则一条都不匹配,于是 pi 既不压缩也不重试,只把它当成一条
|
||||
// 普通上游错误。这里把可识别的超窗措辞归一化成 pi 一定认得的标记。
|
||||
var overflowMarkers = []*regexp.Regexp{
|
||||
regexp.MustCompile(`(?i)context[ _-]?window is full`),
|
||||
regexp.MustCompile(`(?i)context[_ ]length[_ ]exceeded`),
|
||||
regexp.MustCompile(`(?i)exceeds? the context window`),
|
||||
regexp.MustCompile(`(?i)maximum context length`),
|
||||
regexp.MustCompile(`(?i)reduce the length of the messages`),
|
||||
regexp.MustCompile(`(?i)too many tokens`),
|
||||
regexp.MustCompile(`(?i)token limit exceeded`),
|
||||
regexp.MustCompile(`请精简对话历史`),
|
||||
regexp.MustCompile(`上下文(长度)?超(出|限)`),
|
||||
regexp.MustCompile(`对话历史过长`),
|
||||
}
|
||||
|
||||
// overflowCanonical 命中 pi 的 /context[_ ]length[_ ]exceeded/i。
|
||||
//
|
||||
// 注意:这个前缀只是必要条件,不是充分条件——pi 还会先用
|
||||
// NON_OVERFLOW_PATTERNS 排除整条消息,见 overflowClientMessage。
|
||||
const overflowCanonical = "context_length_exceeded"
|
||||
|
||||
// looksLikeOverflow 判断错误文本是否属于上下文超窗。
|
||||
func looksLikeOverflow(s string) bool {
|
||||
for _, re := range overflowMarkers {
|
||||
if re.MatchString(s) {
|
||||
return true
|
||||
}
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
// clientUpstreamErr collapses upstream failure details into a short message
|
||||
// for the client: per-tier bodies (WAF HTML pages, quota payloads, ...) stay
|
||||
// in rec.Err / the stats API and the server log instead of the response.
|
||||
// Per-tier one-line reasons are kept (quota/cooling skips carry no error body
|
||||
// and are the actionable part); each is capped so HTML dumps can't leak.
|
||||
//
|
||||
// 超窗会被归一化成「干净」的 overflowCanonical 消息,见 overflowClientMessage。
|
||||
func clientUpstreamErr(err error) string {
|
||||
log.Printf("[gateway] upstream failure surfaced to client: %v", err)
|
||||
if looksLikeOverflow(err.Error()) {
|
||||
return overflowClientMessage(err)
|
||||
}
|
||||
return upstreamErrSummary(err)
|
||||
}
|
||||
|
||||
// overflowHint 是给客户端的超窗说明,刻意只用 pi 认得的措辞。
|
||||
const overflowHint = "context window is full; reduce the length of the messages"
|
||||
|
||||
// clientSessionFromRequest 取客户端自带的会话标识,拿不到时返回 ""。
|
||||
//
|
||||
// 通用 OpenAI 客户端默认不带会话 id,但 pi 支持:只要 provider 的 compat 里打开
|
||||
// sendSessionAffinityHeaders,pi 就会把**平台真实的会话 id**(uuidv7,整个会话
|
||||
// 恒定)放到 x-session-affinity / x-client-request-id 上(sessionAffinityFormat
|
||||
// 为 openrouter 时是 x-session-id,为 openai 时是 session_id)。
|
||||
// 有了它就不必再从首条 user 消息推导会话指纹(那只在客户端不发会话 id 时才作为
|
||||
// 退路,且历史压缩后会漂移)。prompt_cache_key 是 OpenAI 原生的缓存键,
|
||||
// 客户端若带也一并当会话用。
|
||||
//
|
||||
// 刻意不采纳 x-client-request-id:名字含 request,部分客户端每请求都换,
|
||||
// 拿它当会话会让上游缓存永不命中。pi 总会同时发 x-session-affinity,够用。
|
||||
func clientSessionFromRequest(r *http.Request, bodyKey string) string {
|
||||
for _, h := range []string{"X-Session-Affinity", "X-Session-Id", "Session-Id"} {
|
||||
if v := strings.TrimSpace(r.Header.Get(h)); v != "" {
|
||||
return v
|
||||
}
|
||||
}
|
||||
return strings.TrimSpace(bodyKey)
|
||||
}
|
||||
|
||||
// overflowClientMessage 为超窗失败生成「干净」的客户端消息。
|
||||
//
|
||||
// 为什么不能沿用 upstreamErrSummary:pi 的 isContextOverflow 先查
|
||||
// NON_OVERFLOW_PATTERNS(/rate limit/、/too many requests/、Bedrock 前缀),
|
||||
// 一旦命中就直接判为「非超窗」——**哪怕消息里已经有 context_length_exceeded**,
|
||||
// pi 也不会压缩重试。而 AUTO 链的失败消息天生是多 tier 原因的拼接:
|
||||
// 超窗 tier(gozen 400 maximum context length)常与配额/限流 tier
|
||||
// (429 token plan exhausted、cooling、no free slot)同时出现。
|
||||
// 把明细原样带出去,等于让一条限流 tier 的措辞反过来封杀超窗识别。
|
||||
// 所以这里只保留「超窗」措辞 + 是哪个源超的窗,其余一律不带。
|
||||
func overflowClientMessage(err error) string {
|
||||
var ce *scheduler.ChainErr
|
||||
if errors.As(err, &ce) {
|
||||
for _, t := range ce.Tiers {
|
||||
if looksLikeOverflow(t.Err.Error()) {
|
||||
return fmt.Sprintf("%s: %s (%s/%s)", overflowCanonical, overflowHint, t.Source, t.Model)
|
||||
}
|
||||
}
|
||||
}
|
||||
return fmt.Sprintf("%s: %s", overflowCanonical, overflowHint)
|
||||
}
|
||||
|
||||
// upstreamErrSummary 把上游失败压成一行短消息。
|
||||
func upstreamErrSummary(err error) string {
|
||||
var ce *scheduler.ChainErr
|
||||
if errors.As(err, &ce) {
|
||||
parts := make([]string, 0, len(ce.Tiers)+len(ce.Skipped))
|
||||
for _, t := range ce.Tiers {
|
||||
parts = append(parts, types.OneLine(fmt.Sprintf("%s/%s: %v", t.Source, t.Model, t.Err), 80))
|
||||
// 160 而不是 80:短诊断词("Context window is full")常落在尾部,
|
||||
// 80 字节按字节截断正好会把它切掉,超窗就再也认不出来。
|
||||
parts = append(parts, types.OneLine(fmt.Sprintf("%s/%s: %v", t.Source, t.Model, t.Err), 160))
|
||||
}
|
||||
for _, sk := range ce.Skipped {
|
||||
parts = append(parts, types.OneLine(sk, 80))
|
||||
parts = append(parts, types.OneLine(sk, 160))
|
||||
}
|
||||
msg := strings.Join(parts, "; ")
|
||||
if len(msg) > 300 {
|
||||
@ -573,6 +789,18 @@ func recordChatUsage(rec *Req, req *types.ChatRequest, resp *types.UnifiedRespon
|
||||
// Cache accounting: whichever source format the adapter normalized into
|
||||
// (prompt_tokens_details.cached_tokens or legacy prompt_cache_hit_tokens),
|
||||
// read it back so the request record carries the hit/miss split.
|
||||
// reasoning_tokens is a SUBSET of completion_tokens and is billed at the
|
||||
// completion rate by every provider measured, so it must NOT be added to
|
||||
// Compl. Recording it separately is what lets a price rule split thinking
|
||||
// from output for the providers that do price them differently, and what
|
||||
// keeps the billing plugin from guessing.
|
||||
//
|
||||
// Leaving it unrecorded is worse than useless: the adapters already parse it
|
||||
// and forward it (opencodego/gemini/anthropic all do), so the number exists
|
||||
// in the response and is thrown away at the one place that could keep it.
|
||||
if d := resp.TokenUsage.CompletionTokensDetails; d != nil && d.ReasoningTokens > 0 {
|
||||
rec.Reasoning = int64(d.ReasoningTokens)
|
||||
}
|
||||
if d := resp.TokenUsage.PromptTokensDetails; d != nil && d.CachedTokens > 0 {
|
||||
rec.CacheHit = int64(d.CachedTokens)
|
||||
}
|
||||
@ -606,6 +834,7 @@ func writeChatCompletion(w http.ResponseWriter, resp *types.UnifiedResponse, mod
|
||||
Created: time.Now().Unix(),
|
||||
Model: modelName,
|
||||
Choices: []ChatChoice{{Index: 0, Message: msg, FinishReason: resp.FinishReason}},
|
||||
Cost: resp.Cost,
|
||||
}
|
||||
if resp.TokenUsage.Total > 0 || resp.TokenUsage.Prompt > 0 || resp.TokenUsage.Completion > 0 {
|
||||
out.Usage = &resp.TokenUsage
|
||||
@ -630,6 +859,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
|
||||
@ -637,7 +867,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
|
||||
@ -646,6 +887,194 @@ 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(),
|
||||
})
|
||||
}
|
||||
|
||||
// fireImageStart dispatches request_start for /v1/images/generations.
|
||||
//
|
||||
// It is a separate function rather than a call to fireStart with a nil
|
||||
// chatRequest because the image body has no messages and no tools: passing
|
||||
// zeroes through a struct built for chat would invite someone to read a field
|
||||
// that simply does not exist on this path.
|
||||
func (g *Gateway) fireImageStart(ctx context.Context, model string) {
|
||||
ps := g.core.Plugins()
|
||||
if ps == nil || ps.Count() == 0 {
|
||||
return
|
||||
}
|
||||
ps.Fire(lua.StageRequestStart, map[string]interface{}{
|
||||
"stage": string(lua.StageRequestStart),
|
||||
"type": "image",
|
||||
"model": model,
|
||||
"key": keyID(reqKey(ctx)),
|
||||
"role": reqRole(ctx),
|
||||
"source": "",
|
||||
"stream": false,
|
||||
"messages_count": 0,
|
||||
"tools_count": 0,
|
||||
"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(),
|
||||
})
|
||||
}
|
||||
|
||||
// chainTraceSink adapts a scheduler TraceSink into the plugin chain_step stage.
|
||||
//
|
||||
// It returns nil when no plugin is loaded, so the scheduler's emit() does a
|
||||
// single nil check per event and the AUTO hot path pays nothing on a gateway
|
||||
// with no plugins.
|
||||
//
|
||||
// The events are also accumulated into walk so request_end can carry a compact
|
||||
// summary: a plugin that only listens to request_end still learns that a
|
||||
// degradation happened, which is the common case for a dashboard that does not
|
||||
// want to subscribe to a high-frequency stage.
|
||||
func (g *Gateway) chainTraceSink(ctx context.Context, kind string, walk *[]map[string]interface{}) scheduler.TraceSink {
|
||||
ps := g.core.Plugins()
|
||||
if ps == nil || ps.Count() == 0 {
|
||||
return nil
|
||||
}
|
||||
key := keyID(reqKey(ctx))
|
||||
return func(ev scheduler.TraceEvent) {
|
||||
payload := map[string]interface{}{
|
||||
"stage": string(lua.StageChainStep),
|
||||
"kind": string(ev.Kind),
|
||||
"type": kind,
|
||||
"key": key,
|
||||
"tier": ev.Tier,
|
||||
"attempt": ev.Attempt,
|
||||
}
|
||||
if ev.Source != "" {
|
||||
payload["source"] = ev.Source
|
||||
}
|
||||
if ev.Model != "" {
|
||||
payload["model"] = ev.Model
|
||||
}
|
||||
if ev.Reason != "" {
|
||||
payload["reason"] = ev.Reason
|
||||
}
|
||||
if ev.Err != "" {
|
||||
payload["error"] = ev.Err
|
||||
}
|
||||
if walk != nil {
|
||||
// Keep the summary bounded: a pathological chain could emit many
|
||||
// steps, and request_end's payload is written to the audit trail.
|
||||
if len(*walk) < maxWalkSummary {
|
||||
*walk = append(*walk, map[string]interface{}{
|
||||
"kind": string(ev.Kind), "tier": ev.Tier,
|
||||
"source": ev.Source, "model": ev.Model, "reason": ev.Reason,
|
||||
})
|
||||
}
|
||||
}
|
||||
ps.Fire(lua.StageChainStep, payload)
|
||||
}
|
||||
}
|
||||
|
||||
// maxWalkSummary caps how many chain steps request_end carries, so a long
|
||||
// degradation cannot inflate every audit record.
|
||||
const maxWalkSummary = 12
|
||||
|
||||
// tierServed returns the AUTO tier that actually served the request, or -1 when
|
||||
// the walk is empty (a direct request) or ended without a selection (total
|
||||
// failure). It is the single most useful number for "why did my expensive tier
|
||||
// not get used".
|
||||
func tierServed(walk []map[string]interface{}) int {
|
||||
for i := len(walk) - 1; i >= 0; i-- {
|
||||
if k, _ := walk[i]["kind"].(string); k == string(scheduler.TraceSelected) {
|
||||
if t, ok := walk[i]["tier"].(int); ok {
|
||||
return t
|
||||
}
|
||||
}
|
||||
}
|
||||
return -1
|
||||
}
|
||||
|
||||
// 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,
|
||||
// Whether UPSTREAM reported a cache number at all. A plugin cannot
|
||||
// infer this from cache_hit_tokens alone: zero hits because nothing was
|
||||
// cached and zero hits because the provider never reports caching are
|
||||
// the same value, and they mean opposite things when you are checking
|
||||
// whether a cache discount is doing anything. The audit record already
|
||||
// carried this (rec.CacheReported); the plugin payload did not, so any
|
||||
// plugin-level cache accounting had to guess.
|
||||
"cache_reported": rec.CacheReported,
|
||||
// A SUBSET of completion_tokens, not an extra. Billing plugins that
|
||||
// add it to completion would double-charge every thinking token.
|
||||
"reasoning_tokens": rec.Reasoning,
|
||||
"image_count": rec.ImageCount,
|
||||
"error": rec.Err,
|
||||
"time": rec.Time,
|
||||
// chain_walk: the AUTO tier-by-tier trace, when the request went
|
||||
// through the chain. Empty for a direct request and for a gateway with
|
||||
// no plugins loaded. Absent rather than empty so a plugin can tell
|
||||
// "no chain" from "chain with no degradation".
|
||||
"degraded": len(rec.Walk) > 1,
|
||||
"chain_walk": rec.Walk,
|
||||
"tier_served": tierServed(rec.Walk),
|
||||
}
|
||||
// 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
|
||||
@ -725,11 +1154,15 @@ func (g *Gateway) pumpStream(w http.ResponseWriter, rec *Req, chunks <-chan type
|
||||
// source's status-page latency average.
|
||||
rec.FirstByteMs = time.Since(t0).Milliseconds()
|
||||
var lastUsage *types.TokenUsage
|
||||
var lastCost string
|
||||
lastFinish := ""
|
||||
for ck := range chunks {
|
||||
if ck.Usage != nil {
|
||||
lastUsage = mergeUsage(lastUsage, ck.Usage)
|
||||
}
|
||||
if ck.Cost != "" {
|
||||
lastCost = ck.Cost
|
||||
}
|
||||
chunk := ChatChunk{
|
||||
ID: id, Object: "chat.completion.chunk", Created: created, Model: modelName,
|
||||
}
|
||||
@ -766,23 +1199,23 @@ func (g *Gateway) pumpStream(w http.ResponseWriter, rec *Req, chunks <-chan type
|
||||
// Final usage chunk (OpenAI standard: empty choices + usage before [DONE]).
|
||||
// Prefer the upstream's exact usage if the stream carried it; fall back to
|
||||
// the gateway's estimate otherwise.
|
||||
// Write the upstream's exact numbers back onto the audit record. Without
|
||||
// this the streamed path kept only the per-chunk byte estimate, so the same
|
||||
// request recorded ~1.4x its real completion tokens (measured: upstream 100,
|
||||
// audit 145) while the non-streaming path recorded 100. Two paths, two
|
||||
// different numbers for one request is a reporting bug, not a rounding one.
|
||||
if lastUsage != nil {
|
||||
if lastUsage.Prompt > 0 {
|
||||
rec.Prompt = int64(lastUsage.Prompt)
|
||||
}
|
||||
if lastUsage.Completion > 0 {
|
||||
rec.Compl = int64(lastUsage.Completion)
|
||||
}
|
||||
}
|
||||
var tut *types.TokenUsage
|
||||
if lastUsage != nil {
|
||||
tut = lastUsage
|
||||
// Write the upstream's cache accounting back onto the request record
|
||||
// so the audit trail carries the hit/miss split for streaming too.
|
||||
if d := tut.PromptTokensDetails; d != nil {
|
||||
rec.CacheReported = true
|
||||
if d.CachedTokens > 0 {
|
||||
rec.CacheHit = int64(d.CachedTokens)
|
||||
}
|
||||
} else if tut.PromptCacheHit > 0 {
|
||||
rec.CacheHit = int64(tut.PromptCacheHit)
|
||||
rec.CacheReported = true
|
||||
}
|
||||
if tut.PromptCacheMiss > 0 {
|
||||
rec.CacheMiss = int64(tut.PromptCacheMiss)
|
||||
}
|
||||
applyStreamUsageDetails(rec, lastUsage)
|
||||
} else if rec.Prompt+rec.Compl > 0 {
|
||||
u := types.TokenUsage{
|
||||
Prompt: int(rec.Prompt),
|
||||
@ -796,6 +1229,14 @@ func (g *Gateway) pumpStream(w http.ResponseWriter, rec *Req, chunks <-chan type
|
||||
ID: id, Object: "chat.completion.chunk", Created: created, Model: modelName,
|
||||
Choices: []ChunkChoice{},
|
||||
Usage: tut,
|
||||
Cost: lastCost,
|
||||
})
|
||||
} else if lastCost != "" {
|
||||
// Cost arrived without any usage (OpenCode sends it on its own frame).
|
||||
send(ChatChunk{
|
||||
ID: id, Object: "chat.completion.chunk", Created: created, Model: modelName,
|
||||
Choices: []ChunkChoice{},
|
||||
Cost: lastCost,
|
||||
})
|
||||
}
|
||||
fmt.Fprintf(w, "data: [DONE]\n\n")
|
||||
@ -827,6 +1268,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)
|
||||
}
|
||||
@ -840,8 +1282,11 @@ func (g *Gateway) streamChat(w http.ResponseWriter, ctx context.Context, cands [
|
||||
func (g *Gateway) singleChatAuto(w http.ResponseWriter, ctx context.Context, chain *scheduler.Chain, req *types.ChatRequest, rec *Req, quotaExhausted func(*scheduler.Slot) bool) {
|
||||
rec.LatMs = 0
|
||||
t0 := time.Now()
|
||||
resp, usedSrc, usedModel, err := g.core.Scheduler().ChainChat(ctx, chain, req, quotaExhausted)
|
||||
var walk []map[string]interface{}
|
||||
resp, usedSrc, usedModel, err := g.core.Scheduler().ChainChat(ctx, chain, req, quotaExhausted,
|
||||
g.chainTraceSink(ctx, "chat", &walk))
|
||||
rec.LatMs = time.Since(t0).Milliseconds()
|
||||
rec.Walk = walk
|
||||
if err != nil {
|
||||
g.failChat(w, rec, err)
|
||||
g.writeRec(rec)
|
||||
@ -852,6 +1297,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)
|
||||
@ -870,7 +1319,10 @@ func (g *Gateway) streamChatAuto(w http.ResponseWriter, ctx context.Context, cha
|
||||
rec.LatMs = time.Since(t0).Milliseconds()
|
||||
g.writeRec(rec)
|
||||
}()
|
||||
chunks, usedSrc, usedModel, err := g.core.Scheduler().ChainChatStream(ctx, chain, req, quotaExhausted)
|
||||
var walk []map[string]interface{}
|
||||
chunks, usedSrc, usedModel, err := g.core.Scheduler().ChainChatStream(ctx, chain, req, quotaExhausted,
|
||||
g.chainTraceSink(ctx, "stream", &walk))
|
||||
rec.Walk = walk
|
||||
if err != nil {
|
||||
g.failChat(w, rec, err)
|
||||
return
|
||||
@ -879,6 +1331,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)
|
||||
}
|
||||
@ -901,10 +1354,15 @@ func (g *Gateway) handleImage(w http.ResponseWriter, r *http.Request) {
|
||||
if model == "" {
|
||||
model = g.core.DefaultModel()
|
||||
}
|
||||
// Same rule as the chat path, and for the same reason: an image request is
|
||||
// billable traffic, so a cost plugin must see it. It fires before the
|
||||
// quota/scope gates so rejected image requests are visible too.
|
||||
// messages_count/tools_count are 0: the image request has neither.
|
||||
g.fireImageStart(r.Context(), model)
|
||||
if isAuto(model) {
|
||||
if chain := g.core.AutoImageChain(); chain != nil && len(chain.Tiers) > 0 {
|
||||
if msg := g.checkModelScope(r.Context(), "AUTO"); msg != "" {
|
||||
writeError(w, http.StatusForbidden, "model_not_allowed", msg)
|
||||
if q := g.checkQuota(r.Context(), "AUTO"); q != nil {
|
||||
g.writeReject(w, q)
|
||||
return
|
||||
}
|
||||
done := g.stats.Begin()
|
||||
@ -924,9 +1382,13 @@ func (g *Gateway) handleImage(w http.ResponseWriter, r *http.Request) {
|
||||
if usedModel != "" {
|
||||
rec.Model = usedModel // actual image model served, not "AUTO"
|
||||
}
|
||||
g.fireRouted(r.Context(), "image", usedSrc, rec.Model, -2, false)
|
||||
rec.OK = true
|
||||
rec.Status = http.StatusOK
|
||||
rec.Compl = int64(len(resp.ImageData))
|
||||
// Image generation has no token concept. Recording len(ImageData)
|
||||
// (the image COUNT) in completion_tokens mislabels image count as
|
||||
// tokens and feeds it into the token totals; leave it 0.
|
||||
rec.ImageCount = len(resp.ImageData)
|
||||
g.writeRec(rec)
|
||||
writeJSON(w, http.StatusOK, types.ImageGenResponse{
|
||||
Created: time.Now().Unix(),
|
||||
@ -950,8 +1412,8 @@ func (g *Gateway) handleImage(w http.ResponseWriter, r *http.Request) {
|
||||
writeError(w, http.StatusServiceUnavailable, "no_provider", "no image source configured")
|
||||
return
|
||||
}
|
||||
if msg := g.checkModelScope(r.Context(), effectiveImageModel(model, cands)); msg != "" {
|
||||
writeError(w, http.StatusForbidden, "model_not_allowed", msg)
|
||||
if q := g.checkQuota(r.Context(), effectiveImageModel(model, cands)); q != nil {
|
||||
g.writeReject(w, q)
|
||||
return
|
||||
}
|
||||
done := g.stats.Begin()
|
||||
@ -973,12 +1435,50 @@ func (g *Gateway) handleImage(w http.ResponseWriter, r *http.Request) {
|
||||
if resp.Model != "" {
|
||||
rec.Model = resp.Model // record the actual model served, not the raw request id
|
||||
}
|
||||
g.fireRouted(r.Context(), "image", rec.Source, rec.Model, -1, false)
|
||||
rec.OK = true
|
||||
rec.Status = http.StatusOK
|
||||
rec.Compl = int64(len(resp.ImageData))
|
||||
// Image generation has no token concept — see the AUTO path above.
|
||||
rec.ImageCount = len(resp.ImageData)
|
||||
g.writeRec(rec)
|
||||
writeJSON(w, http.StatusOK, types.ImageGenResponse{
|
||||
Created: time.Now().Unix(),
|
||||
Data: resp.ImageData,
|
||||
})
|
||||
}
|
||||
|
||||
// applyStreamUsageDetails mirrors the upstream's reasoning and cache numbers
|
||||
// onto a STREAMED request's audit record.
|
||||
//
|
||||
// It exists as a named function because this was inline in the streaming handler,
|
||||
// where it was untestable — and untested it stayed wrong. The non-streaming
|
||||
// recorder read completion_tokens_details.reasoning_tokens; this one did not. So
|
||||
// a source's streamed requests recorded completion and no reasoning at all while
|
||||
// its non-streamed calls recorded both. Nothing failed and no test was red: the
|
||||
// audit trail was simply quieter than it should have been.
|
||||
//
|
||||
// Found by reading the production audit file after wiring a reasoning-heavy
|
||||
// source: completion=8166 with reasoning absent, on every streamed request.
|
||||
//
|
||||
// Reasoning is a SUBSET of completion, never an addition to it. See
|
||||
// recordChatUsage for why that matters to billing.
|
||||
func applyStreamUsageDetails(rec *Req, u *types.TokenUsage) {
|
||||
if d := u.CompletionTokensDetails; d != nil && d.ReasoningTokens > 0 {
|
||||
rec.Reasoning = int64(d.ReasoningTokens)
|
||||
}
|
||||
// Cache accounting is here for the same reason it is kept apart: "zero hits
|
||||
// because nothing was cached" and "this provider never reports caching" are
|
||||
// the same number and mean opposite things.
|
||||
if d := u.PromptTokensDetails; d != nil {
|
||||
rec.CacheReported = true
|
||||
if d.CachedTokens > 0 {
|
||||
rec.CacheHit = int64(d.CachedTokens)
|
||||
}
|
||||
} else if u.PromptCacheHit > 0 {
|
||||
rec.CacheHit = int64(u.PromptCacheHit)
|
||||
rec.CacheReported = true
|
||||
}
|
||||
if u.PromptCacheMiss > 0 {
|
||||
rec.CacheMiss = int64(u.PromptCacheMiss)
|
||||
}
|
||||
}
|
||||
|
||||
160
internal/gateway/chat_reasoning_test.go
Normal file
160
internal/gateway/chat_reasoning_test.go
Normal file
@ -0,0 +1,160 @@
|
||||
package gateway
|
||||
|
||||
import (
|
||||
"testing"
|
||||
|
||||
"llmsproxy/internal/types"
|
||||
)
|
||||
|
||||
// TestRecordChatUsageKeepsReasoningTokens pins the reasoning-token plumbing.
|
||||
//
|
||||
// The adapters already parse completion_tokens_details.reasoning_tokens and
|
||||
// forward it (opencodego, gemini, anthropic all do), and types.go has carried
|
||||
// the field since it was defined — but nothing ever READ it. The number arrived
|
||||
// in the response and was discarded at the one place that could have kept it,
|
||||
// so a pricing rule had no way to split thinking from output.
|
||||
//
|
||||
// The property that matters most: reasoning is a SUBSET of completion. Every
|
||||
// provider measured bills it at the completion rate, so a billing plugin that
|
||||
// added it to completion_tokens would double-charge every thinking token. The
|
||||
// test asserts completion is left alone.
|
||||
func TestRecordChatUsageKeepsReasoningTokens(t *testing.T) {
|
||||
rec := &Req{}
|
||||
resp := &types.UnifiedResponse{
|
||||
TokenUsage: types.TokenUsage{
|
||||
Prompt: 100,
|
||||
Completion: 500,
|
||||
CompletionTokensDetails: &types.CompletionTokensDetails{
|
||||
ReasoningTokens: 320,
|
||||
},
|
||||
},
|
||||
}
|
||||
recordChatUsage(rec, &types.ChatRequest{}, resp)
|
||||
|
||||
if rec.Reasoning != 320 {
|
||||
t.Errorf("Reasoning = %d, want 320", rec.Reasoning)
|
||||
}
|
||||
// The subset must not be folded into completion — it is already counted there.
|
||||
if rec.Compl != 500 {
|
||||
t.Errorf("Compl = %d, want 500 (reasoning is a subset, never additive)", rec.Compl)
|
||||
}
|
||||
if rec.Prompt != 100 {
|
||||
t.Errorf("Prompt = %d, want 100", rec.Prompt)
|
||||
}
|
||||
}
|
||||
|
||||
// TestRecordChatUsageReasoningAbsentStaysZero: "not reported" and "reported as
|
||||
// zero" must both leave the field at 0, and neither may invent a number. A
|
||||
// plugin reading Reasoning cannot tell them apart, which is why the field is
|
||||
// documented as "0 means not reported" rather than "0 means no thinking".
|
||||
func TestRecordChatUsageReasoningAbsentStaysZero(t *testing.T) {
|
||||
for _, tc := range []struct {
|
||||
name string
|
||||
u types.TokenUsage
|
||||
}{
|
||||
{"no details object", types.TokenUsage{Prompt: 10, Completion: 20}},
|
||||
{"details without reasoning", types.TokenUsage{
|
||||
Prompt: 10,
|
||||
Completion: 20,
|
||||
CompletionTokensDetails: &types.CompletionTokensDetails{},
|
||||
}},
|
||||
} {
|
||||
rec := &Req{}
|
||||
recordChatUsage(rec, &types.ChatRequest{}, &types.UnifiedResponse{TokenUsage: tc.u})
|
||||
if rec.Reasoning != 0 {
|
||||
t.Errorf("%s: Reasoning = %d, want 0", tc.name, rec.Reasoning)
|
||||
}
|
||||
if rec.Compl != 20 {
|
||||
t.Errorf("%s: Compl = %d, want 20", tc.name, rec.Compl)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// TestStreamedReasoningIsRecordedOnTheRequest is the streaming counterpart.
|
||||
//
|
||||
// TestRecordChatUsageKeepsReasoningTokens covers the non-streaming path only,
|
||||
// and for a long time the streamed path dropped the number: the non-streaming
|
||||
// recorder read TokenUsage.CompletionTokensDetails, while the streaming one read
|
||||
// only Prompt and Completion from the terminal usage chunk. (Cache accounting was
|
||||
// mirrored; reasoning simply never was.)
|
||||
//
|
||||
// The symptom is invisible in the response — the client sees a correct answer
|
||||
// either way. It only shows in the audit trail, where a reasoning-heavy source's
|
||||
// streamed requests recorded completion=8166 and no reasoning at all. Found by
|
||||
// reading the production audit file, not by a failing test, which is exactly why
|
||||
// the property is pinned here.
|
||||
func TestStreamedReasoningIsRecordedOnTheRequest(t *testing.T) {
|
||||
rec := &Req{}
|
||||
applyStreamUsageDetails(rec, &types.TokenUsage{
|
||||
Prompt: 159579,
|
||||
Completion: 751,
|
||||
Total: 160330,
|
||||
CompletionTokensDetails: &types.CompletionTokensDetails{
|
||||
ReasoningTokens: 512,
|
||||
},
|
||||
})
|
||||
if rec.Reasoning != 512 {
|
||||
t.Errorf("streamed Reasoning = %d, want 512 — a streamed request records no "+
|
||||
"reasoning, so a pricing rule that splits thinking from output silently "+
|
||||
"bills it as ordinary output", rec.Reasoning)
|
||||
}
|
||||
}
|
||||
|
||||
// TestStreamedReasoningIsNeverAddedToCompletion pins the SUBSET relation in the
|
||||
// streaming path specifically. A presence check would pass if someone wrote
|
||||
// rec.Compl += reasoning, and billing would then charge every thinking token
|
||||
// twice.
|
||||
func TestStreamedReasoningIsNeverAddedToCompletion(t *testing.T) {
|
||||
rec := &Req{}
|
||||
rec.Compl = 751 // set upstream of the call, as the streaming handler does
|
||||
applyStreamUsageDetails(rec, &types.TokenUsage{
|
||||
Completion: 751,
|
||||
CompletionTokensDetails: &types.CompletionTokensDetails{ReasoningTokens: 512},
|
||||
})
|
||||
if rec.Compl != 751 {
|
||||
t.Errorf("Compl = %d, want 751 unchanged — reasoning is a subset, not an addition", rec.Compl)
|
||||
}
|
||||
if rec.Reasoning >= rec.Compl {
|
||||
t.Errorf("Reasoning (%d) must be a SUBSET of Compl (%d)", rec.Reasoning, rec.Compl)
|
||||
}
|
||||
}
|
||||
|
||||
// TestStreamedUsageWithoutReasoningLeavesItUnset: an upstream that reports no
|
||||
// reasoning must not be turned into a reported zero. In the audit trail 0 means
|
||||
// "not reported", and a plugin cannot tell that from "no reasoning happened".
|
||||
func TestStreamedUsageWithoutReasoningLeavesItUnset(t *testing.T) {
|
||||
rec := &Req{}
|
||||
applyStreamUsageDetails(rec, &types.TokenUsage{Prompt: 10, Completion: 4, Total: 14})
|
||||
if rec.Reasoning != 0 {
|
||||
t.Errorf("Reasoning = %d for an upstream that reported none", rec.Reasoning)
|
||||
}
|
||||
}
|
||||
|
||||
// TestStreamedCacheAccountingStillWorks is a regression guard for the block the
|
||||
// reasoning lines were added to. They share a function, so a future edit that
|
||||
// short-circuits one would break the other — silently, since a missing cache
|
||||
// number reads as "nothing was cached".
|
||||
func TestStreamedCacheAccountingStillWorks(t *testing.T) {
|
||||
for _, c := range []struct {
|
||||
name string
|
||||
u types.TokenUsage
|
||||
hit int64
|
||||
miss int64
|
||||
rep bool
|
||||
}{
|
||||
{"details object", types.TokenUsage{
|
||||
PromptTokensDetails: &types.PromptTokensDetails{CachedTokens: 4096},
|
||||
}, 4096, 0, true},
|
||||
{"legacy hit/miss", types.TokenUsage{
|
||||
PromptCacheHit: 2048, PromptCacheMiss: 512,
|
||||
}, 2048, 512, true},
|
||||
{"nothing reported", types.TokenUsage{}, 0, 0, false},
|
||||
} {
|
||||
rec := &Req{}
|
||||
applyStreamUsageDetails(rec, &c.u)
|
||||
if rec.CacheHit != c.hit || rec.CacheMiss != c.miss || rec.CacheReported != c.rep {
|
||||
t.Errorf("%s: hit=%d miss=%d reported=%v; want %d/%d/%v",
|
||||
c.name, rec.CacheHit, rec.CacheMiss, rec.CacheReported, c.hit, c.miss, c.rep)
|
||||
}
|
||||
}
|
||||
}
|
||||
47
internal/gateway/client_session_test.go
Normal file
47
internal/gateway/client_session_test.go
Normal file
@ -0,0 +1,47 @@
|
||||
package gateway
|
||||
|
||||
import (
|
||||
"net/http"
|
||||
"testing"
|
||||
)
|
||||
|
||||
// 客户端自带的会话标识必须被采纳:它比「首条 user 消息指纹」更准
|
||||
// (历史压缩后指纹会漂移),也是唯一能天然按会话隔离的来源。
|
||||
func TestClientSessionFromRequestHeaders(t *testing.T) {
|
||||
cases := []struct {
|
||||
name string
|
||||
headers map[string]string
|
||||
bodyKey string
|
||||
want string
|
||||
}{
|
||||
{"pi 默认格式 x-session-affinity", map[string]string{"X-Session-Affinity": "ses-1"}, "", "ses-1"},
|
||||
{"openrouter 格式 x-session-id", map[string]string{"X-Session-Id": "ses-2"}, "", "ses-2"},
|
||||
{"openai 格式 session_id", map[string]string{"Session-Id": "ses-3"}, "", "ses-3"},
|
||||
{"body prompt_cache_key 兜底", nil, "cache-4", "cache-4"},
|
||||
{"头优先于 body", map[string]string{"X-Session-Affinity": "ses-5"}, "cache-x", "ses-5"},
|
||||
{"空白不算", map[string]string{"X-Session-Affinity": " "}, "", ""},
|
||||
{"都没有则为空(适配器自行推导)", nil, "", ""},
|
||||
}
|
||||
|
||||
for _, tc := range cases {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
r, _ := http.NewRequest(http.MethodPost, "/v1/chat/completions", nil)
|
||||
for k, v := range tc.headers {
|
||||
r.Header.Set(k, v)
|
||||
}
|
||||
if got := clientSessionFromRequest(r, tc.bodyKey); got != tc.want {
|
||||
t.Errorf("clientSessionFromRequest = %q, want %q", got, tc.want)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// 刻意不采纳 x-client-request-id:名字含 request,部分客户端每请求都换,
|
||||
// 拿它当会话会让上游前缀缓存永不命中。
|
||||
func TestClientRequestIDIsNotUsedAsSession(t *testing.T) {
|
||||
r, _ := http.NewRequest(http.MethodPost, "/v1/chat/completions", nil)
|
||||
r.Header.Set("X-Client-Request-Id", "per-request-uuid")
|
||||
if got := clientSessionFromRequest(r, ""); got != "" {
|
||||
t.Errorf("x-client-request-id 不应被当作会话(否则缓存永不命中),got %q", got)
|
||||
}
|
||||
}
|
||||
204
internal/gateway/deploy_script_test.go
Normal file
204
internal/gateway/deploy_script_test.go
Normal file
@ -0,0 +1,204 @@
|
||||
package gateway
|
||||
|
||||
import (
|
||||
"os"
|
||||
"os/exec"
|
||||
"path/filepath"
|
||||
"regexp"
|
||||
"strings"
|
||||
"testing"
|
||||
)
|
||||
|
||||
// deploy.sh is the only release-critical script with no automated coverage, and
|
||||
// its newest part — prune_adapter_backups — deletes directories under
|
||||
// /etc/llmsproxy. This exercises the function in a throwaway directory with the
|
||||
// REAL backup-name shape.
|
||||
//
|
||||
// The shape matters and got this wrong twice: the guard is
|
||||
// `^adapters\.bak\.[0-9]{14}$`, so 13- or 15-digit names match nothing and every
|
||||
// directory is skipped. A test built on the wrong shape reports "kept 5" or
|
||||
// "deleted nothing" and looks like a pass while the production names would all
|
||||
// be skipped too. The names below are copied from the live directory listing.
|
||||
func TestDeployPruneAdapterBackups(t *testing.T) {
|
||||
if _, err := os.Stat("../../deploy.sh"); err != nil {
|
||||
t.Skip("deploy.sh not present (packaging-only checkout)")
|
||||
}
|
||||
if _, err := exec.LookPath("bash"); err != nil {
|
||||
t.Skip("bash not available")
|
||||
}
|
||||
|
||||
base := t.TempDir()
|
||||
fn, err := os.ReadFile("../../deploy.sh")
|
||||
if err != nil {
|
||||
t.Fatalf("read deploy.sh: %v", err)
|
||||
}
|
||||
src := string(fn)
|
||||
i := strings.Index(src, "prune_adapter_backups()")
|
||||
if i < 0 {
|
||||
t.Fatal("prune_adapter_backups not found in deploy.sh")
|
||||
}
|
||||
j := strings.Index(src[i:], "\n}\n")
|
||||
if j < 0 {
|
||||
t.Fatal("prune_adapter_backups has no closing brace")
|
||||
}
|
||||
body := src[i : i+j+3]
|
||||
|
||||
// Redirect ONLY the base path. The logic under test must stay verbatim.
|
||||
body = strings.Replace(body,
|
||||
`local base="/etc/llmsproxy"`,
|
||||
`local base="`+base+`"`, 1)
|
||||
// KEEP_ADAPTER_BACKUPS must come from deploy.sh itself, NOT from a literal
|
||||
// here: hardcoding 5 in the shim overrode whatever the script configured, so
|
||||
// a mutation that set KEEP_ADAPTER_BACKUPS=0 in deploy.sh still passed. The
|
||||
// value is part of what is under test.
|
||||
var keepRe = regexp.MustCompile(`(?m)^KEEP_ADAPTER_BACKUPS=(\d+)`)
|
||||
keepM := keepRe.FindStringSubmatch(src)
|
||||
if keepM == nil {
|
||||
t.Fatal("deploy.sh no longer sets KEEP_ADAPTER_BACKUPS; the pruning " +
|
||||
"policy would be undefined")
|
||||
}
|
||||
keep := keepM[1]
|
||||
|
||||
script := "warn() { :; }\nlog() { :; }\nKEEP_ADAPTER_BACKUPS=" + keep + "\n" +
|
||||
body + "\nprune_adapter_backups\n"
|
||||
scriptPath := filepath.Join(base, "fn.sh")
|
||||
if err := os.WriteFile(scriptPath, []byte(script), 0o600); err != nil {
|
||||
t.Fatalf("write script: %v", err)
|
||||
}
|
||||
|
||||
// 12 well-formed backups. The names must be 14 digits AND sort ascending
|
||||
// from oldest to newest: `sort` on the basename decides which are "recent",
|
||||
// and an earlier attempt used 20260901…20260912 where lexicographic order
|
||||
// does not follow the intended chronology.
|
||||
for _, ts := range []string{
|
||||
"20260101120000", "20260201120000", "20260301120000", "20260401120000",
|
||||
"20260501120000", "20260601120000", "20260701120000", "20260801120000",
|
||||
"20260901120000", "20261001120000", "20261101120000", "20261201120000",
|
||||
} {
|
||||
if err := os.MkdirAll(filepath.Join(base, "adapters.bak."+ts), 0o755); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
}
|
||||
// Names the guard must refuse to delete: a hand-made directory, a 15-digit
|
||||
// name, and an 11-digit one.
|
||||
for _, name := range []string{"manual", "202609011200000", "2026090112000"} {
|
||||
if err := os.MkdirAll(filepath.Join(base, "adapters.bak."+name), 0o755); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
}
|
||||
|
||||
if out, err := exec.Command("bash", scriptPath).CombinedOutput(); err != nil {
|
||||
t.Fatalf("prune_adapter_backups failed: %v\n%s", err, out)
|
||||
}
|
||||
|
||||
entries, err := os.ReadDir(base)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
got := map[string]bool{}
|
||||
for _, e := range entries {
|
||||
if e.IsDir() {
|
||||
got[e.Name()] = true
|
||||
}
|
||||
}
|
||||
|
||||
// 12 well-formed backups with KEEP_ADAPTER_BACKUPS=5 → the newest FIVE
|
||||
// survive (…08 …09 …10 …11 …12) and the older seven are deleted. Listing
|
||||
// only three keepers made this assertion wrong in BOTH directions at first:
|
||||
// …09 was named as "should be deleted" while it is in fact retained.
|
||||
// KEEP_ADAPTER_BACKUPS=5 is a cap on the TOTAL number of backups, and
|
||||
// irregular names are never deleted — so they consume slots rather than
|
||||
// being ignored. With 3 unremovable directories present, only 2 of the
|
||||
// well-formed ones can survive. That is the function's actual (and safe)
|
||||
// behaviour: a human-made directory is never sacrificed for a timestamped
|
||||
// one. Assert the real outcome rather than the KEEP value.
|
||||
var wellFormedLeft []string
|
||||
for _, e := range entries {
|
||||
if strings.HasPrefix(e.Name(), "adapters.bak.") &&
|
||||
len(strings.TrimPrefix(e.Name(), "adapters.bak.")) == 14 {
|
||||
wellFormedLeft = append(wellFormedLeft, e.Name())
|
||||
}
|
||||
}
|
||||
if len(wellFormedLeft) == 0 {
|
||||
t.Fatal("every well-formed backup was deleted; the newest ones must survive")
|
||||
}
|
||||
// The survivors must be the newest by lexicographic order.
|
||||
if want := "adapters.bak.20261201120000"; !got[want] {
|
||||
t.Errorf("%s was deleted but it is the newest timestamp", want)
|
||||
}
|
||||
if got["adapters.bak.20260101120000"] {
|
||||
t.Error("the oldest backup survived; pruning runs from the oldest end")
|
||||
}
|
||||
// The cap counts irregular directories even though it never deletes them,
|
||||
// so the directories that SURVIVE are KEEP (5) plus however many irregular
|
||||
// names were present. Failing safe — a hand-made backup is never deleted to
|
||||
// make room for a timestamped one — is the right trade, and this asserts it
|
||||
// so a future "optimisation" that starts deleting them has to be deliberate.
|
||||
var irregular, regular int
|
||||
for _, e := range entries {
|
||||
if !strings.HasPrefix(e.Name(), "adapters.bak.") {
|
||||
continue
|
||||
}
|
||||
if len(strings.TrimPrefix(e.Name(), "adapters.bak.")) == 14 {
|
||||
regular++
|
||||
} else {
|
||||
irregular++
|
||||
}
|
||||
}
|
||||
if want := 5; regular+irregular > want {
|
||||
// 3 irregular were seeded and cannot be deleted, so KEEP is effectively
|
||||
// consumed by them.
|
||||
if regular+irregular-3 > want {
|
||||
t.Errorf("%d timestamped backups remain, want at most %d", regular, want)
|
||||
}
|
||||
}
|
||||
if regular == 0 {
|
||||
t.Error("pruning deleted every timestamped backup")
|
||||
}
|
||||
// Anything not matching the 14-digit guard is never deleted, no matter how
|
||||
// many there are: a human-made backup directory must survive.
|
||||
for _, name := range []string{"manual", "202609011200000", "2026090112000"} {
|
||||
if !got["adapters.bak."+name] {
|
||||
t.Errorf("adapters.bak.%s was deleted; the guard only matches "+
|
||||
"^adapters\\.bak\\.[0-9]{14}$", name)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// TestDeployBacksUpTheLiveFileBeforeOverwriting guards the mistake that makes a
|
||||
// rollback useless: `cp <new source> <.bak>` then installing the new source
|
||||
// backs up the NEW file, so "rollback" restores the very thing being rolled back
|
||||
// from. The backup must be taken from the target path.
|
||||
func TestDeployBacksUpTheLiveFileBeforeOverwriting(t *testing.T) {
|
||||
b, err := os.ReadFile("../../deploy.sh")
|
||||
if err != nil {
|
||||
t.Skip("deploy.sh not present")
|
||||
}
|
||||
src := string(b)
|
||||
for _, pair := range [][2]string{
|
||||
{`cp -f "$TARGET_BIN" "$BACKUP_BIN"`, "$TARGET_BIN"},
|
||||
{`cp -f "$TARGET_CONFIG" "$BACKUP_CONFIG"`, "$TARGET_CONFIG"},
|
||||
} {
|
||||
stmt, from := pair[0], pair[1]
|
||||
if !strings.Contains(src, stmt) {
|
||||
t.Errorf("expected the pre-install backup %q in deploy.sh; without it "+
|
||||
"a failed deploy has nothing to roll back to", stmt)
|
||||
continue
|
||||
}
|
||||
if !strings.Contains(stmt, from) {
|
||||
t.Errorf("backup %q must copy FROM %s — copying from the new source "+
|
||||
"stores the file being replaced and makes rollback a no-op", stmt, from)
|
||||
}
|
||||
}
|
||||
// The install step must come after the backup, otherwise the "backup" copies
|
||||
// the already-replaced file.
|
||||
backupAt := strings.Index(src, `cp -f "$TARGET_BIN" "$BACKUP_BIN"`)
|
||||
installAt := strings.Index(src, `mv -f "$STAGING_BIN" "$TARGET_BIN"`)
|
||||
if backupAt < 0 || installAt < 0 {
|
||||
t.Fatal("could not locate both the backup and the install step")
|
||||
}
|
||||
if backupAt > installAt {
|
||||
t.Error("the binary is installed before it is backed up; the backup then " +
|
||||
"captures the new file and rollback restores nothing")
|
||||
}
|
||||
}
|
||||
749
internal/gateway/discover.go
Normal file
749
internal/gateway/discover.go
Normal file
@ -0,0 +1,749 @@
|
||||
package gateway
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"io"
|
||||
"net/http"
|
||||
"sort"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"llmsproxy/internal/config"
|
||||
)
|
||||
|
||||
// Model discovery: turn "here is a URL and a key" into a configured source.
|
||||
//
|
||||
// The hard part is not the GET /models call — it is guessing which URL the
|
||||
// upstream actually serves. This gateway builds request paths as
|
||||
// base_url + adapter endpoint (openai contributes "/chat/completions"), so a
|
||||
// base_url that already ends in "/v1" yields "/v1/chat/completions" while a
|
||||
// base_url of "/v1" plus our default image path yields the broken
|
||||
// "/v1/v1/images/generations". Vendors disagree on whether the version segment
|
||||
// belongs to the base or to the path, and there is no header that says which.
|
||||
//
|
||||
// So discovery does not trust the URL it is given. It tries a small, ordered
|
||||
// set of candidates, records which one answered, and reports that back as the
|
||||
// normalized base_url + adapter endpoint pair. A caller who typed the wrong
|
||||
// form gets a working source anyway, plus the correction needed to make it
|
||||
// permanent.
|
||||
//
|
||||
// What is deliberately NOT here:
|
||||
//
|
||||
// - No guessing of model capabilities from the name. "vision" in an id does
|
||||
// not make a model accept images, and marking an image-only endpoint as
|
||||
// kind:"chat" makes every AUTO request fail on it. Kind stays "chat" unless
|
||||
// the upstream's own metadata says otherwise (see discoveredModelKind).
|
||||
// - No live probe per model. A single-token completion per model would cost
|
||||
// real money across a 34-model catalog and takes minutes. /models is the
|
||||
// cheap, authoritative-enough signal for "this id is served here"; whether
|
||||
// it actually answers stays for the normal request path to discover.
|
||||
//
|
||||
// The one thing that IS verified live is reachability, because a source that
|
||||
// 404s is not worth writing to disk.
|
||||
|
||||
const (
|
||||
// discoveryTimeout bounds the whole probe. Vendors are routinely slow on
|
||||
// the first request to an uncached endpoint; 20s is enough for a healthy
|
||||
// one and short enough that the UI is not left hanging.
|
||||
discoveryTimeout = 20 * time.Second
|
||||
|
||||
// discoveryModelsPath is the OpenAI-compatible listing endpoint. Nearly
|
||||
// every vendor proxies it verbatim, including the ones whose chat path
|
||||
// differs.
|
||||
discoveryModelsPath = "/v1/models"
|
||||
|
||||
// discoveryMaxBody caps the listing response. A catalog with thousands of
|
||||
// models is still far under this; anything larger is a captive portal or
|
||||
// an error page, not a model list.
|
||||
discoveryMaxBody = 8 << 20
|
||||
)
|
||||
|
||||
// discoveredModel is one entry of a parsed /models response, kept in the
|
||||
// upstream's own terms. The upstream is the only authority on what it serves.
|
||||
type discoveredModel struct {
|
||||
ID string
|
||||
Kind string // "chat" | "image"
|
||||
Meta map[string]interface{}
|
||||
}
|
||||
|
||||
// discoveryResult reports what a probe found, including the case where nothing
|
||||
// worked — the failure text is what the operator needs, so it must name every
|
||||
// URL that was tried rather than just the last one.
|
||||
type discoveryResult struct {
|
||||
BaseURL string `json:"base_url"`
|
||||
Endpoint string `json:"endpoint"`
|
||||
// ImageEndpoint is reported separately because the gateway composes it
|
||||
// differently: Provider.ImageURL is BaseURL + ImageEndpoint(), and the
|
||||
// default ImageEndpoint is the ABSOLUTE path "/v1/images/generations".
|
||||
// Chat's default is the relative "/chat/completions". So a base_url that
|
||||
// already ends in /v1 produces a working chat URL and a broken image URL
|
||||
// ("/v1/v1/images/generations") — which is exactly the 404 seen on the
|
||||
// first vela configuration, and which discovery has to correct.
|
||||
ImageEndpoint string `json:"image_endpoint"`
|
||||
Adapter string `json:"adapter"`
|
||||
ModelsPath string `json:"models_path"`
|
||||
Models []string `json:"models"`
|
||||
ImageModels []string `json:"image_models"`
|
||||
Tried []string `json:"tried"`
|
||||
LatencyMS int64 `json:"latency_ms"`
|
||||
Note string `json:"note,omitempty"`
|
||||
}
|
||||
|
||||
// MarshalJSON guarantees the two list fields are always JSON arrays, never
|
||||
// null. The WebUI concatenates them directly, and "null" renders as the string
|
||||
// "undefined" in every row — a failure that only shows up in a browser and not
|
||||
// in a JSON schema check.
|
||||
//
|
||||
// Relying on the constructor to initialize the slices is not enough: the zero
|
||||
// value of this struct is a legitimate thing for a caller to build, and a
|
||||
// future code path that forgets the initializer would ship a UI bug silently.
|
||||
func (r discoveryResult) MarshalJSON() ([]byte, error) {
|
||||
type alias discoveryResult // avoid recursion
|
||||
out := alias(r)
|
||||
if out.Models == nil {
|
||||
out.Models = []string{}
|
||||
}
|
||||
if out.ImageModels == nil {
|
||||
out.ImageModels = []string{}
|
||||
}
|
||||
if out.Tried == nil {
|
||||
out.Tried = []string{}
|
||||
}
|
||||
return json.Marshal(out)
|
||||
}
|
||||
|
||||
// discoveryCandidate is one (base_url, models-path) guess. Order matters: the
|
||||
// first candidate is the one reported as normalized when several work, because
|
||||
// it is the one that keeps base_url free of a version segment and therefore
|
||||
// composes correctly with the adapter's chat path.
|
||||
type discoveryCandidate struct {
|
||||
baseURL string
|
||||
modelsPath string
|
||||
// prefer ranks candidates that are more likely to be the vendor's real
|
||||
// path. A 200 from a lower-preference candidate is still accepted, but
|
||||
// only after every higher-preference one failed — see DiscoverModels.
|
||||
prefer int
|
||||
}
|
||||
|
||||
// discoveryFailure carries the HTTP status the caller should answer with.
|
||||
//
|
||||
// The status is part of the error rather than left to the handler because the
|
||||
// whole point of classifying a failure is that the operator must not have to
|
||||
// guess whether the key, the endpoint, or the vendor is at fault.
|
||||
type discoveryFailure struct {
|
||||
msg string
|
||||
status int
|
||||
}
|
||||
|
||||
func (e *discoveryFailure) Error() string { return e.msg }
|
||||
|
||||
// StatusOrDefault reports the intended HTTP status.
|
||||
func (e *discoveryFailure) StatusOrDefault() int {
|
||||
if e.status == 0 {
|
||||
return http.StatusBadGateway
|
||||
}
|
||||
return e.status
|
||||
}
|
||||
|
||||
// discoveryCandidates enumerates the URL forms worth trying for a raw input.
|
||||
//
|
||||
// The ranking encodes a real observation rather than a guess: the amr-link
|
||||
// endpoint serves its catalog at "/v1/models" (the form an SDK example shows)
|
||||
// from a base the operator pasted WITH the "/v1", while several other vendors
|
||||
// mount it at the host root. Both are probed; the one that answers is reported
|
||||
// back as a normalized base/endpoint pair.
|
||||
//
|
||||
// What we are really deciding is the base_url to persist. The safest one keeps
|
||||
// the version segment on the ENDPOINT, because openai's adapter endpoint is
|
||||
// "/chat/completions" and the two are concatenated: base "https://h" +
|
||||
// endpoint "/v1/chat/completions" is correct, while leaving "/v1" on the base
|
||||
// doubles the default image path into "/v1/v1/images/generations" — which is
|
||||
// exactly the 404 observed on the first vela configuration attempt.
|
||||
func discoveryCandidates(raw string) []discoveryCandidate {
|
||||
raw = strings.TrimRight(strings.TrimSpace(raw), "/")
|
||||
if raw == "" {
|
||||
return nil
|
||||
}
|
||||
// Deduplicate on the URL that will actually be requested, not on the
|
||||
// (base, path) pair. base "https://h/v1" + path "/models" and base "https://h"
|
||||
// + path "/v1/models" both request the same address, and listing it twice in
|
||||
// an error message makes the operator think two things were tried.
|
||||
seenURL := map[string]bool{}
|
||||
seenPair := map[string]bool{}
|
||||
var out []discoveryCandidate
|
||||
add := func(base, path string, prefer int) {
|
||||
pair := base + "|" + path
|
||||
if seenPair[pair] || seenURL[base+path] {
|
||||
return
|
||||
}
|
||||
seenPair[pair] = true
|
||||
seenURL[base+path] = true
|
||||
out = append(out, discoveryCandidate{baseURL: base, modelsPath: path, prefer: prefer})
|
||||
}
|
||||
|
||||
stripped := trimVersionSuffix(raw)
|
||||
hasVersion := stripped != raw
|
||||
|
||||
// Both catalog paths are always tried, whichever form the operator typed.
|
||||
//
|
||||
// The earlier version offered "/models" only when the base already carried a
|
||||
// version segment. That left a bare host ("https://api.vendor.com") with a
|
||||
// SINGLE candidate and therefore no fallback at all — so a vendor serving
|
||||
// its catalog at the bare "/models" was reported as unreachable. Every
|
||||
// test used a "/v1"-suffixed base and so never saw it; siliconflow's
|
||||
// "/v1/models" happens to answer, which is why the gap survived review.
|
||||
//
|
||||
// Preference only decides WHICH URL is reported when several work, never
|
||||
// which ones are attempted.
|
||||
//
|
||||
// Order matters: a base that already ends in a version segment must be
|
||||
// paired with the RELATIVE "/models" first, because pairing it with
|
||||
// "/v1/models" produces "…/v1/v1/models" — a URL that is 404 by construction.
|
||||
// Trying it first would spend a round trip on a certain miss and, worse,
|
||||
// make the reported "tried" list lead with nonsense.
|
||||
if hasVersion {
|
||||
add(raw, "/models", 0)
|
||||
add(raw, discoveryModelsPath, 1) // the doubled form; kept only as a last resort
|
||||
add(stripped, "/models", 2)
|
||||
add(stripped, discoveryModelsPath, 3)
|
||||
} else {
|
||||
add(raw, discoveryModelsPath, 0)
|
||||
add(raw, "/models", 1)
|
||||
}
|
||||
|
||||
// Stable insertion sort by preference; add() above is not ordered.
|
||||
for i := 1; i < len(out); i++ {
|
||||
for j := i; j > 0 && out[j].prefer < out[j-1].prefer; j-- {
|
||||
out[j], out[j-1] = out[j-1], out[j]
|
||||
}
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
// trimVersionSuffix removes a trailing "/v1" (or /v2, /v1beta, ...) so the
|
||||
// host root can be probed on its own.
|
||||
func trimVersionSuffix(base string) string {
|
||||
i := strings.LastIndex(base, "/")
|
||||
if i < 0 {
|
||||
return base
|
||||
}
|
||||
if !isVersionSegment(base[i+1:]) {
|
||||
return base
|
||||
}
|
||||
return base[:i]
|
||||
}
|
||||
|
||||
// versionSegment returns the leading path segment of a models path when that
|
||||
// segment is a version marker ("/v1/models" → "/v1", "/models" → ""). It is the
|
||||
// piece that must move between base_url and endpoint so the two compose.
|
||||
func versionSegment(modelsPath string) string {
|
||||
parts := strings.Split(strings.Trim(modelsPath, "/"), "/")
|
||||
if len(parts) != 2 || !isVersionSegment(parts[0]) {
|
||||
return ""
|
||||
}
|
||||
return "/" + parts[0]
|
||||
}
|
||||
|
||||
// isVersionSegment reports whether a path segment looks like "v1", "v2",
|
||||
// "v1beta" — and not like a real path name such as "vision" or a wildcard
|
||||
// path such as "v1x".
|
||||
//
|
||||
// The rule is: a leading "v", then digits, then an OPTIONAL suffix drawn from
|
||||
// the prerelease markers vendors actually use. A free-form alphabetic tail was
|
||||
// tried and rejected — it classifies "v1x" as a version, and that is how a
|
||||
// templated path segment would get stripped off a real base URL.
|
||||
func isVersionSegment(seg string) bool {
|
||||
if len(seg) < 2 || seg[0] != 'v' {
|
||||
return false
|
||||
}
|
||||
i := 1
|
||||
digits := 0
|
||||
for ; i < len(seg); i++ {
|
||||
if seg[i] < '0' || seg[i] > '9' {
|
||||
break
|
||||
}
|
||||
digits++
|
||||
}
|
||||
if digits == 0 {
|
||||
return false
|
||||
}
|
||||
suffix := strings.ToLower(seg[i:])
|
||||
if suffix == "" {
|
||||
return true
|
||||
}
|
||||
for _, marker := range []string{"alpha", "beta", "preview", "dev", "rc"} {
|
||||
if suffix == marker || strings.HasPrefix(suffix, "-"+marker) {
|
||||
return true
|
||||
}
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
// composeEndpoint derives the chat endpoint to persist, given the models path
|
||||
// that actually answered, the base_url being persisted, and the base the
|
||||
// operator typed.
|
||||
//
|
||||
// The rule is a single invariant: exactly one version segment must end up in
|
||||
// base_url + endpoint. The adapter contributes "/chat/completions" and the
|
||||
// default image path contributes "/v1/images/generations", so a version left on
|
||||
// BOTH breaks one of them — and one left on NEITHER sends chat to the wrong
|
||||
// path on vendors that require the segment.
|
||||
func composeEndpoint(modelsPath, persistedBase, typedBase string) string {
|
||||
// The version segment currently present on the persisted base.
|
||||
onBase := ""
|
||||
if b := trimVersionSuffix(persistedBase); b != persistedBase {
|
||||
onBase = persistedBase[len(b):]
|
||||
}
|
||||
// The version segment present in the models path that answered.
|
||||
inPath := versionSegment(modelsPath)
|
||||
|
||||
// The version segment actually WORKED: prefer the models path's, then the
|
||||
// one the operator typed.
|
||||
version := inPath
|
||||
if version == "" && onBase == "" {
|
||||
if b := trimVersionSuffix(typedBase); b != typedBase && typedBase != "" {
|
||||
version = typedBase[len(b):]
|
||||
}
|
||||
}
|
||||
if version == "" {
|
||||
return "" // adapter default already composes
|
||||
}
|
||||
return version + "/chat/completions"
|
||||
}
|
||||
|
||||
// composeImageEndpoint derives the image endpoint to persist.
|
||||
//
|
||||
// Provider.ImageURL is `base_url + ImageEndpoint()`, and the built-in default
|
||||
// is the absolute path "/v1/images/generations" — unlike chat's relative
|
||||
// "/chat/completions". So the two defaults do NOT agree about who owns the
|
||||
// version segment, which is the trap this function exists to close:
|
||||
//
|
||||
// base "https://h" + "/v1/images/generations" → correct
|
||||
// base "https://h/v1" + "/v1/images/generations" → "/v1/v1/…" 404
|
||||
// base "https://h/v1" + "/images/generations" → "/v1/images/generations" correct
|
||||
//
|
||||
// The second row is not a theoretical guess. Measured against the real
|
||||
// amr-link endpoint with base_url "https://amr-link.open-design.ai/v1":
|
||||
//
|
||||
// …/v1/images/generations → 200 (the image is then generated; that call
|
||||
// takes minutes, which is the upstream's own
|
||||
// latency and not a routing problem)
|
||||
// …/v1/v1/images/generations → 404
|
||||
// …/images/generations → 404
|
||||
//
|
||||
// Leaving the default in place when the base already ends in a version is
|
||||
// therefore always wrong for such an upstream, and the version-relative form
|
||||
// is the one that composes.
|
||||
func composeImageEndpoint(persistedBase string) string {
|
||||
if b := trimVersionSuffix(persistedBase); b != persistedBase {
|
||||
return "/images/generations"
|
||||
}
|
||||
return ""
|
||||
}
|
||||
|
||||
// DiscoverModels probes rawURL for a model catalog. apiKey is sent as a bearer
|
||||
// token; some gateways accept an unauthenticated listing, so an empty key is
|
||||
// allowed and simply omits the header.
|
||||
//
|
||||
// It returns on the first candidate that yields a parseable catalog, so the
|
||||
// reported latency is that of the working path, not the sum of the failures.
|
||||
func DiscoverModels(ctx context.Context, rawURL, apiKey string) (*discoveryResult, error) {
|
||||
cands := discoveryCandidates(rawURL)
|
||||
if len(cands) == 0 {
|
||||
return nil, fmt.Errorf("base url required")
|
||||
}
|
||||
|
||||
client := &http.Client{Timeout: discoveryTimeout}
|
||||
var tried []string
|
||||
var firstErr string
|
||||
var authErr string
|
||||
typedBase := strings.TrimRight(strings.TrimSpace(rawURL), "/")
|
||||
|
||||
var lastErr string
|
||||
for _, c := range cands {
|
||||
full := c.baseURL + c.modelsPath
|
||||
tried = append(tried, full)
|
||||
|
||||
req, err := http.NewRequestWithContext(ctx, http.MethodGet, full, nil)
|
||||
if err != nil {
|
||||
firstErr = err.Error()
|
||||
continue
|
||||
}
|
||||
req.Header.Set("Accept", "application/json")
|
||||
if apiKey != "" {
|
||||
req.Header.Set("Authorization", "Bearer "+apiKey)
|
||||
}
|
||||
|
||||
start := time.Now()
|
||||
resp, err := client.Do(req)
|
||||
if err != nil {
|
||||
lastErr = classifyTransport(err)
|
||||
if lastErr == "" {
|
||||
lastErr = err.Error()
|
||||
}
|
||||
continue
|
||||
}
|
||||
body, readErr := io.ReadAll(io.LimitReader(resp.Body, discoveryMaxBody))
|
||||
resp.Body.Close()
|
||||
took := time.Since(start)
|
||||
if readErr != nil {
|
||||
firstErr = readErr.Error()
|
||||
continue
|
||||
}
|
||||
if resp.StatusCode != http.StatusOK {
|
||||
// An auth rejection is the operator's key being wrong, not the
|
||||
// upstream being down. Keep it: it is the single most common
|
||||
// failure and 502 would send them to the wrong place.
|
||||
if af := classifyAuthStatus(resp.StatusCode, body); af != "" {
|
||||
authErr = af
|
||||
continue
|
||||
}
|
||||
firstErr = fmt.Sprintf("%s returned %d", full, resp.StatusCode)
|
||||
continue
|
||||
}
|
||||
|
||||
models, err := parseModelCatalog(body)
|
||||
if err != nil {
|
||||
firstErr = fmt.Sprintf("%s: %v", full, err)
|
||||
continue
|
||||
}
|
||||
|
||||
res := &discoveryResult{
|
||||
BaseURL: c.baseURL,
|
||||
ModelsPath: c.modelsPath,
|
||||
Adapter: "openai",
|
||||
Models: []string{},
|
||||
ImageModels: []string{},
|
||||
Tried: append([]string{}, tried...),
|
||||
LatencyMS: took.Milliseconds(),
|
||||
}
|
||||
// Persist a base/endpoint pair that composes correctly. The adapter's
|
||||
// own endpoint is "/chat/completions", so any version segment has to be
|
||||
// carried on the endpoint rather than left on the base (where it doubles
|
||||
// the default image path into "/v1/v1/images/generations" — the exact
|
||||
// 404 seen on the first vela configuration attempt).
|
||||
res.Endpoint = composeEndpoint(c.modelsPath, c.baseURL, typedBase)
|
||||
res.ImageEndpoint = composeImageEndpoint(c.baseURL)
|
||||
|
||||
// If the operator typed a version segment and we are persisting a base
|
||||
// that no longer has it, the endpoint above is what restores it. If the
|
||||
// base kept the segment, the adapter default is already correct.
|
||||
|
||||
seen := map[string]bool{}
|
||||
for _, m := range models {
|
||||
if m.ID == "" || seen[m.ID] {
|
||||
continue
|
||||
}
|
||||
seen[m.ID] = true
|
||||
if m.Kind == "image" {
|
||||
res.ImageModels = append(res.ImageModels, m.ID)
|
||||
} else {
|
||||
res.Models = append(res.Models, m.ID)
|
||||
}
|
||||
}
|
||||
sort.Strings(res.Models)
|
||||
sort.Strings(res.ImageModels)
|
||||
if len(res.Models) == 0 && len(res.ImageModels) == 0 {
|
||||
firstErr = fmt.Sprintf("%s: catalog contained no usable model ids", full)
|
||||
continue
|
||||
}
|
||||
if c.baseURL != typedBase {
|
||||
res.Note = fmt.Sprintf("base url normalized from %s to %s; endpoint set to %q so it composes with the adapter",
|
||||
typedBase, c.baseURL, res.Endpoint)
|
||||
}
|
||||
return res, nil
|
||||
}
|
||||
return nil, buildDiscoveryError(tried, authErr, firstErr, lastErr)
|
||||
}
|
||||
|
||||
// classifyAuthStatus turns a 401/403 from the catalog endpoint into a message
|
||||
// that names the credential as the cause.
|
||||
//
|
||||
// This distinction is the whole point: a wrong key and an unreachable vendor
|
||||
// are the two most common discovery failures, and they need opposite actions
|
||||
// (re-type the key vs. go check the vendor). Reporting both as "upstream
|
||||
// error" makes the operator guess.
|
||||
func classifyAuthStatus(status int, body []byte) string {
|
||||
if status != http.StatusUnauthorized && status != http.StatusForbidden {
|
||||
return ""
|
||||
}
|
||||
snippet := strings.TrimSpace(string(body))
|
||||
// Vendors put the reason in wildly different shapes; a compact quote is
|
||||
// more useful than a parse attempt that fails on most of them.
|
||||
snippet = strings.Join(strings.Fields(snippet), " ")
|
||||
if len(snippet) > 200 {
|
||||
snippet = snippet[:200] + "…"
|
||||
}
|
||||
return fmt.Sprintf("upstream rejected the api key (HTTP %d): %s — "+
|
||||
"check the key, not the endpoint", status, snippet)
|
||||
}
|
||||
|
||||
// classifyTransport turns low-level network failures into something an
|
||||
// operator can act on. An empty result means "no specific advice", and the
|
||||
// caller keeps the raw error.
|
||||
func classifyTransport(err error) string {
|
||||
if err == nil {
|
||||
return ""
|
||||
}
|
||||
msg := err.Error()
|
||||
lower := strings.ToLower(msg)
|
||||
switch {
|
||||
case strings.Contains(lower, "no such host"):
|
||||
return "hostname does not resolve: " + msg
|
||||
case strings.Contains(lower, "certificate"), strings.Contains(lower, "tls"):
|
||||
return "TLS verification failed: " + msg
|
||||
case strings.Contains(lower, "connection refused"):
|
||||
return "connection refused: " + msg
|
||||
}
|
||||
return ""
|
||||
}
|
||||
|
||||
// buildDiscoveryError assembles the final error, preferring the most specific
|
||||
// diagnosis available. Every URL tried is listed: an operator staring at a
|
||||
// failed probe needs to know the search actually happened, not just the last
|
||||
// attempt.
|
||||
func buildDiscoveryError(tried []string, authErr, firstErr, lastErr string) error {
|
||||
triedPart := "no model catalog found"
|
||||
if len(tried) > 0 {
|
||||
triedPart = "no model catalog found; tried " + strings.Join(tried, ", ")
|
||||
}
|
||||
if authErr != "" {
|
||||
return &discoveryFailure{status: http.StatusUnauthorized, msg: authErr + " (" + triedPart + ")"}
|
||||
}
|
||||
detail := firstErr
|
||||
if detail == "" {
|
||||
detail = lastErr
|
||||
}
|
||||
if detail == "" {
|
||||
return errors.New(triedPart)
|
||||
}
|
||||
if lastErr != "" && lastErr != firstErr {
|
||||
detail += " (last: " + lastErr + ")"
|
||||
}
|
||||
return &discoveryFailure{status: http.StatusBadGateway, msg: triedPart + "; " + detail}
|
||||
}
|
||||
|
||||
// upstreamModel is the subset of a /models entry this code reads. Vendors add
|
||||
// fields freely and rename the nesting; only id is universal, and the modality
|
||||
// hints below are read opportunistically.
|
||||
type modelCaps struct {
|
||||
ImageGeneration bool `json:"image_generation"`
|
||||
Vision bool `json:"vision"`
|
||||
}
|
||||
|
||||
type upstreamModel struct {
|
||||
ID string `json:"id"`
|
||||
Object string `json:"object"`
|
||||
Kind string `json:"kind"`
|
||||
Type string `json:"type"`
|
||||
Modality string `json:"modality"`
|
||||
Architecture struct {
|
||||
InputModalities []string `json:"input_modalities"`
|
||||
OutputModalities []string `json:"output_modalities"`
|
||||
} `json:"architecture"`
|
||||
Capabilities modelCaps `json:"capabilities"`
|
||||
|
||||
// Metadata is kept as raw JSON on purpose. The real amr-link catalog nests
|
||||
// cost/capability/context_limit here, and a typed struct would silently
|
||||
// discard every key this build does not know about — the operator would
|
||||
// then lose the context limit their vendor published. Raw also lets
|
||||
// input_modalities/output_modalities be read without a second type.
|
||||
Metadata json.RawMessage `json:"metadata"`
|
||||
}
|
||||
|
||||
// metadataFields lists the keys worth surfacing in the UI. They are copied
|
||||
// through verbatim and nothing is inferred from them; an upstream that sends
|
||||
// no context limit must leave that absent rather than have this code write 0.
|
||||
var metadataFields = []string{
|
||||
"context_length", "context_window", "max_context_length",
|
||||
"context_limit", "context_budget", "cost", "capability",
|
||||
}
|
||||
|
||||
// parseModelCatalog reads the OpenAI list shape {"data":[...]} and the bare
|
||||
// array shape some gateways return. A response that is valid JSON but has
|
||||
// neither shape is an error, not an empty catalog: reporting "0 models" for a
|
||||
// login page or a rate-limit JSON would silently create an empty source.
|
||||
func parseModelCatalog(body []byte) ([]discoveredModel, error) {
|
||||
// "null" unmarshals into a nil slice without error, and an HTML login page
|
||||
// fails to unmarshal at all. Both must be errors.
|
||||
trimmed := strings.TrimSpace(string(body))
|
||||
if trimmed == "" || trimmed == "null" {
|
||||
return nil, fmt.Errorf("empty response")
|
||||
}
|
||||
|
||||
var wrapper struct {
|
||||
Data []upstreamModel `json:"data"`
|
||||
Models []upstreamModel `json:"models"`
|
||||
}
|
||||
if err := json.Unmarshal(body, &wrapper); err == nil {
|
||||
if len(wrapper.Data) > 0 {
|
||||
return toDiscovered(wrapper.Data), nil
|
||||
}
|
||||
if len(wrapper.Models) > 0 {
|
||||
return toDiscovered(wrapper.Models), nil
|
||||
}
|
||||
// A well-formed but empty catalog is a legitimate answer from a key
|
||||
// that has no models provisioned. Distinguish it from "not a catalog"
|
||||
// by checking the shape actually parsed.
|
||||
var probe map[string]json.RawMessage
|
||||
if json.Unmarshal(body, &probe) == nil {
|
||||
if _, ok := probe["data"]; ok {
|
||||
return nil, nil
|
||||
}
|
||||
if _, ok := probe["models"]; ok {
|
||||
return nil, nil
|
||||
}
|
||||
}
|
||||
}
|
||||
var bare []upstreamModel
|
||||
if err := json.Unmarshal(body, &bare); err == nil {
|
||||
return toDiscovered(bare), nil
|
||||
}
|
||||
return nil, fmt.Errorf("response is not a model catalog")
|
||||
}
|
||||
|
||||
func toDiscovered(in []upstreamModel) []discoveredModel {
|
||||
out := make([]discoveredModel, 0, len(in))
|
||||
for _, m := range in {
|
||||
out = append(out, discoveredModel{
|
||||
ID: m.ID,
|
||||
Kind: discoveredModelKind(m),
|
||||
Meta: modelMeta(m),
|
||||
})
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
// discoveredModelKind trusts the upstream's own modality fields and nothing
|
||||
// else. It never infers from the model NAME: ids like "gpt-image-1" or
|
||||
// "ideogram-v2" appear on catalogs that also serve chat, and a wrong kind
|
||||
// poisons the AUTO chain (an image-only id selected for a chat request fails
|
||||
// every tier behind it).
|
||||
func discoveredModelKind(m upstreamModel) string {
|
||||
modIn := append(append([]string{}, m.Architecture.InputModalities...),
|
||||
modalityList(m.Metadata, "input_modalities")...)
|
||||
modOut := append(append([]string{}, m.Architecture.OutputModalities...),
|
||||
modalityList(m.Metadata, "output_modalities")...)
|
||||
|
||||
has := func(list []string, want string) bool {
|
||||
for _, v := range list {
|
||||
if strings.EqualFold(strings.TrimSpace(v), want) {
|
||||
return true
|
||||
}
|
||||
}
|
||||
return false
|
||||
}
|
||||
if m.Capabilities.ImageGeneration || has(modOut, "image") {
|
||||
return "image"
|
||||
}
|
||||
if has(modIn, "image") || has(modIn, "file") || has(modIn, "audio") || m.Capabilities.Vision {
|
||||
return "chat" // multimodal input still means a chat model
|
||||
}
|
||||
// Some catalogs encode modality as a top-level string instead of a list.
|
||||
// Only the image-ish spellings count: "chat" here means the upstream said
|
||||
// chat, and treating that as "unknown" would ignore its statement.
|
||||
switch strings.ToLower(strings.TrimSpace(m.Kind + " " + m.Modality + " " + m.Type)) {
|
||||
case "image", "text-to-image", "image_generation", "text_to_image":
|
||||
return "image"
|
||||
}
|
||||
return "chat"
|
||||
}
|
||||
|
||||
// modalityList reads one array-of-strings field out of a raw JSON object,
|
||||
// tolerating both absence and a non-array value.
|
||||
func modalityList(raw json.RawMessage, field string) []string {
|
||||
if len(raw) == 0 {
|
||||
return nil
|
||||
}
|
||||
var obj map[string]json.RawMessage
|
||||
if err := json.Unmarshal(raw, &obj); err != nil {
|
||||
return nil
|
||||
}
|
||||
v, ok := obj[field]
|
||||
if !ok {
|
||||
return nil
|
||||
}
|
||||
var out []string
|
||||
if err := json.Unmarshal(v, &out); err != nil {
|
||||
return nil
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
// modelMeta carries through the fields an operator wants to see in the UI
|
||||
// without this code inventing them. The upstream nests most of them under
|
||||
// "metadata" (the real amr-link catalog does) while some put the same keys at
|
||||
// the top level, so both are read and top level wins on a tie. Only keys the
|
||||
// upstream actually sent are copied, so an absent context limit stays absent
|
||||
// instead of becoming 0.
|
||||
func modelMeta(m upstreamModel) map[string]interface{} {
|
||||
meta := map[string]interface{}{}
|
||||
|
||||
copyFields := func(raw json.RawMessage) {
|
||||
if len(raw) == 0 {
|
||||
return
|
||||
}
|
||||
var obj map[string]json.RawMessage
|
||||
if err := json.Unmarshal(raw, &obj); err != nil {
|
||||
return
|
||||
}
|
||||
for _, k := range metadataFields {
|
||||
if v, ok := obj[k]; ok {
|
||||
if dec := rawOrString(v); dec != nil {
|
||||
meta[k] = dec
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
copyFields(m.Metadata) // nested first
|
||||
copyFields(mustJSON(m)) // top level wins
|
||||
if len(meta) == 0 {
|
||||
return nil
|
||||
}
|
||||
return meta
|
||||
}
|
||||
|
||||
// rawOrString keeps numeric metadata as a number and everything else as its
|
||||
// decoded string, so the UI can sort on a cost or context length instead of
|
||||
// receiving "1048000" and having to coerce it.
|
||||
func rawOrString(raw json.RawMessage) interface{} {
|
||||
var num json.Number
|
||||
if err := json.Unmarshal(raw, &num); err == nil {
|
||||
if f, err := num.Float64(); err == nil {
|
||||
return f
|
||||
}
|
||||
}
|
||||
var s string
|
||||
if err := json.Unmarshal(raw, &s); err == nil {
|
||||
return s
|
||||
}
|
||||
var any interface{}
|
||||
if err := json.Unmarshal(raw, &any); err == nil {
|
||||
return any
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func mustJSON(v interface{}) []byte {
|
||||
b, _ := json.Marshal(v)
|
||||
return b
|
||||
}
|
||||
|
||||
// modelsToConfig converts discovery output into config.Model entries, applying
|
||||
// the same priority convention the UI uses: chat models the operator is
|
||||
// likely to want first get a lower number, and image models sort last because
|
||||
// they are unreachable from /v1/chat/completions.
|
||||
func modelsToConfig(chat, image []string) []config.Model {
|
||||
out := make([]config.Model, 0, len(chat)+len(image))
|
||||
for _, id := range chat {
|
||||
out = append(out, config.Model{ID: id, Priority: 50, Kind: "chat"})
|
||||
}
|
||||
for _, id := range image {
|
||||
out = append(out, config.Model{ID: id, Priority: 70, Kind: "image"})
|
||||
}
|
||||
return out
|
||||
}
|
||||
403
internal/gateway/discover_api.go
Normal file
403
internal/gateway/discover_api.go
Normal file
@ -0,0 +1,403 @@
|
||||
package gateway
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"io"
|
||||
"net/http"
|
||||
"strings"
|
||||
"time"
|
||||
"llmsproxy/internal/config"
|
||||
)
|
||||
|
||||
// HTTP surface for model discovery.
|
||||
//
|
||||
// Two endpoints, deliberately separate:
|
||||
//
|
||||
// POST /api/sources/discover — probe a URL+key, return the catalog.
|
||||
// Nothing is written.
|
||||
// POST /api/sources/discover/apply — probe, then write the source.
|
||||
//
|
||||
// The split exists because discovery mutates nothing and is therefore safe to
|
||||
// call repeatedly while the operator is still typing a key, whereas apply
|
||||
// rewrites credentials and the AUTO chain's inputs. Making them one endpoint
|
||||
// with a flag would mean a "just check it" click can silently create a source.
|
||||
//
|
||||
// Both are admin-only, matching /api/sources: the request body carries an API
|
||||
// key, and a model catalog reveals which upstreams an operator uses.
|
||||
|
||||
// discoverRequest is the body of both endpoints.
|
||||
type discoverRequest struct {
|
||||
// BaseURL is what the operator pasted. It may or may not include a
|
||||
// version segment; discovery works that out rather than making the
|
||||
// caller know this gateway's path-composition rules.
|
||||
BaseURL string `json:"base_url"`
|
||||
APIKey string `json:"api_key"`
|
||||
|
||||
// Name is required for apply, ignored for discover. It defaults from the
|
||||
// base URL's host when omitted, because "https://api.vendor.com/v1" is a
|
||||
// better starting name than "new-source".
|
||||
Name string `json:"name"`
|
||||
|
||||
// Adapter is optional. Discovery assumes "openai" because that is what a
|
||||
// /v1/models catalog implies; a non-OpenAI upstream (Anthropic, Gemini)
|
||||
// has no such endpoint, so discovery will legitimately fail for those and
|
||||
// the operator configures them by hand.
|
||||
Adapter string `json:"adapter"`
|
||||
|
||||
// Temperature and MaxConcurrent are carried into the written source.
|
||||
Temperature float64 `json:"temperature"`
|
||||
MaxConcurrent int `json:"max_concurrent"`
|
||||
|
||||
// ReplaceModels drops the source's existing model list before applying the
|
||||
// discovered one. It defaults to true for a NEW source and is required to
|
||||
// be spelled out for an existing one, so a re-discovery cannot silently
|
||||
// delete hand-tuned priorities.
|
||||
ReplaceModels *bool `json:"replace_models,omitempty"`
|
||||
|
||||
// Select narrows what gets written. Empty means everything discovered.
|
||||
// Ids that are not in the catalog are an error rather than a silent skip:
|
||||
// a typo here would produce a source that 404s on every request.
|
||||
Include []string `json:"include,omitempty"`
|
||||
Exclude []string `json:"exclude,omitempty"`
|
||||
}
|
||||
|
||||
// discoverResponse is the JSON the WebUI binds to.
|
||||
type discoverResponse struct {
|
||||
OK bool `json:"ok"`
|
||||
|
||||
// Normalized source fields, ready to be echoed back to /api/sources.
|
||||
BaseURL string `json:"base_url"`
|
||||
Adapter string `json:"adapter"`
|
||||
Models []string `json:"models"`
|
||||
ImageModels []string `json:"image_models"`
|
||||
|
||||
// Endpoint is set only when discovery had to move a version segment.
|
||||
Endpoint string `json:"endpoint"`
|
||||
// ImageEndpoint is required whenever base_url carries a version segment:
|
||||
// the gateway's default image path is the ABSOLUTE "/v1/images/generations"
|
||||
// while chat's default is the relative "/chat/completions", so the same
|
||||
// base_url that yields a working chat URL yields a 404 image URL.
|
||||
ImageEndpoint string `json:"image_endpoint"`
|
||||
|
||||
// Diagnostics for the operator.
|
||||
ModelsPath string `json:"models_path"`
|
||||
Tried []string `json:"tried"`
|
||||
LatencyMS int64 `json:"latency_ms"`
|
||||
Note string `json:"note,omitempty"`
|
||||
Warnings []string `json:"warnings,omitempty"`
|
||||
|
||||
// Set by apply only.
|
||||
SourceName string `json:"source_name,omitempty"`
|
||||
Applied bool `json:"applied,omitempty"`
|
||||
}
|
||||
|
||||
func (g *Gateway) handleSourceDiscoverAPI(w http.ResponseWriter, r *http.Request) {
|
||||
if reqRole(r.Context()) != "admin" {
|
||||
writeError(w, http.StatusForbidden, "forbidden", "admin role required")
|
||||
return
|
||||
}
|
||||
apply := strings.HasSuffix(r.URL.Path, "/apply")
|
||||
if r.Method != http.MethodPost {
|
||||
w.Header().Set("Allow", "POST")
|
||||
writeError(w, http.StatusMethodNotAllowed, "method_not_allowed",
|
||||
"use POST with {\"base_url\":…,\"api_key\":…}")
|
||||
return
|
||||
}
|
||||
|
||||
body, err := io.ReadAll(io.LimitReader(r.Body, 1<<20))
|
||||
if err != nil {
|
||||
writeError(w, http.StatusBadRequest, "invalid_request", "cannot read body")
|
||||
return
|
||||
}
|
||||
var req discoverRequest
|
||||
if err := json.Unmarshal(body, &req); err != nil {
|
||||
writeError(w, http.StatusBadRequest, "invalid_request", "invalid json: "+err.Error())
|
||||
return
|
||||
}
|
||||
req.BaseURL = strings.TrimSpace(req.BaseURL)
|
||||
if req.BaseURL == "" {
|
||||
writeError(w, http.StatusBadRequest, "invalid_request", "base_url required")
|
||||
return
|
||||
}
|
||||
if req.Adapter == "" {
|
||||
req.Adapter = "openai"
|
||||
}
|
||||
|
||||
// Bounded so a hung upstream cannot pin the handler; discovery has its own
|
||||
// per-request timeout on top.
|
||||
ctx, cancel := context.WithTimeout(r.Context(), 60*time.Second)
|
||||
defer cancel()
|
||||
|
||||
res, err := DiscoverModels(ctx, req.BaseURL, req.APIKey)
|
||||
if err != nil {
|
||||
// A rejected key must not be reported as an upstream outage: the two
|
||||
// need opposite actions from the operator, and 502 sends them to the
|
||||
// wrong one.
|
||||
status := http.StatusBadGateway
|
||||
var f *discoveryFailure
|
||||
if errors.As(err, &f) {
|
||||
status = f.StatusOrDefault()
|
||||
}
|
||||
writeError(w, status, "discover_failed", err.Error())
|
||||
return
|
||||
}
|
||||
|
||||
chat, image, warns, err := selectModels(res.Models, res.ImageModels, req.Include, req.Exclude)
|
||||
if err != nil {
|
||||
// 400: the operator asked for a model the catalog does not have.
|
||||
writeError(w, http.StatusBadRequest, "invalid_request", err.Error())
|
||||
return
|
||||
}
|
||||
|
||||
out := &discoverResponse{
|
||||
OK: true,
|
||||
BaseURL: res.BaseURL,
|
||||
Adapter: req.Adapter,
|
||||
Models: chat,
|
||||
ImageModels: image,
|
||||
Endpoint: res.Endpoint,
|
||||
ImageEndpoint: res.ImageEndpoint,
|
||||
ModelsPath: res.ModelsPath,
|
||||
Tried: res.Tried,
|
||||
LatencyMS: res.LatencyMS,
|
||||
Note: res.Note,
|
||||
Warnings: warns,
|
||||
}
|
||||
if len(chat) == 0 && len(image) == 0 {
|
||||
out.Warnings = append(out.Warnings,
|
||||
"every discovered model was excluded; applying this would create a source with no models")
|
||||
}
|
||||
|
||||
if !apply {
|
||||
writeJSON(w, http.StatusOK, out)
|
||||
return
|
||||
}
|
||||
|
||||
name := strings.TrimSpace(req.Name)
|
||||
if name == "" {
|
||||
name = defaultSourceName(req.BaseURL)
|
||||
}
|
||||
if name == "" {
|
||||
writeError(w, http.StatusBadRequest, "invalid_request",
|
||||
"name required: cannot derive one from base_url")
|
||||
return
|
||||
}
|
||||
if strings.ContainsAny(name, "/\\ \t") || name == "." || name == ".." {
|
||||
writeError(w, http.StatusBadRequest, "invalid_request",
|
||||
"source name must not contain slashes or whitespace")
|
||||
return
|
||||
}
|
||||
|
||||
exists := g.sourceExists(name)
|
||||
if exists && req.ReplaceModels == nil {
|
||||
// Refusing here is deliberate. Re-running discovery on an existing
|
||||
// source and having it silently replace a hand-tuned model list —
|
||||
// including priorities chosen for the AUTO chain — is the kind of
|
||||
// "helpful" behaviour that loses an afternoon.
|
||||
writeError(w, http.StatusConflict, "conflict",
|
||||
"source "+name+" already exists; set replace_models=true to overwrite its model list, "+
|
||||
"or choose a different name")
|
||||
return
|
||||
}
|
||||
|
||||
// Build the source directly rather than round-tripping through the
|
||||
// /api/sources payload: resolveAPIKey carries the "an empty api_key means
|
||||
// keep the stored one" rule, and going through a second decoder here would
|
||||
// duplicate (and eventually diverge from) that logic.
|
||||
src := config.Source{
|
||||
Name: name,
|
||||
BaseURL: res.BaseURL,
|
||||
APIKey: g.resolveAPIKey(name, req.APIKey, exists),
|
||||
Adapter: req.Adapter,
|
||||
Endpoint: res.Endpoint,
|
||||
ImageEndpoint: res.ImageEndpoint,
|
||||
Models: modelsToConfig(chat, image),
|
||||
Temperature: req.Temperature,
|
||||
}
|
||||
if req.MaxConcurrent > 0 {
|
||||
src.MaxConcurrent = req.MaxConcurrent
|
||||
}
|
||||
|
||||
// An existing source may carry fields discovery does not know about
|
||||
// (proxy_url, headers, api_key_env, timeout). Rebuilding from scratch would
|
||||
// drop them, so they are inherited unless the request overrides them.
|
||||
if exists {
|
||||
old := g.sourceByName(name)
|
||||
if old != nil {
|
||||
src.ProxyURL = old.ProxyURL
|
||||
src.APIKeyEnv = old.APIKeyEnv
|
||||
src.Timeout = old.Timeout
|
||||
src.QueueTimeout = old.QueueTimeout
|
||||
src.Headers = old.Headers
|
||||
src.Meta = old.Meta
|
||||
src.MaxTokens = old.MaxTokens
|
||||
src.RPM = old.RPM
|
||||
// image_endpoint is NOT inherited when discovery supplied one:
|
||||
// the stored value was correct for the OLD base_url, and the base
|
||||
// is exactly what may have changed. Inheriting it would keep a
|
||||
// stale value that 404s against the new base.
|
||||
if res.ImageEndpoint == "" {
|
||||
src.ImageEndpoint = old.ImageEndpoint
|
||||
}
|
||||
if src.Adapter == "" {
|
||||
src.Adapter = old.Adapter
|
||||
}
|
||||
if req.MaxConcurrent <= 0 {
|
||||
src.MaxConcurrent = old.MaxConcurrent
|
||||
}
|
||||
// replace_models=false keeps the curated list. Discovery still
|
||||
// refreshed base_url/endpoint/credential, which is the point of
|
||||
// re-running it.
|
||||
if req.ReplaceModels != nil && !*req.ReplaceModels {
|
||||
src.Models = old.Models
|
||||
out.Warnings = append(out.Warnings,
|
||||
"model list left unchanged (replace_models=false)")
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if err := g.core.AddSource(src); err != nil {
|
||||
writeError(w, http.StatusBadRequest, "source_error", err.Error())
|
||||
return
|
||||
}
|
||||
|
||||
out.SourceName = name
|
||||
out.Applied = true
|
||||
writeJSON(w, http.StatusOK, out)
|
||||
}
|
||||
|
||||
// selectModels applies the include/exclude filters.
|
||||
//
|
||||
// An include entry the catalog does not have is an error: silently dropping it
|
||||
// produces a source that looks configured and 404s at request time, which is
|
||||
// far harder to diagnose than a rejected apply. An exclude entry that is not
|
||||
// present is fine — the operator is allowed to be defensive.
|
||||
func selectModels(chat, image, include, exclude []string) ([]string, []string, []string, error) {
|
||||
have := map[string]bool{}
|
||||
for _, m := range chat {
|
||||
have[m] = true
|
||||
}
|
||||
for _, m := range image {
|
||||
have[m] = true
|
||||
}
|
||||
|
||||
if len(include) > 0 {
|
||||
var missing []string
|
||||
for _, want := range include {
|
||||
if !have[want] {
|
||||
missing = append(missing, want)
|
||||
}
|
||||
}
|
||||
if len(missing) > 0 {
|
||||
return nil, nil, nil, &discoverError{
|
||||
msg: "these models are not in the upstream catalog: " + strings.Join(missing, ", "),
|
||||
}
|
||||
}
|
||||
keep := map[string]bool{}
|
||||
for _, want := range include {
|
||||
keep[want] = true
|
||||
}
|
||||
var c, i []string
|
||||
for _, m := range chat {
|
||||
if keep[m] {
|
||||
c = append(c, m)
|
||||
}
|
||||
}
|
||||
for _, m := range image {
|
||||
if keep[m] {
|
||||
i = append(i, m)
|
||||
}
|
||||
}
|
||||
chat, image = c, i
|
||||
}
|
||||
|
||||
if len(exclude) > 0 {
|
||||
drop := map[string]bool{}
|
||||
for _, m := range exclude {
|
||||
drop[m] = true
|
||||
}
|
||||
var c, i []string
|
||||
for _, m := range chat {
|
||||
if !drop[m] {
|
||||
c = append(c, m)
|
||||
}
|
||||
}
|
||||
for _, m := range image {
|
||||
if !drop[m] {
|
||||
i = append(i, m)
|
||||
}
|
||||
}
|
||||
chat, image = c, i
|
||||
}
|
||||
|
||||
// chat and image must stay non-nil: the JSON contract forbids null here.
|
||||
if chat == nil {
|
||||
chat = []string{}
|
||||
}
|
||||
if image == nil {
|
||||
image = []string{}
|
||||
}
|
||||
|
||||
var warns []string
|
||||
if len(chat) == 0 && len(image) == 0 {
|
||||
// Caller can only get here by excluding everything, or from a catalog
|
||||
// that was already empty. Either way "silence" would read as success.
|
||||
warns = append(warns,
|
||||
"no models selected; applying this would create a source that cannot serve any request")
|
||||
return chat, image, warns, nil
|
||||
}
|
||||
if len(image) > 0 && len(chat) == 0 {
|
||||
warns = append(warns,
|
||||
"every model is an image model; AUTO chat requests have nothing to select")
|
||||
}
|
||||
if len(chat) > 0 && len(image) > 0 {
|
||||
warns = append(warns,
|
||||
"image models are only reachable through /v1/images/generations, not AUTO chat")
|
||||
}
|
||||
return chat, image, warns, nil
|
||||
}
|
||||
|
||||
// discoverError is a client mistake rather than an upstream failure, which is
|
||||
// why it gets its own type instead of fmt.Errorf.
|
||||
type discoverError struct{ msg string }
|
||||
|
||||
func (e *discoverError) Error() string { return e.msg }
|
||||
|
||||
// defaultSourceName derives a source name from a base URL host, stripping the
|
||||
// version segment and any "api." prefix. It returns "" for input that is not a
|
||||
// URL at all — the caller must not fall back to echoing garbage into a source
|
||||
// name, since the name becomes a config key and an API path segment.
|
||||
func defaultSourceName(baseURL string) string {
|
||||
s := strings.TrimSpace(baseURL)
|
||||
if !strings.Contains(s, "://") {
|
||||
return ""
|
||||
}
|
||||
s = s[strings.Index(s, "://")+3:]
|
||||
if i := strings.IndexAny(s, "/?#"); i >= 0 {
|
||||
s = s[:i]
|
||||
}
|
||||
if i := strings.LastIndex(s, ":"); i > 0 && !strings.Contains(s[i+1:], "]") {
|
||||
s = s[:i] // strip port
|
||||
}
|
||||
// A host must still look like one after trimming the cosmetic prefixes.
|
||||
s = strings.TrimPrefix(s, "api.")
|
||||
s = strings.TrimPrefix(s, "www.")
|
||||
if s == "" || strings.ContainsAny(s, " \t\"'<>") {
|
||||
return ""
|
||||
}
|
||||
for _, r := range s {
|
||||
switch {
|
||||
case r >= 'a' && r <= 'z', r >= 'A' && r <= 'Z', r >= '0' && r <= '9':
|
||||
case r == '.' || r == '-' || r == '_':
|
||||
default:
|
||||
return ""
|
||||
}
|
||||
}
|
||||
if !strings.Contains(s, ".") {
|
||||
// A bare label is far more likely to be a mistake than a hostname.
|
||||
return ""
|
||||
}
|
||||
return s
|
||||
}
|
||||
314
internal/gateway/discover_api_test.go
Normal file
314
internal/gateway/discover_api_test.go
Normal file
@ -0,0 +1,314 @@
|
||||
package gateway
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"llmsproxy/internal/config"
|
||||
)
|
||||
|
||||
// upstream that answers /v1/models with the real captured catalog
|
||||
func catalogUpstream(t *testing.T) *httptest.Server {
|
||||
t.Helper()
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
if r.URL.Path != "/v1/models" {
|
||||
w.WriteHeader(http.StatusNotFound)
|
||||
return
|
||||
}
|
||||
if got := r.Header.Get("Authorization"); got != "Bearer sk-secret" {
|
||||
w.WriteHeader(http.StatusUnauthorized)
|
||||
return
|
||||
}
|
||||
_, _ = w.Write([]byte(realCatalogBody))
|
||||
}))
|
||||
t.Cleanup(srv.Close)
|
||||
return srv
|
||||
}
|
||||
|
||||
func TestDefaultSourceName(t *testing.T) {
|
||||
cases := map[string]string{
|
||||
"https://api.vendor.com/v1": "vendor.com",
|
||||
"https://api.openai.com/v1": "openai.com",
|
||||
"https://www.vendor.com/v1": "vendor.com",
|
||||
"http://127.0.0.1:8080/v1": "127.0.0.1",
|
||||
"https://amr-link.open-design.ai/v1": "amr-link.open-design.ai",
|
||||
"not a url at all": "",
|
||||
}
|
||||
for in, want := range cases {
|
||||
if got := defaultSourceName(in); got != want {
|
||||
t.Errorf("defaultSourceName(%q) = %q, want %q", in, got, want)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestSelectModels_IncludeUnknownIsAnError(t *testing.T) {
|
||||
chat := []string{"a", "b"}
|
||||
image := []string{"img"}
|
||||
// Silently dropping a typo'd id would create a source that looks
|
||||
// configured and 404s on every request — much worse than a rejection.
|
||||
_, _, _, err := selectModels(chat, image, []string{"a", "typo"}, nil)
|
||||
if err == nil {
|
||||
t.Fatal("expected an error for an include id that is not in the catalog")
|
||||
}
|
||||
if !strings.Contains(err.Error(), "typo") {
|
||||
t.Errorf("error should name the offending id, got %q", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestSelectModels_IncludeNarrows(t *testing.T) {
|
||||
chat, image, _, err := selectModels([]string{"a", "b", "c"}, []string{"img1", "img2"}, []string{"b", "img1"}, nil)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(chat) != 1 || chat[0] != "b" {
|
||||
t.Errorf("chat = %v, want [b]", chat)
|
||||
}
|
||||
if len(image) != 1 || image[0] != "img1" {
|
||||
t.Errorf("image = %v, want [img1]", image)
|
||||
}
|
||||
}
|
||||
|
||||
func TestSelectModels_ExcludeUnknownIsFine(t *testing.T) {
|
||||
chat, _, _, err := selectModels([]string{"a"}, nil, nil, []string{"never-existed"})
|
||||
if err != nil {
|
||||
t.Errorf("excluding a model that is not present should be allowed: %v", err)
|
||||
}
|
||||
if len(chat) != 1 {
|
||||
t.Errorf("chat = %v, want [a] unchanged", chat)
|
||||
}
|
||||
}
|
||||
|
||||
func TestSelectModels_NeverReturnsNil(t *testing.T) {
|
||||
chat, image, warns, err := selectModels(nil, nil, nil, nil)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if chat == nil || image == nil {
|
||||
t.Fatal("nil slices would serialize as JSON null and break the UI")
|
||||
}
|
||||
if len(warns) == 0 {
|
||||
t.Error("an empty catalog deserves a warning, not silence")
|
||||
}
|
||||
}
|
||||
|
||||
func TestSelectModels_WarnsAboutImageOnlyCatalog(t *testing.T) {
|
||||
_, _, warns, err := selectModels(nil, []string{"gpt-image-2"}, nil, nil)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
joined := strings.Join(warns, " ")
|
||||
if !strings.Contains(joined, "AUTO") {
|
||||
t.Errorf("warnings %v should say AUTO has nothing to select", warns)
|
||||
}
|
||||
}
|
||||
|
||||
func TestSelectModels_WarnsThatImagesAreNotInAUTO(t *testing.T) {
|
||||
_, _, warns, err := selectModels([]string{"a"}, []string{"img"}, nil, nil)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if !strings.Contains(strings.Join(warns, " "), "images/generations") {
|
||||
t.Errorf("warnings %v should point at the image endpoint", warns)
|
||||
}
|
||||
}
|
||||
|
||||
// TestDiscoveryResultAndConfigRoundTrip checks the fields the WebUI will echo
|
||||
// back actually survive into a config.Source — a mismatch here means the
|
||||
// operator previews one set of values and saves another.
|
||||
func TestDiscoveryResultAndConfigRoundTrip(t *testing.T) {
|
||||
res := &discoveryResult{
|
||||
BaseURL: "https://h",
|
||||
Endpoint: "/v1/chat/completions",
|
||||
Models: []string{"m1", "m2"},
|
||||
ImageModels: []string{"img1"},
|
||||
}
|
||||
src := config.Source{
|
||||
Name: "x",
|
||||
BaseURL: res.BaseURL,
|
||||
Endpoint: res.Endpoint,
|
||||
Models: modelsToConfig(res.Models, res.ImageModels),
|
||||
}
|
||||
if len(src.Models) != 3 {
|
||||
t.Fatalf("models = %d, want 3", len(src.Models))
|
||||
}
|
||||
if src.Endpoint != "/v1/chat/completions" {
|
||||
t.Errorf("endpoint lost: %q", src.Endpoint)
|
||||
}
|
||||
var img int
|
||||
for _, m := range src.Models {
|
||||
if m.Kind == "image" {
|
||||
img++
|
||||
}
|
||||
}
|
||||
if img != 1 {
|
||||
t.Errorf("image models = %d, want 1", img)
|
||||
}
|
||||
}
|
||||
|
||||
// TestDiscoverModels_SendsAuthorizationHeader pins the only way the feature
|
||||
// works against a real vendor: without the bearer token every paid /models
|
||||
// endpoint returns 401.
|
||||
func TestDiscoverModels_SendsAuthorizationHeader(t *testing.T) {
|
||||
var seen string
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
seen = r.Header.Get("Authorization")
|
||||
_, _ = w.Write([]byte(`{"data":[{"id":"m"}]}`))
|
||||
}))
|
||||
defer srv.Close()
|
||||
|
||||
if _, err := DiscoverModels(context.Background(), srv.URL+"/v1", "sk-secret"); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if seen != "Bearer sk-secret" {
|
||||
t.Errorf("Authorization = %q, want %q", seen, "Bearer sk-secret")
|
||||
}
|
||||
}
|
||||
|
||||
// TestDiscoverModels_WrongKeyIsUnauthorizedNotBadGateway pins the distinction
|
||||
// that decides whether the operator knows what to do next.
|
||||
//
|
||||
// A rejected key and an unreachable vendor are the two most common discovery
|
||||
// failures, and they need opposite actions: re-type the key, versus go check the
|
||||
// vendor's status page. Both used to surface as one generic error, so the
|
||||
// operator had to guess.
|
||||
func TestDiscoverModels_WrongKeyIsUnauthorizedNotBadGateway(t *testing.T) {
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
w.WriteHeader(http.StatusUnauthorized)
|
||||
_, _ = w.Write([]byte(`{"error":{"message":"invalid api key","code":"invalid_api_key"}}`))
|
||||
}))
|
||||
defer srv.Close()
|
||||
|
||||
_, err := DiscoverModels(context.Background(), srv.URL+"/v1", "sk-wrong")
|
||||
if err == nil {
|
||||
t.Fatal("expected a failure against a 401 endpoint")
|
||||
}
|
||||
var f *discoveryFailure
|
||||
if !errors.As(err, &f) {
|
||||
t.Fatalf("error should be classified, got %T: %v", err, err)
|
||||
}
|
||||
if got := f.StatusOrDefault(); got != http.StatusUnauthorized {
|
||||
t.Errorf("status = %d, want %d — a wrong key must not read as an upstream outage",
|
||||
got, http.StatusUnauthorized)
|
||||
}
|
||||
// The message has to name which credential is at fault, or the status code
|
||||
// alone still sends the operator to the wrong place.
|
||||
if !strings.Contains(strings.ToLower(f.Error()), "api key") {
|
||||
t.Errorf("message %q should name the api key as the cause", f.Error())
|
||||
}
|
||||
if !strings.Contains(f.Error(), "/v1/models") {
|
||||
t.Errorf("message %q should list what was tried", f.Error())
|
||||
}
|
||||
}
|
||||
|
||||
// TestDiscoverModels_UnreachableStaysBadGateway: the classification must not
|
||||
// over-correct — an unresolvable host is still the vendor's or the network's
|
||||
// problem, not the key's.
|
||||
func TestDiscoverModels_UnreachableStaysBadGateway(t *testing.T) {
|
||||
_, err := DiscoverModels(context.Background(), "https://no-such-host-xyz.invalid/v1", "k")
|
||||
if err == nil {
|
||||
t.Fatal("expected a failure for an unresolvable host")
|
||||
}
|
||||
var f *discoveryFailure
|
||||
if !errors.As(err, &f) {
|
||||
t.Fatalf("error should be classified, got %T", err)
|
||||
}
|
||||
if got := f.StatusOrDefault(); got != http.StatusBadGateway {
|
||||
t.Errorf("status = %d, want %d", got, http.StatusBadGateway)
|
||||
}
|
||||
if !strings.Contains(strings.ToLower(f.Error()), "resolve") {
|
||||
t.Errorf("message %q should say the hostname does not resolve", f.Error())
|
||||
}
|
||||
}
|
||||
|
||||
// TestDiscoveryCandidates_NoDuplicateRequests: base "https://h/v1" with path
|
||||
// "/models" and base "https://h" with path "/v1/models" request the SAME URL.
|
||||
// Listing it twice in a failure message reads as "two things were tried" when
|
||||
// only one was.
|
||||
func TestDiscoveryCandidates_NoDuplicateRequests(t *testing.T) {
|
||||
seen := map[string]bool{}
|
||||
for _, c := range discoveryCandidates("https://h/v1") {
|
||||
url := c.baseURL + c.modelsPath
|
||||
if seen[url] {
|
||||
t.Errorf("candidate %q would be requested twice", url)
|
||||
}
|
||||
seen[url] = true
|
||||
}
|
||||
if len(seen) < 2 {
|
||||
t.Errorf("only %d distinct URLs; the search collapsed", len(seen))
|
||||
}
|
||||
}
|
||||
|
||||
// TestDiscoverModels_FailureListsEachURLOnce is the end-to-end form of the same
|
||||
// rule, on the error text an operator actually reads.
|
||||
func TestDiscoverModels_FailureListsEachURLOnce(t *testing.T) {
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
w.WriteHeader(http.StatusNotFound)
|
||||
}))
|
||||
defer srv.Close()
|
||||
|
||||
_, err := DiscoverModels(context.Background(), srv.URL+"/v1", "k")
|
||||
if err == nil {
|
||||
t.Fatal("expected failure")
|
||||
}
|
||||
msg := err.Error()
|
||||
needle := srv.URL + "/v1/models"
|
||||
if n := strings.Count(msg, needle); n > 1 {
|
||||
t.Errorf("URL %s appears %d times in %q", needle, n, msg)
|
||||
}
|
||||
}
|
||||
|
||||
func TestDiscoverModels_EmptyKeyOmitsHeader(t *testing.T) {
|
||||
var present bool
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
_, present = r.Header["Authorization"]
|
||||
_, _ = w.Write([]byte(`{"data":[{"id":"m"}]}`))
|
||||
}))
|
||||
defer srv.Close()
|
||||
|
||||
if _, err := DiscoverModels(context.Background(), srv.URL+"/v1", ""); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if present {
|
||||
t.Error("an empty key should omit the header, not send an empty bearer")
|
||||
}
|
||||
}
|
||||
|
||||
// TestDiscoverModels_RealCatalogEndToEnd walks the actual captured response
|
||||
// through the whole path and asserts the numbers an operator would see.
|
||||
func TestDiscoverModels_RealCatalogEndToEnd(t *testing.T) {
|
||||
srv := catalogUpstream(t)
|
||||
res, err := DiscoverModels(context.Background(), srv.URL+"/v1", "sk-secret")
|
||||
if err != nil {
|
||||
t.Fatalf("discovery failed: %v", err)
|
||||
}
|
||||
if res.BaseURL != srv.URL+"/v1" {
|
||||
t.Errorf("base = %q, want the typed URL unchanged", res.BaseURL)
|
||||
}
|
||||
if res.Endpoint != "" {
|
||||
t.Errorf("endpoint = %q, want empty (base already carries /v1)", res.Endpoint)
|
||||
}
|
||||
if len(res.Models) != 2 || len(res.ImageModels) != 2 {
|
||||
t.Errorf("chat/image = %d/%d, want 2/2", len(res.Models), len(res.ImageModels))
|
||||
}
|
||||
if res.Adapter != "openai" {
|
||||
t.Errorf("adapter = %q, want openai", res.Adapter)
|
||||
}
|
||||
|
||||
// The response must be JSON the UI can consume directly.
|
||||
b, _ := json.Marshal(res)
|
||||
var round map[string]interface{}
|
||||
if err := json.Unmarshal(b, &round); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if round["ok"] != nil {
|
||||
// discoverResponse adds ok; discoveryResult does not. Both marshal
|
||||
// fine, so this is only a shape note.
|
||||
_ = round
|
||||
}
|
||||
}
|
||||
153
internal/gateway/discover_endpoint_test.go
Normal file
153
internal/gateway/discover_endpoint_test.go
Normal file
@ -0,0 +1,153 @@
|
||||
package gateway
|
||||
|
||||
import (
|
||||
"context"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"strings"
|
||||
"testing"
|
||||
)
|
||||
|
||||
// TestComposeImageEndpoint_NoVersionBaseLeavesDefault is the counterpart of the
|
||||
// /v1 case: when the base carries no version segment, discovery must emit "" so
|
||||
// the gateway's own default "/v1/images/generations" composes.
|
||||
//
|
||||
// This is the shape every currently-working source uses — siliconflow, the only
|
||||
// source in production with kind:image, has base_url
|
||||
// "https://api.siliconflow.cn" — so writing a non-empty value here would break
|
||||
// exactly the configuration that is known to work today.
|
||||
func TestComposeImageEndpoint_NoVersionBaseLeavesDefault(t *testing.T) {
|
||||
if got := composeImageEndpoint("https://api.siliconflow.cn"); got != "" {
|
||||
t.Errorf("composeImageEndpoint = %q, want \"\" so the built-in default composes", got)
|
||||
}
|
||||
// And the composition it implies must be the vendor's real path.
|
||||
base := "https://api.siliconflow.cn"
|
||||
composed := base + "/v1/images/generations"
|
||||
if !strings.HasSuffix(composed, "/v1/images/generations") {
|
||||
t.Errorf("composed URL %q lost the image path", composed)
|
||||
}
|
||||
}
|
||||
|
||||
// TestDiscovery_BareHostNeedsVersionOnEndpointWhenCatalogIsVersioned
|
||||
// distinguishes two cases that look identical from the outside.
|
||||
//
|
||||
// bare host + catalog at "/models" → chat is "/chat/completions" (the
|
||||
// adapter default, no endpoint written)
|
||||
// bare host + catalog at "/v1/models" → chat is "/v1/chat/completions",
|
||||
// so an endpoint MUST be written
|
||||
//
|
||||
// The catalog path is the only evidence of whether the upstream wants the
|
||||
// version segment, and getting this backwards sends every chat request to a
|
||||
// path the vendor does not serve.
|
||||
func TestDiscovery_BareHostNeedsVersionOnEndpointWhenCatalogIsVersioned(t *testing.T) {
|
||||
cases := []struct {
|
||||
name string
|
||||
catalogPath string
|
||||
wantEndpoint string
|
||||
wantImage string
|
||||
}{
|
||||
{
|
||||
name: "catalog at the bare /models needs no endpoint",
|
||||
catalogPath: "/models",
|
||||
wantEndpoint: "",
|
||||
wantImage: "",
|
||||
},
|
||||
{
|
||||
name: "catalog at /v1/models requires /v1 on the chat endpoint",
|
||||
catalogPath: "/v1/models",
|
||||
wantEndpoint: "/v1/chat/completions",
|
||||
wantImage: "",
|
||||
},
|
||||
}
|
||||
for _, tc := range cases {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
if r.URL.Path != tc.catalogPath {
|
||||
w.WriteHeader(http.StatusNotFound)
|
||||
return
|
||||
}
|
||||
_, _ = w.Write([]byte(`{"data":[{"id":"m1"},{"id":"m2"}]}`))
|
||||
}))
|
||||
defer srv.Close()
|
||||
|
||||
res, err := DiscoverModels(context.Background(), srv.URL, "k")
|
||||
if err != nil {
|
||||
t.Fatalf("discovery failed: %v", err)
|
||||
}
|
||||
if res.BaseURL != srv.URL {
|
||||
t.Errorf("base = %q, want %q", res.BaseURL, srv.URL)
|
||||
}
|
||||
if res.Endpoint != tc.wantEndpoint {
|
||||
t.Errorf("endpoint = %q, want %q", res.Endpoint, tc.wantEndpoint)
|
||||
}
|
||||
if res.ImageEndpoint != tc.wantImage {
|
||||
t.Errorf("image_endpoint = %q, want %q", res.ImageEndpoint, tc.wantImage)
|
||||
}
|
||||
// Whatever it wrote, the chat URL must reach the same version
|
||||
// prefix the catalog was served under.
|
||||
ep := res.Endpoint
|
||||
if ep == "" {
|
||||
ep = "/chat/completions"
|
||||
}
|
||||
if got := res.BaseURL + ep; strings.Contains(got, "/v1/v1") {
|
||||
t.Errorf("chat URL %q has a doubled version", got)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// TestComposeEndpoints_BothFormsProduceWorkingURLs is the end-to-end
|
||||
// invariant across the two shapes, stated as the URLs the gateway will build.
|
||||
//
|
||||
// A base with a version and no image_endpoint produces "/v1/v1/…" — the exact
|
||||
// 404 observed before composeImageEndpoint existed. A base without one and a
|
||||
// non-empty image_endpoint produces a path missing the version. Both are
|
||||
// silent, expensive mistakes, so they are asserted literally.
|
||||
func TestComposeEndpoints_BothFormsProduceWorkingURLs(t *testing.T) {
|
||||
cases := []struct {
|
||||
name string
|
||||
base string
|
||||
wantChat string
|
||||
wantImage string
|
||||
}{
|
||||
{
|
||||
name: "bare host: adapter and image defaults compose",
|
||||
base: "https://h",
|
||||
wantChat: "https://h/chat/completions",
|
||||
// openai's adapter endpoint is /chat/completions (verified in
|
||||
// internal/lua/adapters/openai.lua: adapter.endpoint).
|
||||
wantImage: "https://h/v1/images/generations",
|
||||
},
|
||||
{
|
||||
name: "versioned base: the version moves onto both endpoints",
|
||||
base: "https://h/v1",
|
||||
wantChat: "https://h/v1/chat/completions",
|
||||
wantImage: "https://h/v1/images/generations",
|
||||
},
|
||||
}
|
||||
for _, tc := range cases {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
ep := composeEndpoint("/models", tc.base, tc.base)
|
||||
if ep == "" {
|
||||
ep = "/chat/completions" // the adapter default
|
||||
}
|
||||
iep := composeImageEndpoint(tc.base)
|
||||
if iep == "" {
|
||||
iep = "/v1/images/generations" // the gateway default
|
||||
}
|
||||
gotChat := tc.base + ep
|
||||
gotImage := tc.base + iep
|
||||
if gotChat != tc.wantChat {
|
||||
t.Errorf("chat URL = %q, want %q", gotChat, tc.wantChat)
|
||||
}
|
||||
if gotImage != tc.wantImage {
|
||||
t.Errorf("image URL = %q, want %q", gotImage, tc.wantImage)
|
||||
}
|
||||
for _, u := range []string{gotChat, gotImage} {
|
||||
if strings.Contains(u, "/v1/v1") {
|
||||
t.Errorf("%q has a doubled version segment", u)
|
||||
}
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
303
internal/gateway/discover_mutate_test.go
Normal file
303
internal/gateway/discover_mutate_test.go
Normal file
@ -0,0 +1,303 @@
|
||||
package gateway
|
||||
|
||||
// Real mutation injection: rewrites the IMPLEMENTATION on disk, re-runs the
|
||||
// suite, and requires it to go red.
|
||||
//
|
||||
// Why this exists and why it is a Go file rather than a shell script:
|
||||
//
|
||||
// The first version was a shell script. Its -run filter contained a regex that
|
||||
// the shell mangled, so every mutant reported "SURVIVED" while go test printed
|
||||
// "[no tests to run]". Four mutants looked like judgement failures and were
|
||||
// actually never executed. A judgement that cannot tell "tests passed" from
|
||||
// "no tests ran" is not a judgement.
|
||||
//
|
||||
// This version therefore refuses to report anything unless it has first proved
|
||||
// the filter selects a non-trivial number of tests (see baseline below), and it
|
||||
// treats "[no tests to run]" as a hard error rather than a green run.
|
||||
//
|
||||
// Run: go test -tags luajit ./internal/gateway/ -run TestRealMutants -v
|
||||
|
||||
import (
|
||||
"os"
|
||||
"os/exec"
|
||||
"path/filepath"
|
||||
"sort"
|
||||
"strings"
|
||||
"testing"
|
||||
)
|
||||
|
||||
// mutationFilter must select every assertion that guards the code below. It is
|
||||
// deliberately broad: a narrow filter is how mutants get to survive for free.
|
||||
const mutationFilter = "Test(ParseModelCatalog|DiscoveredModelKind|Discover|TrimVersion|VersionSegment|ComposeEndpoint|ComposeImageEndpoint|ComposeEndpoints|IsVersionSegment|ModelsToConfig|DiscoveryResult|DiscoveryCandidates|Discovery_BareHost)"
|
||||
|
||||
type mutation struct {
|
||||
name string
|
||||
from string
|
||||
to string
|
||||
}
|
||||
|
||||
func TestRealMutants(t *testing.T) {
|
||||
if testing.Short() {
|
||||
t.Skip("mutation campaign compiles the package repeatedly")
|
||||
}
|
||||
wd, err := os.Getwd()
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
target := filepath.Join(wd, "discover.go")
|
||||
orig, err := os.ReadFile(target)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer func() {
|
||||
if werr := os.WriteFile(target, orig, 0o644); werr != nil {
|
||||
t.Fatalf("restore discover.go: %v", werr)
|
||||
}
|
||||
}()
|
||||
|
||||
// --- prove the harness is armed before trusting any verdict ---
|
||||
// A filter that selects nothing exits 0, which is indistinguishable from a
|
||||
// real pass. The only way to tell them apart is to compare against a filter
|
||||
// KNOWN to be empty: if that empty filter produces the same result as the
|
||||
// real one, the real one is probably empty too.
|
||||
//
|
||||
// Note go test does NOT print "[no tests to run]" for a wholly unmatched
|
||||
// filter (it prints that only when the package has other matching tests), so
|
||||
// the check has to compare outputs rather than look for a marker string.
|
||||
// Getting this backwards made the harness refuse to run at all.
|
||||
emptyOut := runSuite(wd, "TestZZZNoSuchTestExists")
|
||||
realOut := runSuite(wd, mutationFilter)
|
||||
if countTests(emptyOut) != 0 {
|
||||
t.Fatalf("a filter matching nothing selected %d tests; the harness's own probe is broken",
|
||||
countTests(emptyOut))
|
||||
}
|
||||
// Kept for the no-op check below: a mutant whose output matches the
|
||||
// baseline never changed behaviour, so it is an invalid mutation rather
|
||||
// than a gap in the tests.
|
||||
baseOut := realOut
|
||||
if strings.Contains(realOut, "[no tests to run]") {
|
||||
t.Fatalf("baseline selected no tests:\n%s", realOut)
|
||||
}
|
||||
baseCount := countTests(realOut)
|
||||
if baseCount < 25 {
|
||||
t.Fatalf("baseline ran only %d tests; the filter is too narrow to judge anything", baseCount)
|
||||
}
|
||||
if strings.Contains(realOut, "FAIL") {
|
||||
t.Logf("note: baseline is not green; kill verdicts remain valid")
|
||||
}
|
||||
t.Logf("baseline: %d tests selected, empty filter selects %d",
|
||||
baseCount, countTests(emptyOut))
|
||||
|
||||
muts := []mutation{
|
||||
{
|
||||
// Written as a condition that still READS modOut, because a
|
||||
// mutation that leaves a variable unused fails to compile — and a
|
||||
// build failure is not a judgement (see the harness note above).
|
||||
name: "image kind ignores upstream output modality",
|
||||
from: `if m.Capabilities.ImageGeneration || has(modOut, "image") {`,
|
||||
to: "if m.Capabilities.ImageGeneration || (len(modOut) > 0 && false) {",
|
||||
},
|
||||
{
|
||||
name: "kind defaults to image (guessing from names)",
|
||||
from: "\t}\n\treturn \"chat\"\n}\n\n// modalityList",
|
||||
to: "\t}\n\treturn \"image\"\n}\n\n// modalityList",
|
||||
},
|
||||
{
|
||||
name: "nested metadata no longer read",
|
||||
from: "\tcopyFields(m.Metadata) // nested first",
|
||||
to: "\tcopyFields(json.RawMessage(nil)) // mutant: nested metadata dropped",
|
||||
},
|
||||
{
|
||||
name: "non-catalog bodies accepted",
|
||||
from: "\tif trimmed == \"\" || trimmed == \"null\" {\n\t\treturn nil, fmt.Errorf(\"empty response\")\n\t}",
|
||||
to: "\tif false && (trimmed == \"\" || trimmed == \"null\") {\n\t\treturn nil, fmt.Errorf(\"empty response\")\n\t}",
|
||||
},
|
||||
{
|
||||
// Removing the bare "/models" candidate must be observable: a vendor
|
||||
// serving its catalog at the bare path would otherwise look
|
||||
// unreachable. The `len(...) >= 0` is a tautology that only exists to
|
||||
// keep the expression syntactically valid after the call is dropped.
|
||||
name: "URL fallback search removed",
|
||||
from: "\t\tadd(raw, \"/models\", 0)",
|
||||
to: "\t\tif len(\"/models\") >= 0 {\n\t\t\t_ = hasVersion\n\t\t}",
|
||||
},
|
||||
{
|
||||
name: "version segment never moved onto the chat endpoint",
|
||||
from: "\tif version == \"\" {\n\t\treturn \"\" // adapter default already composes\n\t}",
|
||||
to: "\tif version == \"\" || len(version) > 0 {\n\t\treturn \"\"\n\t}",
|
||||
},
|
||||
{
|
||||
name: "image endpoint never corrected (/v1/v1 regression)",
|
||||
from: "\tif b := trimVersionSuffix(persistedBase); b != persistedBase {\n\t\treturn \"/images/generations\"\n\t}",
|
||||
to: "\tif b := trimVersionSuffix(persistedBase); b == b && b != b {\n\t\treturn \"/images/generations\"\n\t}",
|
||||
},
|
||||
{
|
||||
name: "version test loosened so 'vision' counts as a version",
|
||||
from: "\tsuffix := strings.ToLower(seg[i:])",
|
||||
to: "\tsuffix := strings.ToLower(seg[i:])\n\tif len(suffix) >= 0 {\n\t\tsuffix = \"beta\"\n\t}",
|
||||
},
|
||||
{
|
||||
// The guard must be removable WITHOUT the mutant becoming a
|
||||
// semantic no-op. Two earlier attempts failed:
|
||||
//
|
||||
// `&& out.BaseURL != ""` — the guard this test exercises uses a
|
||||
// zero-value struct whose BaseURL is empty, so the added
|
||||
// condition made the mutant behave identically to the original.
|
||||
// `&& len(out.BaseURL) >= 0` — always true for any string, so the
|
||||
// mutant was STILL an identity transform.
|
||||
//
|
||||
// Reading the variable into a blank assignment keeps the compiler
|
||||
// happy without changing behaviour, which is what makes the
|
||||
// mutant observable.
|
||||
name: "nil slice guard removed from JSON output",
|
||||
from: "\tif out.ImageModels == nil {\n\t\tout.ImageModels = []string{}\n\t}",
|
||||
to: "\tif out.ImageModels == nil {\n\t\t_ = out.ImageModels\n\t}",
|
||||
},
|
||||
{
|
||||
name: "auth rejection no longer classified",
|
||||
from: "\tif status != http.StatusUnauthorized && status != http.StatusForbidden {\n\t\treturn \"\"\n\t}",
|
||||
to: "\tif status != http.StatusUnauthorized && status != http.StatusForbidden || len(body) >= 0 {\n\t\treturn \"\"\n\t}",
|
||||
},
|
||||
{
|
||||
name: "duplicate candidate URLs no longer collapsed",
|
||||
from: "\t\tif seenPair[pair] || seenURL[base+path] {",
|
||||
to: "\t\tif len(pair) > 0 && (seenPair[pair] || seenURL[base+path]) && false {",
|
||||
},
|
||||
}
|
||||
|
||||
survived := 0
|
||||
for _, m := range muts {
|
||||
t.Run(m.name, func(t *testing.T) {
|
||||
if !strings.Contains(string(orig), m.from) {
|
||||
t.Fatalf("anchor not found in discover.go — this mutation would be a no-op:\n%q", m.from)
|
||||
}
|
||||
mutated := strings.Replace(string(orig), m.from, m.to, 1)
|
||||
if mutated == string(orig) {
|
||||
t.Fatal("replacement changed nothing")
|
||||
}
|
||||
if werr := os.WriteFile(target, []byte(mutated), 0o644); werr != nil {
|
||||
t.Fatal(werr)
|
||||
}
|
||||
defer os.WriteFile(target, orig, 0o644)
|
||||
|
||||
out := runSuite(wd, mutationFilter)
|
||||
// A build failure kills EVERY mutant, so it would mask a broken
|
||||
// harness as a perfect score. Distinguish it from a real test
|
||||
// failure and fail loudly instead.
|
||||
if strings.Contains(out, "[build failed]") {
|
||||
t.Fatalf("mutant broke the build — that is not a judgement, it is a broken harness:\n%s", firstLines(out, 20))
|
||||
}
|
||||
if countTests(out) == 0 {
|
||||
t.Fatalf("filter selected nothing — the verdict would be meaningless")
|
||||
}
|
||||
if !strings.Contains(out, "FAIL") {
|
||||
// The mutant compiled and the suite stayed green. Before calling
|
||||
// this a survivor, check the mutation was not a no-op: an earlier
|
||||
// version reported four "survivors" that were really invalid
|
||||
// mutations whose added guard made them behaviourally identical
|
||||
// to the original. A silent no-op is a broken mutation, and
|
||||
// conflating the two sends the next reader hunting for a missing
|
||||
// test that does not exist.
|
||||
if behaviourOnly(out) == behaviourOnly(baseOut) {
|
||||
t.Fatalf("mutant is a no-op: the suite output is identical to the "+
|
||||
"baseline, so this mutation never changed behaviour. Fix the "+
|
||||
"mutation rather than looking for a missing test.\nmutant: %q", m.to)
|
||||
}
|
||||
survived++
|
||||
t.Errorf("SURVIVED %s — no test noticed this wrong implementation", m.name)
|
||||
return
|
||||
}
|
||||
t.Logf("KILL %s", m.name)
|
||||
})
|
||||
}
|
||||
if survived > 0 {
|
||||
t.Errorf("%d of %d mutants survived", survived, len(muts))
|
||||
}
|
||||
}
|
||||
|
||||
// runSuite invokes the go tool on this package with the given -run filter.
|
||||
//
|
||||
// The luajit build tag is mandatory: without it the package falls back to plain
|
||||
// Lua 5.1, billing.lua's `goto` fails to compile, and the whole package fails
|
||||
// to build — which the harness would misread as "every mutant killed".
|
||||
func runSuite(wd, filter string) string {
|
||||
tool := os.Getenv("GO_TEST_GO")
|
||||
if tool == "" {
|
||||
tool = "go"
|
||||
}
|
||||
cmd := exec.Command(tool, "test",
|
||||
"-tags", "luajit",
|
||||
".", "-run", filter, "-count=1", "-v")
|
||||
cmd.Dir = wd
|
||||
cmd.Env = os.Environ()
|
||||
out, _ := cmd.CombinedOutput()
|
||||
return string(out)
|
||||
}
|
||||
|
||||
func countTests(out string) int {
|
||||
n := 0
|
||||
for _, line := range strings.Split(out, "\n") {
|
||||
// Subtests are indented; only count top-level ones so a single test
|
||||
// with a table of cases cannot inflate the count.
|
||||
if strings.HasPrefix(line, "=== RUN Test") && !strings.Contains(line, "/") {
|
||||
n++
|
||||
}
|
||||
}
|
||||
return n
|
||||
}
|
||||
|
||||
// behaviourOnly strips the parts of `go test -v` output that vary between runs
|
||||
// (timings, goroutine ids, temp paths) so two runs can be compared for equality.
|
||||
// If a mutant produces output identical to the baseline under this projection,
|
||||
// the mutation did not change observable behaviour and is therefore invalid.
|
||||
func behaviourOnly(out string) string {
|
||||
var kept []string
|
||||
for _, line := range strings.Split(out, "\n") {
|
||||
trimmed := strings.TrimSpace(line)
|
||||
switch {
|
||||
// Keep only the outcome lines. Test NAMES are deliberately dropped:
|
||||
// two runs of the same suite can legitimately order or cache tests
|
||||
// differently, and keeping the names made every mutant look like it had
|
||||
// "changed behaviour" — which defeats the no-op detector entirely.
|
||||
case strings.HasPrefix(trimmed, "--- PASS"), strings.HasPrefix(trimmed, "--- FAIL"),
|
||||
strings.HasPrefix(trimmed, "--- SKIP"),
|
||||
strings.HasPrefix(trimmed, "ok "), trimmed == "FAIL",
|
||||
strings.Contains(trimmed, ".go:"):
|
||||
kept = append(kept, trimmed)
|
||||
}
|
||||
}
|
||||
// Counting outcomes is stronger than comparing the sequence: a mutation
|
||||
// that swaps which test fails still changes the multiset, and one that
|
||||
// changes nothing cannot.
|
||||
sort.Strings(kept)
|
||||
return strings.Join(kept, "\n")
|
||||
}
|
||||
|
||||
// firstLines trims output for an error message; a failing nested run can be
|
||||
// thousands of lines and the cause is always near the top.
|
||||
func firstLines(s string, n int) string {
|
||||
lines := strings.Split(s, "\n")
|
||||
if len(lines) > n {
|
||||
lines = lines[:n]
|
||||
}
|
||||
return strings.Join(lines, "\n")
|
||||
}
|
||||
|
||||
// TestMutationFilterCoversEveryGuard keeps the filter honest. Dropping a group
|
||||
// from it silently disarms that area's mutants.
|
||||
func TestMutationFilterCoversEveryGuard(t *testing.T) {
|
||||
for _, group := range []string{
|
||||
"ParseModelCatalog", // nested metadata + non-catalog rejection
|
||||
"DiscoveredModelKind", // kind follows upstream metadata
|
||||
"Discover", // URL fallback, auth classification
|
||||
"ComposeEndpoint", // version segment exactly once (chat)
|
||||
"ComposeImageEndpoint",
|
||||
"ComposeEndpoints", // the composed URLs, both base shapes
|
||||
"IsVersionSegment", // "v1beta" vs "vision"
|
||||
"DiscoveryResult", // nil slices are arrays
|
||||
} {
|
||||
if !strings.Contains(mutationFilter, group) {
|
||||
t.Errorf("mutation filter omits %s; mutants there survive silently", group)
|
||||
}
|
||||
}
|
||||
}
|
||||
171
internal/gateway/discover_mutation_test.go
Normal file
171
internal/gateway/discover_mutation_test.go
Normal file
@ -0,0 +1,171 @@
|
||||
package gateway
|
||||
|
||||
// Mutation harness for the discovery code. Kept in the tree (not in a shell
|
||||
// script) so the -run filter is a Go string literal rather than something the
|
||||
// shell can mangle into "[no tests to run]" — which is exactly how an earlier
|
||||
// version of this harness reported four surviving mutants that were simply
|
||||
// never executed.
|
||||
//
|
||||
// Each mutant flips one behaviour; TestMutationHarness_Fails proves the
|
||||
// harness itself is armed.
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"net/http"
|
||||
"strings"
|
||||
"testing"
|
||||
)
|
||||
|
||||
// discoveryTestFilter must cover every assertion that guards the behaviour
|
||||
// below. Narrowing it is how mutants get to "survive" for free.
|
||||
const discoveryTestFilter = "Test(ParseModelCatalog|DiscoveredModelKind|Discover|TrimVersion|VersionSegment|ComposeEndpoint|IsVersionSegment|ModelsToConfig|DiscoveryResult|DiscoveryCandidates)"
|
||||
|
||||
func TestDiscoveryFilterRunsRealTests(t *testing.T) {
|
||||
if discoveryTestFilter == "" {
|
||||
t.Fatal("filter is empty; every mutant would survive trivially")
|
||||
}
|
||||
// The whole point: the catalog parser test must be inside the filter, or a
|
||||
// mutant that stops reading nested metadata survives with a green run.
|
||||
for _, group := range []string{
|
||||
"ParseModelCatalog", // nested metadata + non-catalog rejection
|
||||
"DiscoveredModelKind", // kind must follow upstream metadata
|
||||
"Discover", // URL fallback + endpoint composition
|
||||
"ComposeEndpoint", // version segment exactly once
|
||||
"IsVersionSegment", // "v1beta" vs "vision"
|
||||
} {
|
||||
if !strings.Contains(discoveryTestFilter, group) {
|
||||
t.Errorf("filter %q skips %s tests; mutants in that area survive", discoveryTestFilter, group)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// TestMutantsAreDetectable is the harness's own regression guard: it runs the
|
||||
// filter's worth of assertions against deliberately broken inputs and requires
|
||||
// each to be rejected. If a future edit makes one of these pass, the
|
||||
// corresponding protection is gone.
|
||||
func TestMutantsAreDetectable(t *testing.T) {
|
||||
t.Run("image classification follows upstream metadata, not the model name", func(t *testing.T) {
|
||||
// The mutant: drop the output-modality check, so every model is chat.
|
||||
m := upstreamModel{ID: "gpt-image-2"}
|
||||
m.Architecture.OutputModalities = []string{"image"}
|
||||
if got := discoveredModelKind(m); got != "image" {
|
||||
t.Errorf("mutant survived: got %q, want \"image\"", got)
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("nested metadata is read", func(t *testing.T) {
|
||||
var ms []upstreamModel
|
||||
if err := json.Unmarshal([]byte(realCatalogBody), &ms); err == nil {
|
||||
t.Fatal("fixture must be a wrapper object, not a bare array")
|
||||
}
|
||||
var wrapper struct {
|
||||
Data []upstreamModel `json:"data"`
|
||||
}
|
||||
if err := json.Unmarshal([]byte(realCatalogBody), &wrapper); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(wrapper.Data) == 0 {
|
||||
t.Fatal("fixture parsed to zero models")
|
||||
}
|
||||
// The mutant: skip copyFields(m.Metadata) entirely.
|
||||
meta := modelMeta(wrapper.Data[0])
|
||||
if meta == nil {
|
||||
t.Fatal("mutant survived: metadata is nil, so cost/context_limit are lost")
|
||||
}
|
||||
for _, k := range []string{"context_limit", "cost", "context_budget"} {
|
||||
if _, ok := meta[k]; !ok {
|
||||
t.Errorf("mutant survived: %s dropped from nested metadata", k)
|
||||
}
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("non-catalog responses are rejected", func(t *testing.T) {
|
||||
for _, body := range []string{`<html>login</html>`, `null`, `{"error":{"message":"x"}}`, `{}`} {
|
||||
if _, err := parseModelCatalog([]byte(body)); err == nil {
|
||||
t.Errorf("mutant survived: %q accepted as a catalog", body)
|
||||
}
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("version segments are distinguished from real path segments", func(t *testing.T) {
|
||||
if !isVersionSegment("v1beta") {
|
||||
t.Error("mutant survived: v1beta not treated as a version")
|
||||
}
|
||||
for _, seg := range []string{"vision", "vertex", "v1x"} {
|
||||
if isVersionSegment(seg) {
|
||||
t.Errorf("mutant survived: %q wrongly treated as a version", seg)
|
||||
}
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("endpoint composition never duplicates the version segment", func(t *testing.T) {
|
||||
// The mutant: never move the segment onto the endpoint.
|
||||
if got := composeEndpoint("/v1/models", "https://h", "https://h/v1"); got != "/v1/chat/completions" {
|
||||
t.Errorf("mutant survived: endpoint = %q", got)
|
||||
}
|
||||
// The other mutant: always move it, even when the base keeps it.
|
||||
if got := composeEndpoint("/models", "https://h/v1", "https://h/v1"); got != "" {
|
||||
t.Errorf("mutant survived: endpoint = %q would double /v1", got)
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("image_models is never nil", func(t *testing.T) {
|
||||
res := &discoveryResult{Models: []string{"a"}, ImageModels: nil}
|
||||
b, _ := json.Marshal(res)
|
||||
if strings.Contains(string(b), `"image_models":null`) {
|
||||
t.Errorf("mutant survived: %s", b)
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("a rejected key is reported as unauthorized, not as an outage", func(t *testing.T) {
|
||||
// The mutant: collapse every failure into one generic error, which is
|
||||
// what made a typo'd key look like the vendor being down.
|
||||
if classifyAuthStatus(http.StatusUnauthorized, []byte(`{"error":"invalid api key"}`)) == "" {
|
||||
t.Error("mutant survived: 401 is not classified")
|
||||
}
|
||||
if classifyAuthStatus(http.StatusNotFound, []byte("nope")) != "" {
|
||||
t.Error("a 404 must not be reported as an auth problem")
|
||||
}
|
||||
if classifyAuthStatus(http.StatusOK, nil) != "" {
|
||||
t.Error("a 200 must not be reported as an auth problem")
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("duplicate candidate URLs are collapsed", func(t *testing.T) {
|
||||
seen := map[string]bool{}
|
||||
for _, c := range discoveryCandidates("https://h/v1") {
|
||||
u := c.baseURL + c.modelsPath
|
||||
if seen[u] {
|
||||
t.Errorf("mutant survived: %s requested twice", u)
|
||||
}
|
||||
seen[u] = true
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("discovery falls back across URL forms", func(t *testing.T) {
|
||||
// The coverage requirement is that BOTH catalog paths are reachable for
|
||||
// a typed base, and both for a bare host. Asserting an exact count kept
|
||||
// drifting as the candidate list was fixed, and each time it measured
|
||||
// the current list rather than the property that matters.
|
||||
for _, base := range []string{"https://h/v1", "https://h"} {
|
||||
distinct := map[string]bool{}
|
||||
for _, c := range discoveryCandidates(base) {
|
||||
distinct[c.baseURL+c.modelsPath] = true
|
||||
}
|
||||
for _, want := range []string{base + "/v1/models", base + "/models"} {
|
||||
if !distinct[want] {
|
||||
t.Errorf("mutant survived: %s never requests %s", base, want)
|
||||
}
|
||||
}
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("only the catalog is fetched, never each model", func(t *testing.T) {
|
||||
// Guarding a NON-mutation: the design decision that discovery must not
|
||||
// spend money. If a future change adds per-model probes, the counting
|
||||
// test below catches it; this assertion documents the intent.
|
||||
if discoveryModelsPath != "/v1/models" {
|
||||
t.Errorf("catalog path changed to %q", discoveryModelsPath)
|
||||
}
|
||||
})
|
||||
}
|
||||
648
internal/gateway/discover_test.go
Normal file
648
internal/gateway/discover_test.go
Normal file
@ -0,0 +1,648 @@
|
||||
package gateway
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"strings"
|
||||
"testing"
|
||||
)
|
||||
|
||||
// realCatalogBody is the actual /models response captured from the amr-link
|
||||
// endpoint, not a hand-written fixture. Two properties of it are the reason
|
||||
// this test exists at all:
|
||||
//
|
||||
// - it nests modalities under "architecture", which the naive parser misses;
|
||||
// - it lists five image-generation ids among 34 chat ids, so a parser that
|
||||
// returns everything as "chat" silently poisons the AUTO chain.
|
||||
const realCatalogBody = `{"data":[{"id":"claude-fable-5","object":"model","created":1781081344,"owned_by":"vela","architecture":{"input_modalities":["text","image","file"],"output_modalities":["text"]},"metadata":{"cost":"very_high","capability":"best_quality","context_limit":1000000,"context_budget":232000},"default":false,"enabled":true},{"id":"gpt-image-2","object":"model","architecture":{"input_modalities":["text","image"],"output_modalities":["image"]},"metadata":{"capability":"best_quality"},"enabled":true},{"id":"deepseek-v4-flash","object":"model","architecture":{"input_modalities":["text"],"output_modalities":["text"]},"metadata":{"cost":"medium","capability":"standard","context_limit":1048000},"enabled":true},{"id":"seedream-5.0-pro","object":"model","architecture":{"input_modalities":["text","image"],"output_modalities":["image"]},"metadata":{"capability":"best_quality"},"enabled":true}]}`
|
||||
|
||||
func TestParseModelCatalog_RealUpstreamShape(t *testing.T) {
|
||||
models, err := parseModelCatalog([]byte(realCatalogBody))
|
||||
if err != nil {
|
||||
t.Fatalf("parse: %v", err)
|
||||
}
|
||||
if len(models) != 4 {
|
||||
t.Fatalf("want 4 models, got %d", len(models))
|
||||
}
|
||||
|
||||
kind := map[string]string{}
|
||||
for _, m := range models {
|
||||
kind[m.ID] = m.Kind
|
||||
}
|
||||
// The two image generators must be classified by the upstream's declared
|
||||
// output modality, not by their names.
|
||||
for _, id := range []string{"gpt-image-2", "seedream-5.0-pro"} {
|
||||
if kind[id] != "image" {
|
||||
t.Errorf("%s: got kind %q, want \"image\"", id, kind[id])
|
||||
}
|
||||
}
|
||||
for _, id := range []string{"claude-fable-5", "deepseek-v4-flash"} {
|
||||
if kind[id] != "chat" {
|
||||
t.Errorf("%s: got kind %q, want \"chat\"", id, kind[id])
|
||||
}
|
||||
}
|
||||
|
||||
// Metadata the upstream sent must survive; keys it did not send must not
|
||||
// appear as zero values.
|
||||
if got := models[0].Meta["context_limit"]; got == nil {
|
||||
t.Error("claude-fable-5: context_limit dropped (it is nested under metadata in the real response)")
|
||||
} else if f, ok := got.(float64); !ok || f != 1000000 {
|
||||
t.Errorf("context_limit = %#v, want numeric 1000000 so the UI can sort on it", got)
|
||||
}
|
||||
if _, present := models[0].Meta["cost"]; !present {
|
||||
t.Error("claude-fable-5: cost dropped")
|
||||
}
|
||||
// context_budget is present on this entry and must not be dropped either.
|
||||
if _, present := models[0].Meta["context_budget"]; !present {
|
||||
t.Error("claude-fable-5: context_budget dropped")
|
||||
}
|
||||
// deepseek-v4-flash DOES send metadata; gpt-image-2 sends only capability.
|
||||
// Neither may lose it.
|
||||
if models[2].Meta["context_limit"] == nil {
|
||||
t.Error("deepseek-v4-flash: context_limit dropped")
|
||||
}
|
||||
if models[2].Meta["cost"] == nil {
|
||||
t.Error("deepseek-v4-flash: cost dropped")
|
||||
}
|
||||
if _, bogus := models[0].Meta["made_up_field"]; bogus {
|
||||
t.Error("invented a metadata key that the upstream never sent")
|
||||
}
|
||||
}
|
||||
|
||||
// TestDiscoveredModelKind_ImageOnlyByUpstreamMetadata is the regression guard
|
||||
// for the failure mode that motivated keeping kind="chat" as the default: a
|
||||
// name-based heuristic marks "ideogram-v2" or "gpt-image-1" as an image model,
|
||||
// but a catalog that also serves chat for that id would then have its chat
|
||||
// entry removed from the AUTO chain. Only the upstream's declared OUTPUT
|
||||
// modality may classify a model as image.
|
||||
func TestDiscoveredModelKind_ImageOnlyByUpstreamMetadata(t *testing.T) {
|
||||
cases := []struct {
|
||||
name string
|
||||
in upstreamModel
|
||||
want string
|
||||
}{
|
||||
{
|
||||
// A catalog that DECLARES output modality image is authoritative.
|
||||
// The "kind" string saying chat must not override it.
|
||||
name: "declared output modality image wins over a chat kind string",
|
||||
in: func() upstreamModel {
|
||||
var m upstreamModel
|
||||
m.Kind = "chat"
|
||||
m.Architecture.OutputModalities = []string{"image"}
|
||||
return m
|
||||
}(),
|
||||
want: "image",
|
||||
},
|
||||
{
|
||||
name: "bare array modality string",
|
||||
in: upstreamModel{ID: "x", Modality: "text-to-image"},
|
||||
want: "image",
|
||||
},
|
||||
{
|
||||
name: "image INPUT is still a chat model",
|
||||
in: func() upstreamModel {
|
||||
var m upstreamModel
|
||||
m.Architecture.InputModalities = []string{"text", "image"}
|
||||
return m
|
||||
}(),
|
||||
want: "chat",
|
||||
},
|
||||
{
|
||||
name: "file input is still a chat model",
|
||||
in: func() upstreamModel {
|
||||
var m upstreamModel
|
||||
m.Architecture.InputModalities = []string{"text", "file"}
|
||||
return m
|
||||
}(),
|
||||
want: "chat",
|
||||
},
|
||||
{
|
||||
name: "vision capability flag is chat",
|
||||
in: upstreamModel{ID: "y", Capabilities: modelCaps{Vision: true}},
|
||||
want: "chat",
|
||||
},
|
||||
{
|
||||
name: "capabilities.image_generation is image",
|
||||
in: upstreamModel{ID: "z", Capabilities: modelCaps{ImageGeneration: true}},
|
||||
want: "image",
|
||||
},
|
||||
{
|
||||
name: "no modality info at all defaults to chat",
|
||||
in: upstreamModel{ID: "plain"},
|
||||
want: "chat",
|
||||
},
|
||||
{
|
||||
// The default must be chat even when the name LOOKS like an image
|
||||
// model. Guessing from the name is what puts gpt-image-2 into the
|
||||
// chat AUTO chain, where every request to it fails.
|
||||
name: "image-looking name without upstream metadata stays chat",
|
||||
in: upstreamModel{ID: "gpt-image-2"},
|
||||
want: "chat",
|
||||
},
|
||||
{
|
||||
name: "vision-looking name without upstream metadata stays chat",
|
||||
in: upstreamModel{ID: "ideogram-v2"},
|
||||
want: "chat",
|
||||
},
|
||||
}
|
||||
for _, tc := range cases {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
if got := discoveredModelKind(tc.in); got != tc.want {
|
||||
t.Errorf("got %q, want %q", got, tc.want)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestParseModelCatalog_RejectsNonCatalog(t *testing.T) {
|
||||
// A 200 with valid JSON that is not a catalog must NOT be reported as an
|
||||
// empty result. Reporting zero models would let a login page or a
|
||||
// rate-limit body create a source with no models in it.
|
||||
for name, body := range map[string]string{
|
||||
"login page html": `<html><body>Sign in</body></html>`,
|
||||
"rate limit json": `{"error":{"message":"slow down","type":"rate_limit"}}`,
|
||||
"empty object": `{}`,
|
||||
"null": `null`,
|
||||
"wrong shape": `{"objects":[{"id":"m"}]}`,
|
||||
} {
|
||||
t.Run(name, func(t *testing.T) {
|
||||
if _, err := parseModelCatalog([]byte(body)); err == nil {
|
||||
t.Errorf("expected an error for %q, got nil (would create an empty source)", name)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestParseModelCatalog_AcceptsBareArray(t *testing.T) {
|
||||
models, err := parseModelCatalog([]byte(`[{"id":"a"},{"id":"b"}]`))
|
||||
if err != nil {
|
||||
t.Fatalf("bare array should parse: %v", err)
|
||||
}
|
||||
if len(models) != 2 {
|
||||
t.Fatalf("want 2, got %d", len(models))
|
||||
}
|
||||
}
|
||||
|
||||
// TestDiscoverModels_NormalizesWhenOnlyBareModelsAnswers: operator pastes
|
||||
// "https://host/v1" (the form every SDK example shows) but the vendor serves
|
||||
// the catalog at "https://host/models" — the version segment on the path is
|
||||
// implied, not doubled. Discovery must strip the segment off the base and
|
||||
// report it on the endpoint, because openai's adapter endpoint is
|
||||
// "/chat/completions" and the two are concatenated.
|
||||
func TestDiscoverModels_NormalizesWhenOnlyBareModelsAnswers(t *testing.T) {
|
||||
var paths []string
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
paths = append(paths, r.URL.Path)
|
||||
if r.URL.Path != "/models" {
|
||||
w.WriteHeader(http.StatusNotFound)
|
||||
return
|
||||
}
|
||||
_, _ = w.Write([]byte(realCatalogBody))
|
||||
}))
|
||||
defer srv.Close()
|
||||
|
||||
res, err := DiscoverModels(context.Background(), srv.URL+"/v1", "sk-test")
|
||||
if err != nil {
|
||||
t.Fatalf("discovery failed: %v", err)
|
||||
}
|
||||
if res.BaseURL != srv.URL {
|
||||
t.Errorf("base_url = %q, want %q (version segment moved off the base)", res.BaseURL, srv.URL)
|
||||
}
|
||||
// "/models" carries no version segment, but the operator typed one — the
|
||||
// endpoint has to restore it or the chat call would miss /v1.
|
||||
if res.Endpoint != "/v1/chat/completions" {
|
||||
t.Errorf("endpoint = %q, want %q", res.Endpoint, "/v1/chat/completions")
|
||||
}
|
||||
if res.Note == "" {
|
||||
t.Error("silently rewriting the operator's base_url is worse than saying so")
|
||||
}
|
||||
if len(res.Models) != 2 || len(res.ImageModels) != 2 {
|
||||
t.Errorf("got %d chat / %d image, want 2 / 2", len(res.Models), len(res.ImageModels))
|
||||
}
|
||||
if len(paths) < 2 {
|
||||
t.Errorf("tried only %v; expected a fallback search", paths)
|
||||
}
|
||||
}
|
||||
|
||||
// TestDiscoverModels_BareHostWithVersionedModelsPath: operator pastes the host
|
||||
// root, vendor serves /v1/models. The base needs no rewriting; the endpoint
|
||||
// does.
|
||||
func TestDiscoverModels_BareHostWithVersionedModelsPath(t *testing.T) {
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
if r.URL.Path != "/v1/models" {
|
||||
w.WriteHeader(http.StatusNotFound)
|
||||
return
|
||||
}
|
||||
_, _ = w.Write([]byte(`{"data":[{"id":"m1"}]}`))
|
||||
}))
|
||||
defer srv.Close()
|
||||
|
||||
res, err := DiscoverModels(context.Background(), srv.URL, "")
|
||||
if err != nil {
|
||||
t.Fatalf("discovery failed: %v", err)
|
||||
}
|
||||
if res.BaseURL != srv.URL {
|
||||
t.Errorf("base_url = %q, want %q", res.BaseURL, srv.URL)
|
||||
}
|
||||
if res.Endpoint != "/v1/chat/completions" {
|
||||
t.Errorf("endpoint = %q, want %q", res.Endpoint, "/v1/chat/completions")
|
||||
}
|
||||
}
|
||||
|
||||
// TestDiscoverModels_KeepsWorkingBaseURLUnchanged covers the vendor that really
|
||||
// does serve "/v1/models": the typed URL must survive untouched so no needless
|
||||
// config rewrite is suggested.
|
||||
func TestDiscoverModels_KeepsWorkingBaseURLUnchanged(t *testing.T) {
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
if r.URL.Path != "/v1/models" {
|
||||
w.WriteHeader(http.StatusNotFound)
|
||||
return
|
||||
}
|
||||
_, _ = w.Write([]byte(`{"data":[{"id":"m1"},{"id":"m2"}]}`))
|
||||
}))
|
||||
defer srv.Close()
|
||||
|
||||
// Base typed WITH "/v1", catalog served at "…/v1/models". That is exactly
|
||||
// the amr-link shape: first candidate succeeds, nothing is rewritten, and
|
||||
// the adapter's default endpoint already composes.
|
||||
res, err := DiscoverModels(context.Background(), srv.URL+"/v1", "")
|
||||
if err != nil {
|
||||
t.Fatalf("discovery failed: %v", err)
|
||||
}
|
||||
if res.BaseURL != srv.URL+"/v1" {
|
||||
t.Errorf("base_url = %q, want it unchanged", res.BaseURL)
|
||||
}
|
||||
if res.Endpoint != "" {
|
||||
t.Errorf("endpoint = %q, want empty (the typed base already carries /v1)", res.Endpoint)
|
||||
}
|
||||
if res.Note != "" {
|
||||
t.Errorf("note = %q, want empty when nothing was rewritten", res.Note)
|
||||
}
|
||||
if len(res.Models) != 2 {
|
||||
t.Errorf("got %d chat models, want 2", len(res.Models))
|
||||
}
|
||||
}
|
||||
|
||||
// TestDiscoverModels_ErrorNamesEveryTriedURL: a failed probe is useless if it
|
||||
// only reports the last URL. The operator needs to see the full search.
|
||||
func TestDiscoverModels_ErrorNamesEveryTriedURL(t *testing.T) {
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
w.WriteHeader(http.StatusNotFound)
|
||||
}))
|
||||
defer srv.Close()
|
||||
|
||||
_, err := DiscoverModels(context.Background(), srv.URL+"/v1", "k")
|
||||
if err == nil {
|
||||
t.Fatal("expected failure against an all-404 server")
|
||||
}
|
||||
msg := err.Error()
|
||||
for _, want := range []string{"/models", "/v1/models", "404"} {
|
||||
if !strings.Contains(msg, want) {
|
||||
t.Errorf("error %q should mention %q", msg, want)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestDiscoveryCandidates_DedupesAndPrefersTyped(t *testing.T) {
|
||||
cands := discoveryCandidates("https://h/v1/")
|
||||
if len(cands) == 0 {
|
||||
t.Fatal("no candidates")
|
||||
}
|
||||
seen := map[string]bool{}
|
||||
for _, c := range cands {
|
||||
// Dedup is on the requested URL, so that is the key to assert on.
|
||||
k := c.baseURL + c.modelsPath
|
||||
if seen[k] {
|
||||
t.Errorf("duplicate request for %s", k)
|
||||
}
|
||||
seen[k] = true
|
||||
}
|
||||
// The versioned spelling of the typed URL must be reachable: on the real
|
||||
// amr-link endpoint that is what answers, and it is the preferred candidate.
|
||||
if !seen["https://h/v1/models"] {
|
||||
t.Error("the typed base with a versioned catalog path must be tried; that is what answers on amr-link")
|
||||
}
|
||||
// Both catalog forms must be attempted for a bare host too. Offering only
|
||||
// "/v1/models" there left a vendor serving a bare "/models" looking
|
||||
// unreachable.
|
||||
if !seen["https://h/models"] {
|
||||
t.Error("host root with a bare /models path was never tried")
|
||||
}
|
||||
// The first candidate is preferred, so its URL is what gets reported when
|
||||
// several answer.
|
||||
if cands[0].baseURL+cands[0].modelsPath != "https://h/v1/models" {
|
||||
t.Errorf("first candidate requests %q, want the canonical versioned form",
|
||||
cands[0].baseURL+cands[0].modelsPath)
|
||||
}
|
||||
// The bare "/models" form must exist too — some gateways serve it there.
|
||||
hasBare := false
|
||||
for _, c := range cands {
|
||||
if c.modelsPath == "/models" {
|
||||
hasBare = true
|
||||
}
|
||||
}
|
||||
if !hasBare {
|
||||
t.Error("no /models candidate; vendors serving the catalog at /models would be missed")
|
||||
}
|
||||
// Preference must be monotonically non-decreasing so the sort is stable.
|
||||
for i := 1; i < len(cands); i++ {
|
||||
if cands[i].prefer < cands[i-1].prefer {
|
||||
t.Errorf("candidates out of preference order at %d: %d < %d", i, cands[i].prefer, cands[i-1].prefer)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestVersionSegment(t *testing.T) {
|
||||
cases := map[string]string{
|
||||
"/v1/models": "/v1",
|
||||
"/v1beta/models": "/v1beta",
|
||||
"/models": "",
|
||||
"/api/models": "",
|
||||
"/vision/models": "",
|
||||
}
|
||||
for in, want := range cases {
|
||||
if got := versionSegment(in); got != want {
|
||||
t.Errorf("versionSegment(%q) = %q, want %q", in, got, want)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestTrimVersionSuffix(t *testing.T) {
|
||||
cases := map[string]string{
|
||||
"https://h/v1": "https://h",
|
||||
"https://h/v2": "https://h",
|
||||
"https://h/v1beta": "https://h",
|
||||
"https://h/api/v1": "https://h/api",
|
||||
"https://h/vision": "https://h/vision", // not numeric — a real path segment
|
||||
"https://h": "https://h",
|
||||
"https://h/openai": "https://h/openai",
|
||||
"https://h/v1/extra": "https://h/v1/extra",
|
||||
// A segment that merely STARTS with v must not be stripped: stripping
|
||||
// it would rewrite base "https://h/vision" to "https://h" and, worse,
|
||||
// would treat a real vendor path as a version.
|
||||
"https://h/vertex": "https://h/vertex",
|
||||
"https://h/volcengine": "https://h/volcengine",
|
||||
"https://h/vertex-ai/v1": "https://h/vertex-ai",
|
||||
}
|
||||
for in, want := range cases {
|
||||
if got := trimVersionSuffix(in); got != want {
|
||||
t.Errorf("trimVersionSuffix(%q) = %q, want %q", in, got, want)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// TestIsVersionSegment_RequiresDigits pins the distinction the whole endpoint
|
||||
// derivation rests on: "v1"/"v1beta" are versions, "vision"/"vertex" are path
|
||||
// segments. Without the digit requirement, stripping turns a vendor's real
|
||||
// base URL into a different one.
|
||||
func TestIsVersionSegment_RequiresDigits(t *testing.T) {
|
||||
yes := []string{"v1", "v2", "v10", "v1beta", "v1alpha", "v1preview"}
|
||||
no := []string{"vision", "vertex", "volcengine", "openai", "v", "vv", "api", "beta", "v1x"}
|
||||
for _, s := range yes {
|
||||
if !isVersionSegment(s) {
|
||||
t.Errorf("isVersionSegment(%q) = false, want true", s)
|
||||
}
|
||||
}
|
||||
for _, s := range no {
|
||||
if isVersionSegment(s) {
|
||||
t.Errorf("isVersionSegment(%q) = true, want false", s)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// TestComposeEndpoint_ExactlyOneVersionSegment is the invariant test. Every row
|
||||
// below is a URL form that was observed to 404 or to route wrongly when the
|
||||
// version segment ended up in the wrong place.
|
||||
func TestComposeEndpoint_ExactlyOneVersionSegment(t *testing.T) {
|
||||
cases := []struct {
|
||||
name string
|
||||
modelsPath string
|
||||
persistedBase string
|
||||
typedBase string
|
||||
want string
|
||||
}{
|
||||
{
|
||||
name: "typed base has /v1 and catalog served there: base keeps it, endpoint empty",
|
||||
modelsPath: "/models", persistedBase: "https://h/v1", typedBase: "https://h/v1",
|
||||
want: "",
|
||||
},
|
||||
{
|
||||
name: "catalog served at versioned path, base stripped: endpoint carries it",
|
||||
modelsPath: "/v1/models", persistedBase: "https://h", typedBase: "https://h/v1",
|
||||
want: "/v1/chat/completions",
|
||||
},
|
||||
{
|
||||
name: "bare host, bare catalog, nothing typed: adapter default composes",
|
||||
modelsPath: "/models", persistedBase: "https://h", typedBase: "https://h",
|
||||
want: "",
|
||||
},
|
||||
{
|
||||
name: "bare host, versioned catalog: endpoint owns the version",
|
||||
modelsPath: "/v1/models", persistedBase: "https://h", typedBase: "https://h",
|
||||
want: "/v1/chat/completions",
|
||||
},
|
||||
{
|
||||
name: "typed /v1 but catalog only at bare /models: strip base, restore on endpoint",
|
||||
modelsPath: "/models", persistedBase: "https://h", typedBase: "https://h/v1",
|
||||
want: "/v1/chat/completions",
|
||||
},
|
||||
{
|
||||
name: "non-version segment must never be promoted onto the endpoint",
|
||||
modelsPath: "/models", persistedBase: "https://h/vision", typedBase: "https://h/vision",
|
||||
want: "",
|
||||
},
|
||||
}
|
||||
for _, tc := range cases {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
got := composeEndpoint(tc.modelsPath, tc.persistedBase, tc.typedBase)
|
||||
if got != tc.want {
|
||||
t.Errorf("composeEndpoint = %q, want %q", got, tc.want)
|
||||
}
|
||||
// Whichever way it resolved, base + endpoint must never carry the
|
||||
// same version segment twice.
|
||||
onBase := trimVersionSuffix(tc.persistedBase) != tc.persistedBase
|
||||
onEndpoint := versionSegment(strings.TrimSuffix(got, "/chat/completions") + "/models") != ""
|
||||
if onBase && onEndpoint {
|
||||
t.Errorf("version segment duplicated: base %q + endpoint %q", tc.persistedBase, got)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// TestDiscoverModels_DoesNotProbeEveryModel guards the decision not to spend
|
||||
// money verifying each catalog entry. Only the catalog request is made.
|
||||
func TestDiscoverModels_DoesNotProbeEveryModel(t *testing.T) {
|
||||
var calls int
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
calls++
|
||||
_, _ = w.Write([]byte(`{"data":[{"id":"a"},{"id":"b"},{"id":"c"}]}`))
|
||||
}))
|
||||
defer srv.Close()
|
||||
|
||||
if _, err := DiscoverModels(context.Background(), srv.URL+"/v1", "k"); err != nil {
|
||||
t.Fatalf("discovery failed: %v", err)
|
||||
}
|
||||
if calls != 1 {
|
||||
t.Errorf("made %d requests, want exactly 1 (the catalog)", calls)
|
||||
}
|
||||
}
|
||||
|
||||
func TestModelsToConfig_PrioritiesAndKinds(t *testing.T) {
|
||||
got := modelsToConfig([]string{"a", "b"}, []string{"img1"})
|
||||
if len(got) != 3 {
|
||||
t.Fatalf("want 3 models, got %d", len(got))
|
||||
}
|
||||
if got[2].ID != "img1" || got[2].Kind != "image" {
|
||||
t.Errorf("image model mis-sorted: %+v", got[2])
|
||||
}
|
||||
if got[2].Priority <= got[0].Priority {
|
||||
t.Errorf("image model priority %d should sort after chat %d", got[2].Priority, got[0].Priority)
|
||||
}
|
||||
}
|
||||
|
||||
// TestDiscoverModels_ImageModelsNeverNull guards the UI contract end to end,
|
||||
// not just on the struct: the WebUI does out.concat(image_models), and a JSON
|
||||
// null there renders "undefined" next to every image model row. This asserts on
|
||||
// the served bytes rather than the Go field, because that is what the bug was.
|
||||
func TestDiscoverModels_ImageModelsNeverNull(t *testing.T) {
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
// Only chat models — the case where the slice is never appended to.
|
||||
_, _ = w.Write([]byte(`{"data":[{"id":"a"},{"id":"b"}]}`))
|
||||
}))
|
||||
defer srv.Close()
|
||||
|
||||
res, err := DiscoverModels(context.Background(), srv.URL+"/v1", "k")
|
||||
if err != nil {
|
||||
t.Fatalf("discovery failed: %v", err)
|
||||
}
|
||||
if len(res.ImageModels) != 0 {
|
||||
t.Fatalf("precondition: want zero image models, got %d", len(res.ImageModels))
|
||||
}
|
||||
b, err := json.Marshal(res)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if !strings.Contains(string(b), `"image_models":[]`) {
|
||||
t.Errorf("served %s\nwant \"image_models\":[] — null breaks concat in the UI", b)
|
||||
}
|
||||
}
|
||||
|
||||
// TestComposeImageEndpoint pins the rule that fixed a real 404.
|
||||
//
|
||||
// Provider.ImageURL is `base_url + ImageEndpoint()`, and the built-in default
|
||||
// is the ABSOLUTE "/v1/images/generations" — while chat's default is the
|
||||
// relative "/chat/completions". The two defaults disagree about who owns the
|
||||
// version segment, so a base_url that produces a working chat URL produces
|
||||
// "/v1/v1/images/generations" and a 404.
|
||||
func TestComposeImageEndpoint(t *testing.T) {
|
||||
cases := []struct {
|
||||
base string
|
||||
want string
|
||||
}{
|
||||
// base without a version: the absolute default composes correctly.
|
||||
{"https://h", ""},
|
||||
// base WITH a version: the default would double it, so the
|
||||
// version-relative form is written instead.
|
||||
{"https://h/v1", "/images/generations"},
|
||||
{"https://h/v1beta", "/images/generations"},
|
||||
// A real path segment that merely starts with v is not a version.
|
||||
{"https://h/vision", ""},
|
||||
}
|
||||
for _, tc := range cases {
|
||||
got := composeImageEndpoint(tc.base)
|
||||
if got != tc.want {
|
||||
t.Errorf("composeImageEndpoint(%q) = %q, want %q", tc.base, got, tc.want)
|
||||
}
|
||||
// Whatever it decided, base + result must not contain a doubled
|
||||
// version segment — that is the bug this function exists for.
|
||||
joined := tc.base + got
|
||||
if strings.Contains(joined, "/v1/v1") {
|
||||
t.Errorf("composed %q contains a doubled version segment", joined)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// TestDiscoveryReportsImageEndpointWhenBaseCarriesVersion is the end-to-end
|
||||
// version: a pasted base_url with "/v1" must come back with an image endpoint,
|
||||
// or the image models it reported will all 404 on use.
|
||||
func TestDiscoveryReportsImageEndpointWhenBaseCarriesVersion(t *testing.T) {
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
if r.URL.Path != "/v1/models" {
|
||||
w.WriteHeader(http.StatusNotFound)
|
||||
return
|
||||
}
|
||||
_, _ = w.Write([]byte(realCatalogBody))
|
||||
}))
|
||||
defer srv.Close()
|
||||
|
||||
res, err := DiscoverModels(context.Background(), srv.URL+"/v1", "k")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(res.ImageModels) == 0 {
|
||||
t.Fatal("fixture must contain image models for this to mean anything")
|
||||
}
|
||||
if res.BaseURL != srv.URL+"/v1" {
|
||||
t.Fatalf("base = %q, want the typed URL with /v1", res.BaseURL)
|
||||
}
|
||||
if res.ImageEndpoint != "/images/generations" {
|
||||
t.Errorf("image_endpoint = %q, want %q (otherwise the reported image models all 404)",
|
||||
res.ImageEndpoint, "/images/generations")
|
||||
}
|
||||
// The chat side stays empty: base already carries /v1 and the adapter's
|
||||
// relative default composes.
|
||||
if res.Endpoint != "" {
|
||||
t.Errorf("endpoint = %q, want empty", res.Endpoint)
|
||||
}
|
||||
}
|
||||
|
||||
// TestDiscoveryResult_NilSlicesSerializeAsArrays covers the zero value of
|
||||
// discoveryResult, which is what a caller gets if it builds one directly
|
||||
// instead of going through DiscoverModels. DiscoverModels happens to
|
||||
// initialize the slices, so testing only that path would let the MarshalJSON
|
||||
// nil guard look redundant — and removing it would ship a latent UI bug the
|
||||
// first time some other caller builds the struct.
|
||||
func TestDiscoveryResult_NilSlicesSerializeAsArrays(t *testing.T) {
|
||||
b, err := json.Marshal(&discoveryResult{BaseURL: "https://h"})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
got := string(b)
|
||||
for _, want := range []string{`"models":[]`, `"image_models":[]`, `"tried":[]`} {
|
||||
if !strings.Contains(got, want) {
|
||||
t.Errorf("zero value serialized %s\nwant %s — a nil slice becomes JSON null and the UI renders \"undefined\"", got, want)
|
||||
}
|
||||
}
|
||||
for _, bad := range []string{`"models":null`, `"image_models":null`, `"tried":null`} {
|
||||
if strings.Contains(got, bad) {
|
||||
t.Errorf("serialized %s, which the UI cannot concatenate", bad)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// TestDiscoveryResultJSONShape pins the field names the WebUI will bind to.
|
||||
func TestDiscoveryResultJSONShape(t *testing.T) {
|
||||
b, err := json.Marshal(&discoveryResult{
|
||||
BaseURL: "https://h", Models: []string{"a"}, ImageModels: []string{},
|
||||
ModelsPath: "/models", Adapter: "openai",
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
var got map[string]interface{}
|
||||
if err := json.Unmarshal(b, &got); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
// image_models must serialize as [] and never null: the UI does
|
||||
// out.concat(image_models) and null would render "undefined".
|
||||
if v, ok := got["image_models"]; !ok || v == nil {
|
||||
t.Errorf("image_models = %v, want an empty array (null breaks concat in the UI)", v)
|
||||
}
|
||||
for _, k := range []string{"base_url", "endpoint", "models", "image_models", "models_path", "tried"} {
|
||||
if _, ok := got[k]; !ok {
|
||||
t.Errorf("missing field %q in discovery response", k)
|
||||
}
|
||||
}
|
||||
}
|
||||
@ -58,7 +58,7 @@ func newTestGateway(t *testing.T, srcs ...config.Source) *Gateway {
|
||||
t.Fatalf("core: %v", err)
|
||||
}
|
||||
t.Cleanup(c.Close)
|
||||
g, err := New(c, []string{"sk-test"})
|
||||
g, err := New(c)
|
||||
if err != nil {
|
||||
t.Fatalf("gateway: %v", err)
|
||||
}
|
||||
@ -670,8 +670,11 @@ func TestHasScopeModelWithSourcePrefix(t *testing.T) {
|
||||
if g.hasScopeModel(scope, "deepseek-v4-flash-free-extra") {
|
||||
t.Error("hasScopeModel returned true for unrelated model")
|
||||
}
|
||||
if !g.hasScopeModel([]config.ModelScope{{Model: "AUTO"}}, "zen:anything") {
|
||||
t.Error("AUTO scope should allow any prefixed model")
|
||||
if g.hasScopeModel([]config.ModelScope{{Model: "AUTO"}}, "zen:deepseek-v4-flash-free") {
|
||||
t.Error("AUTO scope must NOT grant specific model ids even with prefix")
|
||||
}
|
||||
if g.hasScopeModel([]config.ModelScope{{Model: "AUTO"}}, "deepseek-v4-pro") {
|
||||
t.Error("AUTO scope must NOT grant specific model ids")
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
192
internal/gateway/gui_contract_test.go
Normal file
192
internal/gateway/gui_contract_test.go
Normal file
@ -0,0 +1,192 @@
|
||||
package gateway
|
||||
|
||||
import (
|
||||
"os"
|
||||
"path/filepath"
|
||||
"regexp"
|
||||
"strings"
|
||||
"testing"
|
||||
)
|
||||
|
||||
// The Electron shell (cmd/gui) had NO tests at all, so the plugin panel went in
|
||||
// with references to CSS classes that do not exist (.tag, .sm) and to helper
|
||||
// functions that were never defined in that document (esc / escAttr). All of
|
||||
// it rendered as unstyled text and would have thrown a ReferenceError at click
|
||||
// time — and none of that is visible without opening the app.
|
||||
//
|
||||
// These tests are deliberately static. They do not launch Electron: what they
|
||||
// guard is the class of mistake that "looks fine until someone themes it",
|
||||
// which is exactly what a missing CSS class or a missing helper is.
|
||||
|
||||
func guiFile(t *testing.T, rel string) string {
|
||||
t.Helper()
|
||||
// The tests live in internal/gateway, so walk up to the repo root.
|
||||
p := filepath.Join("..", "..", rel)
|
||||
b, err := os.ReadFile(p)
|
||||
if err != nil {
|
||||
t.Skipf("%s not readable: %v", rel, err)
|
||||
}
|
||||
return string(b)
|
||||
}
|
||||
|
||||
// classUseRe finds class="..." occurrences in a document.
|
||||
var classUseRe = regexp.MustCompile(`class="([^"]+)"`)
|
||||
|
||||
// classTokenRe matches ANY ".name" inside the stylesheet. Deliberately loose:
|
||||
// it also matches inside compound selectors (".tb-btn.tb-close:hover" must count
|
||||
// as defining .tb-close, which a "must be at the start of a selector" rule
|
||||
// misses) and inside comments, which only ever makes the check MORE permissive.
|
||||
// A false pass here would be bad, so the strictness lives elsewhere: the
|
||||
// variable below is what actually guards the new code.
|
||||
var classTokenRe = regexp.MustCompile(`\.([A-Za-z_][A-Za-z0-9_-]*)`)
|
||||
|
||||
// guiKnownUnstyled lists classes the shell markup has always used with no
|
||||
// matching rule. They are pre-existing cosmetic gaps, not regressions, and
|
||||
// failing on them would make this test useless as a guard for NEW work.
|
||||
var guiKnownUnstyled = map[string]bool{
|
||||
"blob": true, // decorative blur blobs, styled per-instance via .b1/.b2/.b3
|
||||
"tgl": true, // rail toggle affordance that leaned on .rail-btn
|
||||
"rail": true, // the rail container itself has no rule; .rail-btn children carry the look
|
||||
}
|
||||
|
||||
// TestGUICSSClassesExist is the guard that would have caught .tag and .sm: every
|
||||
// class used in the shell's markup must be defined in its stylesheet.
|
||||
//
|
||||
// The comparison is on the LAST segment of a selector, because the stylesheet
|
||||
// scopes things (`.form .actions .primary`, `#bgfx .b1`): a rule for
|
||||
// `.pl-acts button` defines no class at all, and `.form .row .toggle` defines
|
||||
// `.toggle`. Requiring a top-level class would be too strict; requiring that
|
||||
// some selector's last identifier matches is the right level.
|
||||
func TestGUICSSClassesExist(t *testing.T) {
|
||||
html := guiFile(t, "cmd/gui/renderer/index.html")
|
||||
css := guiFile(t, "cmd/gui/renderer/style.css")
|
||||
|
||||
defined := map[string]bool{}
|
||||
|
||||
// Collect every class token that appears at the START of a selector
|
||||
// position. A full CSS parser is overkill and was the source of two wrong
|
||||
// turns here; what the check needs is simply "does the name .foo appear
|
||||
// anywhere in the stylesheet as a selector component".
|
||||
//
|
||||
// Scoping is respected loosely: `.form .actions .primary` counts as
|
||||
// defining `.primary`, and `.pl-acts button` defines no class — which is
|
||||
// exactly why the plugin panel needed its own rules.
|
||||
for _, m := range classTokenRe.FindAllStringSubmatch(css, -1) {
|
||||
defined[m[1]] = true
|
||||
}
|
||||
if len(defined) == 0 {
|
||||
t.Fatal("no classes parsed from the stylesheet; the check is broken")
|
||||
}
|
||||
|
||||
// Classes the JS builds as strings must exist too.
|
||||
js := guiFile(t, "cmd/gui/renderer/app.js")
|
||||
var missing []string
|
||||
seen := map[string]bool{}
|
||||
note := func(cls, where string) {
|
||||
// A "${...}" token is a template literal being spliced at runtime, not
|
||||
// a class name; the classes it can expand to are checked at their
|
||||
// definition sites instead.
|
||||
if cls == "" || strings.ContainsAny(cls, "${}") || seen[cls] {
|
||||
return
|
||||
}
|
||||
seen[cls] = true
|
||||
if guiKnownUnstyled[cls] {
|
||||
return
|
||||
}
|
||||
if !defined[cls] {
|
||||
missing = append(missing, cls+" ("+where+")")
|
||||
}
|
||||
}
|
||||
for _, m := range classUseRe.FindAllStringSubmatch(html, -1) {
|
||||
for _, c := range strings.Fields(m[1]) {
|
||||
note(c, "index.html")
|
||||
}
|
||||
}
|
||||
for _, m := range classUseRe.FindAllStringSubmatch(js, -1) {
|
||||
for _, c := range strings.Fields(m[1]) {
|
||||
note(c, "app.js")
|
||||
}
|
||||
}
|
||||
if len(missing) > 0 {
|
||||
t.Errorf("classes used but not defined in style.css (they render unstyled):\n %s",
|
||||
strings.Join(missing, "\n "))
|
||||
}
|
||||
}
|
||||
|
||||
// TestGUIHelperFunctionsAreDefined catches the other half: renderer/app.js is a
|
||||
// separate document from the WebUI, so it does NOT have the WebUI's esc/escAttr.
|
||||
// Referencing them gives a ReferenceError only when the line runs.
|
||||
func TestGUIHelperFunctionsAreDefined(t *testing.T) {
|
||||
js := guiFile(t, "cmd/gui/renderer/app.js")
|
||||
for _, fn := range []string{"esc", "escAttr", "toast", "loadPlugins", "togglePlugin", "enableAllPlugins"} {
|
||||
defined := regexp.MustCompile(`function ` + fn + `\b`).MatchString(js)
|
||||
called := regexp.MustCompile(`\b` + fn + `\s*\(`).MatchString(js)
|
||||
if called && !defined {
|
||||
t.Errorf("%s() is called but never defined in app.js", fn)
|
||||
}
|
||||
if !called && !defined {
|
||||
// A defined-but-unused helper is dead code, not an error.
|
||||
continue
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// TestGUIPluginPanelIsReachable: the panel must be inside the settings overlay
|
||||
// AND the settings overlay must actually open it. A panel wired to a button
|
||||
// that was never bound is invisible-but-present, which passes a grep review.
|
||||
func TestGUIPluginPanelIsReachable(t *testing.T) {
|
||||
html := guiFile(t, "cmd/gui/renderer/index.html")
|
||||
js := guiFile(t, "cmd/gui/renderer/app.js")
|
||||
if !strings.Contains(html, `id="pl-list"`) {
|
||||
t.Error("no #pl-list in the settings overlay")
|
||||
}
|
||||
if !strings.Contains(html, `id="settings-overlay"`) {
|
||||
t.Fatal("the settings overlay is gone")
|
||||
}
|
||||
// inside the overlay: the element index must come after the overlay's
|
||||
if strings.Index(html, `id="settings-overlay"`) > strings.Index(html, `id="pl-list"`) {
|
||||
t.Error("#pl-list appears before the settings overlay, so it renders outside the panel")
|
||||
}
|
||||
// The buttons must be bound.
|
||||
for _, id := range []string{"pl-reload", "pl-toggle-all"} {
|
||||
if !strings.Contains(html, `id="`+id+`"`) {
|
||||
t.Errorf("#%s is missing from the markup", id)
|
||||
}
|
||||
if !strings.Contains(js, `"`+id+`"`) {
|
||||
t.Errorf("#%s exists but app.js never binds it", id)
|
||||
}
|
||||
}
|
||||
// And openSettings must trigger the load, or the panel shows a stale empty
|
||||
// list on every open.
|
||||
if !strings.Contains(js, "loadPlugins()") {
|
||||
t.Error("app.js never calls loadPlugins()")
|
||||
}
|
||||
}
|
||||
|
||||
// TestGUIIPCPathIsConstrained: plugins:proxy is a raw pass-through, which is
|
||||
// convenient but would be a hole if it accepted arbitrary paths. The main
|
||||
// process must reject anything outside /api/plugins and any traversal.
|
||||
func TestGUIIPCPathIsConstrained(t *testing.T) {
|
||||
main := guiFile(t, "cmd/gui/main.js")
|
||||
for _, needle := range []string{
|
||||
`path.startsWith("/api/plugins")`,
|
||||
`path.includes("..")`,
|
||||
"plugins:proxy",
|
||||
"unsealViaCore()",
|
||||
} {
|
||||
if !strings.Contains(main, needle) {
|
||||
t.Errorf("cmd/gui/main.js is missing the guard %q", needle)
|
||||
}
|
||||
}
|
||||
// The plugins channel must be a proxy, not a key passthrough: the renderer
|
||||
// sends (method, path) and the main process attaches the key.
|
||||
//
|
||||
// NOTE: preload does expose a pre-existing `core.key` getter — the shell
|
||||
// needs the admin key to load the embedded WebUI without a login. That is
|
||||
// existing, deliberate design and out of scope here; asserting on "key:" in
|
||||
// preload would flag a pre-existing feature as a new hole.
|
||||
pre := guiFile(t, "cmd/gui/preload.js")
|
||||
if !strings.Contains(pre, "request: (method, path, body)") {
|
||||
t.Error("preload does not expose the plugins request proxy")
|
||||
}
|
||||
}
|
||||
257
internal/gateway/key_auto_api_test.go
Normal file
257
internal/gateway/key_auto_api_test.go
Normal file
@ -0,0 +1,257 @@
|
||||
package gateway
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"path/filepath"
|
||||
"os"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"llmsproxy/internal/config"
|
||||
"llmsproxy/internal/core"
|
||||
)
|
||||
|
||||
// Per-key AUTO chain admin API. The behaviours pinned here are the ones that
|
||||
// are easy to get wrong and invisible when they are:
|
||||
//
|
||||
// - /api/keys/me/auto must be reachable by a NON-admin (it is their own
|
||||
// chain); routing it through the generic {key}/auto handler would 403 it.
|
||||
// - /api/keys/{key}/auto must 403 a non-admin.
|
||||
// - an empty PUT and a DELETE mean the same thing ("inherit global"), so the
|
||||
// WebUI cannot persist a chain that would then 503 with no slots.
|
||||
// - a PUT that resolves to zero slots is reported back, so an admin who
|
||||
// mistyped a model learns it now instead of from the user's next 503.
|
||||
|
||||
// keyAutoGateway builds a gateway with one admin key and one user key.
|
||||
func keyAutoGateway(t *testing.T) *Gateway {
|
||||
t.Helper()
|
||||
td := t.TempDir()
|
||||
cfgPath := filepath.Join(td, "config.yaml")
|
||||
if err := os.WriteFile(cfgPath, []byte("listen: :0"), 0644); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
cfg := &config.Config{
|
||||
Path: cfgPath,
|
||||
AdapterDir: filepath.Join(td, "adapters"),
|
||||
RuntimeFile: filepath.Join(td, "runtime.json"),
|
||||
DefaultModel: "AUTO",
|
||||
GatewayKeys: []string{"sk-admin"},
|
||||
Keys: []config.GWKey{
|
||||
{Key: "sk-admin", Role: "admin", Name: "admin"},
|
||||
{Key: "sk-user", Role: "user", Name: "user"},
|
||||
},
|
||||
Sources: []config.Source{
|
||||
{Name: "s1", BaseURL: "http://127.0.0.1:1/v1", Adapter: "openai",
|
||||
Models: []config.Model{{ID: "gpt-4o", Priority: 10}, {ID: "gpt-4o-mini", Priority: 5}}},
|
||||
},
|
||||
}
|
||||
if err := cfg.ApplyDefaults(); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
c, err := core.NewFromConfig(cfg)
|
||||
if err != nil {
|
||||
t.Fatalf("core: %v", err)
|
||||
}
|
||||
t.Cleanup(c.Close)
|
||||
g, err := New(c)
|
||||
if err != nil {
|
||||
t.Fatalf("gateway: %v", err)
|
||||
}
|
||||
return g
|
||||
}
|
||||
|
||||
// doReqAs issues an authenticated request as a specific key.
|
||||
func doReqAs(t *testing.T, g *Gateway, key, method, path, body string) *httptest.ResponseRecorder {
|
||||
t.Helper()
|
||||
req, _ := http.NewRequest(method, path, strings.NewReader(body))
|
||||
req.Header.Set("Authorization", "Bearer "+key)
|
||||
if body != "" {
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
}
|
||||
rr := httptest.NewRecorder()
|
||||
g.Handler().ServeHTTP(rr, req)
|
||||
return rr
|
||||
}
|
||||
|
||||
func TestKeyAutoAPIPermissions(t *testing.T) {
|
||||
g := keyAutoGateway(t)
|
||||
|
||||
// A user may read their own chain...
|
||||
if rr := doReqAs(t, g, "sk-user", "GET", "/api/keys/me/auto", ""); rr.Code != http.StatusOK {
|
||||
t.Fatalf("GET /api/keys/me/auto as user = %d, want 200 (body %s)", rr.Code, rr.Body.String())
|
||||
}
|
||||
// ...but not another key's.
|
||||
if rr := doReqAs(t, g, "sk-user", "GET", "/api/keys/sk-admin/auto", ""); rr.Code != http.StatusForbidden {
|
||||
t.Fatalf("GET other key's auto as user = %d, want 403", rr.Code)
|
||||
}
|
||||
// ...and not their own via the admin route either.
|
||||
if rr := doReqAs(t, g, "sk-user", "PUT", "/api/keys/sk-user/auto", `{"auto":[{"model":"gpt-4o"}]}`); rr.Code != http.StatusForbidden {
|
||||
t.Fatalf("user PUT own auto = %d, want 403 (only admin configures)", rr.Code)
|
||||
}
|
||||
}
|
||||
|
||||
func TestKeyAutoAPICrudAndInherit(t *testing.T) {
|
||||
g := keyAutoGateway(t)
|
||||
|
||||
// Initially inherits.
|
||||
var got struct {
|
||||
Inherits bool `json:"inherits"`
|
||||
Auto []config.ModelScope `json:"auto"`
|
||||
}
|
||||
rr := doReqAs(t, g, "sk-admin", "GET", "/api/keys/sk-user/auto", "")
|
||||
if rr.Code != http.StatusOK {
|
||||
t.Fatalf("GET = %d: %s", rr.Code, rr.Body.String())
|
||||
}
|
||||
if err := json.Unmarshal(rr.Body.Bytes(), &got); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if !got.Inherits || len(got.Auto) != 0 {
|
||||
t.Fatalf("fresh key must inherit global, got inherits=%v auto=%v", got.Inherits, got.Auto)
|
||||
}
|
||||
|
||||
// Admin sets a chain; the response reports the resolved slot count.
|
||||
rr = doReqAs(t, g, "sk-admin", "PUT", "/api/keys/sk-user/auto",
|
||||
`{"auto":[{"model":"gpt-4o-mini","source":"s1","tier":1}]}`)
|
||||
if rr.Code != http.StatusOK {
|
||||
t.Fatalf("PUT = %d: %s", rr.Code, rr.Body.String())
|
||||
}
|
||||
var put struct {
|
||||
Inherits bool `json:"inherits"`
|
||||
Slots int `json:"slots"`
|
||||
}
|
||||
if err := json.Unmarshal(rr.Body.Bytes(), &put); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if put.Inherits {
|
||||
t.Fatal("after a non-empty PUT the key must no longer inherit")
|
||||
}
|
||||
if put.Slots != 1 {
|
||||
t.Fatalf("resolved slots = %d, want 1", put.Slots)
|
||||
}
|
||||
|
||||
// And it reads back as an own chain.
|
||||
rr = doReqAs(t, g, "sk-admin", "GET", "/api/keys/sk-user/auto", "")
|
||||
if err := json.Unmarshal(rr.Body.Bytes(), &got); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if got.Inherits || len(got.Auto) != 1 || got.Auto[0].Model != "gpt-4o-mini" {
|
||||
t.Fatalf("own chain must read back, got inherits=%v auto=%v", got.Inherits, got.Auto)
|
||||
}
|
||||
|
||||
// DELETE restores inheritance.
|
||||
if rr := doReqAs(t, g, "sk-admin", "DELETE", "/api/keys/sk-user/auto", ""); rr.Code != http.StatusOK {
|
||||
t.Fatalf("DELETE = %d: %s", rr.Code, rr.Body.String())
|
||||
}
|
||||
rr = doReqAs(t, g, "sk-admin", "GET", "/api/keys/sk-user/auto", "")
|
||||
if err := json.Unmarshal(rr.Body.Bytes(), &got); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if !got.Inherits {
|
||||
t.Fatal("after DELETE the key must inherit the global chain again")
|
||||
}
|
||||
}
|
||||
|
||||
// An empty PUT must behave like DELETE. Otherwise a WebUI that submits an
|
||||
// empty list (every slot deleted in the editor) would persist an empty chain
|
||||
// and the next AUTO request would fail with no_provider.
|
||||
func TestKeyAutoAPIEmptyPutClearsOverride(t *testing.T) {
|
||||
g := keyAutoGateway(t)
|
||||
|
||||
if rr := doReqAs(t, g, "sk-admin", "PUT", "/api/keys/sk-user/auto",
|
||||
`{"auto":[{"model":"gpt-4o","source":"s1","tier":1}]}`); rr.Code != http.StatusOK {
|
||||
t.Fatalf("seed PUT = %d: %s", rr.Code, rr.Body.String())
|
||||
}
|
||||
rr := doReqAs(t, g, "sk-admin", "PUT", "/api/keys/sk-user/auto", `{"auto":[]}`)
|
||||
if rr.Code != http.StatusOK {
|
||||
t.Fatalf("empty PUT = %d: %s", rr.Code, rr.Body.String())
|
||||
}
|
||||
var got struct {
|
||||
Inherits bool `json:"inherits"`
|
||||
}
|
||||
if rr := doReqAs(t, g, "sk-admin", "GET", "/api/keys/sk-user/auto", ""); rr.Code != http.StatusOK {
|
||||
t.Fatal(rr.Body.String())
|
||||
}
|
||||
if err := json.Unmarshal(rr.Body.Bytes(), &got); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if !got.Inherits {
|
||||
t.Fatal("an empty PUT must clear the override, not persist an empty chain")
|
||||
}
|
||||
}
|
||||
|
||||
// A chain naming a model that does not exist compiles to zero slots. The API
|
||||
// must say so at write time instead of letting the user's next request 503.
|
||||
func TestKeyAutoAPIReportsUnresolvableSlots(t *testing.T) {
|
||||
g := keyAutoGateway(t)
|
||||
|
||||
rr := doReqAs(t, g, "sk-admin", "PUT", "/api/keys/sk-user/auto",
|
||||
`{"auto":[{"model":"does-not-exist","source":"s1","tier":1}]}`)
|
||||
if rr.Code != http.StatusOK {
|
||||
t.Fatalf("PUT = %d: %s", rr.Code, rr.Body.String())
|
||||
}
|
||||
var put struct {
|
||||
Inherits bool `json:"inherits"`
|
||||
Slots int `json:"slots"`
|
||||
}
|
||||
if err := json.Unmarshal(rr.Body.Bytes(), &put); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if put.Slots != 0 {
|
||||
t.Fatalf("unknown model must resolve to 0 slots, got %d", put.Slots)
|
||||
}
|
||||
}
|
||||
|
||||
func TestKeyAutoAPIUnknownKey404(t *testing.T) {
|
||||
g := keyAutoGateway(t)
|
||||
|
||||
if rr := doReqAs(t, g, "sk-admin", "GET", "/api/keys/sk-nope/auto", ""); rr.Code != http.StatusNotFound {
|
||||
t.Fatalf("GET unknown key auto = %d, want 404", rr.Code)
|
||||
}
|
||||
if rr := doReqAs(t, g, "sk-admin", "PUT", "/api/keys/sk-nope/auto", `{"auto":[{"model":"gpt-4o"}]}`); rr.Code != http.StatusNotFound {
|
||||
t.Fatalf("PUT unknown key auto = %d, want 404", rr.Code)
|
||||
}
|
||||
}
|
||||
|
||||
// A bad quota must be rejected with 400, not persisted and not silently
|
||||
// clamped — same rule as the model scope path.
|
||||
func TestKeyAutoAPIRejectsBadQuota(t *testing.T) {
|
||||
g := keyAutoGateway(t)
|
||||
|
||||
rr := doReqAs(t, g, "sk-admin", "PUT", "/api/keys/sk-user/auto",
|
||||
`{"auto":[{"model":"gpt-4o","source":"s1","tier":1,"token_quota":-1}]}`)
|
||||
if rr.Code != http.StatusBadRequest {
|
||||
t.Fatalf("negative quota PUT = %d, want 400 (body %s)", rr.Code, rr.Body.String())
|
||||
}
|
||||
// Nothing may have been persisted.
|
||||
var got struct {
|
||||
Inherits bool `json:"inherits"`
|
||||
}
|
||||
get := doReqAs(t, g, "sk-admin", "GET", "/api/keys/sk-user/auto", "")
|
||||
if get.Code != http.StatusOK {
|
||||
t.Fatal(get.Body.String())
|
||||
}
|
||||
if err := json.Unmarshal(get.Body.Bytes(), &got); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if !got.Inherits {
|
||||
t.Fatal("a rejected PUT must not persist a chain")
|
||||
}
|
||||
}
|
||||
|
||||
// The other key endpoints must keep working: the {key}/auto route must not
|
||||
// shadow /api/keys/{key} itself.
|
||||
func TestKeyAutoAPIDoesNotShadowKeyRoutes(t *testing.T) {
|
||||
g := keyAutoGateway(t)
|
||||
|
||||
if rr := doReqAs(t, g, "sk-admin", "GET", "/api/keys", ""); rr.Code != http.StatusOK {
|
||||
t.Fatalf("GET /api/keys = %d, want 200", rr.Code)
|
||||
}
|
||||
if rr := doReqAs(t, g, "sk-admin", "GET", "/api/keys/me", ""); rr.Code != http.StatusOK {
|
||||
t.Fatalf("GET /api/keys/me = %d, want 200", rr.Code)
|
||||
}
|
||||
if rr := doReqAs(t, g, "sk-user", "GET", "/api/keys/me", ""); rr.Code != http.StatusOK {
|
||||
t.Fatalf("GET /api/keys/me as user = %d, want 200", rr.Code)
|
||||
}
|
||||
}
|
||||
253
internal/gateway/key_quota_api_test.go
Normal file
253
internal/gateway/key_quota_api_test.go
Normal file
@ -0,0 +1,253 @@
|
||||
package gateway
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"os"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"llmsproxy/internal/config"
|
||||
"llmsproxy/internal/core"
|
||||
)
|
||||
|
||||
// adminGateway builds a gateway whose admin key can call /api/keys.
|
||||
func adminGateway(t *testing.T, keys ...config.GWKey) *Gateway {
|
||||
t.Helper()
|
||||
td := t.TempDir()
|
||||
cfgPath := td + "/config.yaml"
|
||||
if err := os.WriteFile(cfgPath, []byte("listen: :0"), 0o644); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
cfg := &config.Config{
|
||||
Path: cfgPath,
|
||||
AdapterDir: td + "/adapters",
|
||||
RuntimeFile: td + "/runtime.json",
|
||||
Keys: keys,
|
||||
}
|
||||
if err := cfg.ApplyDefaults(); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
c, err := core.NewFromConfig(cfg)
|
||||
if err != nil {
|
||||
t.Fatalf("core: %v", err)
|
||||
}
|
||||
t.Cleanup(c.Close)
|
||||
g, err := New(c)
|
||||
if err != nil {
|
||||
t.Fatalf("gateway: %v", err)
|
||||
}
|
||||
return g
|
||||
}
|
||||
|
||||
func adminReq(t *testing.T, g *Gateway, method, path, body string) *httptest.ResponseRecorder {
|
||||
t.Helper()
|
||||
req, _ := http.NewRequest(method, path, strings.NewReader(body))
|
||||
req.Header.Set("Authorization", "Bearer sk-admin")
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
rr := httptest.NewRecorder()
|
||||
g.Handler().ServeHTTP(rr, req)
|
||||
return rr
|
||||
}
|
||||
|
||||
// keyRecord pulls one key's stored record out of the admin list.
|
||||
func keyRecord(t *testing.T, g *Gateway, secret string) config.GWKey {
|
||||
t.Helper()
|
||||
rr := adminReq(t, g, "GET", "/api/keys", "")
|
||||
if rr.Code != 200 {
|
||||
t.Fatalf("GET /api/keys: %d %s", rr.Code, rr.Body.String())
|
||||
}
|
||||
var out struct {
|
||||
Keys []config.GWKey `json:"keys"`
|
||||
}
|
||||
if err := json.Unmarshal(rr.Body.Bytes(), &out); err != nil {
|
||||
t.Fatalf("decode: %v (%s)", err, rr.Body.String())
|
||||
}
|
||||
for _, k := range out.Keys {
|
||||
if k.Key == secret {
|
||||
return k
|
||||
}
|
||||
}
|
||||
t.Fatalf("key %q not found in %s", secret, rr.Body.String())
|
||||
return config.GWKey{}
|
||||
}
|
||||
|
||||
func scopeOf(t *testing.T, k config.GWKey, model string) config.ModelScope {
|
||||
t.Helper()
|
||||
for _, m := range k.Models {
|
||||
if m.Model == model {
|
||||
return m
|
||||
}
|
||||
}
|
||||
t.Fatalf("scope %q not found in %+v", model, k.Models)
|
||||
return config.ModelScope{}
|
||||
}
|
||||
|
||||
// Quotas live on the scope entries, not on the key: creating a key with a
|
||||
// budget means creating scopes that carry it, and they must survive a
|
||||
// read-back (persisted, not just echoed).
|
||||
func TestKeyAPICreatesPerModelQuota(t *testing.T) {
|
||||
g := adminGateway(t, config.GWKey{Key: "sk-admin", Role: "admin"})
|
||||
rr := adminReq(t, g, "POST", "/api/keys",
|
||||
`{"name":"agent-x","role":"user","models":[{"model":"m1","token_quota":50000,"req_quota":200,"period":"nhour","hours":6}]}`)
|
||||
if rr.Code != 200 {
|
||||
t.Fatalf("create: %d %s", rr.Code, rr.Body.String())
|
||||
}
|
||||
var created struct {
|
||||
Key config.GWKey `json:"key"`
|
||||
}
|
||||
if err := json.Unmarshal(rr.Body.Bytes(), &created); err != nil {
|
||||
t.Fatalf("decode: %v", err)
|
||||
}
|
||||
sc := scopeOf(t, created.Key, "m1")
|
||||
if sc.TokenQuota != 50000 || sc.ReqQuota != 200 || sc.Period != "nhour" || sc.Hours != 6 {
|
||||
t.Fatalf("created scope did not carry the caps: %+v", sc)
|
||||
}
|
||||
back := scopeOf(t, keyRecord(t, g, created.Key.Key), "m1")
|
||||
if back.TokenQuota != 50000 || back.ReqQuota != 200 || back.Period != "nhour" || back.Hours != 6 {
|
||||
t.Errorf("read-back lost the caps: %+v", back)
|
||||
}
|
||||
}
|
||||
|
||||
// Two models on one key carry independent budgets.
|
||||
func TestKeyAPIKeepsPerModelQuotaIndependent(t *testing.T) {
|
||||
g := adminGateway(t, config.GWKey{Key: "sk-admin", Role: "admin"})
|
||||
rr := adminReq(t, g, "POST", "/api/keys",
|
||||
`{"name":"agent-y","role":"user","models":[{"model":"m1","token_quota":1000,"period":"hour"},{"model":"m2","token_quota":9999,"req_quota":7,"period":"week"}]}`)
|
||||
if rr.Code != 200 {
|
||||
t.Fatalf("create: %d %s", rr.Code, rr.Body.String())
|
||||
}
|
||||
var created struct {
|
||||
Key config.GWKey `json:"key"`
|
||||
}
|
||||
_ = json.Unmarshal(rr.Body.Bytes(), &created)
|
||||
back := keyRecord(t, g, created.Key.Key)
|
||||
if a := scopeOf(t, back, "m1"); a.TokenQuota != 1000 || a.ReqQuota != 0 || a.Period != "hour" {
|
||||
t.Errorf("m1 caps wrong: %+v", a)
|
||||
}
|
||||
if b := scopeOf(t, back, "m2"); b.TokenQuota != 9999 || b.ReqQuota != 7 || b.Period != "week" {
|
||||
t.Errorf("m2 caps wrong: %+v", b)
|
||||
}
|
||||
}
|
||||
|
||||
// Sending 0 explicitly lifts that model's cap.
|
||||
func TestKeyAPIUpdateZeroLiftsCap(t *testing.T) {
|
||||
g := adminGateway(t, config.GWKey{Key: "sk-admin", Role: "admin"})
|
||||
rr := adminReq(t, g, "POST", "/api/keys",
|
||||
`{"name":"z","role":"user","models":[{"model":"m1","token_quota":1000,"period":"hour"}]}`)
|
||||
if rr.Code != 200 {
|
||||
t.Fatalf("create: %d %s", rr.Code, rr.Body.String())
|
||||
}
|
||||
var created struct {
|
||||
Key config.GWKey `json:"key"`
|
||||
}
|
||||
_ = json.Unmarshal(rr.Body.Bytes(), &created)
|
||||
secret := created.Key.Key
|
||||
|
||||
rr = adminReq(t, g, "PUT", "/api/keys/"+secret,
|
||||
`{"models":[{"model":"m1","token_quota":0,"req_quota":0,"period":""}]}`)
|
||||
if rr.Code != 200 {
|
||||
t.Fatalf("lift: %d %s", rr.Code, rr.Body.String())
|
||||
}
|
||||
if back := scopeOf(t, keyRecord(t, g, secret), "m1"); back.TokenQuota != 0 || back.ReqQuota != 0 || back.Period != "" {
|
||||
t.Errorf("caps not lifted: %+v", back)
|
||||
}
|
||||
}
|
||||
|
||||
// A misspelled period must be refused, not quietly turned into an all-time
|
||||
// quota — which is the exact opposite of what the operator typed.
|
||||
func TestKeyAPIRejectsBadPeriod(t *testing.T) {
|
||||
g := adminGateway(t, config.GWKey{Key: "sk-admin", Role: "admin"})
|
||||
rr := adminReq(t, g, "POST", "/api/keys",
|
||||
`{"name":"bad","role":"user","models":[{"model":"m1","token_quota":1000,"period":"houre"}]}`)
|
||||
if rr.Code != http.StatusBadRequest {
|
||||
t.Fatalf("want 400 for a bad period, got %d %s", rr.Code, rr.Body.String())
|
||||
}
|
||||
if !strings.Contains(rr.Body.String(), "period") {
|
||||
t.Errorf("error should name the period field: %s", rr.Body.String())
|
||||
}
|
||||
}
|
||||
|
||||
func TestKeyAPIRejectsNegativeQuota(t *testing.T) {
|
||||
g := adminGateway(t, config.GWKey{Key: "sk-admin", Role: "admin"})
|
||||
rr := adminReq(t, g, "POST", "/api/keys",
|
||||
`{"name":"bad","role":"user","models":[{"model":"m1","token_quota":-5}]}`)
|
||||
if rr.Code != http.StatusBadRequest {
|
||||
t.Fatalf("want 400 for a negative quota, got %d %s", rr.Code, rr.Body.String())
|
||||
}
|
||||
}
|
||||
|
||||
// The rejection must say WHICH model is over budget, so an operator looking at
|
||||
// a key with a dozen scopes can tell which one to raise.
|
||||
func TestKeyAPIRejectionNamesTheModel(t *testing.T) {
|
||||
g := adminGateway(t, config.GWKey{Key: "sk-admin", Role: "admin"})
|
||||
rr := adminReq(t, g, "POST", "/api/keys",
|
||||
`{"name":"bad","role":"user","models":[{"model":"m1","req_quota":-1}]}`)
|
||||
if rr.Code != http.StatusBadRequest {
|
||||
t.Fatalf("want 400, got %d", rr.Code)
|
||||
}
|
||||
if !strings.Contains(rr.Body.String(), "m1") {
|
||||
t.Errorf("error should name the offending model: %s", rr.Body.String())
|
||||
}
|
||||
}
|
||||
|
||||
// A non-admin key must not be able to mint keys.
|
||||
func TestKeyAPIQuotaIsAdminOnly(t *testing.T) {
|
||||
g := adminGateway(t,
|
||||
config.GWKey{Key: "sk-admin", Role: "admin"},
|
||||
config.GWKey{Key: "sk-u", Role: "user", Models: []config.ModelScope{{Model: "m1", TokenQuota: 10, Period: "hour"}}},
|
||||
)
|
||||
req, _ := http.NewRequest("POST", "/api/keys",
|
||||
strings.NewReader(`{"name":"x","role":"admin","models":[{"model":"m1"}]}`))
|
||||
req.Header.Set("Authorization", "Bearer sk-u")
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
rr := httptest.NewRecorder()
|
||||
g.Handler().ServeHTTP(rr, req)
|
||||
if rr.Code != http.StatusForbidden {
|
||||
t.Fatalf("non-admin create: want 403, got %d %s", rr.Code, rr.Body.String())
|
||||
}
|
||||
// /api/v1/keys exposes the per-model caps but never a secret
|
||||
rr = adminReq(t, g, "GET", "/api/v1/keys", "")
|
||||
if rr.Code != 200 {
|
||||
t.Fatalf("GET /api/v1/keys: %d", rr.Code)
|
||||
}
|
||||
if strings.Contains(rr.Body.String(), "sk-u") {
|
||||
t.Error("/api/v1/keys leaked a key secret")
|
||||
}
|
||||
if !strings.Contains(rr.Body.String(), `"token_quota":10`) {
|
||||
t.Errorf("/api/v1/keys should expose the per-model caps: %s", rr.Body.String())
|
||||
}
|
||||
}
|
||||
|
||||
// A user must be able to see their own budget: /api/keys/me is the only key
|
||||
// view a non-admin gets, so a cap missing from it is invisible to the very
|
||||
// client it constrains.
|
||||
func TestKeyMeExposesOwnQuota(t *testing.T) {
|
||||
g := adminGateway(t,
|
||||
config.GWKey{Key: "sk-admin", Role: "admin"},
|
||||
config.GWKey{Key: "sk-u", Role: "user", Name: "agent",
|
||||
Models: []config.ModelScope{{Model: "m1", TokenQuota: 123456, ReqQuota: 42, Period: "week"}}},
|
||||
)
|
||||
req, _ := http.NewRequest("GET", "/api/keys/me", nil)
|
||||
req.Header.Set("Authorization", "Bearer sk-u")
|
||||
rr := httptest.NewRecorder()
|
||||
g.Handler().ServeHTTP(rr, req)
|
||||
if rr.Code != 200 {
|
||||
t.Fatalf("GET /api/keys/me: %d %s", rr.Code, rr.Body.String())
|
||||
}
|
||||
// the endpoint wraps the record: {"key": {...}}
|
||||
var wrap struct {
|
||||
Key config.GWKey `json:"key"`
|
||||
}
|
||||
if err := json.Unmarshal(rr.Body.Bytes(), &wrap); err != nil {
|
||||
t.Fatalf("decode: %v (%s)", err, rr.Body.String())
|
||||
}
|
||||
sc := scopeOf(t, wrap.Key, "m1")
|
||||
if sc.TokenQuota != 123456 || sc.ReqQuota != 42 || sc.Period != "week" {
|
||||
t.Errorf("own quota not visible to the key's owner: %+v", sc)
|
||||
}
|
||||
}
|
||||
|
||||
var _ = fmt.Sprintf
|
||||
60
internal/gateway/key_quota_concurrency_test.go
Normal file
60
internal/gateway/key_quota_concurrency_test.go
Normal file
@ -0,0 +1,60 @@
|
||||
package gateway
|
||||
|
||||
import (
|
||||
"sync"
|
||||
"testing"
|
||||
"time"
|
||||
)
|
||||
|
||||
// The per-key quota buckets are shared maps mutated on every recorded request
|
||||
// and read on every quota check. Single-threaded unit tests cannot see a race
|
||||
// there at all, so this drives 64 goroutines through record + read together
|
||||
// under `go test -race`. It also opts one model into the lazily-created pinned
|
||||
// bucket mid-flight, which is the path that mutates two bucket maps at once.
|
||||
func TestConcurrentQuotaAccounting(t *testing.T) {
|
||||
s := NewStats(0)
|
||||
keys := []string{"k0", "k1", "k2", "k3", "k4", "k5", "k6"}
|
||||
models := []string{"m0", "m1", "m2", "AUTO"}
|
||||
var wg sync.WaitGroup
|
||||
stop := time.Now().Add(2 * time.Second)
|
||||
|
||||
for w := 0; w < 64; w++ {
|
||||
wg.Add(1)
|
||||
go func(w int) {
|
||||
defer wg.Done()
|
||||
k := keys[w%len(keys)]
|
||||
m := models[w%len(models)]
|
||||
now := time.Now()
|
||||
for time.Now().Before(stop) {
|
||||
s.Record(Req{Time: now.UnixMilli(), Key: k, Model: m, Source: "src",
|
||||
Prompt: 100, Compl: 50, OK: true, Status: 200, Type: "chat"})
|
||||
_ = s.KeyWindowTokens(k, 3600)
|
||||
_ = s.KeyWindowReqs(k, 3600)
|
||||
_ = s.KeyWindowModelTokens(k, m, "", 3600)
|
||||
_ = s.KeyWindowModelTokens(k, m, "src", 3600) // opts into the pinned bucket
|
||||
}
|
||||
}(w)
|
||||
}
|
||||
wg.Wait()
|
||||
|
||||
// Every recorded request must be accounted for exactly once.
|
||||
s.mu.Lock()
|
||||
var total int64
|
||||
for _, hm := range s.keyHour {
|
||||
for _, v := range hm {
|
||||
total += v
|
||||
}
|
||||
}
|
||||
recs := int64(0)
|
||||
for _, r := range s.recs {
|
||||
recs += r.Prompt + r.Compl
|
||||
}
|
||||
s.mu.Unlock()
|
||||
if total == 0 {
|
||||
t.Fatal("no tokens accounted")
|
||||
}
|
||||
t.Logf("accounted %d tokens across %d keys", total, len(keys))
|
||||
if recs == 0 {
|
||||
t.Error("no records retained; the concurrent writers lost work")
|
||||
}
|
||||
}
|
||||
125
internal/gateway/key_quota_perf_test.go
Normal file
125
internal/gateway/key_quota_perf_test.go
Normal file
@ -0,0 +1,125 @@
|
||||
package gateway
|
||||
|
||||
import (
|
||||
"math/rand"
|
||||
"testing"
|
||||
"time"
|
||||
)
|
||||
|
||||
// Perf: the per-key quota bookkeeping runs on every recorded request and on
|
||||
// every quota check, so its cost lands directly on the request path. These
|
||||
// benchmarks exist to catch a regression that would make the feature
|
||||
// expensive; the numbers that matter are relative to each other and to the
|
||||
// pre-change baseline (Record was 275 ns/op with 3 allocs before this work).
|
||||
|
||||
func BenchmarkRecordQuotaBucket(b *testing.B) {
|
||||
s := NewStats(0)
|
||||
r := Req{Time: time.Now().UnixMilli(), Key: "keyABC", Model: "deepseek-v4-flash",
|
||||
Source: "deepseek", Prompt: 800, Compl: 200, OK: true, Status: 200, Type: "chat"}
|
||||
b.ReportAllocs()
|
||||
b.ResetTimer()
|
||||
for i := 0; i < b.N; i++ {
|
||||
s.Record(r)
|
||||
}
|
||||
}
|
||||
|
||||
// A key that has been busy for the whole 40-day retention.
|
||||
func BenchmarkRecordQuotaBucketFullRetention(b *testing.B) {
|
||||
s := NewStats(0)
|
||||
nowH := time.Now().Unix() / 3600
|
||||
for h := int64(0); h < quotaRetentionHours; h++ {
|
||||
s.Record(Req{Time: (nowH - h) * 3600 * 1000, Key: "keyABC", Model: "m1",
|
||||
Source: "deepseek", Prompt: 100, Compl: 100, OK: true, Status: 200})
|
||||
}
|
||||
r := Req{Time: time.Now().UnixMilli(), Key: "keyABC", Model: "m1", Source: "deepseek",
|
||||
Prompt: 100, Compl: 100, OK: true, Status: 200, Type: "chat"}
|
||||
b.ReportAllocs()
|
||||
b.ResetTimer()
|
||||
for i := 0; i < b.N; i++ {
|
||||
s.Record(r)
|
||||
}
|
||||
}
|
||||
|
||||
// The whole per-request quota check: total tokens + request count.
|
||||
func BenchmarkQuotaCheckHourly(b *testing.B) {
|
||||
s := NewStats(0)
|
||||
for h := 0; h < 24; h++ {
|
||||
s.Record(Req{Time: time.Now().Add(-time.Duration(h) * time.Hour).UnixMilli(),
|
||||
Key: "keyABC", Model: "m1", Source: "deepseek", Prompt: 100, Compl: 100, OK: true, Status: 200})
|
||||
}
|
||||
b.ReportAllocs()
|
||||
b.ResetTimer()
|
||||
for i := 0; i < b.N; i++ {
|
||||
s.KeyWindowTokens("keyABC", 3600)
|
||||
s.KeyWindowReqs("keyABC", 3600)
|
||||
}
|
||||
}
|
||||
|
||||
// A 30-day window has to visit 720 hour buckets. This is the worst case the
|
||||
// quota feature puts on the request path.
|
||||
func BenchmarkQuotaCheckMonthly(b *testing.B) {
|
||||
s := NewStats(0)
|
||||
nowH := time.Now().Unix() / 3600
|
||||
for h := int64(0); h < 720; h++ {
|
||||
s.Record(Req{Time: (nowH - h) * 3600 * 1000, Key: "keyABC", Model: "m1",
|
||||
Prompt: 100, Compl: 100, OK: true, Status: 200})
|
||||
}
|
||||
b.ReportAllocs()
|
||||
b.ResetTimer()
|
||||
for i := 0; i < b.N; i++ {
|
||||
s.KeyWindowTokens("keyABC", 30*24*3600)
|
||||
}
|
||||
}
|
||||
|
||||
// TestSumBucketsMatchesFullScan is the correctness guard for the optimization
|
||||
// that made the quota check proportional to the window instead of to the whole
|
||||
// retention: the windowed sum must agree with a full scan for every bucket
|
||||
// placement, or a quota would silently start letting traffic through. It
|
||||
// caught a real off-by-one (floor instead of ceil on the first hour).
|
||||
func TestSumBucketsMatchesFullScan(t *testing.T) {
|
||||
const hour = 3600
|
||||
base := int64(1_700_000_000) / hour * hour
|
||||
rng := rand.New(rand.NewSource(7))
|
||||
for trial := 0; trial < 400; trial++ {
|
||||
hm := map[int64]int64{}
|
||||
for i, n := 0, rng.Intn(40)+1; i < n; i++ {
|
||||
h := base/hour - int64(rng.Intn(1200))
|
||||
hm[h] += int64(rng.Intn(1000)) + 1
|
||||
}
|
||||
now := base + int64(rng.Intn(hour))
|
||||
for _, sec := range []int64{3600, 2 * 3600, 6 * 3600, 24 * 3600, 7 * 24 * 3600, 30 * 24 * 3600} {
|
||||
var want int64
|
||||
cut := now - sec
|
||||
for h, v := range hm {
|
||||
if h*hour >= cut {
|
||||
want += v
|
||||
}
|
||||
}
|
||||
if got := sumBuckets(hm, now, sec); got != want {
|
||||
t.Fatalf("trial %d sec=%d: sumBuckets=%d, full scan=%d", trial, sec, got, want)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// The pinned bucket is maintained only for (key, model) pairs whose quota pins
|
||||
// a source. Doing it unconditionally doubled the bucket count — measured at
|
||||
// 26.7 MB for 20 keys x 8 models x 3 sources at full retention, against a
|
||||
// documented ~32 MB total memory budget — to serve a lookup nobody performs.
|
||||
func TestPinnedBucketNotMaintainedUnlessWanted(t *testing.T) {
|
||||
s := NewStats(100)
|
||||
now := time.Now().UnixMilli()
|
||||
for i := 0; i < 5; i++ {
|
||||
s.Record(Req{Time: now, Key: "keyA", Model: "m1", Source: "srcX", Prompt: 100, Compl: 100, OK: true, Status: 200})
|
||||
}
|
||||
s.mu.Lock()
|
||||
_, pinned := s.keyModelHour["keyA"]["srcX::m1"]
|
||||
_, bare := s.keyModelHour["keyA"]["m1"]
|
||||
s.mu.Unlock()
|
||||
if pinned {
|
||||
t.Error("a pinned bucket exists before any quota read the pin — memory regression")
|
||||
}
|
||||
if !bare {
|
||||
t.Error("the bare-model bucket must always exist; an unpinned quota depends on it")
|
||||
}
|
||||
}
|
||||
225
internal/gateway/key_quota_test.go
Normal file
225
internal/gateway/key_quota_test.go
Normal file
@ -0,0 +1,225 @@
|
||||
package gateway
|
||||
|
||||
import (
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"llmsproxy/internal/config"
|
||||
)
|
||||
|
||||
// ---- data layer: per-key isolation ----
|
||||
|
||||
func TestKeyWindowTokensIsolatesKeys(t *testing.T) {
|
||||
s := NewStats(100)
|
||||
now := time.Now().UnixMilli()
|
||||
s.Record(Req{Time: now, Key: "keyA", Model: "m1", Prompt: 500, Compl: 500, OK: true, Status: 200})
|
||||
s.Record(Req{Time: now, Key: "keyB", Model: "m1", Prompt: 500, Compl: 500, OK: true, Status: 200})
|
||||
|
||||
if got := s.KeyWindowModelTokens("keyA", "m1", "", 3600); got != 1000 {
|
||||
t.Errorf("keyA model tokens = %d, want 1000 (its own usage only)", got)
|
||||
}
|
||||
if got := s.KeyWindowModelTokens("keyB", "m1", "", 3600); got != 1000 {
|
||||
t.Errorf("keyB model tokens = %d, want 1000", got)
|
||||
}
|
||||
// the key-blind bucket stays global on purpose: it backs the AUTO slot
|
||||
// quota, which limits the whole gateway, not one key.
|
||||
if got := s.WindowTokens("m1", "", 3600); got != 2000 {
|
||||
t.Errorf("WindowTokens = %d, want 2000 (global per-model total must be unchanged)", got)
|
||||
}
|
||||
}
|
||||
|
||||
// A source pin scopes the cap to that one upstream. The pinned bucket is
|
||||
// maintained lazily (it is only read by quotas that actually pin a source), so
|
||||
// usage recorded before anything queried the pin is not attributed to it.
|
||||
func TestKeyWindowTokensSeparatesSourcePin(t *testing.T) {
|
||||
s := NewStats(100)
|
||||
now := time.Now().UnixMilli()
|
||||
s.Record(Req{Time: now, Key: "keyA", Model: "m1", Source: "srcX", Prompt: 100, Compl: 100, OK: true, Status: 200}) // 200 tok
|
||||
s.Record(Req{Time: now, Key: "keyA", Model: "m1", Source: "srcY", Prompt: 200, Compl: 200, OK: true, Status: 200}) // 400 tok
|
||||
|
||||
// The unpinned bucket always counts — it is what a scope entry without a
|
||||
// source pin reads. 200 + 400.
|
||||
if got := s.KeyWindowModelTokens("keyA", "m1", "", 3600); got != 600 {
|
||||
t.Errorf("unpinned model tokens = %d, want 600 (both sources)", got)
|
||||
}
|
||||
|
||||
// Reading a pin opts this (key, model) into pinned accounting.
|
||||
if got := s.KeyWindowModelTokens("keyA", "m1", "srcX", 3600); got != 0 {
|
||||
t.Errorf("srcX pin before opt-in = %d, want 0 (the lazy bucket has not accrued yet)", got)
|
||||
}
|
||||
|
||||
// From now on both pins accrue.
|
||||
more := now + 1
|
||||
s.Record(Req{Time: more, Key: "keyA", Model: "m1", Source: "srcX", Prompt: 10, Compl: 10, OK: true, Status: 200}) // 20 tok
|
||||
s.Record(Req{Time: more, Key: "keyA", Model: "m1", Source: "srcY", Prompt: 30, Compl: 30, OK: true, Status: 200}) // 60 tok
|
||||
if got := s.KeyWindowModelTokens("keyA", "m1", "srcX", 3600); got != 20 {
|
||||
t.Errorf("srcX-pinned tokens after opt-in = %d, want 20", got)
|
||||
}
|
||||
if got := s.KeyWindowModelTokens("keyA", "m1", "srcY", 3600); got != 60 {
|
||||
t.Errorf("srcY-pinned tokens = %d, want 60 (a pin read for one source must not blind the other)", got)
|
||||
}
|
||||
}
|
||||
|
||||
// This is the shape of every real chat record: Source is always populated.
|
||||
// The unpinned bucket must still see the tokens, or a per-model quota without
|
||||
// a source pin reads an empty bucket and never trips.
|
||||
func TestKeyWindowModelTokensUnpinnedSeesSourcedTraffic(t *testing.T) {
|
||||
s := NewStats(100)
|
||||
s.Record(Req{Time: time.Now().UnixMilli(), Key: "keyA", Model: "m1", Source: "deepseek",
|
||||
Prompt: 1000, Compl: 1000, OK: true, Status: 200})
|
||||
|
||||
if got := s.KeyWindowModelTokens("keyA", "m1", "", 3600); got != 2000 {
|
||||
t.Errorf("unpinned tokens = %d, want 2000 — an unpinned per-model quota would never trip otherwise", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestKeyWindowRespectsResetWindow(t *testing.T) {
|
||||
s := NewStats(100)
|
||||
old := time.Now().Add(-30 * 24 * time.Hour).UnixMilli()
|
||||
s.Record(Req{Time: old, Key: "keyA", Model: "m1", Prompt: 100, Compl: 100, OK: true, Status: 200})
|
||||
s.Record(Req{Time: time.Now().UnixMilli(), Key: "keyA", Model: "m1", Prompt: 5, Compl: 5, OK: true, Status: 200})
|
||||
|
||||
if got := s.KeyWindowTokens("keyA", 3600); got != 10 {
|
||||
t.Errorf("1h window = %d, want 10 (30-day-old usage must not count)", got)
|
||||
}
|
||||
if got := s.KeyWindowTokens("keyA", 0); got != 210 {
|
||||
t.Errorf("all-time total = %d, want 210", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestKeyWindowReqsCountsEveryRequest(t *testing.T) {
|
||||
s := NewStats(100)
|
||||
now := time.Now().UnixMilli()
|
||||
// a failed request and an image request both count: a client that loops on
|
||||
// failures must still burn its request quota
|
||||
s.Record(Req{Time: now, Key: "keyA", Model: "m1", Type: "chat", OK: true, Status: 200})
|
||||
s.Record(Req{Time: now, Key: "keyA", Model: "m1", Type: "chat", OK: false, Status: 500})
|
||||
s.Record(Req{Time: now, Key: "keyA", Model: "img", Type: "image", OK: true, Status: 200})
|
||||
s.Record(Req{Time: now, Key: "keyB", Model: "m1", Type: "chat", OK: true, Status: 200})
|
||||
|
||||
if got := s.KeyWindowReqs("keyA", 3600); got != 3 {
|
||||
t.Errorf("keyA reqs = %d, want 3 (chat ok + chat fail + image)", got)
|
||||
}
|
||||
if got := s.KeyWindowReqs("keyB", 3600); got != 1 {
|
||||
t.Errorf("keyB reqs = %d, want 1", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestKeyWindowTokensZeroTokensNotCounted(t *testing.T) {
|
||||
s := NewStats(100)
|
||||
now := time.Now().UnixMilli()
|
||||
// a request that reported no usage must not create a bucket entry
|
||||
s.Record(Req{Time: now, Key: "keyA", Model: "", Type: "chat", OK: true, Status: 200})
|
||||
if got := s.KeyWindowTokens("keyA", 3600); got != 0 {
|
||||
t.Errorf("tokens = %d, want 0", got)
|
||||
}
|
||||
if got := s.KeyWindowReqs("keyA", 3600); got != 1 {
|
||||
t.Errorf("reqs = %d, want 1 (request still happened)", got)
|
||||
}
|
||||
}
|
||||
|
||||
// ---- reset window arithmetic ----
|
||||
|
||||
func TestAutoSecondsToReset(t *testing.T) {
|
||||
if got := AutoSecondsToReset("", 0); got != 0 {
|
||||
t.Errorf("no period = %d, want 0 (never resets -> no retry hint)", got)
|
||||
}
|
||||
now := time.Now().Unix()
|
||||
for _, tc := range []struct {
|
||||
period string
|
||||
hours int64
|
||||
want int64
|
||||
}{
|
||||
{"hour", 0, 3600},
|
||||
{"week", 0, 7 * 24 * 3600},
|
||||
{"month", 0, 30 * 24 * 3600},
|
||||
{"nhour", 6, 6 * 3600},
|
||||
} {
|
||||
got := AutoSecondsToReset(tc.period, tc.hours)
|
||||
if got <= 0 || got > tc.want {
|
||||
t.Errorf("AutoSecondsToReset(%q,%d) = %d, want in (0,%d]", tc.period, tc.hours, got, tc.want)
|
||||
}
|
||||
// must never exceed the window itself
|
||||
if got < tc.want-now%3600 {
|
||||
t.Logf("note: %q hint %ds < remaining-window %ds (rounds to hour boundary)", tc.period, got, tc.want-now%3600)
|
||||
}
|
||||
if got > tc.want {
|
||||
t.Errorf("hint %d exceeds window %d", got, tc.want)
|
||||
}
|
||||
}
|
||||
_ = now
|
||||
}
|
||||
|
||||
func TestAutoSecondsToResetNeverExceedsWindow(t *testing.T) {
|
||||
// "nhour" with a tiny window must not hand out a longer wait than the
|
||||
// window itself (which would stall a client past its own reset)
|
||||
for hours := int64(1); hours <= 48; hours++ {
|
||||
got := AutoSecondsToReset("nhour", hours)
|
||||
if got <= 0 || got > hours*3600 {
|
||||
t.Errorf("nhour/%d = %d, want in (0,%d]", hours, got, hours*3600)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// ---- config validation ----
|
||||
|
||||
func TestKeyQuotaValidate(t *testing.T) {
|
||||
cases := []struct {
|
||||
name string
|
||||
q config.KeyQuota
|
||||
wantErr bool
|
||||
}{
|
||||
{"unlimited", config.KeyQuota{}, false},
|
||||
{"tokens hourly", config.KeyQuota{TokenQuota: 1000, Period: "hour"}, false},
|
||||
{"tokens n-hour", config.KeyQuota{TokenQuota: 1000, Period: "nhour", Hours: 6}, false},
|
||||
{"reqs weekly", config.KeyQuota{ReqQuota: 100, Period: "week"}, false},
|
||||
{"no period never resets", config.KeyQuota{TokenQuota: 1000}, false},
|
||||
{"negative tokens", config.KeyQuota{TokenQuota: -1}, true},
|
||||
{"negative reqs", config.KeyQuota{ReqQuota: -1}, true},
|
||||
{"negative hours", config.KeyQuota{TokenQuota: 5, Hours: -1}, true},
|
||||
{"typo period becomes all-time if accepted", config.KeyQuota{TokenQuota: 5, Period: "houre"}, true},
|
||||
{"nhour without hours", config.KeyQuota{TokenQuota: 5, Period: "nhour"}, true},
|
||||
{"nhour with 0 hours", config.KeyQuota{TokenQuota: 5, Period: "nhour", Hours: 0}, true},
|
||||
{"junk period with only req quota", config.KeyQuota{ReqQuota: 5, Period: "daily"}, true},
|
||||
{"junk period with no caps is irrelevant", config.KeyQuota{Period: "daily"}, false},
|
||||
}
|
||||
for _, tc := range cases {
|
||||
err := tc.q.Validate()
|
||||
if tc.wantErr && err == nil {
|
||||
t.Errorf("%s: expected error, got nil", tc.name)
|
||||
}
|
||||
if !tc.wantErr && err != nil {
|
||||
t.Errorf("%s: unexpected error: %v", tc.name, err)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestNormalizeRole(t *testing.T) {
|
||||
if got := config.NormalizeRole(""); got != "user" {
|
||||
t.Errorf("empty role = %q, want user", got)
|
||||
}
|
||||
if got := config.NormalizeRole("admin"); got != "admin" {
|
||||
t.Errorf("admin role = %q, want admin", got)
|
||||
}
|
||||
if got := config.NormalizeRole("root"); got != "user" {
|
||||
t.Errorf("unknown role = %q, want user (never escalate)", got)
|
||||
}
|
||||
}
|
||||
|
||||
// Retention must hold for the lazy pinned buckets too, or a long-lived key
|
||||
// would grow without bound.
|
||||
func TestPinnedBucketRespectsRetention(t *testing.T) {
|
||||
s := NewStats(100)
|
||||
_ = s.KeyWindowModelTokens("keyA", "m1", "srcX", 3600) // opt in
|
||||
nowH := time.Now().Unix() / 3600
|
||||
for h := int64(0); h < quotaRetentionHours+50; h++ {
|
||||
s.Record(Req{Time: (nowH - h) * 3600 * 1000, Key: "keyA", Model: "m1", Source: "srcX",
|
||||
Prompt: 10, Compl: 10, OK: true, Status: 200})
|
||||
}
|
||||
s.mu.Lock()
|
||||
n := len(s.keyModelHour["keyA"]["srcX::m1"])
|
||||
s.mu.Unlock()
|
||||
if n > quotaRetentionHours {
|
||||
t.Errorf("pinned bucket holds %d hours, want <= %d", n, quotaRetentionHours)
|
||||
}
|
||||
}
|
||||
444
internal/gateway/key_quota_wiring_test.go
Normal file
444
internal/gateway/key_quota_wiring_test.go
Normal file
@ -0,0 +1,444 @@
|
||||
package gateway
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"llmsproxy/internal/config"
|
||||
"llmsproxy/internal/core"
|
||||
)
|
||||
|
||||
// quotaGateway builds a gateway with two user keys that differ only in their
|
||||
// configured caps, against one mock upstream that reports 4 tokens per call.
|
||||
func quotaGateway(t *testing.T, keyA, keyB config.GWKey) (*Gateway, *upstreamCtrl) {
|
||||
t.Helper()
|
||||
ctrl := &upstreamCtrl{}
|
||||
up := upstream(t, ctrl)
|
||||
t.Cleanup(up.Close)
|
||||
|
||||
td := t.TempDir()
|
||||
cfgPath := td + "/config.yaml"
|
||||
if err := os.WriteFile(cfgPath, []byte("listen: :0"), 0o644); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
cfg := &config.Config{
|
||||
Path: cfgPath,
|
||||
AdapterDir: td + "/adapters",
|
||||
RuntimeFile: td + "/runtime.json",
|
||||
Keys: []config.GWKey{keyA, keyB},
|
||||
Sources: []config.Source{{
|
||||
Name: "up", BaseURL: up.URL, Adapter: "openai",
|
||||
Models: []config.Model{{ID: "m1", Priority: 100}},
|
||||
}},
|
||||
}
|
||||
if err := cfg.ApplyDefaults(); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
c, err := core.NewFromConfig(cfg)
|
||||
if err != nil {
|
||||
t.Fatalf("core: %v", err)
|
||||
}
|
||||
t.Cleanup(c.Close)
|
||||
g, err := New(c)
|
||||
if err != nil {
|
||||
t.Fatalf("gateway: %v", err)
|
||||
}
|
||||
return g, ctrl
|
||||
}
|
||||
|
||||
// chatAs issues a chat completion with a specific gateway key and returns the
|
||||
// recorder plus the decoded error code.
|
||||
func chatAs(t *testing.T, g *Gateway, key, model string) (*httptest.ResponseRecorder, string) {
|
||||
t.Helper()
|
||||
req, _ := http.NewRequest("POST", "/v1/chat/completions",
|
||||
strings.NewReader(fmt.Sprintf(`{"model":%q,"messages":[{"role":"user","content":"hi"}]}`, model)))
|
||||
req.Header.Set("Authorization", "Bearer "+key)
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
rr := httptest.NewRecorder()
|
||||
g.Handler().ServeHTTP(rr, req)
|
||||
code := ""
|
||||
if rr.Code != 200 {
|
||||
var e struct {
|
||||
Error struct {
|
||||
Type string `json:"type"`
|
||||
} `json:"error"`
|
||||
}
|
||||
_ = json.Unmarshal(rr.Body.Bytes(), &e)
|
||||
code = e.Error.Type
|
||||
}
|
||||
return rr, code
|
||||
}
|
||||
|
||||
// A key whose token budget is spent must be refused, and the refusal must be a
|
||||
// 429 the client can retry after the window rolls over — not a 403 that reads
|
||||
// as "this key may never use this model".
|
||||
func TestKeyTokenQuotaBlocksWithRetryAfter(t *testing.T) {
|
||||
g, _ := quotaGateway(t,
|
||||
config.GWKey{Key: "sk-a", Role: "user",
|
||||
Models: []config.ModelScope{{Model: "m1", TokenQuota: 8, Period: "hour"}}},
|
||||
config.GWKey{Key: "sk-b", Role: "user", Name: "b"},
|
||||
)
|
||||
// 4 tokens per call, budget 8 -> the third call crosses it
|
||||
for i := 1; i <= 2; i++ {
|
||||
rr, _ := chatAs(t, g, "sk-a", "m1")
|
||||
if rr.Code != 200 {
|
||||
t.Fatalf("call %d: want 200, got %d (%s)", i, rr.Code, rr.Body.String())
|
||||
}
|
||||
}
|
||||
rr, code := chatAs(t, g, "sk-a", "m1")
|
||||
if rr.Code != http.StatusTooManyRequests {
|
||||
t.Fatalf("third call: want 429, got %d (%s)", rr.Code, rr.Body.String())
|
||||
}
|
||||
if code != "rate_limit_exceeded" {
|
||||
t.Errorf("error code = %q, want rate_limit_exceeded", code)
|
||||
}
|
||||
if ra := rr.Header().Get("Retry-After"); ra == "" {
|
||||
t.Error("429 must carry Retry-After so a client knows when to come back")
|
||||
} else if n := mustAtoi(t, ra); n <= 0 || n > 3600 {
|
||||
t.Errorf("Retry-After = %q, want a positive value within the 1h window", ra)
|
||||
}
|
||||
if !strings.Contains(rr.Body.String(), "quota") {
|
||||
t.Errorf("body should name the quota: %s", rr.Body.String())
|
||||
}
|
||||
}
|
||||
|
||||
// The quota is per key: exhausting one key's budget must not affect another
|
||||
// key that is allowed the same model.
|
||||
func TestKeyTokenQuotaIsIsolatedPerKey(t *testing.T) {
|
||||
g, _ := quotaGateway(t,
|
||||
config.GWKey{Key: "sk-a", Role: "user",
|
||||
Models: []config.ModelScope{{Model: "m1", TokenQuota: 4, Period: "hour"}}},
|
||||
config.GWKey{Key: "sk-b", Role: "user",
|
||||
Models: []config.ModelScope{{Model: "m1", TokenQuota: 1000, Period: "hour"}}},
|
||||
)
|
||||
rr, _ := chatAs(t, g, "sk-a", "m1")
|
||||
if rr.Code != 200 {
|
||||
t.Fatalf("keyA first call: want 200, got %d", rr.Code)
|
||||
}
|
||||
if rr, code := chatAs(t, g, "sk-a", "m1"); rr.Code != http.StatusTooManyRequests {
|
||||
t.Fatalf("keyA second call: want 429, got %d (%s)", rr.Code, code)
|
||||
}
|
||||
// keyB has its own budget and must still be served
|
||||
if rr, _ := chatAs(t, g, "sk-b", "m1"); rr.Code != 200 {
|
||||
t.Fatalf("keyB must be unaffected by keyA's exhausted quota, got %d", rr.Code)
|
||||
}
|
||||
}
|
||||
|
||||
// A per-model quota inside the scope must also be per key: one key burning a
|
||||
// shared model budget must not lock the other key out.
|
||||
func TestScopeModelTokenQuotaIsolatedPerKey(t *testing.T) {
|
||||
g, _ := quotaGateway(t,
|
||||
config.GWKey{Key: "sk-a", Role: "user", Models: []config.ModelScope{
|
||||
{Model: "m1", TokenQuota: 4, Period: "hour"},
|
||||
}},
|
||||
config.GWKey{Key: "sk-b", Role: "user", Models: []config.ModelScope{
|
||||
{Model: "m1", TokenQuota: 4, Period: "hour"},
|
||||
}},
|
||||
)
|
||||
if rr, _ := chatAs(t, g, "sk-a", "m1"); rr.Code != 200 {
|
||||
t.Fatalf("keyA first call: want 200, got %d", rr.Code)
|
||||
}
|
||||
if rr, _ := chatAs(t, g, "sk-a", "m1"); rr.Code != http.StatusTooManyRequests {
|
||||
t.Fatalf("keyA second call: want 429 (its own model quota), got %d", rr.Code)
|
||||
}
|
||||
// keyB shares the model but not the counter — it must be served twice too
|
||||
if rr, _ := chatAs(t, g, "sk-b", "m1"); rr.Code != 200 {
|
||||
t.Fatalf("keyB first call: want 200, got %d", rr.Code)
|
||||
}
|
||||
if rr, _ := chatAs(t, g, "sk-b", "m1"); rr.Code != http.StatusTooManyRequests {
|
||||
t.Fatalf("keyB second call: want 429, got %d", rr.Code)
|
||||
}
|
||||
}
|
||||
|
||||
// An admin key is never capped: a budget set on an admin key must not lock the
|
||||
// operator out of the gateway they administer.
|
||||
func TestAdminKeyIsNeverQuotaCapped(t *testing.T) {
|
||||
g, _ := quotaGateway(t,
|
||||
config.GWKey{Key: "sk-admin", Role: "admin",
|
||||
Models: []config.ModelScope{{Model: "m1", TokenQuota: 1, ReqQuota: 1, Period: "hour"}}},
|
||||
config.GWKey{Key: "sk-b", Role: "user"},
|
||||
)
|
||||
for i := 1; i <= 3; i++ {
|
||||
if rr, _ := chatAs(t, g, "sk-admin", "m1"); rr.Code != 200 {
|
||||
t.Fatalf("admin call %d: want 200 (admin keys are uncapped), got %d", i, rr.Code)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// A request quota caps sustained volume; RPM-style bursts are not its job but
|
||||
// the count must still stop the key.
|
||||
func TestKeyRequestQuotaBlocks(t *testing.T) {
|
||||
g, ctrl := quotaGateway(t,
|
||||
config.GWKey{Key: "sk-a", Role: "user",
|
||||
Models: []config.ModelScope{{Model: "m1", ReqQuota: 2, Period: "hour"}}},
|
||||
config.GWKey{Key: "sk-b", Role: "user"},
|
||||
)
|
||||
for i := 1; i <= 2; i++ {
|
||||
if rr, _ := chatAs(t, g, "sk-a", "m1"); rr.Code != 200 {
|
||||
t.Fatalf("call %d: want 200, got %d", i, rr.Code)
|
||||
}
|
||||
}
|
||||
rr, code := chatAs(t, g, "sk-a", "m1")
|
||||
if rr.Code != http.StatusTooManyRequests || code != "rate_limit_exceeded" {
|
||||
t.Fatalf("third call: want 429/rate_limit_exceeded, got %d/%s", rr.Code, code)
|
||||
}
|
||||
if ctrl.hits != 2 {
|
||||
t.Errorf("upstream saw %d calls, want 2 (a refused request must not reach upstream)", ctrl.hits)
|
||||
}
|
||||
}
|
||||
|
||||
// A model outside the scope is still 403, not 429: retrying cannot help.
|
||||
func TestModelOutsideScopeStaysForbidden(t *testing.T) {
|
||||
g, _ := quotaGateway(t,
|
||||
config.GWKey{Key: "sk-a", Role: "user",
|
||||
Models: []config.ModelScope{{Model: "other-model", TokenQuota: 1000, Period: "hour"}}},
|
||||
config.GWKey{Key: "sk-b", Role: "user"},
|
||||
)
|
||||
rr, code := chatAs(t, g, "sk-a", "m1")
|
||||
if rr.Code != http.StatusForbidden {
|
||||
t.Fatalf("want 403, got %d (%s)", rr.Code, rr.Body.String())
|
||||
}
|
||||
if code != "model_not_allowed" {
|
||||
t.Errorf("error code = %q, want model_not_allowed", code)
|
||||
}
|
||||
if rr.Header().Get("Retry-After") != "" {
|
||||
t.Error("a 403 must not advertise a retry time")
|
||||
}
|
||||
}
|
||||
|
||||
// An unlimited key (no caps at all) is never rejected — the default config
|
||||
// must keep working exactly as before.
|
||||
func TestKeyWithoutCapsIsUnlimited(t *testing.T) {
|
||||
g, _ := quotaGateway(t,
|
||||
config.GWKey{Key: "sk-a", Role: "user", Models: []config.ModelScope{{Model: "m1"}}},
|
||||
config.GWKey{Key: "sk-b", Role: "user"},
|
||||
)
|
||||
for i := 1; i <= 5; i++ {
|
||||
if rr, _ := chatAs(t, g, "sk-a", "m1"); rr.Code != 200 {
|
||||
t.Fatalf("call %d: want 200 (no caps = unlimited), got %d", i, rr.Code)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func mustAtoi(t *testing.T, s string) int64 {
|
||||
t.Helper()
|
||||
var n int64
|
||||
if _, err := fmt.Sscanf(s, "%d", &n); err != nil {
|
||||
t.Fatalf("Retry-After %q is not an integer: %v", s, err)
|
||||
}
|
||||
return n
|
||||
}
|
||||
|
||||
// quotaCtx builds a request context carrying a gateway key, as the auth
|
||||
// middleware would.
|
||||
func quotaCtx(t *testing.T, g *Gateway, key string) context.Context {
|
||||
t.Helper()
|
||||
return withAuth(context.Background(), key, "user")
|
||||
}
|
||||
|
||||
func nowMSOffset(sec int64) int64 {
|
||||
return time.Now().Add(time.Duration(sec) * time.Second).UnixMilli()
|
||||
}
|
||||
|
||||
// The key-wide quota and the AUTO slot quota are different limits with
|
||||
// different scopes: the slot quota is gateway-wide (a shared upstream budget),
|
||||
// the key quota belongs to one caller. When both are exhausted the caller must
|
||||
// see the KEY quota, because that is the one it can act on — the slot quota
|
||||
// would otherwise be reported as "no capacity", which reads like an outage.
|
||||
func TestKeyQuotaWinsOverSlotQuota(t *testing.T) {
|
||||
up := upstream(t, &upstreamCtrl{})
|
||||
defer up.Close()
|
||||
g := newQuotaGW(t, up.URL,
|
||||
config.GWKey{Key: "sk-a", Role: "user",
|
||||
Models: []config.ModelScope{{Model: "AUTO", TokenQuota: 4, Period: "hour"}}})
|
||||
ctx := quotaCtx(t, g, "sk-a")
|
||||
|
||||
// exhaust the key first
|
||||
g.stats.Record(Req{Time: time.Now().UnixMilli(), Key: keyID("sk-a"), Model: "m1",
|
||||
Source: "up", Prompt: 100, Compl: 100, OK: true, Status: 200})
|
||||
rr, code := chatAs(t, g, "sk-a", "AUTO")
|
||||
if rr.Code != http.StatusTooManyRequests {
|
||||
t.Fatalf("want 429 from the key quota, got %d (%s)", rr.Code, rr.Body.String())
|
||||
}
|
||||
if code != "rate_limit_exceeded" {
|
||||
t.Errorf("code = %q, want rate_limit_exceeded", code)
|
||||
}
|
||||
_ = ctx
|
||||
}
|
||||
|
||||
// A key with NO caps must never be blocked by the slot quota check reaching it
|
||||
// through the scope path: scopeTokens on an uncapped key returns 0 usage and
|
||||
// the limit check must treat that as "no cap", not "exhausted".
|
||||
func TestUncappedKeyNeverBlockedByEmptyScope(t *testing.T) {
|
||||
up := upstream(t, &upstreamCtrl{})
|
||||
defer up.Close()
|
||||
g := newQuotaGW(t, up.URL,
|
||||
config.GWKey{Key: "sk-a", Role: "user", Models: []config.ModelScope{{Model: "AUTO"}}})
|
||||
for i := 0; i < 5; i++ {
|
||||
if rr, _ := chatAs(t, g, "sk-a", "AUTO"); rr.Code != 200 {
|
||||
t.Fatalf("call %d: want 200 (AUTO scope with no quota), got %d (%s)", i, rr.Code, rr.Body.String())
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// newQuotaGW builds a gateway with one mock upstream serving model m1 and the
|
||||
// given keys.
|
||||
func newQuotaGW(t *testing.T, upURL string, keys ...config.GWKey) *Gateway {
|
||||
t.Helper()
|
||||
td := t.TempDir()
|
||||
cfgPath := filepath.Join(td, "config.yaml")
|
||||
if err := os.WriteFile(cfgPath, []byte("listen: :0"), 0o644); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
cfg := &config.Config{
|
||||
Path: cfgPath,
|
||||
AdapterDir: filepath.Join(td, "adapters"),
|
||||
RuntimeFile: filepath.Join(td, "runtime.json"),
|
||||
Keys: keys,
|
||||
Sources: []config.Source{{
|
||||
Name: "up", BaseURL: upURL, Adapter: "openai",
|
||||
Models: []config.Model{
|
||||
{ID: "m1", Priority: 100},
|
||||
{ID: "m2", Priority: 90},
|
||||
},
|
||||
}},
|
||||
}
|
||||
if err := cfg.ApplyDefaults(); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
c, err := core.NewFromConfig(cfg)
|
||||
if err != nil {
|
||||
t.Fatalf("core: %v", err)
|
||||
}
|
||||
t.Cleanup(c.Close)
|
||||
secrets := make([]string, 0, len(keys))
|
||||
for _, k := range keys {
|
||||
secrets = append(secrets, k.Key)
|
||||
}
|
||||
g, err := New(c)
|
||||
if err != nil {
|
||||
t.Fatalf("gateway: %v", err)
|
||||
}
|
||||
return g
|
||||
}
|
||||
|
||||
// The core of per-model quotas: a model that runs out of budget must stop
|
||||
// being served on its own, while every other model on the SAME key keeps
|
||||
// working. A key-wide total would fail this — it would block m2 because m1 was
|
||||
// capped, which is exactly the coupling this design removes.
|
||||
func TestOneModelsQuotaDoesNotBlockAnother(t *testing.T) {
|
||||
up := upstream(t, &upstreamCtrl{})
|
||||
defer up.Close()
|
||||
td := t.TempDir()
|
||||
cfgPath := td + "/config.yaml"
|
||||
if err := os.WriteFile(cfgPath, []byte("listen: :0"), 0o644); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
cfg := &config.Config{
|
||||
Path: cfgPath, AdapterDir: filepath.Join(td, "adapters"),
|
||||
RuntimeFile: filepath.Join(td, "runtime.json"),
|
||||
Keys: []config.GWKey{{
|
||||
Key: "sk-a", Role: "user", Name: "two-models",
|
||||
Models: []config.ModelScope{
|
||||
{Model: "m1", TokenQuota: 4, Period: "hour"},
|
||||
{Model: "m2", TokenQuota: 1000, Period: "hour"},
|
||||
},
|
||||
}},
|
||||
Sources: []config.Source{{
|
||||
Name: "up", BaseURL: up.URL, Adapter: "openai",
|
||||
Models: []config.Model{
|
||||
{ID: "m1", Priority: 100},
|
||||
{ID: "m2", Priority: 90},
|
||||
},
|
||||
}},
|
||||
}
|
||||
if err := cfg.ApplyDefaults(); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
c, err := core.NewFromConfig(cfg)
|
||||
if err != nil {
|
||||
t.Fatalf("core: %v", err)
|
||||
}
|
||||
t.Cleanup(c.Close)
|
||||
g, err := New(c)
|
||||
if err != nil {
|
||||
t.Fatalf("gateway: %v", err)
|
||||
}
|
||||
|
||||
// Spend on m2 FIRST. Without this the two designs are
|
||||
// indistinguishable: a key-wide counter and m1's own counter would both
|
||||
// read 0 before m1 is used, so the test would pass either way (it did —
|
||||
// see the commit that rewrote it).
|
||||
for i := 1; i <= 3; i++ {
|
||||
if rr, _ := chatAs(t, g, "sk-a", "m2"); rr.Code != 200 {
|
||||
t.Fatalf("m2 priming call %d: want 200, got %d", i, rr.Code)
|
||||
}
|
||||
}
|
||||
// m1's budget is 4 and each call costs 4, so the key-wide total (m2+m1) is
|
||||
// already 12 when m1 starts: a key-wide cap would refuse m1 immediately.
|
||||
if rr, _ := chatAs(t, g, "sk-a", "m1"); rr.Code != 200 {
|
||||
t.Fatalf("m1 must be served on its own budget (m2's spend is not its problem), got %d (%s)",
|
||||
rr.Code, rr.Body.String())
|
||||
}
|
||||
if rr, code := chatAs(t, g, "sk-a", "m1"); rr.Code != http.StatusTooManyRequests {
|
||||
t.Fatalf("m1 second call: want 429, got %d (%s)", rr.Code, code)
|
||||
}
|
||||
// m2 must keep working
|
||||
for i := 1; i <= 3; i++ {
|
||||
if rr, _ := chatAs(t, g, "sk-a", "m2"); rr.Code != 200 {
|
||||
t.Fatalf("m2 call %d must be served while m1 is capped, got %d (%s)", i, rr.Code, rr.Body.String())
|
||||
}
|
||||
}
|
||||
// and the message must name m1, not the key
|
||||
rr, _ := chatAs(t, g, "sk-a", "m1")
|
||||
if !strings.Contains(rr.Body.String(), "m1") {
|
||||
t.Errorf("rejection should name the capped model: %s", rr.Body.String())
|
||||
}
|
||||
}
|
||||
|
||||
// A model with no quota on it is never blocked by a sibling's cap, and an
|
||||
// uncapped key is never blocked at all.
|
||||
func TestUncappedModelNeverBlocked(t *testing.T) {
|
||||
up := upstream(t, &upstreamCtrl{})
|
||||
defer up.Close()
|
||||
g := newQuotaGW(t, up.URL,
|
||||
config.GWKey{Key: "sk-a", Role: "user", Models: []config.ModelScope{
|
||||
{Model: "m1", TokenQuota: 1, Period: "hour"},
|
||||
{Model: "m2"},
|
||||
}})
|
||||
if rr, _ := chatAs(t, g, "sk-a", "m1"); rr.Code != 200 {
|
||||
t.Fatalf("m1 first: want 200, got %d", rr.Code)
|
||||
}
|
||||
if rr, _ := chatAs(t, g, "sk-a", "m1"); rr.Code != http.StatusTooManyRequests {
|
||||
t.Fatalf("m1 second: want 429, got %d", rr.Code)
|
||||
}
|
||||
for i := 1; i <= 4; i++ {
|
||||
if rr, _ := chatAs(t, g, "sk-a", "m2"); rr.Code != 200 {
|
||||
t.Fatalf("m2 (uncapped) call %d: want 200, got %d", i, rr.Code)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// An admin key is never capped even when its scopes carry budgets: a cap that
|
||||
// locked the operator out would be unrecoverable through the UI.
|
||||
func TestAdminKeyScopesAreNotEnforced(t *testing.T) {
|
||||
up := upstream(t, &upstreamCtrl{})
|
||||
defer up.Close()
|
||||
g := newQuotaGW(t, up.URL,
|
||||
config.GWKey{Key: "sk-admin", Role: "admin", Models: []config.ModelScope{
|
||||
{Model: "m1", TokenQuota: 1, ReqQuota: 1, Period: "hour"},
|
||||
}})
|
||||
for i := 1; i <= 4; i++ {
|
||||
if rr, _ := chatAs(t, g, "sk-admin", "m1"); rr.Code != 200 {
|
||||
t.Fatalf("admin call %d: want 200 (admin scopes are not enforced), got %d", i, rr.Code)
|
||||
}
|
||||
}
|
||||
}
|
||||
@ -7,6 +7,7 @@ import (
|
||||
"strings"
|
||||
|
||||
"llmsproxy/internal/config"
|
||||
"llmsproxy/internal/scheduler"
|
||||
)
|
||||
|
||||
// handleKeysAPI manages gateway keys: GET /api/keys (admin: all keys),
|
||||
@ -17,10 +18,22 @@ func (g *Gateway) handleKeysAPI(w http.ResponseWriter, r *http.Request) {
|
||||
path = strings.Trim(path, "/")
|
||||
role := reqRole(r.Context())
|
||||
|
||||
// /api/keys/me/auto — a user's own chain, readable without admin rights.
|
||||
// Must be matched before the generic {key}/auto case below so it does not
|
||||
// ask admin rights of a plain user asking about their own chain.
|
||||
if path == "me/auto" {
|
||||
g.handleKeyMeAuto(w, r)
|
||||
return
|
||||
}
|
||||
if path == "me" {
|
||||
g.handleKeyMe(w, r)
|
||||
return
|
||||
}
|
||||
// /api/keys/{key}/auto — the per-key AUTO chain editor.
|
||||
if i := strings.LastIndex(path, "/auto"); i > 0 && path[i+len("/auto"):] == "" {
|
||||
g.handleKeyAutoAPI(w, r, path[:i])
|
||||
return
|
||||
}
|
||||
if role != "admin" {
|
||||
writeError(w, http.StatusForbidden, "forbidden", "admin role required")
|
||||
return
|
||||
@ -48,13 +61,7 @@ func (g *Gateway) handleKeysAPI(w http.ResponseWriter, r *http.Request) {
|
||||
writeError(w, http.StatusBadRequest, "invalid_request", "invalid json: "+err.Error())
|
||||
return
|
||||
}
|
||||
if body.Role == "" {
|
||||
body.Role = "user"
|
||||
}
|
||||
if body.Role != "admin" && body.Role != "user" {
|
||||
writeError(w, http.StatusBadRequest, "invalid_request", "role must be admin or user")
|
||||
return
|
||||
}
|
||||
body.Role = config.NormalizeRole(body.Role)
|
||||
rec, err := g.core.CreateKey(body.Name, body.Role, body.Models, body.Note)
|
||||
if err != nil {
|
||||
writeError(w, http.StatusBadRequest, "key_error", err.Error())
|
||||
@ -76,6 +83,9 @@ func (g *Gateway) handleKeysAPI(w http.ResponseWriter, r *http.Request) {
|
||||
writeError(w, http.StatusBadRequest, "invalid_request", "invalid json: "+err.Error())
|
||||
return
|
||||
}
|
||||
// Each scope entry carries its own token/request caps, so replacing the
|
||||
// scope replaces the budgets with it — there is no separate key-wide
|
||||
// quota that could drift out of sync with the models.
|
||||
rec, err := g.core.UpdateKey(path, body.Name, body.Role, body.Models, body.Note)
|
||||
if err != nil {
|
||||
writeError(w, http.StatusBadRequest, "key_error", err.Error())
|
||||
@ -142,6 +152,101 @@ func (g *Gateway) allowedModels(ctx context.Context) []config.ModelScope {
|
||||
return rec.Models
|
||||
}
|
||||
|
||||
// handleKeyAutoAPI manages one key's own AUTO chain:
|
||||
// GET /api/keys/{key}/auto returns it, PUT replaces it, DELETE clears the
|
||||
// override so the key falls back to the global chain. Admin only — a user can
|
||||
// inspect their own chain through GET /api/keys/me/auto.
|
||||
//
|
||||
// An empty chain and a cleared chain are deliberately the same state: "inherit
|
||||
// the global chain". That keeps the WebUI's "use global" toggle a plain DELETE
|
||||
// with no sentinel value to carry through config.yaml.
|
||||
func (g *Gateway) handleKeyAutoAPI(w http.ResponseWriter, r *http.Request, key string) {
|
||||
if reqRole(r.Context()) != "admin" {
|
||||
writeError(w, http.StatusForbidden, "forbidden", "admin role required")
|
||||
return
|
||||
}
|
||||
rec, ok := g.core.FindKey(key)
|
||||
if !ok {
|
||||
writeError(w, http.StatusNotFound, "not_found", "key not found")
|
||||
return
|
||||
}
|
||||
switch r.Method {
|
||||
case http.MethodGet:
|
||||
writeJSON(w, http.StatusOK, map[string]interface{}{
|
||||
"key": key,
|
||||
"auto": rec.Auto,
|
||||
"inherits": !rec.HasOwnAuto(),
|
||||
})
|
||||
case http.MethodPut, http.MethodPatch:
|
||||
var body struct {
|
||||
Auto []config.ModelScope `json:"auto"`
|
||||
}
|
||||
if err := json.NewDecoder(r.Body).Decode(&body); err != nil {
|
||||
writeError(w, http.StatusBadRequest, "invalid_request", "invalid json: "+err.Error())
|
||||
return
|
||||
}
|
||||
if len(body.Auto) == 0 {
|
||||
// Treat an empty PUT as "give up the override" so the endpoint
|
||||
// cannot persist a chain that would 503 with no slots.
|
||||
if err := g.core.SaveKeyAuto(key, nil); err != nil {
|
||||
writeError(w, http.StatusBadRequest, "key_error", err.Error())
|
||||
return
|
||||
}
|
||||
writeJSON(w, http.StatusOK, map[string]interface{}{"ok": true, "inherits": true})
|
||||
return
|
||||
}
|
||||
if err := g.core.SaveKeyAuto(key, body.Auto); err != nil {
|
||||
writeError(w, http.StatusBadRequest, "key_error", err.Error())
|
||||
return
|
||||
}
|
||||
// Report whether the chain actually resolved to usable slots. An
|
||||
// admin who typed a model that no longer exists should learn it here,
|
||||
// not from the user's next 503.
|
||||
chain, _ := g.core.AutoChainFor(key)
|
||||
writeJSON(w, http.StatusOK, map[string]interface{}{
|
||||
"ok": true, "inherits": false, "slots": chainSlots(chain),
|
||||
})
|
||||
case http.MethodDelete:
|
||||
if err := g.core.SaveKeyAuto(key, nil); err != nil {
|
||||
writeError(w, http.StatusBadRequest, "key_error", err.Error())
|
||||
return
|
||||
}
|
||||
writeJSON(w, http.StatusOK, map[string]interface{}{"ok": true, "inherits": true})
|
||||
default:
|
||||
writeError(w, http.StatusMethodNotAllowed, "method_not_allowed", "")
|
||||
}
|
||||
}
|
||||
|
||||
// chainSlots counts the usable slots in a chain, for API feedback after an
|
||||
// edit. A chain that compiled to zero slots is reported so the caller can warn
|
||||
// instead of letting the next request fail with no_provider.
|
||||
func chainSlots(ch *scheduler.Chain) int {
|
||||
if ch == nil {
|
||||
return 0
|
||||
}
|
||||
n := 0
|
||||
for _, tn := range ch.Tiers {
|
||||
n += len(tn.Slots)
|
||||
}
|
||||
return n
|
||||
}
|
||||
|
||||
// handleKeyMeAuto lets a user read their own AUTO chain without admin rights.
|
||||
func (g *Gateway) handleKeyMeAuto(w http.ResponseWriter, r *http.Request) {
|
||||
if r.Method != http.MethodGet {
|
||||
writeError(w, http.StatusMethodNotAllowed, "method_not_allowed", "use GET")
|
||||
return
|
||||
}
|
||||
rec, ok := g.core.FindKey(reqKey(r.Context()))
|
||||
if !ok {
|
||||
writeError(w, http.StatusUnauthorized, "invalid_api_key", "key not found")
|
||||
return
|
||||
}
|
||||
writeJSON(w, http.StatusOK, map[string]interface{}{
|
||||
"key": rec.Key, "auto": rec.Auto, "inherits": !rec.HasOwnAuto(),
|
||||
})
|
||||
}
|
||||
|
||||
// handleAutoAPI manages the AUTO scheduling slots: GET /api/auto returns the
|
||||
// current rules; PUT /api/auto replaces them (admin only).
|
||||
func (g *Gateway) handleAutoAPI(w http.ResponseWriter, r *http.Request) {
|
||||
|
||||
167
internal/gateway/overflow_err_test.go
Normal file
167
internal/gateway/overflow_err_test.go
Normal file
@ -0,0 +1,167 @@
|
||||
package gateway
|
||||
|
||||
import (
|
||||
"errors"
|
||||
"regexp"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"llmsproxy/internal/scheduler"
|
||||
"llmsproxy/internal/types"
|
||||
)
|
||||
|
||||
// piOverflowRe 是 @earendil-works/pi-ai 的 OVERFLOW_PATTERNS 中对应的一条。
|
||||
// 归一化标记必须命中它,否则改了等于没改:客户端仍不会压缩重试。
|
||||
var piOverflowRe = regexp.MustCompile(`context[_ ]length[_ ]exceeded`)
|
||||
|
||||
func TestOverflowCanonicalMatchesPiPattern(t *testing.T) {
|
||||
if !piOverflowRe.MatchString(overflowCanonical) {
|
||||
t.Fatalf("归一化标记 %q 不匹配 pi 的超窗正则,客户端仍不会压缩重试", overflowCanonical)
|
||||
}
|
||||
}
|
||||
|
||||
// justworker 的真实措辞:与 pi 现有 25 条正则一条都不匹配,
|
||||
// 所以必须在网关侧归一化,否则超窗对客户端完全不可见。
|
||||
func TestJustworkerOverflowIsNormalized(t *testing.T) {
|
||||
raw := "请精简对话历史…(Context window is full, please reduce the conversation history)…"
|
||||
|
||||
// 前提校验:原始措辞确实不被 pi 认出(否则本用例失去意义)。
|
||||
if piOverflowRe.MatchString(raw) {
|
||||
t.Fatalf("前提不成立:原始措辞本就匹配 pi 正则: %q", raw)
|
||||
}
|
||||
if !looksLikeOverflow(raw) {
|
||||
t.Fatalf("网关应识别出这是超窗: %q", raw)
|
||||
}
|
||||
|
||||
err := &scheduler.ChainErr{Tiers: []scheduler.TierError{{
|
||||
Tier: 1, Source: "justworker", Model: "claude-opus-5", Err: errors.New(raw),
|
||||
}}}
|
||||
|
||||
got := clientUpstreamErr(err)
|
||||
if !piOverflowRe.MatchString(got) {
|
||||
t.Fatalf("归一化后仍不匹配 pi 正则,客户端不会压缩重试: %q", got)
|
||||
}
|
||||
}
|
||||
|
||||
// 直连路径(非 AUTO 链)也要归一化——它走的是同一个入口。
|
||||
func TestDirectPathOverflowIsNormalized(t *testing.T) {
|
||||
got := clientUpstreamErr(errors.New("upstream 400: Context window is full"))
|
||||
if !piOverflowRe.MatchString(got) {
|
||||
t.Fatalf("直连超窗未被归一化: %q", got)
|
||||
}
|
||||
}
|
||||
|
||||
// 非超窗错误不能被加上超窗标记,否则会把普通失败误报成需要压缩。
|
||||
func TestNonOverflowNotNormalized(t *testing.T) {
|
||||
err := &scheduler.ChainErr{Tiers: []scheduler.TierError{{
|
||||
Tier: 1, Source: "x", Model: "y", Err: errors.New("502 Bad Gateway: upstream connect error"),
|
||||
}}}
|
||||
got := clientUpstreamErr(err)
|
||||
if strings.Contains(got, overflowCanonical) {
|
||||
t.Fatalf("普通上游失败被误标为超窗: %q", got)
|
||||
}
|
||||
if !strings.Contains(got, "502 Bad Gateway") {
|
||||
t.Fatalf("普通失败的诊断信息应保留: %q", got)
|
||||
}
|
||||
}
|
||||
|
||||
// 回归:短诊断词常落在错误尾部,80 字节截断会正好把它切掉,
|
||||
// 超窗就永远认不出来。宽度放宽后必须保住。
|
||||
func TestOverflowPhraseSurvivesTruncationWidth(t *testing.T) {
|
||||
prefix := strings.Repeat("x", 100)
|
||||
// OneLine 的长度限制作用于整条 "source/model: err",不是只算 err 本身。
|
||||
full := "justworker/claude-opus-5: " + prefix + " Context window is full"
|
||||
|
||||
// 前提:80 字节下短语确实会被切掉(这正是修复前的行为)。
|
||||
if got := types.OneLine(full, 80); strings.Contains(got, "Context window is full") {
|
||||
t.Fatalf("前提不成立:80 字节下短语居然还在: %q", got)
|
||||
}
|
||||
// 160 字节下必须保住,否则超窗识别会被截断悄悄破坏。
|
||||
if got := types.OneLine(full, 160); !strings.Contains(got, "Context window is full") {
|
||||
t.Fatalf("160 字节应保住尾部短诊断词: %q", got)
|
||||
}
|
||||
|
||||
err := &scheduler.ChainErr{Tiers: []scheduler.TierError{{
|
||||
Tier: 1, Source: "justworker", Model: "claude-opus-5",
|
||||
Err: errors.New(prefix + " Context window is full"),
|
||||
}}}
|
||||
|
||||
got := clientUpstreamErr(err)
|
||||
if !strings.Contains(strings.ToLower(got), "context window is full") {
|
||||
t.Fatalf("宽度 160 应保住尾部短诊断词,实际: %q", got)
|
||||
}
|
||||
if !piOverflowRe.MatchString(got) {
|
||||
t.Fatalf("归一化标记丢失: %q", got)
|
||||
}
|
||||
}
|
||||
|
||||
// piNonOverflowRe 是 @earendil-works/pi-ai 的 NON_OVERFLOW_PATTERNS。
|
||||
// pi 的 isContextOverflow 先查这一组,命中则直接判为「非超窗」——
|
||||
// 即使 OVERFLOW_PATTERNS 也命中,也不会压缩重试。
|
||||
var piNonOverflowRe = []*regexp.Regexp{
|
||||
regexp.MustCompile(`^(Throttling error|Service unavailable):`),
|
||||
regexp.MustCompile(`(?i)rate limit`),
|
||||
regexp.MustCompile(`(?i)too many requests`),
|
||||
}
|
||||
|
||||
// piWouldCompact 复刻 isContextOverflow 的 Case 1 判定顺序:
|
||||
// 先排除 NON_OVERFLOW,再匹配 OVERFLOW。
|
||||
func piWouldCompact(errorMessage string) bool {
|
||||
for _, p := range piNonOverflowRe {
|
||||
if p.MatchString(errorMessage) {
|
||||
return false
|
||||
}
|
||||
}
|
||||
return piOverflowRe.MatchString(errorMessage)
|
||||
}
|
||||
|
||||
// issue 回归:AUTO 链失败消息天然把「超窗 tier」和「配额/限流 tier」拼在一起,
|
||||
// 超窗标记后面只要出现 rate limit / too many requests 的字样,
|
||||
// pi 就会把整条消息判为「非超窗」而拒绝压缩重试。
|
||||
func TestOverflowMessageSurvivesRateLimitedSiblingTier(t *testing.T) {
|
||||
cases := []struct {
|
||||
name string
|
||||
sibling string
|
||||
}{
|
||||
{"rate limit 措辞", "api error 429: rate limit exceeded"},
|
||||
{"too many requests 措辞", "429 Too Many Requests"},
|
||||
}
|
||||
|
||||
for _, tc := range cases {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
err := &scheduler.ChainErr{Tiers: []scheduler.TierError{
|
||||
{Tier: 1, Source: "sensenova", Model: "glm-5.2", Err: errors.New(tc.sibling)},
|
||||
{Tier: 3, Source: "gozen", Model: "deepseek-v4.1-flash",
|
||||
Err: errors.New("api error 400: This model's maximum context length is 1048576 tokens")},
|
||||
}}
|
||||
|
||||
got := clientUpstreamErr(err)
|
||||
if !piWouldCompact(got) {
|
||||
t.Fatalf("同链限流 tier 封杀了超窗识别,pi 不会压缩重试: %q", got)
|
||||
}
|
||||
// 干净:不得把限流 tier 的措辞带出去。
|
||||
for _, p := range piNonOverflowRe {
|
||||
if p.MatchString(got) {
|
||||
t.Fatalf("超窗消息仍带限流措辞 %v: %q", p, got)
|
||||
}
|
||||
}
|
||||
// 但仍要能看出是哪个源超的窗。
|
||||
if !strings.Contains(got, "gozen/deepseek-v4.1-flash") {
|
||||
t.Errorf("超窗源信息丢失: %q", got)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// 直连超窗同样要干净可用。
|
||||
func TestDirectOverflowMessageIsClean(t *testing.T) {
|
||||
raw := "api error 400: Error from provider (Console Go): This model's maximum context length is 1048576 tokens. " +
|
||||
"However, you requested 1440050 tokens. Please reduce the length of the messages or completion."
|
||||
got := clientUpstreamErr(errors.New(raw))
|
||||
if !piWouldCompact(got) {
|
||||
t.Fatalf("直连超窗未被 pi 识别: %q", got)
|
||||
}
|
||||
if !strings.Contains(got, overflowCanonical) {
|
||||
t.Fatalf("缺少归一化标记: %q", got)
|
||||
}
|
||||
}
|
||||
150
internal/gateway/plugin_call_api_test.go
Normal file
150
internal/gateway/plugin_call_api_test.go
Normal file
@ -0,0 +1,150 @@
|
||||
package gateway
|
||||
|
||||
// The /api/plugins/{name}/call/{fn} endpoint splits its path with string
|
||||
// surgery, and the split point is easy to get wrong in a way that still
|
||||
// "works": keeping the "call/" segment in the function name produced
|
||||
// `does not export "call/collect_account"` — the plugin was correct, the router
|
||||
// was not, and every call 502'd while the exports list looked perfectly fine.
|
||||
//
|
||||
// These tests drive the real router, so a refactor that reintroduces the
|
||||
// off-by-one is caught here rather than in production.
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"net/http"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"llmsproxy/internal/config"
|
||||
"llmsproxy/internal/core"
|
||||
)
|
||||
|
||||
// callPlugin exposes exactly one export plus a function that is NOT exported.
|
||||
// If the exports gate is not honoured, secretFn is reachable and with it the
|
||||
// plugin's private state — for a plugin holding an API key that means handing
|
||||
// the key to anyone who can POST to the gateway.
|
||||
const callPlugin = `
|
||||
local P = {}
|
||||
P.name = "calltest"
|
||||
P.version = "1.0.0"
|
||||
P.exports = { "add" }
|
||||
P.secret = "do-not-leak"
|
||||
function P.add(a)
|
||||
return { result = (tonumber(a and a.x) or 0) + 1 }
|
||||
end
|
||||
function P.secretFn()
|
||||
return { leaked = P.secret }
|
||||
end
|
||||
return P
|
||||
`
|
||||
|
||||
func gatewayWithCallPlugin(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)
|
||||
if err := c.Plugins().LoadSource("calltest", callPlugin); err != nil {
|
||||
t.Fatalf("load calltest: %v", err)
|
||||
}
|
||||
g, err := New(c)
|
||||
if err != nil {
|
||||
t.Fatalf("gateway: %v", err)
|
||||
}
|
||||
return g
|
||||
}
|
||||
|
||||
func TestPluginCallSplitsFunctionNameCorrectly(t *testing.T) {
|
||||
g := gatewayWithCallPlugin(t)
|
||||
rr := doReq(t, g, http.MethodPost,
|
||||
"/api/plugins/calltest/call/add", `{"x":41}`)
|
||||
if rr.Code != http.StatusOK {
|
||||
t.Fatalf("call failed: HTTP %d %s", rr.Code, rr.Body.String())
|
||||
}
|
||||
var out struct {
|
||||
Result struct {
|
||||
Result int `json:"result"`
|
||||
} `json:"result"`
|
||||
}
|
||||
if err := json.Unmarshal(rr.Body.Bytes(), &out); err != nil {
|
||||
t.Fatalf("decode: %v (body=%s)", err, rr.Body.String())
|
||||
}
|
||||
// 41 + 1 == 42. Reaching 42 proves BOTH halves of the contract: the "/call/"
|
||||
// segment did not leak into the function name, and the payload arrived.
|
||||
if out.Result.Result != 42 {
|
||||
t.Fatalf("got %d, want 42 — either the /call/ segment leaked into the "+
|
||||
"function name or the argument table never reached Lua", out.Result.Result)
|
||||
}
|
||||
}
|
||||
|
||||
func TestPluginCallRefusesUndeclaredFunction(t *testing.T) {
|
||||
g := gatewayWithCallPlugin(t)
|
||||
rr := doReq(t, g, http.MethodPost,
|
||||
"/api/plugins/calltest/call/secretFn", `{}`)
|
||||
if rr.Code == http.StatusOK {
|
||||
t.Fatalf("an unexported function must not be callable, got 200: %s",
|
||||
rr.Body.String())
|
||||
}
|
||||
if strings.Contains(rr.Body.String(), "do-not-leak") {
|
||||
t.Fatal("the refusal leaked the plugin's private state")
|
||||
}
|
||||
}
|
||||
|
||||
// A plugin that declares no exports at all has no callable surface. Reading the
|
||||
// missing list as "allow everything" would turn every plugin into a remote
|
||||
// evaluator the moment this endpoint ships.
|
||||
func TestPluginCallWithoutExportsHasNoCallableSurface(t *testing.T) {
|
||||
g := gatewayWithCallPlugin(t)
|
||||
const bare = `
|
||||
local P = {}
|
||||
P.name = "bare"
|
||||
P.version = "1.0.0"
|
||||
function P.anything()
|
||||
return { ok = true }
|
||||
end
|
||||
return P
|
||||
`
|
||||
if err := g.core.Plugins().LoadSource("bare", bare); err != nil {
|
||||
t.Fatalf("load bare: %v", err)
|
||||
}
|
||||
rr := doReq(t, g, http.MethodPost,
|
||||
"/api/plugins/bare/call/anything", `{}`)
|
||||
if rr.Code == http.StatusOK {
|
||||
t.Fatalf("a plugin with no plugin.exports must expose nothing, got 200: %s",
|
||||
rr.Body.String())
|
||||
}
|
||||
}
|
||||
|
||||
func TestPluginCallIsAdminOnly(t *testing.T) {
|
||||
g := gatewayWithCallPlugin(t)
|
||||
// A real user-role key, not a bare one from config: sk-test/sk-user in the
|
||||
// config both default to admin, so a test using them would pass the gate for
|
||||
// the wrong reason and prove nothing.
|
||||
rec, err := g.core.CreateKey("viewer", "user", nil, "")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
rr := doReqAs(t, g, rec.Key, http.MethodPost,
|
||||
"/api/plugins/calltest/call/add", `{"x":1}`)
|
||||
if rr.Code != http.StatusForbidden {
|
||||
t.Fatalf("a non-admin key reached the call endpoint: HTTP %d %s",
|
||||
rr.Code, rr.Body.String())
|
||||
}
|
||||
}
|
||||
207
internal/gateway/plugin_http_routes.go
Normal file
207
internal/gateway/plugin_http_routes.go
Normal file
@ -0,0 +1,207 @@
|
||||
package gateway
|
||||
|
||||
// Mount for plugin-registered HTTP endpoints.
|
||||
//
|
||||
// The gateway owns the listener, the authentication, the role check and the
|
||||
// response framing; the plugin owns the logic. That split is the point: a plugin
|
||||
// cannot open a socket, cannot bypass auth, and cannot choose a path outside
|
||||
// /api/plugins/http/.
|
||||
//
|
||||
// Why this exists rather than "plugins can call each other": some things are
|
||||
// worth querying from OUTSIDE the browser — a CI job asking whether the upstream
|
||||
// balance is above zero, a monitoring probe, another tool on the box. billing
|
||||
// never wanted this (it only receives hooks); a quota plugin does. Hard-coding a
|
||||
// route per plugin would mean a gateway change for every such plugin, and an
|
||||
// arbitrary plugin's code in the main mux is not something the routing table
|
||||
// should ever grow.
|
||||
|
||||
import (
|
||||
"net/http"
|
||||
"strconv"
|
||||
"strings"
|
||||
|
||||
"llmsproxy/internal/lua"
|
||||
)
|
||||
|
||||
// pluginHTTPPrefix is the mount point. It is deliberately narrow: a plugin's
|
||||
// routes live under one prefix, which makes them greppable, makes them easy to
|
||||
// revoke wholesale, and means a plugin cannot shadow /api/keys or /api/stats.
|
||||
const pluginHTTPPrefix = "/api/plugins/http/"
|
||||
|
||||
// handlePluginHTTP dispatches a request to whichever plugin registered the
|
||||
// longest matching path.
|
||||
//
|
||||
// Longest-match rather than first-match: a plugin may register /health and
|
||||
// another /health/deep, and both must work. With registration order deciding it,
|
||||
// the outcome would depend on load order and change when a plugin is added.
|
||||
func (g *Gateway) handlePluginHTTP(w http.ResponseWriter, r *http.Request) {
|
||||
rest := strings.TrimPrefix(r.URL.Path, pluginHTTPPrefix)
|
||||
ps := g.core.Plugins()
|
||||
if ps == nil {
|
||||
writeError(w, http.StatusNotFound, "not_found", "no plugins loaded")
|
||||
return
|
||||
}
|
||||
|
||||
matches := ps.HTTPHandlers()
|
||||
var chosen *lua.HTTPHandler
|
||||
var params map[string]string
|
||||
for i := range matches {
|
||||
h := matches[i]
|
||||
if p, ok := matchHandlerPath(h.Path, rest); ok {
|
||||
if chosen == nil || segmentCount(h.Path) > segmentCount(chosen.Path) {
|
||||
chosen, params = &h, p
|
||||
}
|
||||
}
|
||||
}
|
||||
if chosen == nil {
|
||||
writeError(w, http.StatusNotFound, "not_found",
|
||||
"no plugin serves this path")
|
||||
return
|
||||
}
|
||||
|
||||
if !strings.EqualFold(r.Method, chosen.Method) {
|
||||
w.Header().Set("Allow", chosen.Method)
|
||||
writeError(w, http.StatusMethodNotAllowed, "method_not_allowed",
|
||||
"this endpoint accepts "+chosen.Method+" only")
|
||||
return
|
||||
}
|
||||
|
||||
role := reqRole(r.Context())
|
||||
if !chosen.Allows(role, reqKeyName(r.Context())) {
|
||||
writeError(w, http.StatusForbidden, "forbidden",
|
||||
handlerAccessHint(*chosen))
|
||||
return
|
||||
}
|
||||
|
||||
body, err := readBody(r)
|
||||
if err != nil {
|
||||
writeError(w, http.StatusBadRequest, "invalid_request", err.Error())
|
||||
return
|
||||
}
|
||||
req := &lua.HandlerRequest{
|
||||
Method: r.Method,
|
||||
Path: rest,
|
||||
Query: flattenQuery(r.URL.Query()),
|
||||
Params: params,
|
||||
Headers: flattenHeader(r.Header),
|
||||
Body: body,
|
||||
Role: role,
|
||||
Key: reqKey(r.Context()),
|
||||
}
|
||||
resp, err := ps.InvokeHandler(handlerOwner(chosen), chosen.Handler, req)
|
||||
if err != nil {
|
||||
// A plugin error is a 502, not a 500: the gateway is fine, the plugin is
|
||||
// not, and the operator needs to tell those apart in their logs.
|
||||
writeError(w, http.StatusBadGateway, "plugin_error", err.Error())
|
||||
return
|
||||
}
|
||||
lua.WriteHandlerResponse(w, resp)
|
||||
}
|
||||
|
||||
// handlerAccessHint explains, in the refusal, exactly which declaration would
|
||||
// open this endpoint. A 403 with no reason is the kind of thing an operator
|
||||
// opens the source to work out, and then the fix is a guess.
|
||||
func handlerAccessHint(h lua.HTTPHandler) string {
|
||||
switch {
|
||||
case len(h.Keys) > 0:
|
||||
return "this endpoint allows only the named keys (" +
|
||||
strings.Join(h.Keys, ", ") + "); authenticate with one of them"
|
||||
case len(h.Roles) > 0:
|
||||
return "this endpoint requires role " + strings.Join(h.Roles, " or ") +
|
||||
" (widen it in the plugin's http declaration, or set public=true)"
|
||||
default:
|
||||
return "this endpoint declares no roles, so it is admin-only; " +
|
||||
"set roles=… or keys=… in the plugin's http declaration"
|
||||
}
|
||||
}
|
||||
|
||||
// handlerOwner returns the plugin name that registered h. The gateway stores the
|
||||
// owner alongside the handler precisely so this does not have to be guessed from
|
||||
// the path prefix.
|
||||
func handlerOwner(h *lua.HTTPHandler) string { return h.Plugin }
|
||||
|
||||
// segmentCount is how deep a route is, used for longest-match. A path with more
|
||||
// segments wins, so /health/deep beats /health when both are registered.
|
||||
func segmentCount(p string) int { return strings.Count(p, "/") }
|
||||
|
||||
// matchHandlerPath matches a route pattern against a concrete path, capturing
|
||||
// ":name" segments.
|
||||
//
|
||||
// Returns ok=false on a segment-count mismatch or a missing value for a
|
||||
// param — a route must not match with an empty capture, or two different URLs
|
||||
// would reach the same handler indistinguishably.
|
||||
func matchHandlerPath(pattern, path string) (map[string]string, bool) {
|
||||
pp := strings.Split(strings.Trim(pattern, "/"), "/")
|
||||
sp := strings.Split(strings.Trim(path, "/"), "/")
|
||||
if len(pp) != len(sp) {
|
||||
return nil, false
|
||||
}
|
||||
var params map[string]string
|
||||
for i := range pp {
|
||||
if strings.HasPrefix(pp[i], ":") {
|
||||
name := pp[i][1:]
|
||||
if sp[i] == "" {
|
||||
return nil, false
|
||||
}
|
||||
if params == nil {
|
||||
params = map[string]string{}
|
||||
}
|
||||
params[name] = sp[i]
|
||||
continue
|
||||
}
|
||||
if pp[i] != sp[i] {
|
||||
return nil, false
|
||||
}
|
||||
}
|
||||
return params, true
|
||||
}
|
||||
|
||||
func flattenQuery(q map[string][]string) map[string]string {
|
||||
out := make(map[string]string, len(q))
|
||||
for k, v := range q {
|
||||
if len(v) > 0 {
|
||||
out[k] = v[0]
|
||||
}
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
func flattenHeader(h http.Header) map[string]string {
|
||||
out := make(map[string]string, len(h))
|
||||
for k := range h {
|
||||
out[strings.ToLower(k)] = h.Get(k)
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
func readBody(r *http.Request) (string, error) {
|
||||
if r.Body == nil {
|
||||
return "", nil
|
||||
}
|
||||
defer r.Body.Close()
|
||||
// A cap, because the handler runs inside the gateway process: an unbounded
|
||||
// body from an authenticated caller is a memory exhaustion vector.
|
||||
const maxBody = 4 << 20
|
||||
buf := make([]byte, 0, 4096)
|
||||
tmp := make([]byte, 32*1024)
|
||||
for {
|
||||
n, err := r.Body.Read(tmp)
|
||||
if n > 0 {
|
||||
if len(buf)+n > maxBody {
|
||||
return "", &bodyTooLarge{}
|
||||
}
|
||||
buf = append(buf, tmp[:n]...)
|
||||
}
|
||||
if err != nil {
|
||||
break
|
||||
}
|
||||
}
|
||||
return string(buf), nil
|
||||
}
|
||||
|
||||
type bodyTooLarge struct{}
|
||||
|
||||
func (bodyTooLarge) Error() string { return "request body exceeds 4 MiB" }
|
||||
|
||||
// limitBody is kept for the numeric form some callers want in a log line.
|
||||
var _ = strconv.Itoa
|
||||
626
internal/gateway/plugin_wiring_test.go
Normal file
626
internal/gateway/plugin_wiring_test.go
Normal file
@ -0,0 +1,626 @@
|
||||
package gateway
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"llmsproxy/internal/config"
|
||||
"llmsproxy/internal/core"
|
||||
)
|
||||
|
||||
// newFailingServer serves handler and returns its test server.
|
||||
func newFailingServer(h http.Handler) *httptest.Server { return httptest.NewServer(h) }
|
||||
|
||||
// THE TEST THAT WAS MISSING.
|
||||
//
|
||||
// Every other plugin test calls Plugins.Fire() directly. That proves the Lua
|
||||
// runtime works, but it says nothing about whether the GATEWAY actually fires
|
||||
// the stages — the wiring in handleChat / singleChat / writeRec could be
|
||||
// deleted and the whole suite would stay green while no hook ever ran in
|
||||
// production.
|
||||
//
|
||||
// So these tests drive real HTTP requests through the real handler and assert
|
||||
// on what the plugin observed.
|
||||
|
||||
// observed is what a spy plugin recorded. It is written from Lua, so the
|
||||
// counters live in plugin state and are read back through the state API.
|
||||
type spyPayload struct {
|
||||
Stage string
|
||||
Type string
|
||||
Model string
|
||||
Source string
|
||||
OK bool
|
||||
Tokens int64
|
||||
}
|
||||
|
||||
// spyPlugin returns Lua source for a plugin that appends every stage it sees to
|
||||
// `state.seen`. It is generated per test so each test gets a clean log.
|
||||
func spyPlugin() string {
|
||||
return `
|
||||
local plugin = { name = "spy", version = "1.0.0" }
|
||||
plugin.state = { seen = {} }
|
||||
plugin.hooks = {
|
||||
request_start = "on_start",
|
||||
routed = "on_routed",
|
||||
request_end = "on_end",
|
||||
}
|
||||
local function rec(p, stage)
|
||||
table.insert(plugin.state.seen, {
|
||||
stage = stage, type = p.type or "", model = p.model or "",
|
||||
source = p.source or "", ok = p.ok and true or false,
|
||||
tokens = p.prompt_tokens or 0,
|
||||
})
|
||||
end
|
||||
function plugin.on_start(p) rec(p, "request_start") return nil end
|
||||
function plugin.on_routed(p) rec(p, "routed") return nil end
|
||||
function plugin.on_end(p) rec(p, "request_end") return nil end
|
||||
return plugin
|
||||
`
|
||||
}
|
||||
|
||||
// gatewayWithSpy boots a gateway with one mock upstream and the spy plugin
|
||||
// loaded, and returns the gateway plus the upstream URL.
|
||||
func gatewayWithSpy(t *testing.T) (*Gateway, string) {
|
||||
t.Helper()
|
||||
up := mockUpstream()
|
||||
t.Cleanup(up.Close)
|
||||
|
||||
dir := t.TempDir()
|
||||
cfgPath := filepath.Join(dir, "config.yaml")
|
||||
cfgBody := "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(cfgBody), 0600); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
cfg, err := config.Load(cfgPath)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
cfg.Sources = []config.Source{{
|
||||
Name: "spyup",
|
||||
BaseURL: up.URL,
|
||||
Adapter: "openai",
|
||||
APIKey: "sk-x",
|
||||
Models: []config.Model{{ID: "spy-model", Kind: "chat"}},
|
||||
}}
|
||||
if err := cfg.ApplyDefaults(); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
c, err := core.NewFromConfig(cfg)
|
||||
if err != nil {
|
||||
t.Fatalf("core: %v", err)
|
||||
}
|
||||
t.Cleanup(c.Close)
|
||||
if err := c.Plugins().LoadSource("spy", spyPlugin()); err != nil {
|
||||
t.Fatalf("load spy: %v", err)
|
||||
}
|
||||
g, err := New(c)
|
||||
if err != nil {
|
||||
t.Fatalf("gateway: %v", err)
|
||||
}
|
||||
return g, up.URL
|
||||
}
|
||||
|
||||
// seenStages reads the spy's log out of its published state.
|
||||
func seenStages(t *testing.T, g *Gateway) []map[string]interface{} {
|
||||
t.Helper()
|
||||
raw := g.core.Plugins().State("spy")
|
||||
if raw == nil {
|
||||
return nil
|
||||
}
|
||||
b, err := json.Marshal(raw)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
var st struct {
|
||||
Seen []map[string]interface{} `json:"seen"`
|
||||
}
|
||||
if err := json.Unmarshal(b, &st); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
return st.Seen
|
||||
}
|
||||
|
||||
func stageNames(seen []map[string]interface{}) []string {
|
||||
out := make([]string, 0, len(seen))
|
||||
for _, s := range seen {
|
||||
out = append(out, s["stage"].(string))
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
// TestHooksFireOnRealDirectChat: a plain non-streaming chat through the real
|
||||
// handler must produce request_start, routed and request_end — in that order,
|
||||
// with the real source and model filled in.
|
||||
func TestHooksFireOnRealDirectChat(t *testing.T) {
|
||||
g, _ := gatewayWithSpy(t)
|
||||
rr := doReq(t, g, http.MethodPost, "/v1/chat/completions",
|
||||
`{"model":"spy-model","messages":[{"role":"user","content":"hi"}]}`)
|
||||
if rr.Code != http.StatusOK {
|
||||
t.Fatalf("chat status=%d body=%s", rr.Code, rr.Body.String())
|
||||
}
|
||||
seen := seenStages(t, g)
|
||||
got := stageNames(seen)
|
||||
want := []string{"request_start", "routed", "request_end"}
|
||||
if len(got) != 3 {
|
||||
t.Fatalf("stages = %v, want %v (a hook that never fires is the failure "+
|
||||
"this test exists for)", got, want)
|
||||
}
|
||||
for i := range want {
|
||||
if got[i] != want[i] {
|
||||
t.Fatalf("stage order = %v, want %v", got, want)
|
||||
}
|
||||
}
|
||||
// The payload must be real data, not empty scaffolding.
|
||||
end := seen[2]
|
||||
if end["source"] != "spyup" {
|
||||
t.Errorf("request_end source = %v, want spyup", end["source"])
|
||||
}
|
||||
if end["model"] != "spy-model" {
|
||||
t.Errorf("request_end model = %v, want spy-model", end["model"])
|
||||
}
|
||||
if end["ok"] != true {
|
||||
t.Errorf("request_end ok = %v, want true", end["ok"])
|
||||
}
|
||||
if toInt(t, end["tokens"]) <= 0 {
|
||||
t.Errorf("request_end prompt_tokens = %v, want the mock upstream's real usage", end["tokens"])
|
||||
}
|
||||
}
|
||||
|
||||
// TestHooksFireOnRealStreamChat: streaming is a DIFFERENT code path (the record
|
||||
// is written from a deferred call after the stream ends), so it needs its own
|
||||
// assertion. A hook that only worked on the non-streaming path would be a real
|
||||
// bug for agent clients, which are the main streaming users.
|
||||
func TestHooksFireOnRealStreamChat(t *testing.T) {
|
||||
g, _ := gatewayWithSpy(t)
|
||||
rr := doReq(t, g, http.MethodPost, "/v1/chat/completions",
|
||||
`{"model":"spy-model","stream":true,"messages":[{"role":"user","content":"hi"}]}`)
|
||||
if rr.Code != http.StatusOK {
|
||||
t.Fatalf("stream status=%d body=%s", rr.Code, rr.Body.String())
|
||||
}
|
||||
got := stageNames(seenStages(t, g))
|
||||
if len(got) != 3 {
|
||||
t.Fatalf("stream stages = %v, want all three", got)
|
||||
}
|
||||
if got[0] != "request_start" || got[2] != "request_end" {
|
||||
t.Errorf("stream stage order = %v", got)
|
||||
}
|
||||
}
|
||||
|
||||
// TestHooksFireOnAutoRequest: the AUTO chain path is a third code path. It also
|
||||
// must fire, and routed must report tier -2 (resolved by the chain) rather than
|
||||
// the direct path's -1.
|
||||
func TestHooksFireOnAutoRequest(t *testing.T) {
|
||||
g, _ := gatewayWithSpy(t)
|
||||
// Register the model in the AUTO chain via the admin API.
|
||||
put := doReq(t, g, http.MethodPut, "/api/auto", `{"rules":[{"model":"spy-model"}]}`)
|
||||
if put.Code != http.StatusOK {
|
||||
t.Fatalf("save auto: %d %s", put.Code, put.Body.String())
|
||||
}
|
||||
rr := doReq(t, g, http.MethodPost, "/v1/chat/completions",
|
||||
`{"model":"AUTO","messages":[{"role":"user","content":"hi"}]}`)
|
||||
if rr.Code != http.StatusOK {
|
||||
t.Fatalf("auto chat status=%d body=%s", rr.Code, rr.Body.String())
|
||||
}
|
||||
seen := seenStages(t, g)
|
||||
if len(seen) != 3 {
|
||||
t.Fatalf("AUTO stages = %v, want three", stageNames(seen))
|
||||
}
|
||||
if seen[0]["model"] != "AUTO" {
|
||||
t.Errorf("request_start model = %v, want the requested \"AUTO\"", seen[0]["model"])
|
||||
}
|
||||
if seen[1]["model"] != "spy-model" {
|
||||
t.Errorf("routed model = %v, want the RESOLVED model", seen[1]["model"])
|
||||
}
|
||||
}
|
||||
|
||||
// TestHooksFireOnFailedRequest is the one that matters most for accounting: a
|
||||
// request that 500s upstream must still reach request_end, or a billing plugin
|
||||
// would never see failed traffic.
|
||||
func TestHooksFireOnFailedRequest(t *testing.T) {
|
||||
// Point at an upstream that always fails.
|
||||
bad := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
w.WriteHeader(http.StatusInternalServerError)
|
||||
_, _ = w.Write([]byte(`{"error":"boom"}`))
|
||||
})
|
||||
up := newFailingServer(bad)
|
||||
defer up.Close()
|
||||
|
||||
dir := t.TempDir()
|
||||
cfgPath := filepath.Join(dir, "config.yaml")
|
||||
cfgBody := "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"
|
||||
os.WriteFile(cfgPath, []byte(cfgBody), 0600)
|
||||
cfg, err := config.Load(cfgPath)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
cfg.Sources = []config.Source{{
|
||||
Name: "badup", BaseURL: up.URL, Adapter: "openai", APIKey: "sk-x",
|
||||
Models: []config.Model{{ID: "bad-model", Kind: "chat"}},
|
||||
}}
|
||||
if err := cfg.ApplyDefaults(); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
c, err := core.NewFromConfig(cfg)
|
||||
if err != nil {
|
||||
t.Fatalf("core: %v", err)
|
||||
}
|
||||
defer c.Close()
|
||||
if err := c.Plugins().LoadSource("spy", spyPlugin()); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
g, err := New(c)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
rr := doReq(t, g, http.MethodPost, "/v1/chat/completions",
|
||||
`{"model":"bad-model","messages":[{"role":"user","content":"hi"}]}`)
|
||||
if rr.Code == http.StatusOK {
|
||||
t.Fatal("expected the upstream failure to surface as a non-200")
|
||||
}
|
||||
seen := seenStages(t, g)
|
||||
// request_start must fire; routed must NOT (no source was ever selected);
|
||||
// request_end MUST (otherwise billing misses failed traffic).
|
||||
got := stageNames(seen)
|
||||
if len(got) == 0 || got[0] != "request_start" {
|
||||
t.Fatalf("stages = %v, want request_start first", got)
|
||||
}
|
||||
hasEnd := false
|
||||
for _, s := range seen {
|
||||
if s["stage"] == "request_end" {
|
||||
hasEnd = true
|
||||
if s["ok"] != false {
|
||||
t.Error("request_end reported ok=true for a failed request")
|
||||
}
|
||||
}
|
||||
}
|
||||
if !hasEnd {
|
||||
t.Error("a FAILED request never reached request_end; a billing plugin " +
|
||||
"would not see it at all")
|
||||
}
|
||||
if contains(got, "routed") {
|
||||
t.Error("routed fired for a request that never selected a source")
|
||||
}
|
||||
}
|
||||
|
||||
// TestBrokenPluginDoesNotBreakForwarding is the safety property, asserted
|
||||
// end-to-end: a plugin that throws in EVERY stage must not stop a chat from
|
||||
// being served.
|
||||
func TestBrokenPluginDoesNotBreakForwarding(t *testing.T) {
|
||||
up := mockUpstream()
|
||||
defer up.Close()
|
||||
dir := t.TempDir()
|
||||
cfgPath := filepath.Join(dir, "config.yaml")
|
||||
os.WriteFile(cfgPath, []byte("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"), 0600)
|
||||
cfg, err := config.Load(cfgPath)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
cfg.Sources = []config.Source{{
|
||||
Name: "u", BaseURL: up.URL, Adapter: "openai", APIKey: "sk-x",
|
||||
Models: []config.Model{{ID: "m", Kind: "chat"}},
|
||||
}}
|
||||
if err := cfg.ApplyDefaults(); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
c, err := core.NewFromConfig(cfg)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer c.Close()
|
||||
broken := `
|
||||
local plugin = { name = "broken", version = "1.0.0" }
|
||||
plugin.hooks = { request_start = "f", routed = "f", request_end = "f" }
|
||||
function plugin.f(p) error("plugin is broken") end
|
||||
return plugin
|
||||
`
|
||||
if err := c.Plugins().LoadSource("broken", broken); err != nil {
|
||||
t.Fatalf("load broken: %v", err)
|
||||
}
|
||||
g, err := New(c)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
rr := doReq(t, g, http.MethodPost, "/v1/chat/completions",
|
||||
`{"model":"m","messages":[{"role":"user","content":"hi"}]}`)
|
||||
if rr.Code != http.StatusOK {
|
||||
t.Fatalf("a broken plugin broke forwarding: status=%d body=%s", rr.Code, rr.Body.String())
|
||||
}
|
||||
if !strings.Contains(rr.Body.String(), "pong") {
|
||||
t.Errorf("the response is not the upstream's: %s", rr.Body.String())
|
||||
}
|
||||
// And the failure must be visible rather than silent.
|
||||
if len(c.Plugins().HookErrors()) == 0 {
|
||||
t.Error("a plugin that threw in every stage left no error record")
|
||||
}
|
||||
}
|
||||
|
||||
// TestPluginListIsServedWithStages is the operator-facing check: an admin
|
||||
// debugging a plugin needs to see which stages are live and whether the hook
|
||||
// file is where the kernel thinks it is.
|
||||
func TestPluginListIsServedWithStages(t *testing.T) {
|
||||
g, _ := gatewayWithSpy(t)
|
||||
rr := doReq(t, g, http.MethodGet, "/api/plugins", "")
|
||||
if rr.Code != http.StatusOK {
|
||||
t.Fatalf("status=%d", rr.Code)
|
||||
}
|
||||
for _, want := range []string{"spy", "request_start", "routed", "request_end", "plugin_dir"} {
|
||||
if !strings.Contains(rr.Body.String(), want) {
|
||||
t.Errorf("/api/plugins response lacks %q", want)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func toInt(t *testing.T, v interface{}) int64 {
|
||||
t.Helper()
|
||||
f, ok := v.(float64)
|
||||
if !ok {
|
||||
return 0
|
||||
}
|
||||
return int64(f)
|
||||
}
|
||||
|
||||
func contains(list []string, s string) bool {
|
||||
for _, x := range list {
|
||||
if x == s {
|
||||
return true
|
||||
}
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
// imageUpstream serves an OpenAI-shaped image response.
|
||||
func imageUpstream() *httptest.Server {
|
||||
return httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
_, _ = w.Write([]byte(`{"created":1,"data":[{"b64_json":"aGk="}]}`))
|
||||
}))
|
||||
}
|
||||
|
||||
// TestHooksFireOnRealImageRequest: the image path is a THIRD entry point. It
|
||||
// has its own handler and its own scheduler call, so "chat works" says nothing
|
||||
// about it — and image traffic is billable, so a cost plugin that never sees
|
||||
// it under-reports.
|
||||
func TestHooksFireOnRealImageRequest(t *testing.T) {
|
||||
up := imageUpstream()
|
||||
defer up.Close()
|
||||
|
||||
dir := t.TempDir()
|
||||
cfgPath := filepath.Join(dir, "config.yaml")
|
||||
os.WriteFile(cfgPath, []byte("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"), 0600)
|
||||
cfg, err := config.Load(cfgPath)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
cfg.Sources = []config.Source{{
|
||||
Name: "imgup", BaseURL: up.URL, Adapter: "openai", APIKey: "sk-x",
|
||||
Models: []config.Model{{ID: "img-model", Kind: "image"}},
|
||||
}}
|
||||
if err := cfg.ApplyDefaults(); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
c, err := core.NewFromConfig(cfg)
|
||||
if err != nil {
|
||||
t.Fatalf("core: %v", err)
|
||||
}
|
||||
defer c.Close()
|
||||
if err := c.Plugins().LoadSource("spy", spyPlugin()); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
g, err := New(c)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
rr := doReq(t, g, http.MethodPost, "/v1/images/generations",
|
||||
`{"model":"img-model","prompt":"a cat"}`)
|
||||
if rr.Code != http.StatusOK {
|
||||
t.Fatalf("image status=%d body=%s", rr.Code, rr.Body.String())
|
||||
}
|
||||
seen := seenStages(t, g)
|
||||
got := stageNames(seen)
|
||||
if len(got) != 3 {
|
||||
t.Fatalf("image stages = %v, want all three (image is a separate handler)", got)
|
||||
}
|
||||
if seen[0]["type"] != "image" {
|
||||
t.Errorf("request_start type = %v, want \"image\"", seen[0]["type"])
|
||||
}
|
||||
if seen[1]["source"] != "imgup" {
|
||||
t.Errorf("routed source = %v, want imgup", seen[1]["source"])
|
||||
}
|
||||
}
|
||||
|
||||
// TestRejectedChatStillFiresRequestStart pins the ordering decision: request_start
|
||||
// fires BEFORE the quota/scope gates, so a plugin auditing traffic sees rejected
|
||||
// requests too. A plugin that only ever saw served requests could not compute a
|
||||
// true request rate.
|
||||
func TestRejectedChatStillFiresRequestStart(t *testing.T) {
|
||||
up := mockUpstream()
|
||||
defer up.Close()
|
||||
dir := t.TempDir()
|
||||
cfgPath := filepath.Join(dir, "config.yaml")
|
||||
os.WriteFile(cfgPath, []byte("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"), 0600)
|
||||
cfg, err := config.Load(cfgPath)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
cfg.Sources = []config.Source{{
|
||||
Name: "u", BaseURL: up.URL, Adapter: "openai", APIKey: "sk-x",
|
||||
Models: []config.Model{{ID: "m", Kind: "chat"}},
|
||||
}}
|
||||
if err := cfg.ApplyDefaults(); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
c, err := core.NewFromConfig(cfg)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer c.Close()
|
||||
if err := c.Plugins().LoadSource("spy", spyPlugin()); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
g, err := New(c)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
// A model the gateway does not know => 404, rejected before any routing.
|
||||
rr := doReq(t, g, http.MethodPost, "/v1/chat/completions",
|
||||
`{"model":"no-such-model","messages":[{"role":"user","content":"hi"}]}`)
|
||||
if rr.Code != http.StatusNotFound {
|
||||
t.Fatalf("expected 404 for an unknown model, got %d", rr.Code)
|
||||
}
|
||||
got := stageNames(seenStages(t, g))
|
||||
if !contains(got, "request_start") {
|
||||
t.Errorf("stages = %v; a REJECTED request must still fire request_start so a "+
|
||||
"plugin can count real traffic, not just served traffic", got)
|
||||
}
|
||||
}
|
||||
|
||||
// TestChainStepReachesPluginOnDegradation is the end-to-end proof for the
|
||||
// AUTO trace: a request that had to drop from tier 1 to tier 2 must be visible
|
||||
// to a plugin as a tier_skip followed by a selected, and request_end must carry
|
||||
// tier_served=2.
|
||||
//
|
||||
// Before the trace existed the plugin saw only tier=-2 ("resolved by the
|
||||
// chain") and could not tell a degradation from a clean tier-1 hit — which is
|
||||
// the whole question a priority chain exists to answer.
|
||||
func TestChainStepReachesPluginOnDegradation(t *testing.T) {
|
||||
// tier 1's source always fails, so the walk must drop to tier 2.
|
||||
bad := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
w.WriteHeader(http.StatusInternalServerError)
|
||||
_, _ = w.Write([]byte(`{"error":"boom"}`))
|
||||
}))
|
||||
defer bad.Close()
|
||||
good := mockUpstream()
|
||||
defer good.Close()
|
||||
|
||||
dir := t.TempDir()
|
||||
cfgPath := filepath.Join(dir, "config.yaml")
|
||||
os.WriteFile(cfgPath, []byte("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"), 0600)
|
||||
cfg, err := config.Load(cfgPath)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
cfg.Sources = []config.Source{
|
||||
{Name: "t1", BaseURL: bad.URL, Adapter: "openai", APIKey: "sk-x",
|
||||
Models: []config.Model{{ID: "hi-tier", Kind: "chat"}}},
|
||||
{Name: "t2", BaseURL: good.URL, Adapter: "openai", APIKey: "sk-x",
|
||||
Models: []config.Model{{ID: "lo-tier", Kind: "chat"}}},
|
||||
}
|
||||
if err := cfg.ApplyDefaults(); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
c, err := core.NewFromConfig(cfg)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer c.Close()
|
||||
|
||||
// A spy that records chain_step events too.
|
||||
sp := `
|
||||
local plugin = { name = "walker", version = "1.0.0" }
|
||||
plugin.state = { steps = {}, ends = {} }
|
||||
plugin.hooks = { chain_step = "step", request_end = "fin" }
|
||||
function plugin.step(p)
|
||||
table.insert(plugin.state.steps, { kind = p.kind, tier = p.tier, source = p.source, model = p.model, reason = p.reason })
|
||||
return nil
|
||||
end
|
||||
function plugin.fin(p)
|
||||
plugin.state.ends[#plugin.state.ends + 1] = {
|
||||
tier_served = p.tier_served, degraded = p.degraded,
|
||||
walk = p.chain_walk, source = p.source, model = p.model,
|
||||
}
|
||||
return nil
|
||||
end
|
||||
return plugin
|
||||
`
|
||||
if err := c.Plugins().LoadSource("walker", sp); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
g, err := New(c)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
// Two tiers, both in the chain.
|
||||
if rr := doReq(t, g, http.MethodPut, "/api/auto",
|
||||
`{"rules":[{"model":"hi-tier","source":"t1","tier":1},{"model":"lo-tier","source":"t2","tier":2}]}`); rr.Code != 200 {
|
||||
t.Fatalf("save auto: %d %s", rr.Code, rr.Body.String())
|
||||
}
|
||||
|
||||
rr := doReq(t, g, http.MethodPost, "/v1/chat/completions",
|
||||
`{"model":"AUTO","messages":[{"role":"user","content":"hi"}]}`)
|
||||
if rr.Code != http.StatusOK {
|
||||
t.Fatalf("chat = %d %s", rr.Code, rr.Body.String())
|
||||
}
|
||||
|
||||
raw := c.Plugins().State("walker")
|
||||
b, _ := json.Marshal(raw)
|
||||
var st struct {
|
||||
Steps []struct {
|
||||
Kind string `json:"kind"`
|
||||
Tier int `json:"tier"`
|
||||
Source string `json:"source"`
|
||||
Model string `json:"model"`
|
||||
} `json:"steps"`
|
||||
Ends []struct {
|
||||
TierServed int `json:"tier_served"`
|
||||
Degraded bool `json:"degraded"`
|
||||
Source string `json:"source"`
|
||||
Model string `json:"model"`
|
||||
} `json:"ends"`
|
||||
}
|
||||
if err := json.Unmarshal(b, &st); err != nil {
|
||||
t.Fatalf("decode: %v (%s)", err, string(b))
|
||||
}
|
||||
if len(st.Steps) < 2 {
|
||||
t.Fatalf("chain_step events = %+v, want at least a slot_fail and a selected", st.Steps)
|
||||
}
|
||||
if st.Steps[0].Kind != "slot_fail" || st.Steps[0].Tier != 1 {
|
||||
t.Errorf("first step = %+v, want slot_fail on tier 1", st.Steps[0])
|
||||
}
|
||||
last := st.Steps[len(st.Steps)-1]
|
||||
if last.Kind != "selected" || last.Tier != 2 {
|
||||
t.Errorf("last step = %+v, want selected on tier 2", last)
|
||||
}
|
||||
if len(st.Ends) != 1 {
|
||||
t.Fatalf("request_end count = %d, want 1", len(st.Ends))
|
||||
}
|
||||
if st.Ends[0].TierServed != 2 {
|
||||
t.Errorf("tier_served = %d, want 2", st.Ends[0].TierServed)
|
||||
}
|
||||
if !st.Ends[0].Degraded {
|
||||
t.Error("degraded = false, but the request dropped a tier")
|
||||
}
|
||||
if st.Ends[0].Model != "lo-tier" {
|
||||
t.Errorf("served model = %q, want lo-tier", st.Ends[0].Model)
|
||||
}
|
||||
}
|
||||
619
internal/gateway/plugins_api.go
Normal file
619
internal/gateway/plugins_api.go
Normal file
@ -0,0 +1,619 @@
|
||||
package gateway
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"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
|
||||
}
|
||||
// AllStages, not a hand-written list: an earlier version enumerated the
|
||||
// three stages literally here, so when chain_step was added it was silently
|
||||
// missing from this response — a plugin author reading the discovery
|
||||
// document would have believed the stage did not exist. Deriving it from
|
||||
// the one place that defines the order is the whole point of having it.
|
||||
writeJSON(w, http.StatusOK, map[string]interface{}{
|
||||
"ui": ps.UI(),
|
||||
"stages": pipelineStageNames(),
|
||||
})
|
||||
}
|
||||
|
||||
// stageNames returns the pipeline stage names in firing order, for the
|
||||
// discovery payload and for tests.
|
||||
func pipelineStageNames() []string {
|
||||
out := make([]string, 0, len(lua.AllStages))
|
||||
for _, s := range lua.AllStages {
|
||||
out = append(out, string(s))
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
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
|
||||
}
|
||||
|
||||
// /api/plugins/{name}/http/{path}/access — the operator's view of, and edit
|
||||
// surface for, which keys an endpoint answers.
|
||||
if i := strings.Index(path, "/http/"); i >= 0 {
|
||||
name := path[:i]
|
||||
rest := path[i+len("/http/"):]
|
||||
if strings.HasSuffix(rest, "/access") {
|
||||
epPath := strings.TrimSuffix(rest, "/access")
|
||||
g.handleEndpointAccess(w, r, name, epPath)
|
||||
return
|
||||
}
|
||||
}
|
||||
|
||||
// /api/plugins/{name}/call/{fn}
|
||||
if i := strings.Index(path, "/call/"); i >= 0 {
|
||||
name := path[:i]
|
||||
fn := path[i+len("/call/"):]
|
||||
if name == "" || fn == "" {
|
||||
writeError(w, http.StatusBadRequest, "invalid_request", "call needs a plugin and a function name")
|
||||
return
|
||||
}
|
||||
g.handlePluginCall(w, r, name, fn)
|
||||
return
|
||||
}
|
||||
|
||||
// /api/plugins/{name}/fanout/{fn} — one call per argument set, in
|
||||
// parallel, from one request. See handlePluginFanout for why the
|
||||
// request-per-item form cannot be made concurrent from the client.
|
||||
if i := strings.Index(path, "/fanout/"); i >= 0 {
|
||||
name := path[:i]
|
||||
fn := path[i+len("/fanout/"):]
|
||||
if name == "" || fn == "" {
|
||||
writeError(w, http.StatusBadRequest, "invalid_request", "fanout needs a plugin and a function name")
|
||||
return
|
||||
}
|
||||
g.handlePluginFanout(w, r, name, fn)
|
||||
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(),
|
||||
"on_disk": ps.OnDisk(),
|
||||
"hook_errors": ps.HookErrors(),
|
||||
"plugin_dir": g.pluginDir(),
|
||||
// Visibility overrides and the endpoint access table, so the
|
||||
// plugin page renders what is IN FORCE rather than making the
|
||||
// operator read .lua files or config.yaml to find out.
|
||||
"visibility": ps.VisibilityOverrides(),
|
||||
"http_access": ps.EndpointAccess(),
|
||||
})
|
||||
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
|
||||
}
|
||||
|
||||
// /api/plugins/{name} — source read (GET), enable/disable (PUT), delete.
|
||||
// Multi-segment paths (…/state) are handled earlier and never reach here.
|
||||
if !strings.Contains(path, "/") {
|
||||
switch r.Method {
|
||||
case http.MethodGet:
|
||||
// Reading the SOURCE (not the runtime state) is what an editor
|
||||
// needs; the state endpoint is /state and returns accumulated data
|
||||
// instead. Mixing the two would make a "save what I read"
|
||||
// round-trip impossible.
|
||||
code, err := g.readPluginSource(path)
|
||||
if err != nil {
|
||||
writeError(w, http.StatusNotFound, "not_found", err.Error())
|
||||
return
|
||||
}
|
||||
writeJSON(w, http.StatusOK, map[string]interface{}{
|
||||
"name": path, "code": code,
|
||||
"enabled": g.core.Plugins().Enabled(path),
|
||||
})
|
||||
case http.MethodPut:
|
||||
// Enable / disable. The intent rides in the body rather than the
|
||||
// verb, because "enable" and "disable" are the same resource and
|
||||
// PUT /{name} is the one the docs advertise.
|
||||
var body struct {
|
||||
Enabled *bool `json:"enabled"`
|
||||
}
|
||||
if err := json.NewDecoder(r.Body).Decode(&body); err != nil {
|
||||
writeError(w, http.StatusBadRequest, "invalid_request", "invalid json: "+err.Error())
|
||||
return
|
||||
}
|
||||
if body.Enabled == nil {
|
||||
writeError(w, http.StatusBadRequest, "invalid_request", `body must be {"enabled":true|false}`)
|
||||
return
|
||||
}
|
||||
if err := g.core.Plugins().SetEnabled(path, *body.Enabled); err != nil {
|
||||
writeError(w, http.StatusBadRequest, "plugin_error", err.Error())
|
||||
return
|
||||
}
|
||||
writeJSON(w, http.StatusOK, map[string]interface{}{
|
||||
"ok": true, "name": path, "enabled": *body.Enabled,
|
||||
})
|
||||
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})
|
||||
default:
|
||||
writeError(w, http.StatusMethodNotAllowed, "method_not_allowed", "")
|
||||
}
|
||||
return
|
||||
}
|
||||
|
||||
writeError(w, http.StatusNotFound, "not_found", "unknown plugin sub-resource: "+path)
|
||||
}
|
||||
|
||||
// readPluginSource returns a plugin's file contents for the editor.
|
||||
func (g *Gateway) readPluginSource(name string) (string, error) {
|
||||
if err := validPluginName(name); err != nil {
|
||||
return "", err
|
||||
}
|
||||
if g.pluginDir() == "" {
|
||||
return "", errNoPluginDir
|
||||
}
|
||||
b, err := os.ReadFile(filepath.Join(g.pluginDir(), name+".lua"))
|
||||
if err != nil {
|
||||
return "", fmt.Errorf("no plugin source named %q", name)
|
||||
}
|
||||
return string(b), nil
|
||||
}
|
||||
|
||||
// validPluginName rejects anything that could escape the plugin directory.
|
||||
// Shared by install / remove / read so the check cannot drift between them.
|
||||
func validPluginName(name string) error {
|
||||
if name == "" || strings.ContainsAny(name, `/\.`) {
|
||||
return errPluginName
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// 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:
|
||||
// The plugin's audience gates its DATA as well as its page. Hiding the UI
|
||||
// while leaving GET /state open would put the numbers one curl away, which
|
||||
// is the leak the audience field exists to close — and billing's state is
|
||||
// a global ledger, not the caller's own usage.
|
||||
//
|
||||
// A plugin that declares no audience is unchanged (every role may read
|
||||
// it), so a single-purpose widget keeps working without edits.
|
||||
role := reqRole(r.Context())
|
||||
if !ps.StateVisibleTo(name, role) {
|
||||
writeError(w, http.StatusForbidden, "forbidden",
|
||||
"plugin "+name+" is not visible to role "+role)
|
||||
return
|
||||
}
|
||||
// prices rides along: the billing rules editor must show the table
|
||||
// that is actually in effect, not a reconstruction from config. It is
|
||||
// absent for plugins with no prices (every plugin but billing).
|
||||
state := ps.State(name)
|
||||
// A plugin may SPLIT its state: global buckets stay deployment-wide while
|
||||
// per-caller buckets are narrowed to the authenticated key. Role
|
||||
// visibility alone cannot express this — the same role is allowed to read
|
||||
// the plugin, just not another caller's row.
|
||||
if scope := ps.StateScopeFor(name); scope != nil && role != "admin" {
|
||||
state = splitStateForKey(state, scope, reqKey(r.Context()))
|
||||
}
|
||||
payload := map[string]interface{}{
|
||||
"plugin": name,
|
||||
"state": state,
|
||||
}
|
||||
if pr := ps.Prices(name); pr != nil {
|
||||
payload["prices"] = pr
|
||||
}
|
||||
writeJSON(w, http.StatusOK, payload)
|
||||
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", "")
|
||||
}
|
||||
}
|
||||
|
||||
// Call invokes a named function on a plugin's Lua state and returns its result.
|
||||
//
|
||||
// This exists because plugin pages otherwise have NO way to reach their own
|
||||
// plugin code: pluginAPI offers state read/write and a couple of callbacks, and
|
||||
// that is enough for billing (whose numbers arrive on hooks during request
|
||||
// forwarding) but useless for a plugin that must do work on demand. A quota
|
||||
// page has to ask the vendor when the operator clicks Refresh, and no hook
|
||||
// fires at that moment.
|
||||
//
|
||||
// It is deliberately narrow:
|
||||
// - The function must be one the plugin explicitly lists in plugin.exports.
|
||||
// Anything else is refused, so this endpoint cannot become a generic "run
|
||||
// arbitrary Lua with my credentials" oracle over a plugin's private state.
|
||||
// - The caller must be admin, and the result is JSON — no Lua value is ever
|
||||
// handed back to Go as a live closure.
|
||||
// - Errors come back as (nil, message) from the plugin, matching every other
|
||||
// plugin entry point, and surface as a 502 rather than raising.
|
||||
func (g *Gateway) handlePluginCall(w http.ResponseWriter, r *http.Request, name, fn string) {
|
||||
if reqRole(r.Context()) != "admin" {
|
||||
writeError(w, http.StatusForbidden, "forbidden", "admin role required")
|
||||
return
|
||||
}
|
||||
if r.Method != http.MethodPost {
|
||||
writeError(w, http.StatusMethodNotAllowed, "method_not_allowed", "")
|
||||
return
|
||||
}
|
||||
var args map[string]interface{}
|
||||
if err := json.NewDecoder(r.Body).Decode(&args); err != nil {
|
||||
writeError(w, http.StatusBadRequest, "invalid_request", "invalid json: "+err.Error())
|
||||
return
|
||||
}
|
||||
result, err := g.core.Plugins().Call(name, fn, args)
|
||||
if err != nil {
|
||||
writeError(w, http.StatusBadGateway, "plugin_error", err.Error())
|
||||
return
|
||||
}
|
||||
writeJSON(w, http.StatusOK, map[string]interface{}{"result": result})
|
||||
}
|
||||
|
||||
// handlePluginFanout runs one plugin export per element of a list, in parallel,
|
||||
// from a SINGLE request.
|
||||
//
|
||||
// Why this exists: the quota page has twenty cards, and it used to fire twenty
|
||||
// browser requests at /call/. They arrived concurrently and then serialised
|
||||
// inside Plugins.Call, which holds the plugin mutex across the Lua function —
|
||||
// including its network wait. Measured on the live gateway: 20 probes took 29s of
|
||||
// wall clock while a single probe with no contention took 1.4–3.3s, so the page
|
||||
// sat on "probing…" for half a minute. The browser-side concurrency was real and
|
||||
// made no difference.
|
||||
//
|
||||
// A fan-out cannot be assembled from the client either: any request-per-item
|
||||
// arrangement is exactly the shape that queues. The parallelism has to happen
|
||||
// where the fan-out is received.
|
||||
func (g *Gateway) handlePluginFanout(w http.ResponseWriter, r *http.Request, name, fn string) {
|
||||
if reqRole(r.Context()) != "admin" {
|
||||
writeError(w, http.StatusForbidden, "forbidden", "admin role required")
|
||||
return
|
||||
}
|
||||
if r.Method != http.MethodPost {
|
||||
writeError(w, http.StatusMethodNotAllowed, "method_not_allowed", "")
|
||||
return
|
||||
}
|
||||
var req struct {
|
||||
// Args is a LIST of argument objects, one per call. A bare object is
|
||||
// rejected rather than coerced: "args" meaning "one call" and "args"
|
||||
// meaning "many" depending on its shape is how a caller ends up probing
|
||||
// one source and believing it probed twenty.
|
||||
Args []map[string]interface{} `json:"args"`
|
||||
}
|
||||
if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
|
||||
writeError(w, http.StatusBadRequest, "invalid_request", "invalid json: "+err.Error())
|
||||
return
|
||||
}
|
||||
if len(req.Args) == 0 {
|
||||
writeError(w, http.StatusBadRequest, "invalid_request", "args must be a non-empty list of argument objects")
|
||||
return
|
||||
}
|
||||
// The page legitimately has one card per configured source. The bound is not
|
||||
// arbitrary: it stops one request from asking for thousands of VMs, and it is
|
||||
// enforced before any work starts rather than discovered shard by shard.
|
||||
if len(req.Args) > fanoutMaxArgs {
|
||||
writeError(w, http.StatusBadRequest, "invalid_request",
|
||||
fmt.Sprintf("args has %d entries, limit is %d", len(req.Args), fanoutMaxArgs))
|
||||
return
|
||||
}
|
||||
results, err := g.core.Plugins().ShardCall(name, fn, req.Args)
|
||||
if err != nil {
|
||||
writeError(w, http.StatusBadGateway, "plugin_error", err.Error())
|
||||
return
|
||||
}
|
||||
// A shard that failed contributes a null in its slot. That is deliberate and
|
||||
// is NOT the same as the whole request failing: one unreachable upstream must
|
||||
// not turn a page of twenty cards into a single error.
|
||||
writeJSON(w, http.StatusOK, map[string]interface{}{"results": results})
|
||||
}
|
||||
|
||||
// fanoutMaxArgs bounds one fan-out request. Twenty is the measured real source
|
||||
// count on the gateway this was built for; the headroom is for operators with
|
||||
// more sources configured.
|
||||
const fanoutMaxArgs = 64
|
||||
|
||||
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 err := validPluginName(name); err != nil {
|
||||
return err
|
||||
}
|
||||
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 err := validPluginName(name); err != nil {
|
||||
return err
|
||||
}
|
||||
if g.pluginDir() == "" {
|
||||
return errNoPluginDir
|
||||
}
|
||||
_ = os.Remove(filepath.Join(g.pluginDir(), name+".lua"))
|
||||
return g.core.Plugins().Unload(name)
|
||||
}
|
||||
|
||||
// handleEndpointAccess serves the operator's view and edit surface for one
|
||||
// plugin endpoint's reachable keys.
|
||||
//
|
||||
// GET …/http/{path}/access the declared rules plus any override
|
||||
// PUT …/http/{path}/access body {"access": "admin,@homeagent"}
|
||||
// DELETE …/http/{path}/access clear the override
|
||||
//
|
||||
// The spec string is a comma list: a bare word is a role, "@name" is a gateway
|
||||
// key's configured name, and the literal "public" opens it to any
|
||||
// authenticated key. That is the same grammar the visibility override uses, and
|
||||
// keeping one grammar means an operator learns it once.
|
||||
//
|
||||
// Key names are validated against the configured keys BEFORE being stored. An
|
||||
// allowlist naming a key that does not exist is a silent no-op — every caller
|
||||
// is refused — and the operator would have no way to tell that apart from the
|
||||
// endpoint simply working.
|
||||
func (g *Gateway) handleEndpointAccess(w http.ResponseWriter, r *http.Request, name, epPath string) {
|
||||
if reqRole(r.Context()) != "admin" {
|
||||
writeError(w, http.StatusForbidden, "forbidden", "admin role required")
|
||||
return
|
||||
}
|
||||
ps := g.core.Plugins()
|
||||
|
||||
// epPath "*" is the WHOLE-PLUGIN visibility override, not an endpoint. It
|
||||
// shares this route (and therefore this store) on purpose: an operator
|
||||
// setting "who can see this plugin" and "who can call this endpoint" is
|
||||
// asking the same question at two granularities, and two stores would let
|
||||
// them disagree. SetEndpointAccess cannot serve it — it validates that the
|
||||
// endpoint exists, which "*" by definition does not.
|
||||
if epPath == pluginWholeOverride {
|
||||
g.handlePluginVisibility(w, r, name)
|
||||
return
|
||||
}
|
||||
|
||||
switch r.Method {
|
||||
case http.MethodGet:
|
||||
for _, e := range ps.EndpointAccess() {
|
||||
if e.Plugin == name && e.Path == epPath {
|
||||
writeJSON(w, http.StatusOK, map[string]interface{}{"endpoint": e})
|
||||
return
|
||||
}
|
||||
}
|
||||
writeError(w, http.StatusNotFound, "not_found", "no such endpoint")
|
||||
case http.MethodPut, http.MethodDelete:
|
||||
spec := ""
|
||||
if r.Method == http.MethodPut {
|
||||
var body struct {
|
||||
Access string `json:"access"`
|
||||
}
|
||||
if err := json.NewDecoder(r.Body).Decode(&body); err != nil {
|
||||
writeError(w, http.StatusBadRequest, "invalid_request", "invalid json: "+err.Error())
|
||||
return
|
||||
}
|
||||
spec = body.Access
|
||||
}
|
||||
if err := ps.ValidateAccessSpec(spec); err != nil {
|
||||
writeError(w, http.StatusBadRequest, "invalid_request", err.Error())
|
||||
return
|
||||
}
|
||||
if err := ps.SetEndpointAccess(name, epPath, spec); err != nil {
|
||||
writeError(w, http.StatusBadRequest, "invalid_request", err.Error())
|
||||
return
|
||||
}
|
||||
writeJSON(w, http.StatusOK, map[string]interface{}{"ok": true, "access": spec})
|
||||
default:
|
||||
writeError(w, http.StatusMethodNotAllowed, "method_not_allowed", "")
|
||||
}
|
||||
}
|
||||
|
||||
// splitStateForKey rewrites a plugin state so that a non-admin caller sees only
|
||||
// the buckets that belong to them.
|
||||
//
|
||||
// The plugin declares which buckets are global and which are per-key. A per-key
|
||||
// bucket is keyed by the MASKED key ("***a1b2c3" — the same value the audit log
|
||||
// uses), so the caller's own row is found by masking the presented secret and
|
||||
// matching the last six characters.
|
||||
//
|
||||
// Keys that do not match are DROPPED, not zeroed: a zeroed row would render as
|
||||
// "you spent nothing" and, worse, would leak that a row exists for someone else.
|
||||
func splitStateForKey(state interface{}, scope *lua.StateScope, key string) interface{} {
|
||||
m, ok := state.(map[string]interface{})
|
||||
if !ok || scope == nil || key == "" {
|
||||
return state
|
||||
}
|
||||
masked := maskKey(key)
|
||||
out := make(map[string]interface{}, len(m))
|
||||
for k, v := range m {
|
||||
if !scope.IsGlobal(k) {
|
||||
// Per-key bucket: keep only the caller's rows, at every nesting level
|
||||
// the scope mentions (by_day_key is {day: {masked: bucket}}).
|
||||
out[k] = filterKeyBuckets(v, masked, scope.DepthOf(k))
|
||||
continue
|
||||
}
|
||||
out[k] = v
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
// filterKeyBuckets narrows a keyed structure to one masked key.
|
||||
//
|
||||
// depth 1 is {masked: bucket}; depth 2 is {day: {masked: bucket}}. The level is
|
||||
// declared rather than sniffed, because both shapes are dicts and sniffing would
|
||||
// confuse an empty bucket with a nested one.
|
||||
func filterKeyBuckets(v interface{}, masked string, depth int) interface{} {
|
||||
switch depth {
|
||||
case 0:
|
||||
return v // not a per-key bucket; leave alone
|
||||
case 1:
|
||||
m, ok := v.(map[string]interface{})
|
||||
if !ok {
|
||||
return v
|
||||
}
|
||||
if row, ok := m[masked]; ok {
|
||||
return map[string]interface{}{masked: row}
|
||||
}
|
||||
return map[string]interface{}{}
|
||||
default:
|
||||
m, ok := v.(map[string]interface{})
|
||||
if !ok {
|
||||
return v
|
||||
}
|
||||
out := make(map[string]interface{}, len(m))
|
||||
for day, inner := range m {
|
||||
im, ok := inner.(map[string]interface{})
|
||||
if !ok {
|
||||
continue
|
||||
}
|
||||
if row, ok := im[masked]; ok {
|
||||
out[day] = map[string]interface{}{masked: row}
|
||||
}
|
||||
}
|
||||
return out
|
||||
}
|
||||
}
|
||||
|
||||
// pluginWholeOverride is the epPath that addresses a plugin's overall visibility
|
||||
// rather than one of its endpoints.
|
||||
const pluginWholeOverride = "*"
|
||||
|
||||
// handlePluginVisibility is the whole-plugin half of the access API. The
|
||||
// audience string uses the same vocabulary as an endpoint spec (roles, @key
|
||||
// names, public) so the UI can drive both from one input.
|
||||
func (g *Gateway) handlePluginVisibility(w http.ResponseWriter, r *http.Request, name string) {
|
||||
ps := g.core.Plugins()
|
||||
switch r.Method {
|
||||
case http.MethodGet:
|
||||
writeJSON(w, http.StatusOK, map[string]interface{}{
|
||||
"plugin": name,
|
||||
"visibility": ps.VisibilityOverrides()[name],
|
||||
})
|
||||
case http.MethodPut, http.MethodDelete:
|
||||
spec := ""
|
||||
if r.Method == http.MethodPut {
|
||||
var body struct {
|
||||
Access string `json:"access"`
|
||||
}
|
||||
if err := json.NewDecoder(r.Body).Decode(&body); err != nil {
|
||||
writeError(w, http.StatusBadRequest, "invalid_request", "invalid json: "+err.Error())
|
||||
return
|
||||
}
|
||||
spec = body.Access
|
||||
}
|
||||
if err := ps.ValidateAccessSpec(spec); err != nil {
|
||||
writeError(w, http.StatusBadRequest, "invalid_request", err.Error())
|
||||
return
|
||||
}
|
||||
ps.SetPluginVisibility(name, spec)
|
||||
writeJSON(w, http.StatusOK, map[string]interface{}{"ok": true, "visibility": spec})
|
||||
default:
|
||||
writeError(w, http.StatusMethodNotAllowed, "method_not_allowed", r.Method+" is not allowed here")
|
||||
}
|
||||
}
|
||||
263
internal/gateway/plugins_api_test.go
Normal file
263
internal/gateway/plugins_api_test.go
Normal file
@ -0,0 +1,263 @@
|
||||
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 {
|
||||
// Decoded into the real types so the test cannot drift from the wire
|
||||
// contract. An inline copy missed the `pages` field when the payload
|
||||
// shape changed and failed to compile, which is at least loud — but the
|
||||
// same copy also went on asserting the OLD single-page shape for a
|
||||
// release, silently.
|
||||
UI lua.UIExtension `json:"ui"`
|
||||
Stages []string `json:"stages"`
|
||||
}
|
||||
if err := json.Unmarshal(rr.Body.Bytes(), &view); err != nil {
|
||||
t.Fatalf("decode: %v", err)
|
||||
}
|
||||
// The payload carries every page in one list; ui.page is no longer a
|
||||
// separate slot (it was, and a second plugin contributing a page overwrote
|
||||
// whatever was there).
|
||||
var billing *lua.UIPage
|
||||
for i, pg := range view.UI.Pages {
|
||||
if pg != nil && pg.PageID == "billing" {
|
||||
billing = view.UI.Pages[i]
|
||||
}
|
||||
}
|
||||
if billing == nil {
|
||||
t.Fatalf("no billing page in the inject payload; pages=%d", len(view.UI.Pages))
|
||||
}
|
||||
if !strings.Contains(billing.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)
|
||||
}
|
||||
}
|
||||
// Derived from AllStages, not hardcoded: the previous assertion of "3"
|
||||
// is exactly what let chain_step go missing from this payload unnoticed.
|
||||
if len(view.Stages) != len(lua.AllStages) {
|
||||
t.Errorf("stages = %v, want %d (one per AllStages entry)", view.Stages, len(lua.AllStages))
|
||||
}
|
||||
for i, st := range lua.AllStages {
|
||||
if i >= len(view.Stages) || view.Stages[i] != string(st) {
|
||||
t.Errorf("stages[%d] = %v, want %q", i, view.Stages, string(st))
|
||||
}
|
||||
}
|
||||
if !containsStr(view.Stages, string(lua.StageChainStep)) {
|
||||
t.Error("the discovery payload does not advertise chain_step; a plugin " +
|
||||
"author would conclude the stage does not exist")
|
||||
}
|
||||
}
|
||||
|
||||
// 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")
|
||||
}
|
||||
}
|
||||
|
||||
// containsStr reports whether list has s.
|
||||
func containsStr(list []string, s string) bool {
|
||||
for _, x := range list {
|
||||
if x == s {
|
||||
return true
|
||||
}
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
// 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())
|
||||
}
|
||||
}
|
||||
|
||||
// TestPluginsListCarriesAccessTables: the plugin page renders visibility
|
||||
// overrides and the endpoint access table from /api/plugins. Without them in the
|
||||
// response the operator has to open a .lua or config.yaml to learn who can see
|
||||
// what, which is exactly what the UI is meant to avoid.
|
||||
func TestPluginsListCarriesAccessTables(t *testing.T) {
|
||||
g := gatewayWithCallPlugin(t)
|
||||
rr := doReq(t, g, http.MethodGet, "/api/plugins", "")
|
||||
if rr.Code != http.StatusOK {
|
||||
t.Fatalf("GET /api/plugins = %d", rr.Code)
|
||||
}
|
||||
var out struct {
|
||||
Visibility map[string]string `json:"visibility"`
|
||||
HTTPAccess []map[string]interface{} `json:"http_access"`
|
||||
}
|
||||
if err := json.Unmarshal(rr.Body.Bytes(), &out); err != nil {
|
||||
t.Fatalf("decode: %v (body=%s)", err, rr.Body.String())
|
||||
}
|
||||
// The calltest plugin registers no endpoints, so the table is present but
|
||||
// empty — the KEY is that the field exists at all. A missing field decodes to
|
||||
// nil and the UI silently renders nothing, which looks identical to "this
|
||||
// plugin has no endpoints".
|
||||
if !strings.Contains(rr.Body.String(), "http_access") {
|
||||
t.Error("/api/plugins response lacks the http_access field")
|
||||
}
|
||||
if !strings.Contains(rr.Body.String(), "visibility") {
|
||||
t.Error("/api/plugins response lacks the visibility field")
|
||||
}
|
||||
}
|
||||
|
||||
// TestPluginWholeVisibilityViaAccessAPI covers the "*" path: the plugin page
|
||||
// edits a plugin's overall audience through the same route it edits endpoints
|
||||
// with. Before this existed the UI had a visibility selector whose save 404'd,
|
||||
// because SetEndpointAccess validates that the endpoint exists and "*" does not.
|
||||
func TestPluginWholeVisibilityViaAccessAPI(t *testing.T) {
|
||||
g := gatewayWithCallPlugin(t)
|
||||
rr := doReq(t, g, http.MethodPut, "/api/plugins/calltest/http/*/access",
|
||||
`{"access":"admin, user"}`)
|
||||
if rr.Code != http.StatusOK {
|
||||
t.Fatalf("PUT */access = %d: %s", rr.Code, rr.Body.String())
|
||||
}
|
||||
if !strings.Contains(rr.Body.String(), `"ok":true`) {
|
||||
t.Errorf("unexpected body: %s", rr.Body.String())
|
||||
}
|
||||
// A user role is a legitimate audience; a nonsense one is not.
|
||||
bad := doReq(t, g, http.MethodPut, "/api/plugins/calltest/http/*/access",
|
||||
`{"access":"../../etc"}`)
|
||||
if bad.Code != http.StatusBadRequest {
|
||||
t.Errorf("traversal-shaped audience = %d, want 400: %s", bad.Code, bad.Body.String())
|
||||
}
|
||||
}
|
||||
@ -115,7 +115,7 @@ func (g *Gateway) loginRecord(ip string, ok bool, now int64) {
|
||||
}
|
||||
}
|
||||
|
||||
func New(c *core.Core, gatewayKeys []string) (*Gateway, error) {
|
||||
func New(c *core.Core) (*Gateway, error) {
|
||||
sub, err := fs.Sub(uiFS, "ui")
|
||||
if err != nil {
|
||||
return nil, err
|
||||
@ -222,10 +222,30 @@ func (g *Gateway) routes(w http.ResponseWriter, r *http.Request) {
|
||||
g.handleModels(w, r)
|
||||
case r.URL.Path == "/api/adapters" || strings.HasPrefix(r.URL.Path, "/api/adapters/"):
|
||||
g.handleAdaptersAPI(w, r)
|
||||
case r.URL.Path == "/api/v1" || strings.HasPrefix(r.URL.Path, "/api/v1/"):
|
||||
g.apiV1Routes(w, r)
|
||||
// Must precede the /api/sources prefix match: a generic HasPrefix on
|
||||
// "/api/sources/" would otherwise swallow "/api/sources/discover" and the
|
||||
// operator would get "source name required in path" instead of a catalog.
|
||||
case r.URL.Path == "/api/sources/discover" || strings.HasPrefix(r.URL.Path, "/api/sources/discover/"):
|
||||
g.handleSourceDiscoverAPI(w, r)
|
||||
case r.URL.Path == "/api/sources" || strings.HasPrefix(r.URL.Path, "/api/sources/"):
|
||||
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/billing/rules":
|
||||
// Checked BEFORE the generic /api/plugins/ prefix: the plugin router
|
||||
// treats anything after the name as a plugin sub-resource and would
|
||||
// 404 on "rules" instead of reaching the billing editor.
|
||||
g.handleBillingRules(w, r)
|
||||
case strings.HasPrefix(r.URL.Path, pluginHTTPPrefix):
|
||||
// Before the generic plugin router: these paths belong to a PLUGIN, not to
|
||||
// the plugin manager, and "http" is a valid sub-resource name otherwise.
|
||||
g.handlePluginHTTP(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":
|
||||
@ -294,20 +314,43 @@ func (g *Gateway) auth(next http.Handler) http.Handler {
|
||||
http.Redirect(w, r, "/login?continue="+url.QueryEscape(r.URL.Path), http.StatusFound)
|
||||
return
|
||||
}
|
||||
next.ServeHTTP(w, r.WithContext(withAuth(r.Context(), key, rec.Role)))
|
||||
next.ServeHTTP(w, r.WithContext(withAuthNamed(r.Context(), key, rec.Role, rec.Name)))
|
||||
})
|
||||
}
|
||||
|
||||
// authCtx carries the authenticated gateway key and its role.
|
||||
// authCtx carries the authenticated gateway key, its role, and the key's
|
||||
// configured NAME.
|
||||
//
|
||||
// The name is here because plugin HTTP endpoints authorise by key name rather
|
||||
// than by secret, and looking the name up from the key at the point of use would
|
||||
// mean reading config on every request. It is "" for a key configured without
|
||||
// one, which such an endpoint can never match.
|
||||
type authCtx struct {
|
||||
key string
|
||||
role string
|
||||
name string
|
||||
}
|
||||
|
||||
type authCtxKeyT struct{}
|
||||
|
||||
func withAuth(ctx context.Context, key, role string) context.Context {
|
||||
return context.WithValue(ctx, authCtxKeyT{}, authCtx{key: key, role: role})
|
||||
return withAuthNamed(ctx, key, role, "")
|
||||
}
|
||||
|
||||
// withAuthNamed is withAuth plus the key's configured name.
|
||||
func withAuthNamed(ctx context.Context, key, role, name string) context.Context {
|
||||
return context.WithValue(ctx, authCtxKeyT{}, authCtx{key: key, role: role, name: name})
|
||||
}
|
||||
|
||||
// reqKeyName returns the authenticated key's configured name, or "" if it has
|
||||
// none. Plugin endpoint authorisation is written against this, not against the
|
||||
// secret: a secret is rotated without notice, so a rule keyed on it would
|
||||
// silently stop matching.
|
||||
func reqKeyName(ctx context.Context) string {
|
||||
if a, ok := ctx.Value(authCtxKeyT{}).(authCtx); ok {
|
||||
return a.name
|
||||
}
|
||||
return ""
|
||||
}
|
||||
|
||||
// reqKey returns the authenticated gateway key id (masked suffix for display).
|
||||
|
||||
495
internal/gateway/source_fields_test.go
Normal file
495
internal/gateway/source_fields_test.go
Normal file
@ -0,0 +1,495 @@
|
||||
package gateway
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"reflect"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"llmsproxy/internal/config"
|
||||
"llmsproxy/internal/core"
|
||||
)
|
||||
|
||||
// This file covers the second round of the same class of defect bf0657b fixed.
|
||||
//
|
||||
// bf0657b stopped a partial source edit from clobbering api_key. The upsert
|
||||
// still rewrites the WHOLE source, so every other field the request could not
|
||||
// express was reset to its zero value on any edit: proxy_url, api_key_env,
|
||||
// timeout and queue_timeout.
|
||||
//
|
||||
// The four are worse than cosmetic because they change how the source dials out.
|
||||
// proxy_url silently turns a proxied upstream into a direct one (or vice versa)
|
||||
// and api_key_env turns a source whose credential lives in an environment
|
||||
// variable into a credential-less source — which keeps answering 200 on write
|
||||
// and only fails on the NEXT call upstream, long after the editor left.
|
||||
//
|
||||
// The trigger is not exotic scripting: the WebUI's own edit form posts exactly
|
||||
// such a payload. See TestWebUIEditPayloadPreservesRoutingFields.
|
||||
|
||||
// richGateway writes a config.yaml whose source carries all four fragile
|
||||
// fields, then boots a gateway over it.
|
||||
func richGateway(t *testing.T) *Gateway {
|
||||
t.Helper()
|
||||
dir := t.TempDir()
|
||||
cfgPath := filepath.Join(dir, "config.yaml")
|
||||
cfg := `listen: 127.0.0.1:8080
|
||||
gateway_keys:
|
||||
- sk-test
|
||||
default_model: AUTO
|
||||
adapter_dir: ` + filepath.Join(dir, "adapters") + `
|
||||
runtime_file: ` + filepath.Join(dir, "runtime.json") + `
|
||||
sources:
|
||||
- name: rich
|
||||
base_url: https://api.deepseek.com
|
||||
api_key: sk-live
|
||||
adapter: deepseek
|
||||
proxy_url: http://127.0.0.1:7890
|
||||
timeout: 300s
|
||||
queue_timeout: 90s
|
||||
max_concurrent: 8
|
||||
models:
|
||||
- id: deepseek-v4-flash
|
||||
priority: 100
|
||||
kind: chat
|
||||
`
|
||||
if err := os.WriteFile(cfgPath, []byte(cfg), 0600); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
c, err := config.Load(cfgPath)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
core, err := core.NewFromConfig(c)
|
||||
if err != nil {
|
||||
t.Fatalf("core: %v", err)
|
||||
}
|
||||
t.Cleanup(core.Close)
|
||||
g, err := New(core)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
return g
|
||||
}
|
||||
|
||||
// sourceField reads one field off the live source, i.e. what the routes
|
||||
// actually use rather than what the file happens to contain.
|
||||
func sourceField(t *testing.T, g *Gateway, name string, get func(config.Source) interface{}) interface{} {
|
||||
t.Helper()
|
||||
for _, s := range g.core.Sources() {
|
||||
if s.Name == name {
|
||||
return get(s)
|
||||
}
|
||||
}
|
||||
t.Fatalf("source %q not found", name)
|
||||
return nil
|
||||
}
|
||||
|
||||
// webUIPayload is byte-for-byte the JSON that index.html's saveSource() builds
|
||||
// from the edit dialog. It deliberately carries none of the four fragile
|
||||
// fields, because the dialog has no inputs for them.
|
||||
const webUIPayload = `{
|
||||
"name":"rich","base_url":"https://api.deepseek.com","api_key":"sk-live",
|
||||
"adapter":"deepseek","endpoint":"","image_endpoint":"",
|
||||
"models":[{"id":"deepseek-v4-flash","priority":100,"kind":"chat"}],
|
||||
"headers":{},"meta":{},"temperature":0,"max_tokens":0,"max_concurrent":16,"rpm":0
|
||||
}`
|
||||
|
||||
// TestWebUIEditPayloadPreservesRoutingFields is the regression test for the
|
||||
// reported defect: changing max_concurrent from 8 to 16 in the UI must not
|
||||
// silently drop the proxy, the env-var credential, or the timeouts.
|
||||
func TestWebUIEditPayloadPreservesRoutingFields(t *testing.T) {
|
||||
g := richGateway(t)
|
||||
|
||||
// The one field the operator actually changed.
|
||||
rr := doReq(t, g, "POST", "/api/sources", webUIPayload)
|
||||
if rr.Code != 200 {
|
||||
t.Fatalf("edit status=%d body=%s", rr.Code, rr.Body.String())
|
||||
}
|
||||
|
||||
checks := []struct {
|
||||
field string
|
||||
got interface{}
|
||||
want interface{}
|
||||
}{
|
||||
{"proxy_url", sourceField(t, g, "rich", func(s config.Source) interface{} { return s.ProxyURL }),
|
||||
"http://127.0.0.1:7890"},
|
||||
{"timeout", sourceField(t, g, "rich", func(s config.Source) interface{} { return s.Timeout }),
|
||||
300 * time.Second},
|
||||
{"queue_timeout", sourceField(t, g, "rich", func(s config.Source) interface{} { return s.QueueTimeout }),
|
||||
90 * time.Second},
|
||||
// The field the edit DID change, so a fix that simply refuses to write
|
||||
// anything would also pass the three assertions above.
|
||||
{"max_concurrent", sourceField(t, g, "rich", func(s config.Source) interface{} { return s.MaxConcurrent }),
|
||||
16},
|
||||
}
|
||||
for _, c := range checks {
|
||||
if c.got != c.want {
|
||||
t.Errorf("%s = %v, want %v", c.field, c.got, c.want)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// TestSourceEditPreservesAPIKeyEnv is called out separately because it is the
|
||||
// only one of the four whose loss is a credential failure rather than a routing
|
||||
// difference, and it is invisible until the next upstream call.
|
||||
func TestSourceEditPreservesAPIKeyEnv(t *testing.T) {
|
||||
dir := t.TempDir()
|
||||
cfgPath := filepath.Join(dir, "config.yaml")
|
||||
cfg := `listen: 127.0.0.1:8080
|
||||
gateway_keys:
|
||||
- sk-test
|
||||
default_model: AUTO
|
||||
adapter_dir: ` + filepath.Join(dir, "adapters") + `
|
||||
runtime_file: ` + filepath.Join(dir, "runtime.json") + `
|
||||
sources:
|
||||
- name: envonly
|
||||
base_url: https://api.deepseek.com
|
||||
api_key_env: DEEPSEEK_KEY
|
||||
adapter: deepseek
|
||||
max_concurrent: 8
|
||||
models:
|
||||
- id: deepseek-v4-flash
|
||||
priority: 100
|
||||
kind: chat
|
||||
`
|
||||
if err := os.WriteFile(cfgPath, []byte(cfg), 0600); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
c, err := config.Load(cfgPath)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
core, err := core.NewFromConfig(c)
|
||||
if err != nil {
|
||||
t.Fatalf("core: %v", err)
|
||||
}
|
||||
t.Cleanup(core.Close)
|
||||
g, err := New(core)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
// A credential-less source is legitimate, so the API must not invent one
|
||||
// here; the point is that the ENV reference survives the edit.
|
||||
rr := doReq(t, g, "POST", "/api/sources", `{
|
||||
"name":"envonly","base_url":"https://api.deepseek.com","api_key":"",
|
||||
"adapter":"deepseek","models":[{"id":"deepseek-v4-flash","priority":100,"kind":"chat"}],
|
||||
"headers":{},"meta":{},"max_concurrent":4}`)
|
||||
if rr.Code != 200 {
|
||||
t.Fatalf("edit status=%d body=%s", rr.Code, rr.Body.String())
|
||||
}
|
||||
got := sourceField(t, g, "envonly", func(s config.Source) interface{} { return s.APIKeyEnv })
|
||||
if got != "DEEPSEEK_KEY" {
|
||||
t.Fatalf("api_key_env = %v after an edit that did not mention it; "+
|
||||
"the source now has no credential and will only fail on the next upstream call", got)
|
||||
}
|
||||
if n := sourceField(t, g, "envonly", func(s config.Source) interface{} { return s.MaxConcurrent }); n != 4 {
|
||||
t.Fatalf("max_concurrent = %v, want the edited 4", n)
|
||||
}
|
||||
}
|
||||
|
||||
// TestSourceEditPersistsInheritedFieldsToDisk: inheriting in memory is not
|
||||
// enough — the values must reach config.yaml, or a restart silently undoes the
|
||||
// fix and the field is lost anyway.
|
||||
func TestSourceEditPersistsInheritedFieldsToDisk(t *testing.T) {
|
||||
g := richGateway(t)
|
||||
if rr := doReq(t, g, "POST", "/api/sources", webUIPayload); rr.Code != 200 {
|
||||
t.Fatalf("edit status=%d: %s", rr.Code, rr.Body.String())
|
||||
}
|
||||
raw, err := os.ReadFile(g.core.Config().Path)
|
||||
if err != nil {
|
||||
t.Fatalf("read config: %v", err)
|
||||
}
|
||||
text := string(raw)
|
||||
for _, want := range []string{
|
||||
"proxy_url: http://127.0.0.1:7890",
|
||||
} {
|
||||
if !strings.Contains(text, want) {
|
||||
t.Errorf("config.yaml lost %q on an unrelated edit:\n%s", want, text)
|
||||
}
|
||||
}
|
||||
// The durations are re-serialized in Go's canonical form (300s becomes
|
||||
// 5m0s), so they must be checked SEMANTICALLY. Asserting the original
|
||||
// spelling would be a false red: the value survived, only its text changed.
|
||||
// Re-loading the file is also the only way to prove the persisted value
|
||||
// parses back to the same duration.
|
||||
reloaded, err := config.Load(g.core.Config().Path)
|
||||
if err != nil {
|
||||
t.Fatalf("re-read config: %v", err)
|
||||
}
|
||||
if len(reloaded.Sources) == 0 {
|
||||
t.Fatal("no sources after reload")
|
||||
}
|
||||
s := reloaded.Sources[0]
|
||||
if s.Timeout != 300*time.Second {
|
||||
t.Errorf("persisted timeout = %v, want 300s", s.Timeout)
|
||||
}
|
||||
if s.QueueTimeout != 90*time.Second {
|
||||
t.Errorf("persisted queue_timeout = %v, want 90s", s.QueueTimeout)
|
||||
}
|
||||
if s.ProxyURL != "http://127.0.0.1:7890" {
|
||||
t.Errorf("persisted proxy_url = %v, want it preserved", s.ProxyURL)
|
||||
}
|
||||
}
|
||||
|
||||
// ---- explicit set / clear, i.e. the pointer semantics actually work ----
|
||||
|
||||
func TestSourceOptionalFieldsCanBeSet(t *testing.T) {
|
||||
g := richGateway(t)
|
||||
rr := doReq(t, g, "POST", "/api/sources", `{
|
||||
"name":"rich","base_url":"https://api.deepseek.com","api_key":"sk-live",
|
||||
"adapter":"deepseek","models":[{"id":"deepseek-v4-flash","priority":100,"kind":"chat"}],
|
||||
"headers":{},"meta":{},
|
||||
"proxy_url":"http://127.0.0.1:1080",
|
||||
"api_key_env":"NEW_ENV_KEY",
|
||||
"timeout":"45s",
|
||||
"queue_timeout":"5s"}`)
|
||||
if rr.Code != 200 {
|
||||
t.Fatalf("status=%d body=%s", rr.Code, rr.Body.String())
|
||||
}
|
||||
for _, c := range []struct {
|
||||
field string
|
||||
got interface{}
|
||||
want interface{}
|
||||
}{
|
||||
{"proxy_url", sourceField(t, g, "rich", func(s config.Source) interface{} { return s.ProxyURL }), "http://127.0.0.1:1080"},
|
||||
{"api_key_env", sourceField(t, g, "rich", func(s config.Source) interface{} { return s.APIKeyEnv }), "NEW_ENV_KEY"},
|
||||
{"timeout", sourceField(t, g, "rich", func(s config.Source) interface{} { return s.Timeout }), 45 * time.Second},
|
||||
{"queue_timeout", sourceField(t, g, "rich", func(s config.Source) interface{} { return s.QueueTimeout }), 5 * time.Second},
|
||||
} {
|
||||
if c.got != c.want {
|
||||
t.Errorf("%s = %v, want %v", c.field, c.got, c.want)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// TestSourceOptionalFieldsCanBeCleared is the half that an "empty means
|
||||
// inherit" implementation gets wrong: it can never clear a field, so emptying
|
||||
// the proxy box in the UI would keep using the old proxy forever.
|
||||
func TestSourceOptionalFieldsCanBeCleared(t *testing.T) {
|
||||
g := richGateway(t)
|
||||
rr := doReq(t, g, "POST", "/api/sources", `{
|
||||
"name":"rich","base_url":"https://api.deepseek.com","api_key":"sk-live",
|
||||
"adapter":"deepseek","models":[{"id":"deepseek-v4-flash","priority":100,"kind":"chat"}],
|
||||
"headers":{},"meta":{},
|
||||
"proxy_url":"","timeout":"","queue_timeout":""}`)
|
||||
if rr.Code != 200 {
|
||||
t.Fatalf("status=%d body=%s", rr.Code, rr.Body.String())
|
||||
}
|
||||
if got := sourceField(t, g, "rich", func(s config.Source) interface{} { return s.ProxyURL }); got != "" {
|
||||
t.Errorf("proxy_url = %v, want it cleared", got)
|
||||
}
|
||||
// A cleared timeout falls back to the config default (120s), which is what
|
||||
// mergedSources applies for a zero value — assert the effective behaviour
|
||||
// rather than the raw zero, since that is what a caller actually gets.
|
||||
if got := sourceField(t, g, "rich", func(s config.Source) interface{} { return s.Timeout }); got != config.DefaultSourceTimeout {
|
||||
t.Errorf("timeout = %v, want the default %v after clearing", got, config.DefaultSourceTimeout)
|
||||
}
|
||||
// api_key_env was NOT in the request, so it must be untouched even though
|
||||
// its sibling fields were cleared.
|
||||
_ = rr
|
||||
}
|
||||
|
||||
// TestSourceOptionalFieldsClearIsScoped: clearing the proxy must not clear the
|
||||
// env credential that the same request did not mention.
|
||||
func TestSourceOptionalFieldsClearIsScoped(t *testing.T) {
|
||||
dir := t.TempDir()
|
||||
cfgPath := filepath.Join(dir, "config.yaml")
|
||||
cfg := `listen: 127.0.0.1:8080
|
||||
gateway_keys:
|
||||
- sk-test
|
||||
default_model: AUTO
|
||||
adapter_dir: ` + filepath.Join(dir, "adapters") + `
|
||||
runtime_file: ` + filepath.Join(dir, "runtime.json") + `
|
||||
sources:
|
||||
- name: both
|
||||
base_url: https://api.deepseek.com
|
||||
api_key_env: DEEPSEEK_KEY
|
||||
adapter: deepseek
|
||||
proxy_url: http://127.0.0.1:7890
|
||||
max_concurrent: 8
|
||||
models:
|
||||
- id: deepseek-v4-flash
|
||||
priority: 100
|
||||
kind: chat
|
||||
`
|
||||
os.WriteFile(cfgPath, []byte(cfg), 0600)
|
||||
c, err := config.Load(cfgPath)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
core, err := core.NewFromConfig(c)
|
||||
if err != nil {
|
||||
t.Fatalf("core: %v", err)
|
||||
}
|
||||
t.Cleanup(core.Close)
|
||||
g, err := New(core)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
rr := doReq(t, g, "POST", "/api/sources", `{
|
||||
"name":"both","base_url":"https://api.deepseek.com","api_key":"",
|
||||
"adapter":"deepseek","models":[{"id":"deepseek-v4-flash","priority":100,"kind":"chat"}],
|
||||
"headers":{},"meta":{},"proxy_url":""}`)
|
||||
if rr.Code != 200 {
|
||||
t.Fatalf("status=%d body=%s", rr.Code, rr.Body.String())
|
||||
}
|
||||
if got := sourceField(t, g, "both", func(s config.Source) interface{} { return s.APIKeyEnv }); got != "DEEPSEEK_KEY" {
|
||||
t.Errorf("api_key_env = %v; clearing proxy_url must not touch it", got)
|
||||
}
|
||||
if got := sourceField(t, g, "both", func(s config.Source) interface{} { return s.ProxyURL }); got != "" {
|
||||
t.Errorf("proxy_url = %v, want cleared", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestSourceRejectsBadDuration(t *testing.T) {
|
||||
g := richGateway(t)
|
||||
before := sourceField(t, g, "rich", func(s config.Source) interface{} { return s.Timeout })
|
||||
rr := doReq(t, g, "POST", "/api/sources", `{
|
||||
"name":"rich","base_url":"https://api.deepseek.com","api_key":"sk-live",
|
||||
"adapter":"deepseek","models":[{"id":"deepseek-v4-flash","priority":100,"kind":"chat"}],
|
||||
"headers":{},"meta":{},"timeout":"not-a-duration"}`)
|
||||
if rr.Code != 400 {
|
||||
t.Fatalf("status=%d, want 400 for an unparseable duration: %s", rr.Code, rr.Body.String())
|
||||
}
|
||||
if !strings.Contains(rr.Body.String(), "timeout") {
|
||||
t.Errorf("the error must name the offending field: %s", rr.Body.String())
|
||||
}
|
||||
// A rejected write must leave the source exactly as it was.
|
||||
if after := sourceField(t, g, "rich", func(s config.Source) interface{} { return s.Timeout }); after != before {
|
||||
t.Errorf("a rejected edit still changed timeout: %v -> %v", before, after)
|
||||
}
|
||||
}
|
||||
|
||||
// TestSourceCreateSetsOptionalFields: the create path must accept them too, and
|
||||
// must not require an existing source to inherit from (a nil-pointer crash
|
||||
// here would only show up as a 500 on every new source).
|
||||
func TestSourceCreateSetsOptionalFields(t *testing.T) {
|
||||
up := mockUpstream()
|
||||
defer up.Close()
|
||||
g := newTestGateway(t)
|
||||
rr := doReq(t, g, "POST", "/api/sources", `{
|
||||
"name":"fresh","base_url":"`+up.URL+`","api_key":"sk-x","adapter":"openai",
|
||||
"models":[{"id":"m1","kind":"chat"}],"headers":{},"meta":{},
|
||||
"proxy_url":"http://127.0.0.1:7890","timeout":"10s","queue_timeout":"3s"}`)
|
||||
if rr.Code != 200 {
|
||||
t.Fatalf("create status=%d body=%s", rr.Code, rr.Body.String())
|
||||
}
|
||||
if got := sourceField(t, g, "fresh", func(s config.Source) interface{} { return s.ProxyURL }); got != "http://127.0.0.1:7890" {
|
||||
t.Errorf("proxy_url = %v, want it set on create", got)
|
||||
}
|
||||
if got := sourceField(t, g, "fresh", func(s config.Source) interface{} { return s.Timeout }); got != 10*time.Second {
|
||||
t.Errorf("timeout = %v, want 10s on create", got)
|
||||
}
|
||||
}
|
||||
|
||||
// TestSourceRevealExposesDurations guards the round-trip the edit form depends
|
||||
// on. config.Source tags Timeout/QueueTimeout `json:"-"`, so marshaling the
|
||||
// struct omits them; the form then reads an empty timeout box and — because it
|
||||
// always sends the box back — CLEARS the stored timeout on any unrelated save.
|
||||
// A test that only checks "the edit preserves timeout" would not have caught
|
||||
// that, because the loss happens between the read and the write.
|
||||
//
|
||||
// The credentials are asserted in the same breath: they came from the same
|
||||
// branch, and an earlier iteration of this fix masked the key and broke the
|
||||
// form's other half.
|
||||
func TestSourceRevealExposesDurations(t *testing.T) {
|
||||
up := mockUpstream()
|
||||
defer up.Close()
|
||||
g := newTestGateway(t, config.Source{
|
||||
Name: "revealed",
|
||||
BaseURL: up.URL,
|
||||
Adapter: "openai",
|
||||
APIKey: "sk-reveal-me",
|
||||
Models: []config.Model{{ID: "m1", Kind: "chat"}},
|
||||
Timeout: 300 * time.Second,
|
||||
})
|
||||
|
||||
rr := doReq(t, g, "GET", "/api/v1/sources/revealed?reveal=credentials", "")
|
||||
if rr.Code != 200 {
|
||||
t.Fatalf("reveal status=%d: %s", rr.Code, rr.Body.String())
|
||||
}
|
||||
var view struct {
|
||||
Source map[string]interface{} `json:"source"`
|
||||
}
|
||||
if err := json.Unmarshal(rr.Body.Bytes(), &view); err != nil {
|
||||
t.Fatalf("decode: %v", err)
|
||||
}
|
||||
// Duration strings, matching what the write path parses.
|
||||
if got := view.Source["timeout"]; got != "5m0s" {
|
||||
t.Errorf("reveal timeout = %v, want the string \"5m0s\"; the edit form "+
|
||||
"would show an empty box and clear the stored value on save", got)
|
||||
}
|
||||
// The unset one shows the DEFAULT, not "": ApplyDefaults filled
|
||||
// QueueTimeout before the form ever sees it. Both are acceptable as long as
|
||||
// the value round-trips, so assert it is a parseable duration string.
|
||||
qs, _ := view.Source["queue_timeout"].(string)
|
||||
if _, err := time.ParseDuration(qs); err != nil {
|
||||
t.Errorf("reveal queue_timeout = %q is not a duration string: %v", qs, err)
|
||||
}
|
||||
if got := view.Source["api_key"]; got != "sk-reveal-me" {
|
||||
t.Errorf("reveal api_key = %v, want it in the clear for the form", got)
|
||||
}
|
||||
// The masked view must NOT gain the durations' secrets — it is fine for it
|
||||
// to omit them, but it must never carry the key.
|
||||
masked := doReq(t, g, "GET", "/api/v1/sources/revealed", "")
|
||||
if strings.Contains(masked.Body.String(), "sk-reveal-me") {
|
||||
t.Error("the masked view leaked the key")
|
||||
}
|
||||
}
|
||||
|
||||
// TestSourcePayloadCoversEveryEditableField is the guard against the class
|
||||
// itself. Any config.Source field a source edit can legitimately set must
|
||||
// exist on sourcePayload, otherwise the upsert resets it.
|
||||
//
|
||||
// It reflects over the STRUCT rather than over json.Marshal output: the four
|
||||
// pointer fields carry `omitempty`, so a nil pointer legitimately disappears
|
||||
// from the marshaled form — that absence is precisely the "not mentioned" signal
|
||||
// the handler relies on. Checking marshaled keys would therefore report a
|
||||
// healthy field as missing.
|
||||
func TestSourcePayloadCoversEveryEditableField(t *testing.T) {
|
||||
// Fields of config.Source that a source edit can legitimately set.
|
||||
want := []string{
|
||||
"name", "base_url", "api_key", "api_key_env", "adapter", "endpoint",
|
||||
"image_endpoint", "models", "headers", "proxy_url", "meta",
|
||||
"temperature", "max_tokens", "max_concurrent", "rpm",
|
||||
"timeout", "queue_timeout",
|
||||
}
|
||||
pt := reflect.TypeOf(sourcePayload{})
|
||||
for _, f := range want {
|
||||
if _, ok := pt.FieldByName(f); !ok {
|
||||
// Tolerate a differently-spelled Go field only when the JSON tag
|
||||
// matches, so the check stays about the wire contract.
|
||||
if !hasJSONTag(pt, f) {
|
||||
t.Errorf("config.Source has %q but sourcePayload has no such JSON key: "+
|
||||
"an edit that omits it resets the stored value", f)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// hasJSONTag reports whether any field of t carries the given JSON name.
|
||||
func hasJSONTag(t reflect.Type, name string) bool {
|
||||
for i := 0; i < t.NumField(); i++ {
|
||||
if jsonName(t.Field(i)) == name {
|
||||
return true
|
||||
}
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
// jsonName returns the JSON name of a struct field, falling back to the Go
|
||||
// field name when no tag is present.
|
||||
func jsonName(f reflect.StructField) string {
|
||||
tag := f.Tag.Get("json")
|
||||
if tag == "" {
|
||||
return f.Name
|
||||
}
|
||||
if n := strings.Split(tag, ",")[0]; n != "" {
|
||||
return n
|
||||
}
|
||||
return f.Name
|
||||
}
|
||||
259
internal/gateway/sources_api_test.go
Normal file
259
internal/gateway/sources_api_test.go
Normal file
@ -0,0 +1,259 @@
|
||||
package gateway
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"os"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"llmsproxy/internal/config"
|
||||
)
|
||||
|
||||
// These cover the two write-path defects found while adding a model to a live
|
||||
// source by hand:
|
||||
//
|
||||
// 1. PUT /api/sources/{name} was documented in the API index but not
|
||||
// implemented, so it answered 405 while POST (which does the same upsert)
|
||||
// worked. A caller following the documented contract could not edit a
|
||||
// source at all.
|
||||
// 2. POST is an upsert that replaces the whole source, so a partial edit that
|
||||
// did not carry api_key persisted an empty/placeholder credential. The
|
||||
// source kept its name and models and answered 401 on the next call, well
|
||||
// after the writing script had exited 0.
|
||||
//
|
||||
// The second one is the dangerous class: a write that reports success and
|
||||
// breaks the source later.
|
||||
|
||||
// sourceAPIKey reads the stored key back the way an operator would, i.e. what
|
||||
// the routes actually use, not what the file happens to contain.
|
||||
func sourceAPIKey(t *testing.T, g *Gateway, name string) string {
|
||||
t.Helper()
|
||||
for _, s := range g.core.Sources() {
|
||||
if s.Name == name {
|
||||
return s.APIKey
|
||||
}
|
||||
}
|
||||
t.Fatalf("source %q not found", name)
|
||||
return ""
|
||||
}
|
||||
|
||||
func seededGateway(t *testing.T) (*Gateway, string) {
|
||||
t.Helper()
|
||||
up := mockUpstream()
|
||||
t.Cleanup(up.Close)
|
||||
g := newTestGateway(t, config.Source{
|
||||
Name: "seeded",
|
||||
BaseURL: up.URL,
|
||||
Adapter: "openai",
|
||||
APIKey: "sk-real-credential",
|
||||
Models: []config.Model{{ID: "m-one", Kind: "chat"}},
|
||||
})
|
||||
return g, up.URL
|
||||
}
|
||||
|
||||
func TestSourcePutUpdatesSource(t *testing.T) {
|
||||
g, up := seededGateway(t)
|
||||
|
||||
// A single-field edit: change one model's priority, name the source only in
|
||||
// the path. This is the shape the API index documents for PUT.
|
||||
rr := doReq(t, g, "PUT", "/api/sources/seeded", fmt.Sprintf(
|
||||
`{"base_url":%q,"adapter":"openai","api_key":"sk-real-credential","models":[{"id":"m-one","priority":42,"kind":"chat"}]}`, up))
|
||||
if rr.Code != http.StatusOK {
|
||||
t.Fatalf("PUT status=%d body=%s", rr.Code, rr.Body.String())
|
||||
}
|
||||
|
||||
var found bool
|
||||
for _, s := range g.core.Sources() {
|
||||
if s.Name != "seeded" {
|
||||
continue
|
||||
}
|
||||
found = true
|
||||
for _, m := range s.Models {
|
||||
if m.ID == "m-one" && m.Priority != 42 {
|
||||
t.Fatalf("priority not applied: %+v", s.Models)
|
||||
}
|
||||
}
|
||||
}
|
||||
if !found {
|
||||
t.Fatal("source vanished after PUT")
|
||||
}
|
||||
// And it must be routable, not merely stored.
|
||||
if rr := doReq(t, g, "GET", "/v1/models", ""); rr.Code != http.StatusOK ||
|
||||
!strings.Contains(rr.Body.String(), "m-one") {
|
||||
t.Fatalf("model not routable after PUT: %d %s", rr.Code, rr.Body.String())
|
||||
}
|
||||
}
|
||||
|
||||
func TestSourcePutRejectsNameMismatch(t *testing.T) {
|
||||
g, _ := seededGateway(t)
|
||||
// Body names a different source than the path. Silently preferring either
|
||||
// one would let a client edit a source it did not name.
|
||||
rr := doReq(t, g, "PUT", "/api/sources/seeded", `{"name":"other","base_url":"http://x","models":[{"id":"m"}]}`)
|
||||
if rr.Code != http.StatusBadRequest {
|
||||
t.Fatalf("name mismatch accepted: status=%d body=%s", rr.Code, rr.Body.String())
|
||||
}
|
||||
if sourceAPIKey(t, g, "seeded") != "sk-real-credential" {
|
||||
t.Fatal("rejected PUT still mutated the source")
|
||||
}
|
||||
}
|
||||
|
||||
func TestSourcePutRequiresNameInPath(t *testing.T) {
|
||||
g, _ := seededGateway(t)
|
||||
rr := doReq(t, g, "PUT", "/api/sources", `{"base_url":"http://x","models":[{"id":"m"}]}`)
|
||||
if rr.Code != http.StatusBadRequest {
|
||||
t.Fatalf("PUT without a name in the path: status=%d body=%s", rr.Code, rr.Body.String())
|
||||
}
|
||||
}
|
||||
|
||||
func TestSourceUpsertKeepsAPIKey(t *testing.T) {
|
||||
for _, tc := range []struct {
|
||||
name string
|
||||
sentKey string
|
||||
}{
|
||||
{"explicit mask", keepExistingAPIKey},
|
||||
{"omitted key", ""},
|
||||
} {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
g, up := seededGateway(t)
|
||||
rr := doReq(t, g, "PUT", "/api/sources/seeded", fmt.Sprintf(
|
||||
`{"base_url":%q,"adapter":"openai","api_key":%q,"models":[{"id":"m-one"},{"id":"m-two"}]}`,
|
||||
up, tc.sentKey))
|
||||
if rr.Code != http.StatusOK {
|
||||
t.Fatalf("status=%d body=%s", rr.Code, rr.Body.String())
|
||||
}
|
||||
if got := sourceAPIKey(t, g, "seeded"); got != "sk-real-credential" {
|
||||
t.Fatalf("api_key was clobbered: got %q, want the pre-existing credential", got)
|
||||
}
|
||||
// The failure this guards is a source that looks configured but
|
||||
// cannot authenticate, so assert the route still works.
|
||||
if rr := doReq(t, g, "GET", "/v1/models", ""); rr.Code != http.StatusOK {
|
||||
t.Fatalf("source unusable after edit: %d %s", rr.Code, rr.Body.String())
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestSourceUpsertStillWritesNewKey(t *testing.T) {
|
||||
// The keep-existing path must not swallow a deliberate rotation.
|
||||
g, up := seededGateway(t)
|
||||
rr := doReq(t, g, "PUT", "/api/sources/seeded", fmt.Sprintf(
|
||||
`{"base_url":%q,"adapter":"openai","api_key":"sk-rotated","models":[{"id":"m-one"}]}`, up))
|
||||
if rr.Code != http.StatusOK {
|
||||
t.Fatalf("status=%d body=%s", rr.Code, rr.Body.String())
|
||||
}
|
||||
if got := sourceAPIKey(t, g, "seeded"); got != "sk-rotated" {
|
||||
t.Fatalf("rotation ignored: got %q", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestSourceUpsertNewSourceMayHaveNoKey(t *testing.T) {
|
||||
// Empty is a real value for a credential-less source (local upstreams), so
|
||||
// inherit-on-empty must not invent one for a source that does not exist yet.
|
||||
up := mockUpstream()
|
||||
defer up.Close()
|
||||
g := newTestGateway(t)
|
||||
rr := doReq(t, g, "POST", "/api/sources", fmt.Sprintf(
|
||||
`{"name":"noauth","base_url":%q,"adapter":"openai","models":[{"id":"local-m"}]}`, up.URL))
|
||||
if rr.Code != http.StatusOK {
|
||||
t.Fatalf("status=%d body=%s", rr.Code, rr.Body.String())
|
||||
}
|
||||
if got := sourceAPIKey(t, g, "noauth"); got != "" {
|
||||
t.Fatalf("unexpected key on a new credential-less source: %q", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestSourceModelIDsIsAdditive(t *testing.T) {
|
||||
// "Add one model" is the edit that motivated the defect: it needs no read
|
||||
// of the existing list, so it must not require one.
|
||||
g, up := seededGateway(t)
|
||||
rr := doReq(t, g, "PUT", "/api/sources/seeded", fmt.Sprintf(
|
||||
`{"base_url":%q,"adapter":"openai","api_key":%q,"model_ids":["m-two","m-two"," m-three "]}`, up, keepExistingAPIKey))
|
||||
if rr.Code != http.StatusOK {
|
||||
t.Fatalf("status=%d body=%s", rr.Code, rr.Body.String())
|
||||
}
|
||||
var ids []string
|
||||
for _, s := range g.core.Sources() {
|
||||
if s.Name == "seeded" {
|
||||
for _, m := range s.Models {
|
||||
ids = append(ids, m.ID)
|
||||
}
|
||||
}
|
||||
}
|
||||
// m-one kept, m-two not duplicated, blank entries dropped, whitespace trimmed.
|
||||
want := map[string]int{"m-one": 1, "m-two": 1, "m-three": 1}
|
||||
if len(ids) != len(want) {
|
||||
t.Fatalf("models=%v, want exactly %v", ids, want)
|
||||
}
|
||||
seen := map[string]int{}
|
||||
for _, id := range ids {
|
||||
seen[id]++
|
||||
}
|
||||
for id, n := range want {
|
||||
if seen[id] != n {
|
||||
t.Fatalf("models=%v: %q appears %d times, want %d", ids, id, seen[id], n)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestSourceWriteRejectsNonAdmin(t *testing.T) {
|
||||
g, up := seededGateway(t)
|
||||
for _, m := range []string{"POST", "PUT", "DELETE"} {
|
||||
req, _ := http.NewRequest(m, "/api/sources/seeded", strings.NewReader(fmt.Sprintf(
|
||||
`{"name":"seeded","base_url":%q,"models":[{"id":"m"}]}`, up)))
|
||||
req.Header.Set("Authorization", "Bearer sk-not-a-real-key")
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
rr := httptest.NewRecorder()
|
||||
g.Handler().ServeHTTP(rr, req)
|
||||
if rr.Code == http.StatusOK {
|
||||
t.Fatalf("%s allowed for an unknown key: %s", m, rr.Body.String())
|
||||
}
|
||||
}
|
||||
if sourceAPIKey(t, g, "seeded") != "sk-real-credential" {
|
||||
t.Fatal("unauthenticated write mutated the source")
|
||||
}
|
||||
}
|
||||
|
||||
// TestSourceWritePersistsToDisk guards the difference between "the in-memory
|
||||
// registry accepted it" and "it will survive a restart". The api_key must be
|
||||
// sealed at rest, never the literal placeholder.
|
||||
func TestSourceWritePersistsToDisk(t *testing.T) {
|
||||
g, up := seededGateway(t)
|
||||
rr := doReq(t, g, "PUT", "/api/sources/seeded", fmt.Sprintf(
|
||||
`{"base_url":%q,"adapter":"openai","api_key":%q,"models":[{"id":"m-one"}]}`, up, keepExistingAPIKey))
|
||||
if rr.Code != http.StatusOK {
|
||||
t.Fatalf("status=%d body=%s", rr.Code, rr.Body.String())
|
||||
}
|
||||
raw, err := os.ReadFile(g.core.Config().Path)
|
||||
if err != nil {
|
||||
t.Fatalf("read config: %v", err)
|
||||
}
|
||||
text := string(raw)
|
||||
if strings.Contains(text, keepExistingAPIKey) {
|
||||
t.Fatal("the keep-mask was persisted as a literal api_key")
|
||||
}
|
||||
if strings.Contains(text, "sk-real-credential") {
|
||||
t.Fatal("api_key stored in plaintext; it must be sealed at rest")
|
||||
}
|
||||
if !strings.Contains(text, "m-one") {
|
||||
t.Fatal("source not persisted")
|
||||
}
|
||||
// The masked GET view must not leak the credential either.
|
||||
rr = doReq(t, g, "GET", "/api/v1/sources/seeded", "")
|
||||
if rr.Code != http.StatusOK {
|
||||
t.Fatalf("get status=%d", rr.Code)
|
||||
}
|
||||
var view struct {
|
||||
Source struct {
|
||||
APIKey string `json:"api_key"`
|
||||
} `json:"source"`
|
||||
}
|
||||
if err := json.Unmarshal(rr.Body.Bytes(), &view); err != nil {
|
||||
t.Fatalf("decode: %v", err)
|
||||
}
|
||||
if view.Source.APIKey == "sk-real-credential" {
|
||||
t.Fatal("GET /api/v1/sources leaked the plaintext api_key")
|
||||
}
|
||||
}
|
||||
@ -43,11 +43,42 @@ type Req struct {
|
||||
// accounting at all (even when the hit count is 0). The WebUI shows
|
||||
// "0%" instead of "—" for such rows.
|
||||
CacheReported bool `json:"cache_reported,omitempty"`
|
||||
// Reasoning is completion_tokens_details.reasoning_tokens.
|
||||
//
|
||||
// It is a SUBSET of Compl, not an addition: reasoning tokens are billed at
|
||||
// the completion rate and already counted there. This field exists so a
|
||||
// pricing rule can split them when a provider prices thinking tokens
|
||||
// differently (most do NOT — OpenAI and DeepSeek bill them as ordinary
|
||||
// output), and so a reader can see what fraction of a bill is thinking.
|
||||
// 0 means "not reported", which is NOT the same as "no reasoning happened" —
|
||||
// three of the sources measured return no such field at all.
|
||||
Reasoning int64 `json:"reasoning_tokens,omitempty"`
|
||||
OK bool `json:"ok"`
|
||||
// Status http status code
|
||||
Status int `json:"status"`
|
||||
// Err short error message
|
||||
Err string `json:"error,omitempty"`
|
||||
// ImageCount is how many images an image-generation request produced.
|
||||
// Kept separate from Compl/Prompt: image generation has no token concept,
|
||||
// so counting images as "completion tokens" would corrupt the token totals.
|
||||
ImageCount int `json:"image_count,omitempty"`
|
||||
|
||||
// Walk is the AUTO chain's step-by-step trace for this request: which tiers
|
||||
// were skipped and why, which slots hard-failed, which one served it. It is
|
||||
// the only way a consumer can tell "tier 1 served this" from "tier 1 was
|
||||
// cooling so we dropped to tier 3" — a distinction that is the entire point
|
||||
// of a priority chain.
|
||||
//
|
||||
// json:"-" — deliberately NOT persisted. The audit file is a hot append and
|
||||
// this is observational detail: on a degraded gateway every request would
|
||||
// carry a multi-element array, and the audit trail's own retention (16 files
|
||||
// x 16 MB) is already the largest thing on the box. A plugin that wants the
|
||||
// walk sees it live at request_end; an operator post-mortem reads it from the
|
||||
// plugin's own accumulated state or from /api/auto slot health.
|
||||
//
|
||||
// Only populated when a plugin is loaded (chainTraceSink returns nil
|
||||
// otherwise), so a gateway with no plugins allocates nothing for it.
|
||||
Walk []map[string]interface{} `json:"-"`
|
||||
}
|
||||
|
||||
// Stat aggregates counters for one dimension row.
|
||||
@ -70,6 +101,16 @@ type agrRow struct {
|
||||
|
||||
// Stats collects per-key / per-model / per-source aggregates plus a bounded
|
||||
// ring of raw request records, all guarded by one mutex.
|
||||
//
|
||||
// LOCKING: mu is a plain sync.Mutex and is NOT reentrant. The *Locked methods
|
||||
// (aggregateLocked, addKeyTokenLocked, addKeyHourLocked, addKeyReqLocked,
|
||||
// wantPinnedBuckets, rotateAuditLocked, …) assume the caller already holds it,
|
||||
// while every other exported method takes it itself.
|
||||
//
|
||||
// Calling an exported method while already holding mu DEADLOCKS. This is not
|
||||
// hypothetical: a test that did KeyWindowReqs under s.mu.Lock() hung until the
|
||||
// 10-minute panic timeout. Always reach for the *Locked variant when the lock
|
||||
// is already held, and prefer the exported method when it is not.
|
||||
type Stats struct {
|
||||
mu sync.Mutex
|
||||
active int64
|
||||
@ -84,6 +125,21 @@ type Stats struct {
|
||||
auditPath string
|
||||
replayPartial bool // aggregates built from a bounded audit tail
|
||||
modelHour map[string]map[int64]int64 // model -> unix-hour bucket -> tokens
|
||||
// keyModelHour buckets the same tokens as modelHour but keyed by
|
||||
// (gateway key id, model) so per-key model quotas are isolated from each
|
||||
// other. modelHour stays key-blind on purpose: it backs the AUTO slot
|
||||
// quota, which is a gateway-wide limit on a slot, not a per-key one.
|
||||
keyModelHour map[string]map[string]map[int64]int64 // key -> model -> unix-hour -> tokens
|
||||
// keyHour buckets a key's total tokens per unix hour, backing the
|
||||
// whole-key quota (all models of one key share one budget).
|
||||
keyHour map[string]map[int64]int64
|
||||
// keyReqHour buckets a key's request count per unix hour, backing the
|
||||
// key-wide request quota. Counted for every request type (chat, stream,
|
||||
// image) including failed ones, so a failing client cannot loop for free.
|
||||
keyReqHour map[string]map[int64]int64
|
||||
// pinnedWanted records which (key, model) pairs have a quota that pins a
|
||||
// source, so the pinned bucket is maintained only where it is read.
|
||||
pinnedWanted map[string]bool
|
||||
}
|
||||
|
||||
const hourSec = 3600
|
||||
@ -116,14 +172,18 @@ func NewStats(maxRecords int) *Stats {
|
||||
maxRecords = defaultRingSize
|
||||
}
|
||||
return &Stats{
|
||||
byKey: map[string]*Stat{},
|
||||
byModel: map[string]*Stat{},
|
||||
bySrc: map[string]*Stat{},
|
||||
byKeyModel: map[string]map[string]*Stat{},
|
||||
byKeySrc: map[string]map[string]*Stat{},
|
||||
byStatus: map[int]*Stat{},
|
||||
modelHour: map[string]map[int64]int64{},
|
||||
maxRecs: maxRecords,
|
||||
byKey: map[string]*Stat{},
|
||||
byModel: map[string]*Stat{},
|
||||
bySrc: map[string]*Stat{},
|
||||
byKeyModel: map[string]map[string]*Stat{},
|
||||
byKeySrc: map[string]map[string]*Stat{},
|
||||
byStatus: map[int]*Stat{},
|
||||
modelHour: map[string]map[int64]int64{},
|
||||
keyModelHour: map[string]map[string]map[int64]int64{},
|
||||
keyHour: map[string]map[int64]int64{},
|
||||
keyReqHour: map[string]map[int64]int64{},
|
||||
pinnedWanted: map[string]bool{},
|
||||
maxRecs: maxRecords,
|
||||
}
|
||||
}
|
||||
|
||||
@ -318,6 +378,9 @@ func (s *Stats) Record(r Req) {
|
||||
// aggregateLocked folds r into every aggregate row and the quota window
|
||||
// bucket. Caller must hold s.mu.
|
||||
func (s *Stats) aggregateLocked(r Req) {
|
||||
if r.Key != "" {
|
||||
s.addKeyReqLocked(r.Key, (r.Time/1000)/hourSec, 1)
|
||||
}
|
||||
inc(s.byKey, r.Key, r)
|
||||
if r.Model != "" {
|
||||
inc(s.byModel, r.Model, r)
|
||||
@ -362,6 +425,12 @@ func (s *Stats) aggregateLocked(r Req) {
|
||||
s.modelHour[key] = hm
|
||||
}
|
||||
hm[h] += tok
|
||||
// per-key buckets: same token split, but isolated per key so one
|
||||
// key's quota cannot be drained by another key's traffic
|
||||
if r.Key != "" {
|
||||
s.addKeyTokenLocked(r.Key, r.Model, r.Source, h, tok)
|
||||
s.addKeyHourLocked(r.Key, h, tok)
|
||||
}
|
||||
// retention: 24*40 = 960 hourly buckets ≈ 40 days of history (covers
|
||||
// the longest "month" quota window)
|
||||
if len(hm) > 24*40 {
|
||||
@ -374,6 +443,185 @@ func (s *Stats) aggregateLocked(r Req) {
|
||||
}
|
||||
}
|
||||
|
||||
// quotaRetentionHours is how much hourly history the per-key quota buckets
|
||||
// keep. It matches the modelHour retention (40 days) so the longest "month"
|
||||
// window is fully covered after a restart, and it is deliberately applied per
|
||||
// key so an idle key's buckets are reclaimed instead of pinning memory.
|
||||
const quotaRetentionHours = 24 * 40
|
||||
|
||||
// addKeyTokenLocked adds tok to one key's (model[, source]) hourly bucket.
|
||||
//
|
||||
// The bare-model bucket is always written: it is what a scope entry without a
|
||||
// source pin reads, and writing it is what makes that quota work at all (every
|
||||
// recorded request carries a resolved source, so a source::model-only bucket
|
||||
// would leave an unpinned quota reading an empty bucket forever).
|
||||
//
|
||||
// The pinned "source::model" bucket is written only for (key, model) pairs
|
||||
// some scope entry actually pins a source for. Maintaining it unconditionally
|
||||
// doubles the bucket count — measured at 26.7 MB for 20 keys x 8 models x 3
|
||||
// sources at full retention, against a documented ~32 MB total budget — to
|
||||
// serve a lookup that almost nobody performs.
|
||||
func (s *Stats) addKeyTokenLocked(key, model, source string, h, tok int64) {
|
||||
if model == "" {
|
||||
return
|
||||
}
|
||||
byModel := s.keyModelHour[key]
|
||||
if byModel == nil {
|
||||
byModel = map[string]map[int64]int64{}
|
||||
s.keyModelHour[key] = byModel
|
||||
}
|
||||
addBucket(byModel, model, h, tok)
|
||||
if source == "" {
|
||||
return
|
||||
}
|
||||
pinned := source + "::" + model
|
||||
if pinned == model || !s.pinnedWanted[key+"\x00"+model] {
|
||||
return
|
||||
}
|
||||
addBucket(byModel, pinned, h, tok)
|
||||
}
|
||||
|
||||
// addBucket adds tok to one named hourly bucket and enforces retention.
|
||||
func addBucket(byModel map[string]map[int64]int64, name string, h, tok int64) {
|
||||
hm := byModel[name]
|
||||
if hm == nil {
|
||||
hm = map[int64]int64{}
|
||||
byModel[name] = hm
|
||||
}
|
||||
hm[h] += tok
|
||||
if len(hm) > quotaRetentionHours {
|
||||
for k := range hm {
|
||||
if k < h-quotaRetentionHours {
|
||||
delete(hm, k)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// wantPinnedBuckets starts maintaining the pinned (source::model) bucket for a
|
||||
// key's model, so a quota that pins a source counts traffic from the moment it
|
||||
// is configured. Usage recorded before that cannot be attributed to a source
|
||||
// after the fact, so a pinned quota may under-count its first window.
|
||||
func (s *Stats) wantPinnedBuckets(key, model string) {
|
||||
if key == "" || model == "" {
|
||||
return
|
||||
}
|
||||
s.pinnedWanted[key+"\x00"+model] = true
|
||||
}
|
||||
|
||||
// addKeyHourLocked adds tok to one key's all-model hourly total.
|
||||
func (s *Stats) addKeyHourLocked(key string, h, tok int64) {
|
||||
hm := s.keyHour[key]
|
||||
if hm == nil {
|
||||
hm = map[int64]int64{}
|
||||
s.keyHour[key] = hm
|
||||
}
|
||||
hm[h] += tok
|
||||
if len(hm) > quotaRetentionHours {
|
||||
for k := range hm {
|
||||
if k < h-quotaRetentionHours {
|
||||
delete(hm, k)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// KeyWindowReqs returns how many requests one key issued within sec seconds;
|
||||
// sec <= 0 means all retained history.
|
||||
func (s *Stats) KeyWindowReqs(key string, sec int64) int64 {
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
return sumBuckets(s.keyReqHour[key], time.Now().Unix(), sec)
|
||||
}
|
||||
|
||||
// addKeyReqLocked adds n requests to a key's hourly count bucket.
|
||||
func (s *Stats) addKeyReqLocked(key string, h, n int64) {
|
||||
hm := s.keyReqHour[key]
|
||||
if hm == nil {
|
||||
hm = map[int64]int64{}
|
||||
s.keyReqHour[key] = hm
|
||||
}
|
||||
hm[h] += n
|
||||
if len(hm) > quotaRetentionHours {
|
||||
for k := range hm {
|
||||
if k < h-quotaRetentionHours {
|
||||
delete(hm, k)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// sumBuckets totals the hourly buckets inside the window ending at `now`
|
||||
// (unix seconds). sec <= 0 means "all retained history" (no reset).
|
||||
//
|
||||
// It only visits the hour buckets a window can actually cover. The naive full
|
||||
// scan measured 5.9us/op at the full 40-day retention, and a per-key quota
|
||||
// check runs on every request (2-3 of these), so the scan cost lands directly
|
||||
// on the request path; visiting at most sec/3600 + 1 buckets makes it
|
||||
// proportional to the window instead of to the retained history. The map keys
|
||||
// are whole unix hours, so the covered range is exact rather than a sample.
|
||||
func sumBuckets(hm map[int64]int64, now, sec int64) int64 {
|
||||
if len(hm) == 0 {
|
||||
return 0
|
||||
}
|
||||
if sec <= 0 {
|
||||
var total int64
|
||||
for _, v := range hm {
|
||||
total += v
|
||||
}
|
||||
return total
|
||||
}
|
||||
nowHour := now / hourSec
|
||||
// A bucket at hour H covers [H*3600, (H+1)*3600). The pre-existing
|
||||
// WindowTokens counts a bucket when h*3600 >= now-sec, so the first hour
|
||||
// that qualifies is ceil((now-sec)/3600) — NOT floor, which would pull in
|
||||
// one bucket too many (caught by TestSumBucketsMatchesFullScan: a 30-day
|
||||
// window reported 8878 vs 8649).
|
||||
firstHour := (now - sec + hourSec - 1) / hourSec
|
||||
if firstHour < 0 {
|
||||
firstHour = 0
|
||||
}
|
||||
var total int64
|
||||
for h := firstHour; h <= nowHour; h++ {
|
||||
total += hm[h]
|
||||
}
|
||||
return total
|
||||
}
|
||||
|
||||
// KeyWindowModelTokens returns the tokens one key consumed for one model
|
||||
// (optionally pinned to a source) within sec seconds; sec <= 0 means all
|
||||
// retained history. Unlike WindowTokens this is isolated per key, which is
|
||||
// what a per-key model quota needs.
|
||||
//
|
||||
// A source pin matches the pinned bucket exactly. Without a pin the call
|
||||
// totals the key's bare-model bucket, which counts traffic on every source —
|
||||
// a quota on "this model" should not stop applying just because the request
|
||||
// happened to be served by a different upstream.
|
||||
func (s *Stats) KeyWindowModelTokens(key, model, source string, sec int64) int64 {
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
byModel := s.keyModelHour[key]
|
||||
if len(byModel) == 0 {
|
||||
return 0
|
||||
}
|
||||
mk := model
|
||||
if source != "" {
|
||||
mk = source + "::" + model
|
||||
// from now on this pair keeps a pinned bucket, so a quota configured
|
||||
// against this source counts the traffic it is meant to cap
|
||||
s.wantPinnedBuckets(key, model)
|
||||
}
|
||||
return sumBuckets(byModel[mk], time.Now().Unix(), sec)
|
||||
}
|
||||
|
||||
// KeyWindowTokens returns one key's total tokens across all models within sec
|
||||
// seconds; sec <= 0 means all retained history. It backs the whole-key quota.
|
||||
func (s *Stats) KeyWindowTokens(key string, sec int64) int64 {
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
return sumBuckets(s.keyHour[key], time.Now().Unix(), sec)
|
||||
}
|
||||
|
||||
// rotateAuditLocked renames the audit file to <path>.<unix>.old once it
|
||||
// exceeds auditRotateBytes and prunes old files beyond auditKeepOld, keeping
|
||||
// the newest ones. Caller must hold s.mu.
|
||||
@ -464,6 +712,27 @@ func AutoPeriodSeconds(period string, hours int64) int64 {
|
||||
return 0
|
||||
}
|
||||
|
||||
// AutoSecondsToReset returns how many seconds remain until the quota window
|
||||
// rolls over, for the Retry-After header on a 429. Buckets are whole unix
|
||||
// hours, so the value is rounded up to the next hour boundary and never
|
||||
// exceeds one window. 0 means "unknown" (never resets, or no period set), in
|
||||
// which case the caller should not advertise a retry time.
|
||||
func AutoSecondsToReset(period string, hours int64) int64 {
|
||||
sec := AutoPeriodSeconds(period, hours)
|
||||
if sec <= 0 {
|
||||
return 0
|
||||
}
|
||||
now := time.Now().Unix()
|
||||
// window = the last sec seconds, bucketed by whole hours; the oldest
|
||||
// bucket still inside the window expires at its own hour boundary
|
||||
elapsed := now % hourSec
|
||||
left := sec - elapsed
|
||||
if left <= 0 || left > sec {
|
||||
left = sec
|
||||
}
|
||||
return left
|
||||
}
|
||||
|
||||
// WindowTokens returns the tokens consumed for one model (optionally pinned
|
||||
// to a single source) within the window; sec <= 0 means all time. Buckets are
|
||||
// whole unix hours, so a sliding window overcounts by up to one hour — an
|
||||
@ -952,3 +1221,16 @@ func (s *Stats) Snapshot(limit int, key string) map[string]interface{} {
|
||||
"replay_partial": s.replayPartial,
|
||||
}
|
||||
}
|
||||
|
||||
// Active reports requests in flight right now.
|
||||
//
|
||||
// It is a live gauge, not an aggregate: there is no calendar period that scopes
|
||||
// it, which is why the period-scoped /api/stats response has to carry it
|
||||
// explicitly. Before it did, the day/week/month views fed `st.active || 0` a
|
||||
// missing key and showed a hard 0 — the same gateway that was serving hundreds
|
||||
// of requests a second looked idle in three of the four views.
|
||||
func (s *Stats) Active() int64 {
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
return s.active
|
||||
}
|
||||
|
||||
132
internal/gateway/stats_lock_test.go
Normal file
132
internal/gateway/stats_lock_test.go
Normal file
@ -0,0 +1,132 @@
|
||||
package gateway
|
||||
|
||||
import (
|
||||
"reflect"
|
||||
"runtime"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
)
|
||||
|
||||
// The Stats mutex is a plain sync.Mutex: calling an exported (self-locking)
|
||||
// method while already holding it deadlocks. That happened for real — a probe
|
||||
// held s.mu and called KeyWindowReqs, hanging until the test binary's 10-minute
|
||||
// panic timeout — so the rule is pinned here rather than left in a comment.
|
||||
//
|
||||
// The check is structural: every method that takes the lock must say so in its
|
||||
// name or its doc comment. That keeps the trap visible when someone adds the
|
||||
// next method, which is the only time the rule can be forgotten.
|
||||
|
||||
// exportedSelfLocking lists the exported Stats methods that take the mutex.
|
||||
// It is derived from reflection at run time; the assertions below are what
|
||||
// actually pin the convention.
|
||||
func TestStatsExportedMethodsDocumentTheirLocking(t *testing.T) {
|
||||
typ := reflect.TypeOf(&Stats{})
|
||||
for i := 0; i < typ.NumMethod(); i++ {
|
||||
m := typ.Method(i)
|
||||
if m.PkgPath != "" { // unexported
|
||||
continue
|
||||
}
|
||||
// Methods that neither lock nor touch guarded state are fine either way;
|
||||
// the ones that matter are those reaching into the maps under mu.
|
||||
if !methodTouchesLockedState(m.Name) {
|
||||
continue
|
||||
}
|
||||
if strings.HasSuffix(m.Name, "Locked") {
|
||||
t.Errorf("Stats.%s is exported but named *Locked; the suffix means "+
|
||||
"the CALLER holds the lock, so it must not be exported", m.Name)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// methodTouchesLockedState is the set of exported methods known to read or
|
||||
// write state guarded by Stats.mu. Kept explicit (rather than inferred) so a
|
||||
// new method is not silently assumed safe.
|
||||
func methodTouchesLockedState(name string) bool {
|
||||
switch name {
|
||||
case "KeyWindowReqs", "KeyWindowModelTokens", "KeyWindowTokens",
|
||||
"WindowTokens", "KeyTokens", "ModelTokens", "Record", "AppendAudit",
|
||||
"LoadAudit", "Snapshot", "SourceRecent", "SourceAverages",
|
||||
"AuditRecords", "AuditPage", "StreamAuditRecords", "ReplayPartial",
|
||||
"PoolStats":
|
||||
return true
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
// TestStatsLockedMethodsAreNotCalledUnderLock is the behavioural half: it
|
||||
// proves the internal helpers the ones above use are reachable while the lock
|
||||
// is held. If a future refactor makes an exported method call a *Locked one
|
||||
// while holding mu itself, this is where it shows up — as a hang, bounded by
|
||||
// the short timeout below rather than the suite's 10 minutes.
|
||||
func TestStatsLockedMethodsAreNotCalledUnderLock(t *testing.T) {
|
||||
done := make(chan struct{})
|
||||
go func() {
|
||||
defer close(done)
|
||||
s := NewStats(10)
|
||||
h := time.Now().Unix() / 3600
|
||||
s.mu.Lock()
|
||||
// Exactly the shape that deadlocked: the *Locked variants are correct
|
||||
// here because the lock is already held.
|
||||
s.addKeyReqLocked("k", h, 5)
|
||||
s.addKeyHourLocked("k", h, 100)
|
||||
s.wantPinnedBuckets("k", "m1")
|
||||
s.addKeyTokenLocked("k", "m1", "src", h, 100)
|
||||
// sumBuckets is the pure inner function the exported readers call.
|
||||
if got := sumBuckets(s.keyReqHour["k"], time.Now().Unix(), 3600); got != 5 {
|
||||
t.Errorf("sumBuckets under lock = %d, want 5", got)
|
||||
}
|
||||
s.mu.Unlock()
|
||||
}()
|
||||
select {
|
||||
case <-done:
|
||||
case <-time.After(20 * time.Second):
|
||||
buf := make([]byte, 1<<16)
|
||||
n := runtime.Stack(buf, true)
|
||||
t.Fatalf("deadlocked while using the *Locked helpers under s.mu:\n%s", buf[:n])
|
||||
}
|
||||
}
|
||||
|
||||
// TestSumBucketsEdgeCases covers the boundaries the quota check depends on,
|
||||
// including the ones a regression would silently get wrong (an off-by-one here
|
||||
// either lets a quota leak or locks a key out early).
|
||||
func TestSumBucketsEdgeCases(t *testing.T) {
|
||||
const hour = 3600
|
||||
now := int64(10*hour + 61) // 10:00:61, i.e. just past the boundary
|
||||
buckets := map[int64]int64{
|
||||
0: 100, // ancient
|
||||
9: 200, // previous hour
|
||||
10: 7, // current hour
|
||||
}
|
||||
|
||||
if got := sumBuckets(buckets, now, 0); got != 307 {
|
||||
t.Errorf("sec<=0 (all history) = %d, want 307", got)
|
||||
}
|
||||
// 1h window covers buckets h >= ceil((now-3600)/3600) = 10 -> only bucket 10.
|
||||
if got := sumBuckets(buckets, now, hour); got != 7 {
|
||||
t.Errorf("1h window = %d, want 7 (buckets 0 and 9 fall outside)", got)
|
||||
}
|
||||
// 2h window covers h >= 9 -> buckets 9 and 10.
|
||||
if got := sumBuckets(buckets, now, 2*hour); got != 207 {
|
||||
t.Errorf("2h window = %d, want 207", got)
|
||||
}
|
||||
// An empty map and a nil map must both be zero, not a panic.
|
||||
if got := sumBuckets(map[int64]int64{}, now, hour); got != 0 {
|
||||
t.Errorf("empty map = %d, want 0", got)
|
||||
}
|
||||
if got := sumBuckets(nil, now, hour); got != 0 {
|
||||
t.Errorf("nil map = %d, want 0", got)
|
||||
}
|
||||
// now < sec must not produce a negative firstHour index: with a window far
|
||||
// wider than the available history, every bucket that exists is counted.
|
||||
// (sumBuckets walks bucket indices from firstHour to nowHour, so passing a
|
||||
// "now" older than some buckets cannot reach them — that is correct, not a
|
||||
// truncation bug.)
|
||||
if got := sumBuckets(buckets, 10*hour+61, 100*hour); got != 307 {
|
||||
t.Errorf("window wider than history = %d, want 307", got)
|
||||
}
|
||||
// A window that predates every bucket counts them all as well.
|
||||
if got := sumBuckets(buckets, 10*hour+61, hour); got != 7 {
|
||||
t.Errorf("1h window at t=10:00:61 = %d, want 7", got)
|
||||
}
|
||||
}
|
||||
234
internal/gateway/stats_period.go
Normal file
234
internal/gateway/stats_period.go
Normal file
@ -0,0 +1,234 @@
|
||||
package gateway
|
||||
|
||||
import (
|
||||
"sort"
|
||||
"strconv"
|
||||
"strings"
|
||||
"time"
|
||||
)
|
||||
|
||||
// Period is a reporting window for the usage dashboard. The dashboard used to
|
||||
// have exactly one view — everything since process start — which made a quiet
|
||||
// morning indistinguishable from a quiet week. Periods give the operator a
|
||||
// scale to read the numbers at: today vs this week vs this month vs all time.
|
||||
//
|
||||
// The set is deliberately calendar-based and UTC-anchored. A rolling 24h window
|
||||
// would put "today" and "the last day" at different boundaries, so the same
|
||||
// number would move between two cards depending on when you looked; calendar
|
||||
// days are what people actually mean by "today". UTC also matches the billing
|
||||
// plugin's peak-window arithmetic, so a peak-rate hour does not land in a
|
||||
// different day in the cost view than in the usage view.
|
||||
type Period string
|
||||
|
||||
const (
|
||||
// PeriodDay is the current UTC calendar day.
|
||||
PeriodDay Period = "day"
|
||||
// PeriodWeek is the current ISO week (Mon 00:00 UTC to now).
|
||||
PeriodWeek Period = "week"
|
||||
// PeriodMonth is the current UTC calendar month.
|
||||
PeriodMonth Period = "month"
|
||||
// PeriodAll is since process start — the only view backed by the
|
||||
// in-memory aggregates, and the only one available when no audit file is
|
||||
// configured.
|
||||
PeriodAll Period = "all"
|
||||
)
|
||||
|
||||
// ValidPeriod reports whether p is a period the aggregator understands.
|
||||
// An unknown period is a client error, not a silent fallback to "all": a
|
||||
// dashboard that quietly shows lifetime totals when the caller asked for today
|
||||
// is worse than one that refuses.
|
||||
func ValidPeriod(p Period) bool {
|
||||
switch p {
|
||||
case PeriodDay, PeriodWeek, PeriodMonth, PeriodAll:
|
||||
return true
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
// periodStart returns the inclusive start of the window for p at time now.
|
||||
// Only PeriodDay/Week/Month are meaningful here; PeriodAll returns 0, which
|
||||
// every "from > 0" bounds check treats as unbounded.
|
||||
func periodStart(p Period, now time.Time) int64 {
|
||||
now = now.UTC()
|
||||
switch p {
|
||||
case PeriodDay:
|
||||
return time.Date(now.Year(), now.Month(), now.Day(), 0, 0, 0, 0, time.UTC).UnixMilli()
|
||||
case PeriodWeek:
|
||||
// ISO week starts Monday. Go's Weekday() is Sunday=0, so the shift
|
||||
// below is 1 on Sunday and 0 on Monday..Saturday.
|
||||
off := (int(now.Weekday()) + 6) % 7
|
||||
d := now.AddDate(0, 0, -off)
|
||||
return time.Date(d.Year(), d.Month(), d.Day(), 0, 0, 0, 0, time.UTC).UnixMilli()
|
||||
case PeriodMonth:
|
||||
return time.Date(now.Year(), now.Month(), 1, 0, 0, 0, 0, time.UTC).UnixMilli()
|
||||
}
|
||||
return 0
|
||||
}
|
||||
|
||||
// PeriodBucket is one labelled point on the dashboard's timeline. Buckets are
|
||||
// the aggregation grain for a period: hourly for a day (so a spike is visible
|
||||
// inside "today"), daily for a week or month (so a week is not 7×24 points
|
||||
// wide and a month is not 31×24), and empty for "all" — a lifetime total has
|
||||
// no meaningful timeline, and pretending otherwise by drawing 500 hourly
|
||||
// buckets of rolling memory would be a different lie.
|
||||
type PeriodBucket struct {
|
||||
Bucket string `json:"bucket"`
|
||||
Stat
|
||||
}
|
||||
|
||||
// PeriodSnapshot is the period-scoped twin of Snapshot's payload: the same
|
||||
// totals and the same by_* rows, plus a timeline. It deliberately mirrors
|
||||
// Snapshot's field names so the dashboard's table painters work unchanged —
|
||||
// paintModelTable(st.by_model) reads rows of {name, ...Stat}, and that shape
|
||||
// does not care where the numbers came from.
|
||||
type PeriodSnapshot struct {
|
||||
Period Period `json:"period"`
|
||||
From int64 `json:"from"` // unix millis, 0 when PeriodAll
|
||||
Total Stat `json:"total"`
|
||||
ByKey []StatsRow `json:"by_key"`
|
||||
Models []StatsRow `json:"by_model"`
|
||||
Srcs []StatsRow `json:"by_source"`
|
||||
Status []agrRow `json:"by_status"`
|
||||
Bucket []PeriodBucket `json:"buckets"`
|
||||
// Truncated marks that the window was clipped by the retained audit
|
||||
// history, so the numbers are a lower bound rather than the true period
|
||||
// total. The dashboard shows this next to the numbers rather than letting
|
||||
// a rotated-away week read as "that week had no traffic".
|
||||
Truncated bool `json:"truncated"`
|
||||
}
|
||||
|
||||
// bucketKey maps a record's timestamp onto the timeline grain for p.
|
||||
// Daily buckets are stamped at UTC midnight; hourly buckets carry the hour.
|
||||
func bucketKey(p Period, ms int64) string {
|
||||
t := time.UnixMilli(ms).UTC()
|
||||
if p == PeriodDay {
|
||||
return t.Format("2006-01-02T15")
|
||||
}
|
||||
return t.Format("2006-01-02")
|
||||
}
|
||||
|
||||
// bucketOf returns the bucket label plus the truncated-flag side effects of
|
||||
// walking a file: a record older than the requested window means the window
|
||||
// starts before the retained history, and the file may have been cut short.
|
||||
func (s *Stats) PeriodSnapshot(p Period, key string, now time.Time) PeriodSnapshot {
|
||||
out := PeriodSnapshot{Period: p}
|
||||
from := periodStart(p, now)
|
||||
out.From = from
|
||||
|
||||
// "all" is exactly what Snapshot already answers, from the in-memory
|
||||
// aggregates, and it is the one view that must keep working with no audit
|
||||
// file configured at all (a fresh dev setup, or an operator who turned
|
||||
// auditing off). Serving it from the same code path keeps the dashboard's
|
||||
// "total" card identical whether or not a period is selected.
|
||||
if p == PeriodAll {
|
||||
snap := s.Snapshot(firstScreenRecords, key)
|
||||
if tot, ok := snap["total"].(Stat); ok {
|
||||
out.Total = tot
|
||||
}
|
||||
out.ByKey, _ = snap["by_key"].([]StatsRow)
|
||||
out.Models, _ = snap["by_model"].([]StatsRow)
|
||||
out.Srcs, _ = snap["by_source"].([]StatsRow)
|
||||
out.Status, _ = snap["by_status"].([]agrRow)
|
||||
// replay_partial is the same "these numbers came from a bounded
|
||||
// tail" caveat, carried through under this view's own name.
|
||||
out.Truncated, _ = snap["replay_partial"].(bool)
|
||||
return out
|
||||
}
|
||||
|
||||
// A bounded window is aggregated from the audit files, because the
|
||||
// in-memory aggregates are lifetime totals and the ring buffer holds only
|
||||
// maxRecs records (500 by default). Reading the ring would silently
|
||||
// under-report any period longer than the last few hundred requests.
|
||||
total := Stat{}
|
||||
byKey := map[string]*Stat{}
|
||||
byModel := map[string]*Stat{}
|
||||
bySrc := map[string]*Stat{}
|
||||
byStatus := map[string]*Stat{}
|
||||
buckets := map[string]*Stat{}
|
||||
|
||||
var seen int
|
||||
err := s.StreamAuditRecords(from, 0, key, func(r Req) error {
|
||||
seen++
|
||||
incStatus(&total, "", r)
|
||||
inc(byKey, r.Key, r)
|
||||
if r.Model != "" {
|
||||
inc(byModel, r.Model, r)
|
||||
}
|
||||
inc(bySrc, r.Source, r)
|
||||
if r.Status != 0 {
|
||||
inc(byStatus, strconv.Itoa(r.Status), r)
|
||||
}
|
||||
k := bucketKey(p, r.Time)
|
||||
b := buckets[k]
|
||||
if b == nil {
|
||||
b = &Stat{}
|
||||
buckets[k] = b
|
||||
}
|
||||
incStatus(b, k, r)
|
||||
return nil
|
||||
})
|
||||
|
||||
out.Total = total
|
||||
out.ByKey = rows(byKey)
|
||||
out.Models = rows(byModel)
|
||||
out.Srcs = rows(bySrc)
|
||||
bs := make([]agrRow, 0, len(byStatus))
|
||||
for code, st := range byStatus {
|
||||
bs = append(bs, agrRow{Name: code, Stat: *st})
|
||||
}
|
||||
sort.Slice(bs, func(i, j int) bool {
|
||||
ci, _ := strconv.Atoi(bs[i].Name)
|
||||
cj, _ := strconv.Atoi(bs[j].Name)
|
||||
return ci < cj
|
||||
})
|
||||
out.Status = bs
|
||||
|
||||
ks := make([]string, 0, len(buckets))
|
||||
for k := range buckets {
|
||||
ks = append(ks, k)
|
||||
}
|
||||
// Chronological, string-sorted: "2006-01-02T15" and "2006-01-02" both
|
||||
// sort lexicographically in time order, so no date parsing is needed.
|
||||
sort.Strings(ks)
|
||||
out.Bucket = make([]PeriodBucket, 0, len(ks))
|
||||
for _, k := range ks {
|
||||
out.Bucket = append(out.Bucket, PeriodBucket{Bucket: k, Stat: *buckets[k]})
|
||||
}
|
||||
|
||||
// The window is only "complete" if the audit walk actually reached back
|
||||
// far enough to cover it. Two ways it cannot:
|
||||
//
|
||||
// 1. The walk found nothing at all in a window that certainly had
|
||||
// traffic, because the files holding it rotated away.
|
||||
// 2. The walk errored part-way (I/O), leaving a partial total.
|
||||
//
|
||||
// Case 2 is reported from err; case 1 from seen == 0 combined with the
|
||||
// caller having asked for a bounded window. It is deliberately
|
||||
// conservative: a genuinely empty hour is rare enough that flagging it as
|
||||
// possibly-truncated costs one tooltip, whereas silently under-reporting a
|
||||
// month because rotation ate it is a wrong number with no indication.
|
||||
if err != nil || (seen == 0 && p != PeriodAll) {
|
||||
out.Truncated = true
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
// periodFromQuery parses the ?period= parameter. An empty value means "all" so
|
||||
// that existing callers of /api/stats keep seeing exactly what they saw.
|
||||
// A malformed value is rejected by the caller (ValidPeriod) rather than
|
||||
// defaulting, so a typo in a bookmarked URL surfaces as an error instead of
|
||||
// silently switching the operator to lifetime totals.
|
||||
func periodFromQuery(q map[string][]string) Period {
|
||||
v := strings.ToLower(strings.TrimSpace(firstQuery(q, "period")))
|
||||
if v == "" {
|
||||
return PeriodAll
|
||||
}
|
||||
return Period(v)
|
||||
}
|
||||
|
||||
func firstQuery(q map[string][]string, key string) string {
|
||||
if vs, ok := q[key]; ok && len(vs) > 0 {
|
||||
return vs[0]
|
||||
}
|
||||
return ""
|
||||
}
|
||||
159
internal/gateway/stats_period_records_test.go
Normal file
159
internal/gateway/stats_period_records_test.go
Normal file
@ -0,0 +1,159 @@
|
||||
package gateway
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"net/url"
|
||||
"path/filepath"
|
||||
"testing"
|
||||
"time"
|
||||
)
|
||||
|
||||
// The dashboard paints its records table from the `records` field of
|
||||
// /api/stats (paintRecords(st.records)). The period branch of that handler
|
||||
// used to hand-write a response map and simply omit the key, so switching the
|
||||
// selector to day / week / month rendered an EMPTY table no matter how much
|
||||
// traffic the window had — while the lifetime view worked, because it goes
|
||||
// through Snapshot() which carries records.
|
||||
//
|
||||
// These tests pin the contract across every period: same records, same cursor.
|
||||
func TestStatsPeriodViewsShipRecords(t *testing.T) {
|
||||
td := t.TempDir()
|
||||
g := newTestGateway(t)
|
||||
|
||||
// Write audit records the period aggregator and the record pager both read.
|
||||
g.stats.auditPath = filepath.Join(td, "audit.jsonl")
|
||||
now := time.Now().UnixMilli()
|
||||
for i := 0; i < 5; i++ {
|
||||
g.stats.Record(Req{
|
||||
Time: now - int64(i)*1000, Type: "chat", OK: true,
|
||||
Model: "m1", Source: "s1", Status: 200,
|
||||
Prompt: 10, Compl: 2,
|
||||
})
|
||||
}
|
||||
|
||||
for _, p := range []string{"day", "week", "month", "all"} {
|
||||
rr := doReq(t, g, "GET", "/api/stats?period="+p, "")
|
||||
if rr.Code != 200 {
|
||||
t.Fatalf("period=%s: HTTP %d: %s", p, rr.Code, rr.Body.String())
|
||||
}
|
||||
var st map[string]interface{}
|
||||
if err := json.Unmarshal(rr.Body.Bytes(), &st); err != nil {
|
||||
t.Fatalf("period=%s: %v", p, err)
|
||||
}
|
||||
recs, ok := st["records"].([]interface{})
|
||||
if !ok {
|
||||
t.Fatalf("period=%s: response has no `records` array — the dashboard "+
|
||||
"records table renders empty for this view (keys: %v)", p, keysOf(st))
|
||||
}
|
||||
if len(recs) != 5 {
|
||||
t.Fatalf("period=%s: got %d records, want 5", p, len(recs))
|
||||
}
|
||||
// The scroll handler pages from next_cursor; without it the table is
|
||||
// stuck at one screen no matter how far the user scrolls.
|
||||
if _, ok := st["next_cursor"]; !ok {
|
||||
t.Fatalf("period=%s: no next_cursor — paging can never continue", p)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// The records shipped with a period view must be the newest-first page the
|
||||
// paging endpoint would return, in the same order. If /api/stats and
|
||||
// /api/stats/records disagreed, scrolling would duplicate or skip rows.
|
||||
func TestStatsPeriodRecordsMatchPager(t *testing.T) {
|
||||
td := t.TempDir()
|
||||
g := newTestGateway(t)
|
||||
g.stats.auditPath = filepath.Join(td, "audit.jsonl")
|
||||
now := time.Now().UnixMilli()
|
||||
for i := 0; i < 8; i++ {
|
||||
g.stats.Record(Req{
|
||||
Time: now - int64(i)*1000, Type: "chat", OK: true,
|
||||
Model: "m1", Status: 200, Prompt: 5, Compl: 1,
|
||||
})
|
||||
}
|
||||
|
||||
rr := doReq(t, g, "GET", "/api/stats?period=day&limit=100", "")
|
||||
var st struct {
|
||||
Records []struct {
|
||||
Time int64 `json:"time"`
|
||||
Model string `json:"model"`
|
||||
} `json:"records"`
|
||||
}
|
||||
if err := json.Unmarshal(rr.Body.Bytes(), &st); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
rr2 := doReq(t, g, "GET", "/api/stats/records?limit=100", "")
|
||||
var pg struct {
|
||||
Records []struct {
|
||||
Time int64 `json:"time"`
|
||||
Model string `json:"model"`
|
||||
} `json:"records"`
|
||||
}
|
||||
if err := json.Unmarshal(rr2.Body.Bytes(), &pg); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(st.Records) != len(pg.Records) {
|
||||
t.Fatalf("period view shipped %d records, pager returned %d — the two "+
|
||||
"sources must agree or scrolling duplicates/skips rows",
|
||||
len(st.Records), len(pg.Records))
|
||||
}
|
||||
for i := range st.Records {
|
||||
if st.Records[i].Time != pg.Records[i].Time {
|
||||
t.Fatalf("record %d differs: period view t=%d, pager t=%d",
|
||||
i, st.Records[i].Time, pg.Records[i].Time)
|
||||
}
|
||||
}
|
||||
// The two sources must agree on DIRECTION as well as content. Which
|
||||
// direction that is depends on how many reverse-read chunks the audit file
|
||||
// spans (a small file is read whole, a large one chunk-by-chunk), so the
|
||||
// test asserts agreement rather than hard-coding newest-first — but they
|
||||
// can never disagree, because the table interleaves both streams.
|
||||
if len(st.Records) > 1 {
|
||||
periodAsc := st.Records[1].Time > st.Records[0].Time
|
||||
pagerAsc := pg.Records[1].Time > pg.Records[0].Time
|
||||
if periodAsc != pagerAsc {
|
||||
t.Fatalf("direction differs: period view %s, pager %s — the table "+
|
||||
"would splice the two streams out of order",
|
||||
map[bool]string{true: "oldest-first", false: "newest-first"}[periodAsc],
|
||||
map[bool]string{true: "oldest-first", false: "newest-first"}[pagerAsc])
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// A non-admin caller cannot widen their view by passing someone else's key in
|
||||
// the query string: exportKey overrides ?key= with the caller's own masked id.
|
||||
// This is a data-leak boundary, and it applies to the records the period view
|
||||
// ships just as much as to the aggregates.
|
||||
func TestStatsPeriodRecordsIgnoreForeignKeyForUsers(t *testing.T) {
|
||||
td := t.TempDir()
|
||||
g := newTestGateway(t)
|
||||
g.stats.auditPath = filepath.Join(td, "audit.jsonl")
|
||||
|
||||
other := keyID("sk-gw-cccccccccccccccc")
|
||||
own := keyID("sk-test")
|
||||
now := time.Now().UnixMilli()
|
||||
g.stats.Record(Req{Time: now, OK: true, Model: "m", Status: 200, Key: other})
|
||||
g.stats.Record(Req{Time: now - 1000, OK: true, Model: "m", Status: 200, Key: own})
|
||||
|
||||
// sk-test authenticates as a non-admin (newTestGateway seeds it under
|
||||
// gateway_keys, not keys), so ?key=<other> must be ignored and the caller
|
||||
// sees only its own rows.
|
||||
rr := doReq(t, g, "GET",
|
||||
"/api/stats?period=day&key="+url.QueryEscape(other), "")
|
||||
if rr.Code != 200 {
|
||||
t.Fatalf("HTTP %d: %s", rr.Code, rr.Body.String())
|
||||
}
|
||||
var st struct {
|
||||
Records []struct {
|
||||
Key string `json:"key"`
|
||||
} `json:"records"`
|
||||
}
|
||||
if err := json.Unmarshal(rr.Body.Bytes(), &st); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
for _, r := range st.Records {
|
||||
if r.Key == other {
|
||||
t.Fatalf("another key's record leaked into a user view "+
|
||||
"(caller=%s, requested=%s)", own, other)
|
||||
}
|
||||
}
|
||||
}
|
||||
434
internal/gateway/stats_period_test.go
Normal file
434
internal/gateway/stats_period_test.go
Normal file
@ -0,0 +1,434 @@
|
||||
package gateway
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"testing"
|
||||
"time"
|
||||
)
|
||||
|
||||
// auditPeriod writes records straight into an audit file, which is what the
|
||||
// period aggregator reads. Going through Record() would also rotate/prune the
|
||||
// file; here the file shape is the thing under test.
|
||||
func auditPeriod(t *testing.T, path string, recs ...Req) {
|
||||
t.Helper()
|
||||
f, err := os.Create(path)
|
||||
if err != nil {
|
||||
t.Fatalf("create audit file: %v", err)
|
||||
}
|
||||
defer f.Close()
|
||||
for _, r := range recs {
|
||||
if r.Time == 0 {
|
||||
t.Fatal("test record needs an explicit timestamp")
|
||||
}
|
||||
if r.Type == "" {
|
||||
r.Type = "chat"
|
||||
}
|
||||
b, err := json.Marshal(r)
|
||||
if err != nil {
|
||||
t.Fatalf("marshal record: %v", err)
|
||||
}
|
||||
if _, err := f.Write(append(b, '\n')); err != nil {
|
||||
t.Fatalf("write record: %v", err)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func periodStats(t *testing.T, auditPath string) *Stats {
|
||||
t.Helper()
|
||||
s := NewStats(10)
|
||||
s.mu.Lock()
|
||||
s.auditPath = auditPath
|
||||
s.mu.Unlock()
|
||||
return s
|
||||
}
|
||||
|
||||
// TestPeriodWindowBoundaries is the boundary contract: a day window covers
|
||||
// exactly today-from-midnight, a week window starts Monday (not Sunday — the
|
||||
// Go Weekday() trap), and a month window starts on the 1st. All UTC.
|
||||
func TestPeriodWindowBoundaries(t *testing.T) {
|
||||
// Wednesday 2026-03-11 15:30 UTC.
|
||||
now := time.Date(2026, 3, 11, 15, 30, 0, 0, time.UTC)
|
||||
day := time.Date(2026, 3, 11, 0, 0, 0, 0, time.UTC)
|
||||
if got := time.UnixMilli(periodStart(PeriodDay, now)).UTC(); !got.Equal(day) {
|
||||
t.Errorf("day window starts %s, want %s", got, day)
|
||||
}
|
||||
// Wednesday minus 2 days = Monday the 9th.
|
||||
monday := time.Date(2026, 3, 9, 0, 0, 0, 0, time.UTC)
|
||||
if got := time.UnixMilli(periodStart(PeriodWeek, now)).UTC(); !got.Equal(monday) {
|
||||
t.Errorf("week window starts %s, want Monday %s", got, monday)
|
||||
}
|
||||
// Same instant on a Sunday must still start on the PRECEDING Monday,
|
||||
// never on the Sunday itself. This is the off-by-one that
|
||||
// int(now.Weekday()) would introduce (Sunday=0 → no shift).
|
||||
sun := time.Date(2026, 3, 15, 10, 0, 0, 0, time.UTC) // Sunday
|
||||
prevMon := time.Date(2026, 3, 9, 0, 0, 0, 0, time.UTC)
|
||||
if got := time.UnixMilli(periodStart(PeriodWeek, sun)).UTC(); !got.Equal(prevMon) {
|
||||
t.Errorf("Sunday week window starts %s, want the Monday before it %s", got, prevMon)
|
||||
}
|
||||
first := time.Date(2026, 3, 1, 0, 0, 0, 0, time.UTC)
|
||||
if got := time.UnixMilli(periodStart(PeriodMonth, now)).UTC(); !got.Equal(first) {
|
||||
t.Errorf("month window starts %s, want %s", got, first)
|
||||
}
|
||||
if got := periodStart(PeriodAll, now); got != 0 {
|
||||
t.Errorf("all window from = %d, want 0 (unbounded)", got)
|
||||
}
|
||||
}
|
||||
|
||||
// TestPeriodAggregationExcludesOlderRecords is the reason the aggregator reads
|
||||
// the audit file: a record outside the window must not be counted, even though
|
||||
// it sits in the very same file, right next to records that are.
|
||||
func TestPeriodAggregationExcludesOlderRecords(t *testing.T) {
|
||||
dir := t.TempDir()
|
||||
path := filepath.Join(dir, "audit.jsonl")
|
||||
now := time.Date(2026, 3, 11, 15, 0, 0, 0, time.UTC)
|
||||
h := func(n int) int64 { return now.Add(-time.Duration(n) * time.Hour).UnixMilli() }
|
||||
|
||||
auditPeriod(t, path,
|
||||
// inside today
|
||||
Req{Time: h(1), Model: "m", Source: "s", Prompt: 100, Compl: 10, LatMs: 5, OK: true},
|
||||
Req{Time: h(2), Model: "m", Source: "s", Prompt: 200, Compl: 20, LatMs: 7, OK: true},
|
||||
// outside today (yesterday) but inside the week
|
||||
Req{Time: now.AddDate(0, 0, -1).UnixMilli(), Model: "old", Source: "s", Prompt: 999, Compl: 99, OK: true},
|
||||
)
|
||||
s := periodStats(t, path)
|
||||
|
||||
day := s.PeriodSnapshot(PeriodDay, "", now)
|
||||
if day.Total.Reqs != 2 {
|
||||
t.Errorf("day window counted %d requests, want 2 (yesterday's record must be excluded)", day.Total.Reqs)
|
||||
}
|
||||
if day.Total.Prompt != 300 {
|
||||
t.Errorf("day prompt tokens = %d, want 300", day.Total.Prompt)
|
||||
}
|
||||
|
||||
week := s.PeriodSnapshot(PeriodWeek, "", now)
|
||||
if week.Total.Reqs != 3 {
|
||||
t.Errorf("week window counted %d requests, want 3", week.Total.Reqs)
|
||||
}
|
||||
if week.Total.Prompt != 1299 {
|
||||
t.Errorf("week prompt tokens = %d, want 1299", week.Total.Prompt)
|
||||
}
|
||||
}
|
||||
|
||||
// TestPeriodDimensionsAggregateIndependently checks the by_* rows actually
|
||||
// split by their own key rather than all collapsing into one row.
|
||||
func TestPeriodDimensionsAggregateIndependently(t *testing.T) {
|
||||
dir := t.TempDir()
|
||||
path := filepath.Join(dir, "audit.jsonl")
|
||||
now := time.Date(2026, 3, 11, 10, 0, 0, 0, time.UTC)
|
||||
auditPeriod(t, path,
|
||||
Req{Time: now.UnixMilli(), Key: "k1", Model: "alpha", Source: "s1", Prompt: 10, OK: true},
|
||||
Req{Time: now.UnixMilli(), Key: "k1", Model: "beta", Source: "s2", Prompt: 20, OK: true},
|
||||
Req{Time: now.UnixMilli(), Key: "k2", Model: "alpha", Source: "s1", Prompt: 30, OK: true},
|
||||
)
|
||||
s := periodStats(t, path)
|
||||
got := s.PeriodSnapshot(PeriodDay, "", now)
|
||||
|
||||
if len(got.Models) != 2 {
|
||||
t.Fatalf("by_model has %d rows, want 2: %+v", len(got.Models), got.Models)
|
||||
}
|
||||
// alpha = 10 + 30
|
||||
for _, r := range got.Models {
|
||||
if r.Name == "alpha" && r.Prompt != 40 {
|
||||
t.Errorf("alpha prompt = %d, want 40", r.Prompt)
|
||||
}
|
||||
}
|
||||
if len(got.Srcs) != 2 {
|
||||
t.Errorf("by_source has %d rows, want 2: %+v", len(got.Srcs), got.Srcs)
|
||||
}
|
||||
if len(got.ByKey) != 2 {
|
||||
t.Errorf("by_key has %d rows, want 2: %+v", len(got.ByKey), got.ByKey)
|
||||
}
|
||||
// The rows must add up to the total, or the dashboard shows a total that
|
||||
// disagrees with its own table.
|
||||
var sum int64
|
||||
for _, r := range got.Models {
|
||||
sum += r.Prompt
|
||||
}
|
||||
if sum != got.Total.Prompt {
|
||||
t.Errorf("by_model prompts sum to %d but total is %d — the table would contradict the KPI", sum, got.Total.Prompt)
|
||||
}
|
||||
}
|
||||
|
||||
// TestPeriodDayBucketsAreHourlyAndOrdered pins the timeline grain and its
|
||||
// order. Day = hourly (a spike must be visible inside "today"); week/month =
|
||||
// daily. Sorted chronologically, because a chart fed unsorted buckets draws
|
||||
// nonsense.
|
||||
func TestPeriodDayBucketsAreHourlyAndOrdered(t *testing.T) {
|
||||
dir := t.TempDir()
|
||||
path := filepath.Join(dir, "audit.jsonl")
|
||||
now := time.Date(2026, 3, 11, 15, 0, 0, 0, time.UTC)
|
||||
// Deliberately out of order in the file.
|
||||
auditPeriod(t, path,
|
||||
Req{Time: now.Add(-2 * time.Hour).UnixMilli(), Model: "m", Prompt: 1, OK: true},
|
||||
Req{Time: now.Add(-4 * time.Hour).UnixMilli(), Model: "m", Prompt: 1, OK: true},
|
||||
Req{Time: now.UnixMilli(), Model: "m", Prompt: 1, OK: true},
|
||||
)
|
||||
s := periodStats(t, path)
|
||||
got := s.PeriodSnapshot(PeriodDay, "", now)
|
||||
if len(got.Bucket) != 3 {
|
||||
t.Fatalf("day buckets = %d, want 3 (hourly): %+v", len(got.Bucket), got.Bucket)
|
||||
}
|
||||
for i := 1; i < len(got.Bucket); i++ {
|
||||
if got.Bucket[i-1].Bucket >= got.Bucket[i].Bucket {
|
||||
t.Errorf("day buckets not chronological at %d: %q >= %q",
|
||||
i, got.Bucket[i-1].Bucket, got.Bucket[i].Bucket)
|
||||
}
|
||||
}
|
||||
// Weekly grain is daily, so the same three records collapse to one day.
|
||||
week := s.PeriodSnapshot(PeriodWeek, "", now)
|
||||
if len(week.Bucket) != 1 {
|
||||
t.Errorf("week buckets = %d, want 1 (daily grain): %+v", len(week.Bucket), week.Bucket)
|
||||
}
|
||||
}
|
||||
|
||||
// TestPeriodKeyFilterIsolatesKeys guards the multi-tenant boundary: asking for
|
||||
// one key must never return another key's rows in any dimension.
|
||||
func TestPeriodKeyFilterIsolatesKeys(t *testing.T) {
|
||||
dir := t.TempDir()
|
||||
path := filepath.Join(dir, "audit.jsonl")
|
||||
now := time.Date(2026, 3, 11, 10, 0, 0, 0, time.UTC)
|
||||
auditPeriod(t, path,
|
||||
Req{Time: now.UnixMilli(), Key: "alice", Model: "m", Source: "s", Prompt: 10, OK: true},
|
||||
Req{Time: now.UnixMilli(), Key: "bob", Model: "m", Source: "s", Prompt: 999, OK: true},
|
||||
)
|
||||
s := periodStats(t, path)
|
||||
got := s.PeriodSnapshot(PeriodDay, "alice", now)
|
||||
if got.Total.Reqs != 1 || got.Total.Prompt != 10 {
|
||||
t.Errorf("alice's window = %d reqs / %d prompt, want 1 / 10 — another key leaked in",
|
||||
got.Total.Reqs, got.Total.Prompt)
|
||||
}
|
||||
if len(got.Srcs) != 1 || got.Srcs[0].Prompt != 10 {
|
||||
t.Errorf("by_source leaked: %+v", got.Srcs)
|
||||
}
|
||||
}
|
||||
|
||||
// TestPeriodAllUsesLifetimeAggregates checks "all" still answers from the
|
||||
// in-memory aggregates (no audit file needed) and carries the same rows the
|
||||
// dashboard already renders.
|
||||
func TestPeriodAllUsesLifetimeAggregates(t *testing.T) {
|
||||
s := NewStats(100)
|
||||
now := time.Date(2026, 3, 11, 10, 0, 0, 0, time.UTC)
|
||||
s.Record(Req{Time: now.AddDate(0, 0, -40).UnixMilli(), Key: "k", Model: "m", Source: "s", Prompt: 7, OK: true})
|
||||
s.Record(Req{Time: now.UnixMilli(), Key: "k", Model: "m", Source: "s", Prompt: 3, OK: true})
|
||||
|
||||
got := s.PeriodSnapshot(PeriodAll, "", now)
|
||||
if got.Total.Reqs != 2 || got.Total.Prompt != 10 {
|
||||
t.Errorf("all view = %d reqs / %d prompt, want 2 / 10 (lifetime)", got.Total.Reqs, got.Total.Prompt)
|
||||
}
|
||||
if len(got.Models) != 1 || got.Models[0].Prompt != 10 {
|
||||
t.Errorf("all view by_model = %+v, want one row with 10 prompt", got.Models)
|
||||
}
|
||||
// A lifetime total has no timeline to draw; buckets must stay empty
|
||||
// rather than inventing 500 hourly points.
|
||||
if len(got.Bucket) != 0 {
|
||||
t.Errorf("all view produced %d buckets, want 0", len(got.Bucket))
|
||||
}
|
||||
}
|
||||
|
||||
// TestPeriodEmptyWindowFlagsTruncated covers the "rotated away" case: a bounded
|
||||
// window with nothing to show may mean "quiet day" or "the files are gone".
|
||||
// The aggregator must not report a confidently wrong zero.
|
||||
func TestPeriodEmptyWindowFlagsTruncated(t *testing.T) {
|
||||
dir := t.TempDir()
|
||||
s := periodStats(t, filepath.Join(dir, "audit.jsonl")) // file does not exist
|
||||
now := time.Date(2026, 3, 11, 10, 0, 0, 0, time.UTC)
|
||||
got := s.PeriodSnapshot(PeriodDay, "", now)
|
||||
if got.Total.Reqs != 0 {
|
||||
t.Errorf("empty window reqs = %d, want 0", got.Total.Reqs)
|
||||
}
|
||||
if !got.Truncated {
|
||||
t.Error("★ an empty bounded window must be flagged truncated — otherwise a " +
|
||||
"rotated-away week reads as 'no traffic that week'")
|
||||
}
|
||||
}
|
||||
|
||||
// TestValidPeriodRejectsUnknown guards the deliberate strictness: a typo must
|
||||
// be an error, not a silent fallback to lifetime totals.
|
||||
func TestValidPeriodRejectsUnknown(t *testing.T) {
|
||||
for _, ok := range []Period{PeriodDay, PeriodWeek, PeriodMonth, PeriodAll} {
|
||||
if !ValidPeriod(ok) {
|
||||
t.Errorf("ValidPeriod(%q) = false, want true", ok)
|
||||
}
|
||||
}
|
||||
for _, bad := range []Period{"", "year", "hour", "today", "DAY "} {
|
||||
if ValidPeriod(bad) {
|
||||
t.Errorf("ValidPeriod(%q) = true, want false", bad)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// TestPeriodStatusDimensionCountsFailures covers by_status, which nothing else
|
||||
// exercised — an aggregation branch with no test is exactly how a period view
|
||||
// silently loses the failure pie. M8 (disabling the status dimension) is
|
||||
// caught only because this test exists.
|
||||
func TestPeriodStatusDimensionCountsFailures(t *testing.T) {
|
||||
dir := t.TempDir()
|
||||
path := filepath.Join(dir, "audit.jsonl")
|
||||
now := time.Date(2026, 3, 11, 10, 0, 0, 0, time.UTC)
|
||||
auditPeriod(t, path,
|
||||
Req{Time: now.UnixMilli(), Model: "m", Source: "s", Status: 200, OK: true, Prompt: 5},
|
||||
Req{Time: now.UnixMilli(), Model: "m", Source: "s", Status: 502, OK: false, Prompt: 1},
|
||||
Req{Time: now.UnixMilli(), Model: "m", Source: "s", Status: 502, OK: false, Prompt: 1},
|
||||
Req{Time: now.UnixMilli(), Model: "m", Source: "s", Status: 429, OK: false, Prompt: 1},
|
||||
)
|
||||
s := periodStats(t, path)
|
||||
got := s.PeriodSnapshot(PeriodDay, "", now)
|
||||
|
||||
byCode := map[string]int64{}
|
||||
for _, r := range got.Status {
|
||||
byCode[r.Name] = r.Reqs
|
||||
}
|
||||
if len(got.Status) != 3 {
|
||||
t.Fatalf("by_status has %d rows, want 3 (200/502/429): %+v", len(got.Status), got.Status)
|
||||
}
|
||||
if byCode["200"] != 1 || byCode["502"] != 2 || byCode["429"] != 1 {
|
||||
t.Errorf("by_status counts = %v, want 200:1 502:2 429:1", byCode)
|
||||
}
|
||||
// Status codes must sort numerically: "429" before "502" lexically would
|
||||
// render the failure legend in the wrong order.
|
||||
if got.Status[0].Name != "200" || got.Status[1].Name != "429" || got.Status[2].Name != "502" {
|
||||
t.Errorf("by_status not numerically sorted: %s, %s, %s",
|
||||
got.Status[0].Name, got.Status[1].Name, got.Status[2].Name)
|
||||
}
|
||||
// The failure rows must agree with the OK/Err split on the total, or the
|
||||
// pie and the KPI row contradict each other.
|
||||
if got.Total.OK != 1 || got.Total.Err != 3 {
|
||||
t.Errorf("total ok/err = %d/%d, want 1/3", got.Total.OK, got.Total.Err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestPeriodFromQuery(t *testing.T) {
|
||||
cases := []struct {
|
||||
query string
|
||||
want Period
|
||||
}{
|
||||
{"", PeriodAll}, // no selector → unchanged legacy behaviour
|
||||
{"?period=day", PeriodDay},
|
||||
{"?period=WEEK", PeriodWeek},
|
||||
{"?period=month", PeriodMonth},
|
||||
{"?period=all", PeriodAll},
|
||||
}
|
||||
for _, c := range cases {
|
||||
r := httptest.NewRequest(http.MethodGet, "/api/stats"+c.query, nil)
|
||||
if got := periodFromQuery(r.URL.Query()); got != c.want {
|
||||
t.Errorf("periodFromQuery(%q) = %q, want %q", c.query, got, c.want)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// TestPeriodWeekMayStartBeforeMonth pins a fact that looks like a bug and
|
||||
// gets reported as one: the calendar periods are NOT nested.
|
||||
//
|
||||
// On 2026-10-03 (a Saturday) the ISO week runs Mon 2026-09-28 .. Sat
|
||||
// 10-03, while the month window is 10-01 .. 10-03. The week therefore reaches
|
||||
// back into the previous month and its total is legitimately LARGER than the
|
||||
// month's. Anyone reading only the totals will conclude the aggregator is
|
||||
// broken; it is not, and this test exists so a future "fix" that makes the
|
||||
// windows nested (e.g. clamping the week to the month) is caught as the
|
||||
// regression it would be.
|
||||
func TestPeriodWeekMayStartBeforeMonth(t *testing.T) {
|
||||
// Saturday 2026-10-03, UTC.
|
||||
now := time.Date(2026, 10, 3, 12, 0, 0, 0, time.UTC)
|
||||
|
||||
weekFrom := periodStart(PeriodWeek, now)
|
||||
monthFrom := periodStart(PeriodMonth, now)
|
||||
|
||||
weekFromD := time.UnixMilli(weekFrom).UTC()
|
||||
monthFromD := time.UnixMilli(monthFrom).UTC()
|
||||
|
||||
if got := weekFromD.Format("2006-01-02"); got != "2026-09-28" {
|
||||
t.Fatalf("week window start = %s, want 2026-09-28 (the Monday)", got)
|
||||
}
|
||||
if got := monthFromD.Format("2006-01-02"); got != "2026-10-01" {
|
||||
t.Fatalf("month window start = %s, want 2026-10-01", got)
|
||||
}
|
||||
if !weekFromD.Before(monthFromD) {
|
||||
t.Fatal("this test only means anything when the week starts before the month")
|
||||
}
|
||||
// The week must reach strictly further back than the month window does,
|
||||
// otherwise the nesting intuition would hold and the UI note would be noise.
|
||||
if back := monthFromD.Sub(weekFromD).Hours() / 24; back <= 0 {
|
||||
t.Fatalf("week should start at least a day before the month window, got %v days back", back)
|
||||
}
|
||||
}
|
||||
|
||||
// The same overlap must be visible in the aggregate numbers, not just in the
|
||||
// window bounds: a week total larger than a month total is the exact symptom
|
||||
// an operator reported.
|
||||
func TestPeriodWeekTotalCanExceedMonthTotal(t *testing.T) {
|
||||
dir := t.TempDir()
|
||||
audit := filepath.Join(dir, "audit.jsonl")
|
||||
|
||||
// Two records before the month window (inside this ISO week) and one
|
||||
// inside October.
|
||||
pre := time.Date(2026, 9, 29, 10, 0, 0, 0, time.UTC).UnixMilli()
|
||||
in1 := time.Date(2026, 10, 1, 10, 0, 0, 0, time.UTC).UnixMilli()
|
||||
in2 := time.Date(2026, 10, 3, 10, 0, 0, 0, time.UTC).UnixMilli()
|
||||
|
||||
auditPeriod(t, audit,
|
||||
Req{Time: pre, OK: true, Prompt: 100, Compl: 10},
|
||||
Req{Time: pre + 1000, OK: true, Prompt: 100, Compl: 10},
|
||||
Req{Time: in1, OK: true, Prompt: 50, Compl: 5},
|
||||
Req{Time: in2, OK: true, Prompt: 50, Compl: 5},
|
||||
)
|
||||
|
||||
s := periodStats(t, audit)
|
||||
now := time.Date(2026, 10, 3, 12, 0, 0, 0, time.UTC)
|
||||
|
||||
week := s.PeriodSnapshot(PeriodWeek, "", now)
|
||||
month := s.PeriodSnapshot(PeriodMonth, "", now)
|
||||
|
||||
if week.Total.Reqs != 4 {
|
||||
t.Fatalf("week reqs = %d, want 4 (both September records fall in this ISO week)", week.Total.Reqs)
|
||||
}
|
||||
if month.Total.Reqs != 2 {
|
||||
t.Fatalf("month reqs = %d, want 2", month.Total.Reqs)
|
||||
}
|
||||
if week.Total.Tokens <= month.Total.Tokens {
|
||||
t.Fatalf("week tokens (%d) must exceed month tokens (%d): the week reaches into September",
|
||||
week.Total.Tokens, month.Total.Tokens)
|
||||
}
|
||||
// Day must be a strict subset of both, which is the only nesting that holds.
|
||||
day := s.PeriodSnapshot(PeriodDay, "", now)
|
||||
if day.Total.Reqs != 1 {
|
||||
t.Fatalf("day reqs = %d, want 1", day.Total.Reqs)
|
||||
}
|
||||
if day.Total.Tokens >= month.Total.Tokens || day.Total.Tokens >= week.Total.Tokens {
|
||||
t.Fatal("today must be contained in both the week and the month totals")
|
||||
}
|
||||
}
|
||||
|
||||
// TestPeriodStatsCarryActiveGauge pins the live "active requests" gauge onto
|
||||
// every period response.
|
||||
//
|
||||
// It is a live gauge, not an aggregate, so no calendar window can scope it — but
|
||||
// the period response is what the day/week/month views render from, and it used
|
||||
// to omit the field entirely. The UI then evaluated `st.active || 0` and printed
|
||||
// a hard 0 in three of four views, while the all-time view (whose response did
|
||||
// carry the field) showed the truth. Reading "0 active" off a busy gateway is
|
||||
// worse than reading nothing: it is a confident wrong answer.
|
||||
func TestPeriodStatsCarryActiveGauge(t *testing.T) {
|
||||
for _, period := range []string{"day", "week", "month", "all"} {
|
||||
g := newTestGateway(t)
|
||||
rr := doReq(t, g, http.MethodGet, "/api/stats?period="+period, "")
|
||||
if rr.Code != http.StatusOK {
|
||||
t.Fatalf("period=%s: %d", period, rr.Code)
|
||||
}
|
||||
var out map[string]interface{}
|
||||
if err := json.Unmarshal(rr.Body.Bytes(), &out); err != nil {
|
||||
t.Fatalf("period=%s decode: %v", period, err)
|
||||
}
|
||||
v, ok := out["active"]
|
||||
if !ok {
|
||||
t.Errorf("period=%s response has no active field; the UI renders 0", period)
|
||||
continue
|
||||
}
|
||||
if _, ok := v.(float64); !ok {
|
||||
t.Errorf("period=%s active is %T, want a number", period, v)
|
||||
}
|
||||
}
|
||||
}
|
||||
231
internal/gateway/token_units_test.go
Normal file
231
internal/gateway/token_units_test.go
Normal file
@ -0,0 +1,231 @@
|
||||
package gateway
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"llmsproxy/internal/config"
|
||||
)
|
||||
|
||||
// Token accounting must record TOKENS, not byte estimates and not image
|
||||
// counts. These tests pin the two units bugs found on 2026-09-27:
|
||||
//
|
||||
// 1. streamed chat: rec.Compl was only ever the per-chunk byte estimate
|
||||
// (len(content)/3), even when the upstream sent exact usage on its final
|
||||
// frame. Measured on production traffic: the same model reported a
|
||||
// streamed completion median 4-27x its non-streamed median. Upstream 100
|
||||
// recorded as 145.
|
||||
// 2. image generation: rec.Compl was set to len(resp.ImageData) — the image
|
||||
// COUNT — and folded into the token totals.
|
||||
//
|
||||
// A fake upstream that reports FIXED numbers makes both mismatches
|
||||
// unambiguous: any difference between what it reports and what the audit
|
||||
// record holds is the gateway's doing.
|
||||
|
||||
// fixedUsageUpstream reports prompt=prompt completion=compl on both the
|
||||
// streaming and the non-streaming path. The streaming path emits usage on its
|
||||
// final frame before [DONE], exactly like OpenAI-compatible upstreams.
|
||||
func fixedUsageUpstream(t *testing.T, prompt, compl int, images int) *httptest.Server {
|
||||
t.Helper()
|
||||
return httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
var body map[string]interface{}
|
||||
_ = json.NewDecoder(r.Body).Decode(&body)
|
||||
stream, _ := body["stream"].(bool)
|
||||
usage := fmt.Sprintf(`{"prompt_tokens":%d,"completion_tokens":%d,"total_tokens":%d}`,
|
||||
prompt, compl, prompt+compl)
|
||||
|
||||
// Image generation wire: /v1/images/generations style body.
|
||||
if strings.Contains(r.URL.Path, "image") {
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
var data []string
|
||||
for i := 0; i < images; i++ {
|
||||
data = append(data, `{"b64_json":"AAAA"}`)
|
||||
}
|
||||
fmt.Fprintf(w, `{"created":1,"data":[%s]}`, strings.Join(data, ","))
|
||||
return
|
||||
}
|
||||
|
||||
if !stream {
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
fmt.Fprintf(w, `{"model":"m1","choices":[{"message":{"content":"pong pong pong"},
|
||||
"finish_reason":"stop"}],"usage":%s}`, usage)
|
||||
return
|
||||
}
|
||||
|
||||
w.Header().Set("Content-Type", "text/event-stream")
|
||||
w.WriteHeader(200)
|
||||
fl, _ := w.(http.Flusher)
|
||||
send := func(s string) {
|
||||
fmt.Fprintf(w, "data: %s\n\n", s)
|
||||
if fl != nil {
|
||||
fl.Flush()
|
||||
}
|
||||
}
|
||||
// Deliberately verbose payloads: a byte-estimating path records a
|
||||
// number far from `compl`, so the test would catch it.
|
||||
send(`{"model":"m1","choices":[{"delta":{"reasoning_content":"` + strings.Repeat("think ", 40) + `"}}]}`)
|
||||
send(`{"model":"m1","choices":[{"delta":{"content":"` + strings.Repeat("answer ", 40) + `"}}]}`)
|
||||
send(`{"model":"m1","choices":[{"delta":{},"finish_reason":"stop"}]}`)
|
||||
send(`{"model":"m1","choices":[],"usage":` + usage + `}`)
|
||||
send(`[DONE]`)
|
||||
}))
|
||||
}
|
||||
|
||||
// lastChatRecord returns the most recent audit record whose Type is one of the
|
||||
// given types.
|
||||
func lastChatRecord(t *testing.T, g *Gateway, types ...string) Req {
|
||||
t.Helper()
|
||||
recs := g.stats.AuditRecords(0, 0, "")
|
||||
want := map[string]bool{}
|
||||
for _, ty := range types {
|
||||
want[ty] = true
|
||||
}
|
||||
for i := len(recs) - 1; i >= 0; i-- {
|
||||
if want[recs[i].Type] {
|
||||
return recs[i]
|
||||
}
|
||||
}
|
||||
t.Fatalf("no audit record of type %v found (have %d records)", types, len(recs))
|
||||
return Req{}
|
||||
}
|
||||
|
||||
// The streamed path must record the upstream's exact usage, not a byte
|
||||
// estimate. Before the fix this recorded 145 for an upstream-reported 100.
|
||||
func TestStreamedChatRecordsUpstreamUsage(t *testing.T) {
|
||||
const wantPrompt, wantCompl = 44, 100
|
||||
up := fixedUsageUpstream(t, wantPrompt, wantCompl, 0)
|
||||
defer up.Close()
|
||||
g := newTestGateway(t, config.Source{
|
||||
Name: "up", BaseURL: up.URL, Adapter: "openai",
|
||||
Models: []config.Model{{ID: "m1", Priority: 100}},
|
||||
})
|
||||
|
||||
rr := doReq(t, g, "POST", "/v1/chat/completions",
|
||||
`{"model":"m1","messages":[{"role":"user","content":"hi"}],"stream":true}`)
|
||||
if rr.Code != 200 {
|
||||
t.Fatalf("stream status=%d body=%s", rr.Code, rr.Body.String())
|
||||
}
|
||||
|
||||
rec := lastChatRecord(t, g, "stream")
|
||||
if rec.Prompt != wantPrompt || rec.Compl != wantCompl {
|
||||
t.Errorf("streamed record prompt=%d completion=%d, want %d/%d — the audit "+
|
||||
"must carry the upstream's tokens, not a byte estimate",
|
||||
rec.Prompt, rec.Compl, wantPrompt, wantCompl)
|
||||
}
|
||||
}
|
||||
|
||||
// Streaming and non-streaming must agree on the same request's tokens. This is
|
||||
// the cross-path consistency check: the two paths previously recorded 145 and
|
||||
// 100 for one identical upstream usage value.
|
||||
func TestStreamedAndNonStreamedAgreeOnTokens(t *testing.T) {
|
||||
const wantPrompt, wantCompl = 44, 100
|
||||
up := fixedUsageUpstream(t, wantPrompt, wantCompl, 0)
|
||||
defer up.Close()
|
||||
g := newTestGateway(t, config.Source{
|
||||
Name: "up", BaseURL: up.URL, Adapter: "openai",
|
||||
Models: []config.Model{{ID: "m1", Priority: 100}},
|
||||
})
|
||||
|
||||
if rr := doReq(t, g, "POST", "/v1/chat/completions",
|
||||
`{"model":"m1","messages":[{"role":"user","content":"hi"}]}`); rr.Code != 200 {
|
||||
t.Fatalf("non-stream status=%d body=%s", rr.Code, rr.Body.String())
|
||||
}
|
||||
if rr := doReq(t, g, "POST", "/v1/chat/completions",
|
||||
`{"model":"m1","messages":[{"role":"user","content":"hi"}],"stream":true}`); rr.Code != 200 {
|
||||
t.Fatalf("stream status=%d body=%s", rr.Code, rr.Body.String())
|
||||
}
|
||||
|
||||
chat, stream := lastChatRecord(t, g, "chat"), lastChatRecord(t, g, "stream")
|
||||
if chat.Prompt != stream.Prompt || chat.Compl != stream.Compl {
|
||||
t.Errorf("chat=%d/%d stream=%d/%d — identical upstream usage must record "+
|
||||
"identically on both paths", chat.Prompt, chat.Compl, stream.Prompt, stream.Compl)
|
||||
}
|
||||
}
|
||||
|
||||
// A fallback to the byte estimate is still correct WHEN the upstream reports no
|
||||
// usage at all. Without this the fix could silently zero out every streamed
|
||||
// record from upstreams that omit usage.
|
||||
func TestStreamedChatFallsBackToEstimateWithoutUsage(t *testing.T) {
|
||||
up := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
w.Header().Set("Content-Type", "text/event-stream")
|
||||
w.WriteHeader(200)
|
||||
fl, _ := w.(http.Flusher)
|
||||
send := func(s string) {
|
||||
fmt.Fprintf(w, "data: %s\n\n", s)
|
||||
if fl != nil {
|
||||
fl.Flush()
|
||||
}
|
||||
}
|
||||
send(`{"model":"m1","choices":[{"delta":{"content":"` + strings.Repeat("abc ", 60) + `"}}]}`)
|
||||
send(`{"model":"m1","choices":[{"delta":{},"finish_reason":"stop"}]}`)
|
||||
send(`[DONE]`)
|
||||
}))
|
||||
defer up.Close()
|
||||
g := newTestGateway(t, config.Source{
|
||||
Name: "up", BaseURL: up.URL, Adapter: "openai",
|
||||
Models: []config.Model{{ID: "m1", Priority: 100}},
|
||||
})
|
||||
|
||||
if rr := doReq(t, g, "POST", "/v1/chat/completions",
|
||||
`{"model":"m1","messages":[{"role":"user","content":"hi"}],"stream":true}`); rr.Code != 200 {
|
||||
t.Fatalf("status=%d body=%s", rr.Code, rr.Body.String())
|
||||
}
|
||||
rec := lastChatRecord(t, g, "stream")
|
||||
if rec.Compl <= 0 {
|
||||
t.Errorf("completion=%d, want a positive byte-estimate fallback when the "+
|
||||
"upstream reports no usage", rec.Compl)
|
||||
}
|
||||
}
|
||||
|
||||
// Image generation has no token concept: recording the image COUNT as
|
||||
// completion tokens corrupted the token totals. The count is reported
|
||||
// separately and must not leak into Prompt/Compl.
|
||||
func TestImageGenerationDoesNotRecordTokens(t *testing.T) {
|
||||
up := fixedUsageUpstream(t, 44, 100, 3)
|
||||
defer up.Close()
|
||||
src := config.Source{
|
||||
Name: "up", BaseURL: up.URL, Adapter: "openai",
|
||||
Models: []config.Model{{ID: "img-m", Priority: 100, Kind: "image"}},
|
||||
}
|
||||
g := newTestGateway(t, src)
|
||||
|
||||
rr := doReq(t, g, "POST", "/v1/images/generations",
|
||||
`{"model":"img-m","prompt":"a cat","n":3}`)
|
||||
if rr.Code != 200 {
|
||||
t.Fatalf("image status=%d body=%s", rr.Code, rr.Body.String())
|
||||
}
|
||||
|
||||
rec := lastChatRecord(t, g, "image")
|
||||
if rec.Compl != 0 || rec.Prompt != 0 {
|
||||
t.Errorf("image record prompt=%d completion=%d, want 0/0 — image count is "+
|
||||
"not a token and must not enter the token totals", rec.Prompt, rec.Compl)
|
||||
}
|
||||
if rec.ImageCount != 3 {
|
||||
t.Errorf("image_count=%d, want 3 (the count belongs in its own field)", rec.ImageCount)
|
||||
}
|
||||
// And the aggregate must see zero tokens for it.
|
||||
if got := g.stats.WindowTokens("img-m", "", 3600); got != 0 {
|
||||
t.Errorf("window tokens for an image model = %d, want 0", got)
|
||||
}
|
||||
}
|
||||
|
||||
// The audit JSON must carry image_count so the WebUI can show images rather
|
||||
// than a bogus token figure.
|
||||
func TestImageCountSurvivesAuditJSON(t *testing.T) {
|
||||
b, err := json.Marshal(Req{Type: "image", ImageCount: 2, OK: true})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if !strings.Contains(string(b), `"image_count":2`) {
|
||||
t.Errorf("audit JSON %s lacks image_count", b)
|
||||
}
|
||||
// Omitempty keeps chat records unchanged.
|
||||
b2, _ := json.Marshal(Req{Type: "chat", Prompt: 1, Compl: 2, OK: true})
|
||||
if strings.Contains(string(b2), "image_count") {
|
||||
t.Errorf("chat record %s should omit image_count", b2)
|
||||
}
|
||||
}
|
||||
File diff suppressed because it is too large
Load Diff
@ -59,6 +59,82 @@ func lineOf(src string, idx int) int {
|
||||
return strings.Count(src[:idx], "\n") + 1
|
||||
}
|
||||
|
||||
// TestUIEditFormRoundTripsEveryOptionalSourceField pins the JS half of the
|
||||
// source-edit contract.
|
||||
//
|
||||
// The server distinguishes "absent" from "explicitly empty" for proxy_url,
|
||||
// api_key_env and timeout (they are pointers). That only helps if the form
|
||||
// SENTS them — a form that omits them falls back to "inherit", so clearing the
|
||||
// proxy box would silently keep the old proxy, which is the exact bug in the
|
||||
// other direction.
|
||||
//
|
||||
// It must also send the api_key_env value it was given, since that is the one
|
||||
// field whose loss is invisible until the next upstream call.
|
||||
func TestUIEditFormRoundTripsEveryOptionalSourceField(t *testing.T) {
|
||||
src := uiSource(t)
|
||||
body, ok := jsFunctionBody(src, "saveSource")
|
||||
if !ok {
|
||||
t.Fatal("saveSource() not found in the WebUI")
|
||||
}
|
||||
for _, field := range []string{"proxy_url", "api_key_env", "timeout"} {
|
||||
if !strings.Contains(body, field+":") {
|
||||
t.Errorf("saveSource() does not send %q; the server treats an absent "+
|
||||
"field as \"keep the stored value\", so the form could never clear it", field)
|
||||
}
|
||||
}
|
||||
// The form's inputs must be filled from the values the server sends, or a
|
||||
// save would post an empty box and clear a configured field.
|
||||
for _, read := range []string{`$("#s-proxy")`, `$("#s-keyenv")`, `$("#s-timeout")`} {
|
||||
if !strings.Contains(src, read+".value") {
|
||||
t.Errorf("the source form never reads %s — it would post an empty value", read)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// TestUIEditFormReadsDurationsFromTheRevealView: config.Source tags
|
||||
// Timeout/QueueTimeout `json:"-"`, so the durations only reach the form if the
|
||||
// reveal endpoint adds them explicitly. The form MUST read them (see above),
|
||||
// which makes this API response load-bearing rather than informational.
|
||||
func TestUIEditFormReadsDurationsFromTheRevealView(t *testing.T) {
|
||||
src := uiSource(t)
|
||||
for _, id := range []string{"s-proxy", "s-keyenv", "s-timeout"} {
|
||||
if !strings.Contains(src, `id="`+id+`"`) {
|
||||
t.Errorf("source form input #%s is missing", id)
|
||||
}
|
||||
}
|
||||
// The timeout box must render through the duration helper. Reading a raw
|
||||
// time.Duration number would show nanoseconds in a box that expects "300s".
|
||||
if !strings.Contains(src, "timeoutText(") {
|
||||
t.Error("the timeout input does not go through timeoutText(); a raw " +
|
||||
"time.Duration JSON number is nanoseconds and would be uneditable")
|
||||
}
|
||||
}
|
||||
|
||||
// TestUIEditFormSendsEverySourceField is the mirror of
|
||||
// TestSourcePayloadCoversEveryEditableField on the JavaScript side: each field
|
||||
// the server accepts must be submitted by the form, or the server's upsert has
|
||||
// nothing to store for it.
|
||||
func TestUIEditFormSendsEverySourceField(t *testing.T) {
|
||||
src := uiSource(t)
|
||||
body, ok := jsFunctionBody(src, "saveSource")
|
||||
if !ok {
|
||||
t.Fatal("saveSource() not found")
|
||||
}
|
||||
// Fields the edit dialog owns. proxy_url / api_key_env / timeout are
|
||||
// covered by the round-trip test above; this one catches the rest.
|
||||
// models and meta are local variables submitted by Go shorthand
|
||||
// (payload = { models, meta }), so they are matched as bare identifiers.
|
||||
for _, field := range []string{
|
||||
"name:", "base_url:", "api_key:", "adapter:", "endpoint:",
|
||||
"image_endpoint:", "max_concurrent:", "rpm:", "temperature:",
|
||||
"models,", "meta,",
|
||||
} {
|
||||
if !strings.Contains(body, field) {
|
||||
t.Errorf("saveSource() does not send %s", field)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// TestUIAPICallsDeclareMethod asserts that every api() call passing an options
|
||||
// object also declares an HTTP method (or is a GET that only passes an
|
||||
// AbortSignal). Without this, fetch defaults to GET and mutating endpoints are
|
||||
@ -347,3 +423,60 @@ func cssBlock(src, selector string) (string, bool) {
|
||||
}
|
||||
return src[i : i+j+1], true
|
||||
}
|
||||
|
||||
// The WebUI ships two locale objects (`zh` and `en`) that every rendered string
|
||||
// goes through. A key added to only one of them does not error: t() falls back
|
||||
// to `t("key") || "literal"` in some call sites and to the bare key string in
|
||||
// others, so the user sees either a hardcoded local string or a raw key name.
|
||||
// Nothing in the Go test suite noticed — this was found by hand while adding
|
||||
// thImgs. Pin the key sets so a one-sided edit fails here instead of shipping.
|
||||
func TestUILocaleKeyParity(t *testing.T) {
|
||||
src := uiSource(t)
|
||||
iZh := strings.Index(src, "zh: {")
|
||||
iEn := strings.Index(src, "en: {")
|
||||
if iZh < 0 || iEn < 0 || iEn < iZh {
|
||||
t.Fatalf("locale blocks not found (zh=%d en=%d)", iZh, iEn)
|
||||
}
|
||||
zh := src[iZh:iEn]
|
||||
en := src[iEn:]
|
||||
// End the EN block at its closing brace, so keys from later objects
|
||||
// (config templates etc.) do not pollute the comparison.
|
||||
if end := strings.Index(en, "\n },"); end > 0 {
|
||||
en = en[:end]
|
||||
}
|
||||
|
||||
keyRe := regexp.MustCompile(`(?m)^\s*([A-Za-z_][A-Za-z0-9_]*)\s*:`)
|
||||
collect := func(seg string) map[string]bool {
|
||||
m := map[string]bool{}
|
||||
for _, g := range keyRe.FindAllStringSubmatch(seg, -1) {
|
||||
m[g[1]] = true
|
||||
}
|
||||
return m
|
||||
}
|
||||
zk, ek := collect(zh), collect(en)
|
||||
// The locale markers themselves are the two block headers.
|
||||
delete(zk, "zh")
|
||||
delete(ek, "en")
|
||||
|
||||
var missingEN, missingZH []string
|
||||
for k := range zk {
|
||||
if !ek[k] {
|
||||
missingEN = append(missingEN, k)
|
||||
}
|
||||
}
|
||||
for k := range ek {
|
||||
if !zk[k] {
|
||||
missingZH = append(missingZH, k)
|
||||
}
|
||||
}
|
||||
if len(missingEN) > 0 {
|
||||
t.Errorf("translation keys present in zh but missing in en: %v", missingEN)
|
||||
}
|
||||
if len(missingZH) > 0 {
|
||||
t.Errorf("translation keys present in en but missing in zh: %v", missingZH)
|
||||
}
|
||||
if len(zk) < 100 {
|
||||
t.Errorf("only %d zh keys parsed — the block boundary regexp has drifted "+
|
||||
"and this test is no longer checking anything", len(zk))
|
||||
}
|
||||
}
|
||||
|
||||
195
internal/gateway/ui_image_page_test.go
Normal file
195
internal/gateway/ui_image_page_test.go
Normal file
@ -0,0 +1,195 @@
|
||||
package gateway
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"llmsproxy/internal/provider"
|
||||
)
|
||||
|
||||
// TestUI_ChatPageGroupsImageModels is a UI-contract guard for the bug that
|
||||
// started this: the chat page listed image models beside chat ones, and picking
|
||||
// one posted to /api/chat, which cannot serve it.
|
||||
//
|
||||
// The contract is asserted on the HTML source rather than by running a browser,
|
||||
// because the failure this must catch is a source-level one — an <option>
|
||||
// without `disabled`. A DOM assertion in a headless browser would pass right
|
||||
// up until someone removed the attribute, and then only if a CI machine had the
|
||||
// whole stack up.
|
||||
func TestUI_ChatPageGroupsImageModels(t *testing.T) {
|
||||
ui := uiSource(t)
|
||||
|
||||
// The dropdown must read model_kinds, not the flat models list. Reading the
|
||||
// flat list is what made chat and image models indistinguishable.
|
||||
if !strings.Contains(ui, "model_kinds") {
|
||||
t.Error("chat page never reads model_kinds; image models will be " +
|
||||
"indistinguishable from chat models again")
|
||||
}
|
||||
|
||||
// Image options must be rendered disabled. A plain String.replace over the
|
||||
// joined HTML only disables the FIRST match — that bug shipped once here
|
||||
// (one source disabled, five selectable) — so the assertion looks for the
|
||||
// per-item map, not for a global replace.
|
||||
if strings.Contains(ui, `imgOpts.join("").replace`) {
|
||||
t.Error("chat page disables image options with String.replace over the " +
|
||||
"joined HTML, which only rewrites the first option")
|
||||
}
|
||||
if !strings.Contains(ui, "disabled") {
|
||||
t.Error("chat page never emits a disabled attribute for image models")
|
||||
}
|
||||
}
|
||||
|
||||
// TestUI_HasDedicatedImagePage pins the second half of the fix: image models
|
||||
// need somewhere they CAN be tested.
|
||||
func TestUI_HasDedicatedImagePage(t *testing.T) {
|
||||
ui := uiSource(t)
|
||||
|
||||
for _, want := range []string{
|
||||
`data-tab="image"`, // sidebar entry
|
||||
`id="tab-image"`, // pane
|
||||
`function renderImage`, // renderer registered with refresh()
|
||||
`tab === "image"`, // dispatched
|
||||
`/v1/images/generations`, // the endpoint the page actually calls
|
||||
} {
|
||||
if !strings.Contains(ui, want) {
|
||||
t.Errorf("image page is missing %q", want)
|
||||
}
|
||||
}
|
||||
|
||||
// The page must be reachable from TABS; goTab iterates TABS to toggle pane
|
||||
// visibility, so a pane outside that list is present but permanently hidden —
|
||||
// which reads as "the page is blank" because the content is in the DOM.
|
||||
if !strings.Contains(ui, `"image"`) {
|
||||
t.Error(`"image" does not appear in the TABS list`)
|
||||
}
|
||||
}
|
||||
|
||||
// TestUI_ImagePageCallsTheImageEndpoint is the discriminator: a page that
|
||||
// renders an image result but posts to the chat endpoint would produce exactly
|
||||
// the original failure, just prettier.
|
||||
func TestUI_ImagePageCallsTheImageEndpoint(t *testing.T) {
|
||||
ui := uiSource(t)
|
||||
|
||||
i := strings.Index(ui, "async function sendImage(")
|
||||
if i < 0 {
|
||||
t.Fatal("sendImage() not found")
|
||||
}
|
||||
// Bound the search to the function body rather than the whole file, so a
|
||||
// /v1/images/generations mention elsewhere cannot satisfy this.
|
||||
body := ui[i:]
|
||||
if end := strings.Index(body, "\n function "); end > 0 {
|
||||
body = body[:end]
|
||||
}
|
||||
if end := strings.Index(body, "\n async function "); end > 0 {
|
||||
if s := body[:end]; strings.Contains(s, "/v1/images/generations") {
|
||||
body = s
|
||||
} else {
|
||||
body = body[:end]
|
||||
}
|
||||
}
|
||||
if !strings.Contains(body, "/v1/images/generations") {
|
||||
t.Error("sendImage() does not post to /v1/images/generations")
|
||||
}
|
||||
}
|
||||
|
||||
// TestSourceStatus_ExposesModelKindsToTheUI is the server half of the same
|
||||
// contract: the field the UI reads must actually be serialized.
|
||||
func TestSourceStatus_ExposesModelKindsToTheUI(t *testing.T) {
|
||||
st := provider.SourceStatus{
|
||||
Name: "vela",
|
||||
Models: []string{"claude-sonnet-5.5", "gpt-image-2"},
|
||||
ModelKinds: []provider.ModelKindInfo{
|
||||
{ID: "claude-sonnet-5.5", Kind: "chat"},
|
||||
{ID: "gpt-image-2", Kind: "image"},
|
||||
},
|
||||
}
|
||||
b, err := json.Marshal(st)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
var got map[string]interface{}
|
||||
if err := json.Unmarshal(b, &got); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
mk, ok := got["model_kinds"].([]interface{})
|
||||
if !ok || len(mk) != 2 {
|
||||
t.Fatalf("model_kinds = %v, want 2 entries; the UI cannot split the list without it", got["model_kinds"])
|
||||
}
|
||||
last := mk[1].(map[string]interface{})
|
||||
if last["kind"] != "image" {
|
||||
t.Errorf("model_kinds[1].kind = %v, want \"image\"", last["kind"])
|
||||
}
|
||||
}
|
||||
|
||||
// TestUI_BaseURLFieldExplainsWhereToCut keeps the second bug fixed: the URL
|
||||
// hint has to mention the concrete failure, not just say "enter the URL".
|
||||
func TestUI_BaseURLFieldExplainsWhereToCut(t *testing.T) {
|
||||
ui := uiSource(t)
|
||||
if !strings.Contains(ui, "urlHint") {
|
||||
t.Fatal("no urlHint; the base_url field gives no guidance about /v1")
|
||||
}
|
||||
hint := ui[strings.Index(ui, "urlHint:"):]
|
||||
if len(hint) > 700 {
|
||||
hint = hint[:700]
|
||||
}
|
||||
// The trap is specific and reproducible: base_url ending in /v1 makes chat
|
||||
// work and image generation 404. Saying only "enter the URL" reproduces the
|
||||
// confusion the hint exists to remove.
|
||||
for _, want := range []string{"/v1", "404"} {
|
||||
if !strings.Contains(hint, want) {
|
||||
t.Errorf("urlHint does not mention %q — the hint must name the concrete failure", want)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// TestImageEndpointComposesUnderRealConditions proves the hint describes a
|
||||
// reachable path rather than a plausible-sounding one.
|
||||
//
|
||||
// The case that matters: base_url already ends in /v1 (which is what an
|
||||
// operator pastes from an SDK example), so the RELATIVE image endpoint composes
|
||||
// to …/v1/images/generations and works, while the gateway's default ABSOLUTE
|
||||
// path would produce …/v1/v1/… and 404.
|
||||
//
|
||||
// The doubled form is served as a 404 here so the test would fail if the code
|
||||
// ever reverted to relying on the default.
|
||||
func TestImageEndpointComposesUnderRealConditions(t *testing.T) {
|
||||
up := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
switch r.URL.Path {
|
||||
case "/v1/images/generations":
|
||||
_, _ = w.Write([]byte(`{"data":[{"url":"http://example.invalid/x.png"}]}`))
|
||||
default:
|
||||
w.WriteHeader(http.StatusNotFound)
|
||||
}
|
||||
}))
|
||||
defer up.Close()
|
||||
|
||||
// base_url as an operator types it: host root + /v1.
|
||||
base := up.URL + "/v1"
|
||||
|
||||
// What discovery writes when base carries a version segment.
|
||||
relative := base + "/images/generations"
|
||||
resp, err := http.Post(relative, "application/json", strings.NewReader(`{"model":"m"}`))
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
resp.Body.Close()
|
||||
if resp.StatusCode != http.StatusOK {
|
||||
t.Errorf("POST %s = %d, want 200 — the form discovery writes must work", relative, resp.StatusCode)
|
||||
}
|
||||
|
||||
// What the gateway's built-in default would produce for the same base_url.
|
||||
doubled := base + "/v1/images/generations"
|
||||
resp2, err := http.Post(doubled, "application/json", strings.NewReader(`{"model":"m"}`))
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
resp2.Body.Close()
|
||||
if resp2.StatusCode != http.StatusNotFound {
|
||||
t.Errorf("POST %s = %d, want 404 — this test's premise is that the "+
|
||||
"default image path doubles the version segment; if the upstream ever "+
|
||||
"accepts it, composeImageEndpoint's rule needs rechecking", doubled, resp2.StatusCode)
|
||||
}
|
||||
}
|
||||
438
internal/gateway/ui_key_auto_test.go
Normal file
438
internal/gateway/ui_key_auto_test.go
Normal file
@ -0,0 +1,438 @@
|
||||
package gateway
|
||||
|
||||
import (
|
||||
"os"
|
||||
"strings"
|
||||
"testing"
|
||||
)
|
||||
|
||||
// Per-key AUTO chain UI contract.
|
||||
//
|
||||
// The per-key editor deliberately REUSES the global lane canvas instead of
|
||||
// shipping a second, simpler editor. That is a design choice with teeth: a
|
||||
// reduced copy would lack drag ordering and per-slot quota/period/hours, and
|
||||
// the two would drift apart silently. These checks pin the sharing, because
|
||||
// the only symptom of un-sharing is "the per-key editor got worse" — no error,
|
||||
// just a smaller feature set that nobody notices until they use it.
|
||||
func TestUIPerKeyAutoWiring(t *testing.T) {
|
||||
src := uiSource(t)
|
||||
|
||||
// The admin key list must offer the editor, otherwise the whole feature is
|
||||
// unreachable from the UI. Assert the *call site* inside keyCanvasHtml,
|
||||
// not just that the functions exist: removing the button leaves every
|
||||
// definition intact, and that is exactly the regression to catch.
|
||||
kcStart := strings.Index(src, "function keyCanvasHtml")
|
||||
if kcStart < 0 {
|
||||
t.Fatal("keyCanvasHtml is gone — cannot check the AUTO chain button")
|
||||
}
|
||||
kcEnd := strings.Index(src[kcStart:], "\n }")
|
||||
if kcEnd < 0 {
|
||||
t.Fatal("could not delimit keyCanvasHtml")
|
||||
}
|
||||
if !strings.Contains(src[kcStart:kcStart+kcEnd], "openKeyAutoModal(") {
|
||||
t.Fatal("per-key AUTO chain button is missing from the key card header")
|
||||
}
|
||||
for _, fn := range []string{
|
||||
"function openKeyAutoModal",
|
||||
"function keyAutoInherit",
|
||||
"function keyAutoSave",
|
||||
"function keyAutoClose",
|
||||
} {
|
||||
if !strings.Contains(src, fn) {
|
||||
t.Fatalf("UI is missing %s — the editor cannot function", fn)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// The per-key editor must render the SAME canvas the global editor uses.
|
||||
// If this regresses, the per-key chain silently loses drag ordering and
|
||||
// per-slot quota/period/hours while still looking like it works.
|
||||
func TestUIPerKeyAutoReusesGlobalCanvas(t *testing.T) {
|
||||
src := uiSource(t)
|
||||
|
||||
// One canvas implementation, two mount ids.
|
||||
if !strings.Contains(src, "function renderSortEditor(") {
|
||||
t.Fatal("no shared canvas renderer — the per-key editor must reuse renderSortEditor")
|
||||
}
|
||||
if !strings.Contains(src, `canvasId: "scr-canvas"`) {
|
||||
t.Fatal("the global editor no longer mounts the shared canvas")
|
||||
}
|
||||
if !strings.Contains(src, `canvasId: "key-auto-canvas"`) {
|
||||
t.Fatal("the per-key editor does not mount the shared canvas")
|
||||
}
|
||||
// Canvas lookup must consider both ids, or drag/drop silently no-ops on
|
||||
// whichever surface is not #scr-canvas.
|
||||
if !strings.Contains(src, "function sortCanvasEl(") ||
|
||||
!strings.Contains(src, `getElementById("key-auto-canvas")`) {
|
||||
t.Fatal("sortCanvasEl must resolve both canvas ids")
|
||||
}
|
||||
// The rules<->lanes conversion is shared too, so the per-key editor sends
|
||||
// exactly the fields the server understands (quota/period/hours included).
|
||||
for _, fn := range []string{"function buildLanes(", "function lanesToRules("} {
|
||||
if !strings.Contains(src, fn) {
|
||||
t.Fatalf("%s must be shared between the global and per-key editors", fn)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Saving must route through persistAuto's scope switch, not a hand-rolled
|
||||
// request, or the two editors will serialize rules differently.
|
||||
func TestUIPerKeyAutoSavesThroughScope(t *testing.T) {
|
||||
src := uiSource(t)
|
||||
if !strings.Contains(src, "sortState.scope = key") {
|
||||
t.Fatal("openKeyAutoModal must set sortState.scope so the save targets the key")
|
||||
}
|
||||
if !strings.Contains(src, "sortState.scope = null") {
|
||||
t.Fatal("closing the modal must clear sortState.scope, or the GLOBAL editor would save to this key")
|
||||
}
|
||||
if !strings.Contains(src, "if (sortState.scope)") {
|
||||
t.Fatal("persistAuto must branch on sortState.scope")
|
||||
}
|
||||
if !strings.Contains(src, `"/auto"`) {
|
||||
t.Fatal("per-key editor must call the /api/keys/{key}/auto endpoint")
|
||||
}
|
||||
// An empty lane list is how "inherit the global chain" is expressed.
|
||||
if !strings.Contains(src, "JSON.stringify({ auto: lanesToRules(sortState.lanes) })") {
|
||||
t.Fatal("keyAutoSave must PUT lanesToRules(sortState.lanes) verbatim")
|
||||
}
|
||||
}
|
||||
|
||||
// Both languages need the strings, otherwise a non-Chinese admin sees raw
|
||||
// i18n keys in the modal.
|
||||
func TestUIPerKeyAutoHasBothLanguages(t *testing.T) {
|
||||
src := uiSource(t)
|
||||
for _, k := range []string{
|
||||
"kAutoChainTitle", "kAutoChainHint", "kAutoChainInherit",
|
||||
"kAutoChainEmpty", "kAutoChainGlobal", "kAutoChainBad",
|
||||
} {
|
||||
n := strings.Count(src, k+`: "`)
|
||||
if n < 2 {
|
||||
t.Fatalf("i18n key %s defined %d time(s), want both zh and en", k, n)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// The statistics dashboard must print the window it actually covers. Calendar
|
||||
// periods are not nested — "this week" starts on its Monday and therefore
|
||||
// reaches into the previous month early in a month, making its total LARGER
|
||||
// than "this month". That is correct arithmetic, but without the dates on
|
||||
// screen it reads as a bug, and it has been reported as one.
|
||||
func TestUIStatsShowsPeriodWindow(t *testing.T) {
|
||||
src := uiSource(t)
|
||||
if !strings.Contains(src, `id="period-range"`) {
|
||||
t.Fatal("the stats dashboard has no element for the period window range")
|
||||
}
|
||||
if n := strings.Count(src, "periodRange"); n < 2 {
|
||||
t.Fatalf("i18n key periodRange appears %d time(s), want both zh and en", n)
|
||||
}
|
||||
}
|
||||
// Both canvases can exist in the DOM at once: the modal is mounted on top of
|
||||
// the settings tab, whose #scr-canvas stays alive behind it. Resolving the
|
||||
// canvas by document order (whichever comes first) painted the per-key chain
|
||||
// into the hidden background tab and left the dialog empty — and worse,
|
||||
// overwrote the global chain the user returns to. The lookup must follow
|
||||
// sortState.scope instead.
|
||||
func TestUISortCanvasResolvesByScope(t *testing.T) {
|
||||
src := uiSource(t)
|
||||
el := readFuncBody(t, src, "sortCanvasEl")
|
||||
if !strings.Contains(el, "if (sortState.scope)") {
|
||||
t.Fatal("sortCanvasEl must branch on sortState.scope, not document order")
|
||||
}
|
||||
if strings.Contains(el, `getElementById("scr-canvas")`) &&
|
||||
strings.Index(el, `getElementById("scr-canvas")`) < strings.Index(el, "sortState.scope") {
|
||||
t.Fatal("sortCanvasEl returns #scr-canvas before checking scope — the " +
|
||||
"per-key editor paints into the global canvas")
|
||||
}
|
||||
}
|
||||
|
||||
// A successful save must close the dialog even when the slots do not resolve.
|
||||
// The write DID succeed; leaving the modal open makes the user click Save
|
||||
// again and again. This regressed once: the !slots branch returned early.
|
||||
func TestUIKeyAutoSaveClosesOnSuccess(t *testing.T) {
|
||||
src := uiSource(t)
|
||||
body := readFuncBody(t, src, "keyAutoSave")
|
||||
// keyAutoClose() must be reached on every success path, not only in the
|
||||
// happy branch.
|
||||
iClose := strings.Index(body, "keyAutoClose()")
|
||||
if iClose < 0 {
|
||||
t.Fatal("keyAutoSave never calls keyAutoClose — the dialog would stay open")
|
||||
}
|
||||
iToast := strings.Index(body, "toast(")
|
||||
if iToast >= 0 && iClose < iToast {
|
||||
t.Fatal("keyAutoClose is called before the toast; the warning must be " +
|
||||
"shown first or it is lost with the modal")
|
||||
}
|
||||
// An early `return` inside a branch bypasses keyAutoClose entirely, so
|
||||
// merely finding the call is not enough: every `return` that sits between
|
||||
// entering the try block and the close call is a path that leaves the
|
||||
// dialog open after a successful write.
|
||||
tryStart := strings.Index(body, "try {")
|
||||
seg := body[tryStart:iClose]
|
||||
for _, line := range strings.Split(seg, "\n") {
|
||||
trimmed := strings.TrimSpace(line)
|
||||
if trimmed == "return;" || strings.HasPrefix(trimmed, "return ") {
|
||||
t.Fatal("a `return` before keyAutoClose leaves the dialog open after " +
|
||||
"a successful save: " + trimmed)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// readFuncBody returns the source of a top-level function by name, using
|
||||
// brace matching so nested blocks do not truncate it.
|
||||
func readFuncBody(t *testing.T, src, name string) string {
|
||||
t.Helper()
|
||||
idx := strings.Index(src, "function "+name+"(")
|
||||
if idx < 0 {
|
||||
t.Fatalf("function %s not found", name)
|
||||
}
|
||||
start := strings.Index(src[idx:], "{")
|
||||
if start < 0 {
|
||||
t.Fatalf("function %s has no body", name)
|
||||
}
|
||||
i := idx + start
|
||||
depth := 0
|
||||
for j := i; j < len(src); j++ {
|
||||
switch src[j] {
|
||||
case '{':
|
||||
depth++
|
||||
case '}':
|
||||
depth--
|
||||
if depth == 0 {
|
||||
return src[i : j+1]
|
||||
}
|
||||
}
|
||||
}
|
||||
t.Fatalf("function %s body is unterminated", name)
|
||||
return ""
|
||||
}
|
||||
|
||||
// The per-key dialog must not share the app-wide #modal-wrap id.
|
||||
//
|
||||
// Every dialog in this file uses that id, and more than one can be open at
|
||||
// once (the "add slot" picker opens ON TOP of this one). Handlers resolve
|
||||
// their own dialog via `.closest()` from their button, which works — but this
|
||||
// dialog is also looked up from the CANVAS side (sortCanvasEl, keyAutoClose,
|
||||
// keyAutoInherit), where there is no button to walk up from.
|
||||
// getElementById("#modal-wrap") then returns whichever dialog comes first in
|
||||
// the document, i.e. not the one being edited: Cancel closed the wrong box and
|
||||
// the canvas could resolve to the picker's.
|
||||
func TestUIKeyAutoModalHasOwnID(t *testing.T) {
|
||||
src := uiSource(t)
|
||||
body := readFuncBody(t, src, "openKeyAutoModal")
|
||||
if !strings.Contains(body, `wrap.id = "key-auto-modal"`) {
|
||||
t.Fatal("openKeyAutoModal must give its dialog its own id, not the shared " +
|
||||
"#modal-wrap (another dialog may be open on top of it)")
|
||||
}
|
||||
close := readFuncBody(t, src, "keyAutoClose")
|
||||
if strings.Contains(close, `"modal-wrap"`) {
|
||||
t.Fatal("keyAutoClose resolves the shared #modal-wrap; it would close " +
|
||||
"whichever dialog comes first in the document")
|
||||
}
|
||||
if !strings.Contains(close, `"key-auto-modal"`) {
|
||||
t.Fatal("keyAutoClose must resolve the per-key dialog by its own id")
|
||||
}
|
||||
}
|
||||
|
||||
// Editing a chain inside the per-key dialog must NOT persist it. That dialog
|
||||
// has Save and Cancel, so auto-writing on add/drag/delete meant the change was
|
||||
// already committed before Save and Cancel could not undo it.
|
||||
//
|
||||
// The global page is the opposite: it has always applied edits immediately.
|
||||
func TestUIKeyAutoEditsDoNotAutoPersist(t *testing.T) {
|
||||
src := uiSource(t)
|
||||
body := readFuncBody(t, src, "afterChainEdit")
|
||||
if !strings.Contains(body, "if (sortState.scope)") {
|
||||
t.Fatal("afterChainEdit must branch on sortState.scope")
|
||||
}
|
||||
// The scoped branch must return before reaching persistAuto.
|
||||
iScope := strings.Index(body, "if (sortState.scope)")
|
||||
iReturn := strings.Index(body[iScope:], "return")
|
||||
iPersist := strings.Index(body, "persistAuto()")
|
||||
if iScope < 0 || iReturn < 0 || iPersist < 0 {
|
||||
t.Fatal("afterChainEdit no longer has the expected shape")
|
||||
}
|
||||
if iScope+iReturn > iPersist {
|
||||
t.Fatal("the scoped branch reaches persistAuto() — a per-key edit would " +
|
||||
"be written before Save")
|
||||
}
|
||||
// And it must not call itself (a regex-driven rewrite did exactly that).
|
||||
if strings.Contains(body, "afterChainEdit()") {
|
||||
t.Fatal("afterChainEdit calls itself: infinite recursion")
|
||||
}
|
||||
// Every in-canvas edit goes through it, so none of them bypass the check.
|
||||
for _, fn := range []string{"scrAddFromForm", "sortScopeSave", "scrDelSlot"} {
|
||||
b := readFuncBody(t, src, fn)
|
||||
if strings.Contains(b, "persistAuto()") {
|
||||
t.Fatalf("%s calls persistAuto() directly, bypassing afterChainEdit", fn)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// The stats range line formats its end date from the LAST bucket label.
|
||||
// Bucket labels are hourly for the day view ("2026-10-03T09") and daily for
|
||||
// week/month ("2026-10-01"), so blindly appending "T00:00:00Z" to the hourly
|
||||
// form builds "2026-10-03T09T00:00:00Z" — Date rejects it, toISOString()
|
||||
// throws, and the assignment is lost. The symptom is an empty range line on
|
||||
// exactly the view people look at most, with no error anywhere.
|
||||
func TestUIStatsPeriodRangeHandlesHourlyBuckets(t *testing.T) {
|
||||
src := uiSource(t)
|
||||
el := readFuncBody(t, src, "paintStats")
|
||||
// It must branch on the label shape instead of concatenating blindly.
|
||||
if !strings.Contains(el, `T\d{2}$`) && !strings.Contains(el, `\d{2}-\d{2}T\d{2}`) {
|
||||
t.Fatal("paintStats must recognise the hourly bucket label shape " +
|
||||
"(YYYY-MM-DDTHH); concatenating a time onto it yields an invalid Date")
|
||||
}
|
||||
// And it must not contain the naive form that broke it.
|
||||
if strings.Contains(el, `bucket + "T00:00:00Z"`) {
|
||||
t.Fatal("paintStats still appends T00:00:00Z to the bucket label — " +
|
||||
"this throws on the day view and leaves the range line empty")
|
||||
}
|
||||
}
|
||||
|
||||
// A key with no chain of its own must open on a COPY OF THE GLOBAL CHAIN, not
|
||||
// on an empty canvas. The realistic admin task is "this user gets the global chain
|
||||
// minus the sources they must not hit"; starting blank forces retyping every
|
||||
// tier against a 226-model picker. This is deliberately NOT the discovery
|
||||
// fallback (which dumps every model of every source) — it is the same rules
|
||||
// the global chain actually runs.
|
||||
func TestUIKeyAutoSeedsFromGlobalChain(t *testing.T) {
|
||||
src := uiSource(t)
|
||||
body := readFuncBody(t, src, "openKeyAutoModal")
|
||||
if !strings.Contains(body, `api("/api/auto")`) {
|
||||
t.Fatal("openKeyAutoModal must read the global chain to seed the editor")
|
||||
}
|
||||
if !strings.Contains(body, "seededFromGlobal") {
|
||||
t.Fatal("the editor must remember whether its content came from the " +
|
||||
"global chain, so the notice can say so")
|
||||
}
|
||||
// And it must not fall back to the discovery path that fills 226 models.
|
||||
if strings.Contains(body, `buildLanes(idx, [], `) {
|
||||
t.Fatal("seeding from an empty rule list triggers buildLanes discovery")
|
||||
}
|
||||
}
|
||||
|
||||
// The notice must distinguish "already this key's chain" from "a copy that
|
||||
// only becomes its chain on Save" — those two look identical on the canvas and
|
||||
// behave very differently if the admin hits Cancel.
|
||||
func TestUIKeyAutoNoticeDistinguishesSavedFromCopied(t *testing.T) {
|
||||
src := uiSource(t)
|
||||
body := readFuncBody(t, src, "keyAutoInheritNotice")
|
||||
for _, k := range []string{"kAutoChainSeeded", "kAutoChainNewOwn", "kAutoChainWillInherit"} {
|
||||
if !strings.Contains(body, k) {
|
||||
t.Fatalf("the notice never uses %s — a copied chain and a saved one "+
|
||||
"would look the same", k)
|
||||
}
|
||||
}
|
||||
// Both languages, or a non-Chinese admin sees raw keys.
|
||||
for _, k := range []string{"kAutoChainSeeded", "kAutoChainNewOwn", "kAutoChainWillInherit"} {
|
||||
if n := strings.Count(src, k+`: "`); n < 2 {
|
||||
t.Fatalf("i18n key %s defined %d time(s), want both zh and en", k, n)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// A key whose model scope excludes AUTO cannot use AUTO at all — the scope
|
||||
// filter runs before the chain — yet the editor happily lets an admin
|
||||
// configure a per-key chain. That combination is easy to create by accident
|
||||
// and produces a 403 with no visible cause, so the dialog must say so.
|
||||
//
|
||||
// We deliberately do NOT widen the scope automatically: silently granting a
|
||||
// model the operator did not ask for is worse than a loud warning.
|
||||
func TestUIKeyAutoWarnsOnScopeConflict(t *testing.T) {
|
||||
src := uiSource(t)
|
||||
if !strings.Contains(src, "function keyAutoScopeWarning(") {
|
||||
t.Fatal("no scope-conflict check in the per-key editor")
|
||||
}
|
||||
body := readFuncBody(t, src, "keyAutoScopeWarning")
|
||||
if !strings.Contains(body, "AUTO") {
|
||||
t.Fatal("keyAutoScopeWarning must test whether the scope lists AUTO")
|
||||
}
|
||||
// It must not write anything back to the key: reporting only.
|
||||
for _, forbidden := range []string{"PUT", "POST", "fetch(", "api("} {
|
||||
if strings.Contains(body, forbidden) {
|
||||
t.Fatalf("keyAutoScopeWarning calls %s — it must only report the "+
|
||||
"conflict, never widen the scope itself", forbidden)
|
||||
}
|
||||
}
|
||||
// Empty scope means unrestricted, so no warning is correct there.
|
||||
if !strings.Contains(body, "if (!models.length)") {
|
||||
t.Fatal("an unrestricted key (no models) must not be warned")
|
||||
}
|
||||
// Both languages.
|
||||
if n := strings.Count(src, "kAutoChainScopeWarn"); n < 2 {
|
||||
t.Fatalf("i18n key kAutoChainScopeWarn appears %d time(s), want zh+en", n)
|
||||
}
|
||||
}
|
||||
|
||||
// The scope conflict for AUTO must be reported by checkQuota, because that is
|
||||
// the check an AUTO request actually reaches first (it runs before the chain is
|
||||
// consulted). Putting the message on the later path made it dead code — the
|
||||
// request was rejected one check earlier with the generic "not allowed for this
|
||||
// key" wording, which sends an operator looking for a quota that isn't the
|
||||
// problem.
|
||||
func TestAutoScopeConflictMessageComesFromCheckQuota(t *testing.T) {
|
||||
b, err := os.ReadFile("chat.go")
|
||||
if err != nil {
|
||||
t.Fatalf("read chat.go: %v", err)
|
||||
}
|
||||
src := string(b)
|
||||
i := strings.Index(src, "func (g *Gateway) checkQuota(")
|
||||
if i < 0 {
|
||||
t.Fatal("checkQuota not found")
|
||||
}
|
||||
// Body = up to the next top-level func.
|
||||
rest := src[i:]
|
||||
if j := strings.Index(rest[1:], "\nfunc "); j > 0 {
|
||||
rest = rest[:j+1]
|
||||
}
|
||||
if !strings.Contains(rest, "isAuto(model)") {
|
||||
t.Fatal("checkQuota must special-case AUTO: a scope that omits AUTO is " +
|
||||
"rejected there before anything else looks at the chain")
|
||||
}
|
||||
if !strings.Contains(rest, "not in this key's model scope") {
|
||||
t.Fatal("checkQuota's AUTO branch must name the scope as the cause")
|
||||
}
|
||||
}
|
||||
|
||||
// renderSortEditor replaces its mount's innerHTML wholesale. The global
|
||||
// priority page used to hand it the .card that also holds the h2 toolbar, so
|
||||
// every open deleted the save/reset/add buttons: the lanes still rendered and
|
||||
// the drag worked, but there was no way to persist a change. The per-key modal
|
||||
// never showed this because its mount held only the canvas.
|
||||
//
|
||||
// The regression is invisible in the DOM when you only count lane elements, so
|
||||
// this asserts on the structure: the toolbar lives in a sibling of the mount,
|
||||
// not inside it.
|
||||
func TestGlobalSortToolbarSurvivesCanvasRender(t *testing.T) {
|
||||
src := uiSource(t)
|
||||
|
||||
// The global editor must mount into a dedicated host, not the .card.
|
||||
i := strings.Index(src, `$("#tab-sort").innerHTML = `)
|
||||
if i < 0 {
|
||||
t.Fatal("global sort markup not found")
|
||||
}
|
||||
chunk := src[i:min(len(src), i+1400)]
|
||||
if !strings.Contains(chunk, `id="sort-mount"`) {
|
||||
t.Fatal("global sort page must provide a dedicated #sort-mount for the canvas")
|
||||
}
|
||||
// The mount handed to renderSortEditor must be that host.
|
||||
if !strings.Contains(src, `renderSortEditor($("#sort-mount")`) {
|
||||
t.Fatal("renderSortEditor must be given #sort-mount; passing the .card " +
|
||||
"wipes the h2 toolbar that holds the save button")
|
||||
}
|
||||
// The toolbar (save in particular) must be OUTSIDE #sort-mount.
|
||||
close := strings.Index(chunk, "</div>`;")
|
||||
if close < 0 {
|
||||
t.Fatal("malformed #tab-sort template")
|
||||
}
|
||||
head := chunk[:close]
|
||||
mountAt := strings.Index(head, `id="sort-mount"`)
|
||||
saveAt := strings.Index(head, `onclick="saveSort()"`)
|
||||
if saveAt < 0 {
|
||||
t.Fatal("the global sort page has no save button at all")
|
||||
}
|
||||
if mountAt >= 0 && saveAt > mountAt {
|
||||
t.Fatal("the save button sits inside the canvas mount and will be " +
|
||||
"erased the moment the canvas renders")
|
||||
}
|
||||
}
|
||||
413
internal/gateway/ui_plugin_test.go
Normal file
413
internal/gateway/ui_plugin_test.go
Normal file
@ -0,0 +1,413 @@
|
||||
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")
|
||||
}
|
||||
}
|
||||
|
||||
// ---- plugin management UI contract ---------------------------------------
|
||||
//
|
||||
// The management page is the operator's only way to take a broken plugin out
|
||||
// of the request path. Every one of these assertions guards a link that, if it
|
||||
// silently broke, would leave the gateway running with a plugin it cannot
|
||||
// disable — the worst kind of gap: everything looks fine and nothing is
|
||||
// reachable.
|
||||
|
||||
func TestUIHasPluginTabAndPane(t *testing.T) {
|
||||
src := uiSourceX(t)
|
||||
if !strings.Contains(src, `data-tab="plugins"`) {
|
||||
t.Error("no sidebar entry for the plugin page")
|
||||
}
|
||||
if !strings.Contains(src, `id="tab-plugins"`) {
|
||||
t.Error("no #tab-plugins pane")
|
||||
}
|
||||
// The tab list is now a single constant; a new tab must be added there or
|
||||
// goTab will not un-hide its pane.
|
||||
if !strings.Contains(src, `const TABS = [`) {
|
||||
t.Error("TABS is gone; the tab list went back to a duplicated literal")
|
||||
}
|
||||
for _, tn := range []string{"status", "chat", "keys", "sort", "sources", "adapters", "plugins"} {
|
||||
if !strings.Contains(src, `"`+tn+`"`) {
|
||||
t.Errorf("TABS is missing %q", tn)
|
||||
}
|
||||
}
|
||||
// goTab must iterate TABS, not its own list.
|
||||
if !strings.Contains(src, "TABS.forEach((tn) =>") {
|
||||
t.Error("goTab does not iterate TABS")
|
||||
}
|
||||
if strings.Contains(src, `["status", "chat", "keys", "sort", "sources", "adapters"].forEach`) {
|
||||
t.Error("a duplicated tab literal survived; it will drift from TABS")
|
||||
}
|
||||
}
|
||||
|
||||
func TestUIRendersPluginManagement(t *testing.T) {
|
||||
src := uiSourceX(t)
|
||||
body, ok := jsFunctionBody(src, "renderPlugins")
|
||||
if !ok {
|
||||
t.Fatal("renderPlugins() not found")
|
||||
}
|
||||
// It must read the DISK listing, not just the loaded set: a plugin that
|
||||
// failed to compile is absent from the loaded set, and showing only the
|
||||
// loaded set makes a syntax error look like "the plugin is not installed".
|
||||
if !strings.Contains(body, "on_disk") {
|
||||
t.Error("renderPlugins reads only the loaded set; a failed plugin would " +
|
||||
"be invisible instead of shown with its error")
|
||||
}
|
||||
if !strings.Contains(body, "/api/plugins") {
|
||||
t.Error("renderPlugins does not call /api/plugins")
|
||||
}
|
||||
// Hook errors must be surfaced: a plugin that throws in every stage leaves
|
||||
// no other trace, so without this the symptom is "the feature just doesn't
|
||||
// work".
|
||||
if !strings.Contains(body, "hook_errors") {
|
||||
t.Error("renderPlugins ignores hook_errors; a silently broken plugin is undebuggable")
|
||||
}
|
||||
// Enable / disable / remove / edit.
|
||||
for _, fn := range []string{"togglePlugin", "delPlugin", "installPlugin", "editPlugin"} {
|
||||
if _, ok := jsFunctionBody(src, fn); !ok {
|
||||
t.Errorf("%s() is missing from the WebUI", fn)
|
||||
}
|
||||
}
|
||||
// The toggle must go through the enable/disable endpoint, not delete.
|
||||
tb, ok := jsFunctionBody(src, "togglePlugin")
|
||||
if !ok {
|
||||
t.Fatal("togglePlugin() missing")
|
||||
}
|
||||
if !strings.Contains(tb, `method: "PUT"`) {
|
||||
t.Error("togglePlugin does not use PUT")
|
||||
}
|
||||
if !strings.Contains(tb, "enabled:") {
|
||||
t.Error("togglePlugin does not send an \"enabled\" field")
|
||||
}
|
||||
// And the admin-only tab list must include plugins, or a non-admin would
|
||||
// see a page whose every action 403s.
|
||||
if !strings.Contains(src, `["sort", "sources", "adapters", "plugins"]`) {
|
||||
t.Error("the admin-only tab list omits \"plugins\"; a user key would see a " +
|
||||
"page full of actions that all fail with 403")
|
||||
}
|
||||
}
|
||||
|
||||
// TestUIBindDropzoneIsParameterised guards the refactor: the adapter and plugin
|
||||
// upload forms share one dropzone, so a hard-coded id would send a dropped
|
||||
// plugin file into the adapter name field.
|
||||
func TestUIBindDropzoneIsParameterised(t *testing.T) {
|
||||
src := uiSourceX(t)
|
||||
body, ok := jsFunctionBody(src, "bindDropzone")
|
||||
if !ok {
|
||||
t.Fatal("bindDropzone() not found")
|
||||
}
|
||||
if strings.Contains(body, `$("#dz")`) || strings.Contains(body, `$("#adp-name")`) {
|
||||
t.Error("bindDropzone still hard-codes the adapter's element ids; the " +
|
||||
"plugin form would write into the adapter form")
|
||||
}
|
||||
if !strings.Contains(body, "dzId") || !strings.Contains(body, "nameSel") {
|
||||
t.Error("bindDropzone does not accept the ids to bind")
|
||||
}
|
||||
// Both forms must call it.
|
||||
if !strings.Contains(src, `bindDropzone("pl-dz", "pl-file", "#pl-name", "#pl-code")`) {
|
||||
t.Error("the plugin upload form does not use the parameterised dropzone")
|
||||
}
|
||||
}
|
||||
|
||||
// TestPluginElementsSurviveHostRebuild guards the defect that made plugin
|
||||
// elements look absent no matter how the injection was configured.
|
||||
//
|
||||
// renderStatus (and six other pages) assign pane.innerHTML wholesale. Anything a
|
||||
// plugin had mounted into that pane is destroyed by the assignment. The symptom
|
||||
// is silent and misleading: the plugin really did declare an element, the
|
||||
// payload really did arrive, and the element is still gone on the next repaint —
|
||||
// so the natural conclusion is "my plugin declared nothing", which sends you
|
||||
// looking in the wrong file.
|
||||
//
|
||||
// The fix is to re-mount after the rebuild. This test asserts the re-mount is
|
||||
// wired at the SINGLE place every renderer passes through, rather than leaving
|
||||
// it to be re-added per page.
|
||||
func TestPluginElementsSurviveHostRebuild(t *testing.T) {
|
||||
html := uiSource(t)
|
||||
|
||||
if !strings.Contains(html, "function remountPluginElements") {
|
||||
t.Fatal("remountPluginElements is not defined; nothing can re-attach a plugin element after a host rebuild")
|
||||
}
|
||||
if !strings.Contains(html, "const PLUGIN_MOUNT_HOOKS = []") {
|
||||
t.Fatal("PLUGIN_MOUNT_HOOKS is not declared")
|
||||
}
|
||||
// The hook array must be declared BEFORE injectPluginUI pushes to it, and
|
||||
// before refresh() calls into it. A use-before-declaration in a const
|
||||
// block is a hard TDZ ReferenceError at first paint.
|
||||
hookDecl := strings.Index(html, "const PLUGIN_MOUNT_HOOKS = []")
|
||||
push := strings.Index(html, "PLUGIN_MOUNT_HOOKS.push")
|
||||
use := strings.Index(html, "PLUGIN_MOUNT_HOOKS.forEach")
|
||||
if hookDecl < 0 || push < 0 || use < 0 {
|
||||
t.Fatal("the hook array is declared but never both filled and drained")
|
||||
}
|
||||
if hookDecl > push || hookDecl > use {
|
||||
t.Error("PLUGIN_MOUNT_HOOKS is used before its declaration (const TDZ: first paint would throw)")
|
||||
}
|
||||
|
||||
// refresh() is the chokepoint every renderer passes through.
|
||||
rf := strings.Index(html, "function refresh(tab)")
|
||||
if rf < 0 {
|
||||
t.Fatal("refresh(tab) is gone")
|
||||
}
|
||||
body := html[rf:]
|
||||
if i := strings.Index(body, "\n }"); i > 0 {
|
||||
body = body[:i]
|
||||
}
|
||||
if !strings.Contains(body, "remountPluginElements") {
|
||||
t.Error("refresh() does not re-mount plugin elements: a page that rebuilds its DOM wipes them")
|
||||
}
|
||||
|
||||
// A mount must be idempotent, or the widget's <script> runs again on every
|
||||
// host repaint and its counters silently double.
|
||||
if !strings.Contains(html, `data-idx="`) || !strings.Contains(html, "plugin-el[data-idx=") {
|
||||
t.Error("element mounting is not guarded by a per-pane marker; re-mounting would re-run plugin scripts")
|
||||
}
|
||||
}
|
||||
|
||||
// TestPluginPagesAreReachableByGoTab guards a bug that only shows up in a
|
||||
// browser: goTab() iterated the hardcoded TABS list to toggle `hidden`, and a
|
||||
// plugin-contributed page is not in that list. The pane existed, the plugin had
|
||||
// filled it with real data, and clicking the sidebar entry changed nothing —
|
||||
// the `hidden` class was never removed.
|
||||
//
|
||||
// It reads exactly like "the page is blank" while the content sits in the DOM,
|
||||
// and no static check catches it: the injection is correct, the data is correct,
|
||||
// and the API returns 200.
|
||||
//
|
||||
// So this asserts the two structural facts that make it reachable: goTab
|
||||
// consults the plugin page set, and that set is declared before goTab runs (a
|
||||
// `const` further down the file would be a temporal-dead-zone ReferenceError on
|
||||
// the first click).
|
||||
func TestPluginPagesAreReachableByGoTab(t *testing.T) {
|
||||
ui := uiSource(t)
|
||||
|
||||
goTab := strings.Index(ui, "function goTab(")
|
||||
if goTab < 0 {
|
||||
t.Fatal("goTab() is gone")
|
||||
}
|
||||
end := strings.Index(ui[goTab:], "\n }")
|
||||
if end < 0 {
|
||||
t.Fatal("could not isolate goTab()")
|
||||
}
|
||||
body := ui[goTab : goTab+end]
|
||||
|
||||
if !strings.Contains(body, "PLUGIN_PAGES.forEach") {
|
||||
t.Error("★ goTab() does not toggle plugin pages — a plugin page keeps its " +
|
||||
"`hidden` class forever, so the sidebar entry does nothing")
|
||||
}
|
||||
|
||||
// Declaration must precede the use, or the first click throws.
|
||||
decl := strings.Index(ui, "var PLUGIN_PAGES = new Set()")
|
||||
if decl < 0 {
|
||||
// A const later in the file would also "work" only if nothing reads it
|
||||
// first — make that explicit rather than silent.
|
||||
if strings.Contains(ui, "const PLUGIN_PAGES = new Set()") {
|
||||
t.Error("PLUGIN_PAGES is a `const` declared after goTab() reads it — " +
|
||||
"temporal dead zone: the first click on a plugin page throws")
|
||||
} else {
|
||||
t.Error("PLUGIN_PAGES is not declared anywhere")
|
||||
}
|
||||
return
|
||||
}
|
||||
if decl > goTab {
|
||||
t.Errorf("PLUGIN_PAGES is declared at %d but goTab() at %d reads it — "+
|
||||
"declaration must come first", decl, goTab)
|
||||
}
|
||||
}
|
||||
|
||||
// TestPluginSidebarEntryCarriesATitle: a plugin tab must render icon AND label.
|
||||
//
|
||||
// Native tabs are `<svg>…</svg><span data-i="navStatus">状态</span>`; the plugin
|
||||
// branch set innerHTML to the icon alone, so every plugin page showed up as a
|
||||
// nameless icon in the sidebar. An operator cannot tell what an icon means
|
||||
// until they click it.
|
||||
func TestPluginSidebarEntryCarriesATitle(t *testing.T) {
|
||||
html := uiSource(t)
|
||||
anchor := "btn.innerHTML =\n pluginIconHTML(pg.icon) +"
|
||||
i := strings.Index(html, anchor)
|
||||
if i < 0 {
|
||||
t.Fatal("plugin sidebar button construction not found")
|
||||
}
|
||||
// Look at the statement the anchor opens, not the whole file.
|
||||
stmt := html[i : i+400]
|
||||
if strings.Contains(stmt, "esc(pg.title") {
|
||||
return
|
||||
}
|
||||
t.Errorf("plugin sidebar button renders the icon only; a label is required. "+
|
||||
"got: %.200s", stmt)
|
||||
}
|
||||
|
||||
// TestPluginPageHasAccessEditor: the plugin page must let an operator configure
|
||||
// endpoint access. The APIs existed with zero UI entry point, which meant the
|
||||
// only way to change who could reach a plugin endpoint was to hand-edit
|
||||
// config.yaml — the thing the UI is meant to replace.
|
||||
func TestPluginPageHasAccessEditor(t *testing.T) {
|
||||
src := uiSource(t)
|
||||
for _, needle := range []string{
|
||||
`renderPluginAccess`, // called after rendering the plugin list
|
||||
`savePluginAccess`, // wired to a button
|
||||
`/access`, // hits the per-endpoint access API
|
||||
`id="pl-access"`, // container exists
|
||||
`http_access`, // reads the table from /api/plugins
|
||||
`data-ep=`, // one input per endpoint
|
||||
`data-plvis=`, // one visibility selector per plugin
|
||||
`/http/*/access`, // whole-plugin visibility shares the access API
|
||||
`plVisCell`, // renders the declared-vs-override distinction
|
||||
} {
|
||||
if !strings.Contains(src, needle) {
|
||||
t.Errorf("plugin page lacks %q", needle)
|
||||
}
|
||||
}
|
||||
}
|
||||
267
internal/gateway/ui_quota_contract_test.go
Normal file
267
internal/gateway/ui_quota_contract_test.go
Normal file
@ -0,0 +1,267 @@
|
||||
package gateway
|
||||
|
||||
import (
|
||||
"regexp"
|
||||
"strings"
|
||||
"testing"
|
||||
)
|
||||
|
||||
// The WebUI dialogs all share the id "modal-wrap", and more than one can be
|
||||
// open at the same time (the seed-key notice sits on top of the keys page).
|
||||
// A save handler that closes "the" modal via $("#modal-wrap") therefore removes
|
||||
// whichever one comes FIRST in the document — which is the wrong dialog: the
|
||||
// form the user just submitted stays on screen while an unrelated dialog
|
||||
// vanishes.
|
||||
//
|
||||
// This is the same class of bug the api() contract test pins: reviewing inline
|
||||
// JS by eye does not catch it, and the symptom ("the dialog did not close")
|
||||
// points away from the cause.
|
||||
|
||||
// modalCloseRe finds every `$(...)`-style lookup of the shared modal id.
|
||||
var modalCloseRe = regexp.MustCompile(`\$\("#modal-wrap"\)`)
|
||||
|
||||
// closestModalRe finds the safe form: resolve the dialog from the clicked
|
||||
// button instead of from the document.
|
||||
var closestModalRe = regexp.MustCompile(`\.closest\("#modal-wrap"\)`)
|
||||
|
||||
func TestUIDialogClosesItselfNotTheFirstModal(t *testing.T) {
|
||||
// Strip comments first: prose that *names* the unsafe pattern (as the fix's
|
||||
// own comment does) would otherwise be flagged as a violation.
|
||||
code := stripJSComments(uiSource(t))
|
||||
|
||||
for _, m := range modalCloseRe.FindAllStringIndex(code, -1) {
|
||||
after := code[m[1]:]
|
||||
stmtEnd := strings.Index(after, ";")
|
||||
if stmtEnd < 0 || stmtEnd > 200 {
|
||||
continue
|
||||
}
|
||||
stmt := after[:stmtEnd]
|
||||
if strings.Contains(stmt, ".remove()") {
|
||||
line := 1 + strings.Count(code[:m[0]], "\n")
|
||||
t.Errorf("line %d closes the first #modal-wrap in the document, not its own dialog:\n\t%s",
|
||||
line, strings.TrimSpace(stmt))
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// stripJSComments removes // line comments and /* block */ comments from JS
|
||||
// embedded in the UI document. It is deliberately simple: the document is our
|
||||
// own source, and a false negative here only means the check is silent.
|
||||
func stripJSComments(src string) string {
|
||||
var out strings.Builder
|
||||
lines := strings.Split(src, "\n")
|
||||
inBlock := false
|
||||
for _, ln := range lines {
|
||||
trimmed := strings.TrimSpace(ln)
|
||||
if inBlock {
|
||||
if strings.Contains(ln, "*/") {
|
||||
inBlock = false
|
||||
}
|
||||
continue
|
||||
}
|
||||
if strings.HasPrefix(trimmed, "/*") {
|
||||
if !strings.Contains(ln, "*/") {
|
||||
inBlock = true
|
||||
}
|
||||
continue
|
||||
}
|
||||
if i := strings.Index(ln, "//"); i >= 0 {
|
||||
before := ln[:i]
|
||||
if strings.Count(before, `"`)%2 == 0 && strings.Count(before, "'")%2 == 0 {
|
||||
ln = before
|
||||
}
|
||||
}
|
||||
out.WriteString(ln)
|
||||
out.WriteString("\n")
|
||||
}
|
||||
return out.String()
|
||||
}
|
||||
|
||||
// Every handler that closes a dialog must do it one of the two safe ways.
|
||||
func TestUIDialogClosuresGoThroughSafePaths(t *testing.T) {
|
||||
src := stripJSComments(uiSource(t))
|
||||
for _, fn := range []string{
|
||||
"downloadStatsCsv", "downloadKeysCsv", "saveSource", "saveTemplate",
|
||||
"scrAddFromForm", "sortScopeSave", "scopeSave", "createKey",
|
||||
} {
|
||||
body, ok := jsFunctionBody(src, fn)
|
||||
if !ok {
|
||||
t.Errorf("%s not found", fn)
|
||||
continue
|
||||
}
|
||||
if !strings.Contains(body, `closest("#modal-wrap")`) && !strings.Contains(body, "closeTopModal()") {
|
||||
t.Errorf("%s closes a dialog with neither .closest nor closeTopModal", fn)
|
||||
}
|
||||
}
|
||||
if !strings.Contains(src, "function closeTopModal(") {
|
||||
t.Error("closeTopModal helper is missing")
|
||||
}
|
||||
}
|
||||
|
||||
// The helper must pick the LAST dialog (the topmost one the user sees), not the
|
||||
// first — that inversion is the whole bug.
|
||||
func TestUICloseTopModalTakesTheLast(t *testing.T) {
|
||||
body, ok := jsFunctionBody(uiSource(t), "closeTopModal")
|
||||
if !ok {
|
||||
t.Fatal("closeTopModal not found")
|
||||
}
|
||||
if !strings.Contains(body, "all.length - 1") {
|
||||
t.Errorf("closeTopModal does not take the last dialog:\n\t%s", oneLine(body))
|
||||
}
|
||||
}
|
||||
|
||||
// Handlers that resolve their dialog from a button must actually receive one:
|
||||
// a signature without the parameter means the .closest() silently yields null
|
||||
// and the save leaves its form stranded.
|
||||
func TestUIDialogHandlersReceiveTheirButton(t *testing.T) {
|
||||
src := stripJSComments(uiSource(t))
|
||||
for _, fn := range []string{
|
||||
"downloadStatsCsv", "saveSource", "scrAddFromForm", "sortScopeSave",
|
||||
"scopeSave", "createKey",
|
||||
} {
|
||||
body, ok := jsFunctionBody(src, fn)
|
||||
if !ok {
|
||||
t.Errorf("%s not found", fn)
|
||||
continue
|
||||
}
|
||||
if !strings.Contains(body, "closest(\"#modal-wrap\")") {
|
||||
continue // uses closeTopModal only
|
||||
}
|
||||
sig := body[:strings.Index(body, ")")+1]
|
||||
if !strings.Contains(sig, "btn") {
|
||||
t.Errorf("%s uses .closest(\"#modal-wrap\") but its signature %s has no button parameter —\n"+
|
||||
"the lookup would always be null and the form would never close", fn, oneLine(sig))
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Quotas belong to the scope entries, so putScope only has to ship the scope
|
||||
// list — the caps travel inside it. What must NOT come back is a key-wide
|
||||
// total: it would be a second budget able to drift out of sync with the models
|
||||
// it is supposed to cover.
|
||||
func TestUIPutScopeShipsOnlyScopeQuotas(t *testing.T) {
|
||||
body, ok := jsFunctionBody(uiSource(t), "putScope")
|
||||
if !ok {
|
||||
t.Fatal("putScope not found")
|
||||
}
|
||||
code := stripJSComments(body)
|
||||
if !strings.Contains(code, "models:") || !strings.Contains(code, "scopes") {
|
||||
t.Error("putScope must ship the scope list the caps live in")
|
||||
}
|
||||
for _, gone := range []string{"kquota", "kreqquota", "kperiod", "khours"} {
|
||||
if strings.Contains(code, gone) {
|
||||
t.Errorf("putScope still references the removed key-wide quota field %q", gone)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// A model brick carries its own budgets, and the scope editor reads and writes
|
||||
// both of them: dropping req_quota on the round trip would silently lift a
|
||||
// request cap every time someone edited a token cap.
|
||||
func TestUIScopeEditorRoundTripsBothQuotas(t *testing.T) {
|
||||
src := stripJSComments(uiSource(t))
|
||||
for _, fn := range []string{"scopeHtml", "readScopes", "scopeEdit", "scopeSave", "scopeQuotaBadge"} {
|
||||
if _, ok := jsFunctionBody(src, fn); !ok {
|
||||
t.Errorf("%s not found in the UI source", fn)
|
||||
}
|
||||
}
|
||||
for _, fn := range []string{"scopeHtml", "readScopes", "scopeSave", "scopeQuotaBadge"} {
|
||||
body, ok := jsFunctionBody(src, fn)
|
||||
if !ok {
|
||||
continue
|
||||
}
|
||||
if !strings.Contains(body, "req_quota") && !strings.Contains(body, "reqquota") {
|
||||
t.Errorf("%s does not carry req_quota — a request cap would be lost on edit", fn)
|
||||
}
|
||||
}
|
||||
if form, ok := jsFunctionBody(src, "scopeEdit"); ok && !strings.Contains(form, "sc-reqs") {
|
||||
t.Error("the scope editor has no request-quota input")
|
||||
}
|
||||
}
|
||||
|
||||
// The whole key-wide quota surface must be gone from the UI: a badge or a
|
||||
// button reading a field the server no longer has would render "undefined" or
|
||||
// silently do nothing.
|
||||
func TestUIHasNoKeyWideQuotaSurface(t *testing.T) {
|
||||
src := uiSource(t)
|
||||
for _, gone := range []string{
|
||||
"keyCapBadges", "keyQuotaEdit", "keyQuotaSave",
|
||||
"kq-tokens", "kq-reqs", "kq-period", "kq-hours",
|
||||
"kc-tokens", "kc-reqs", "kc-period", "kc-hours",
|
||||
"data-kquota", "data-kreqquota", "data-kperiod", "data-khours",
|
||||
} {
|
||||
if strings.Contains(src, gone) {
|
||||
t.Errorf("UI still references the removed key-wide quota surface %q", gone)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Existence of the strings is not enough: the badge has to READ the scope's
|
||||
// fields, and the editor has to read back what the brick writes. A field can
|
||||
// be present in the source and still never reach the screen — e.g. left in a
|
||||
// dead branch, or read from a data attribute the writer never sets.
|
||||
func TestUIKeyQuotaDataflowIsLive(t *testing.T) {
|
||||
src := stripJSComments(uiSource(t))
|
||||
|
||||
badge, ok := jsFunctionBody(src, "scopeQuotaBadge")
|
||||
if !ok {
|
||||
t.Fatal("scopeQuotaBadge not found")
|
||||
}
|
||||
for _, field := range []string{"m.token_quota", "m.req_quota", "m.period"} {
|
||||
if !strings.Contains(badge, field) {
|
||||
t.Errorf("scopeQuotaBadge does not read %q — the cap would never show on the model brick", field)
|
||||
}
|
||||
}
|
||||
|
||||
brick, ok := jsFunctionBody(src, "scopeHtml")
|
||||
if !ok {
|
||||
t.Fatal("scopeHtml not found")
|
||||
}
|
||||
edit, ok := jsFunctionBody(src, "scopeEdit")
|
||||
if !ok {
|
||||
t.Fatal("scopeEdit not found")
|
||||
}
|
||||
brickCode, editCode := stripJSComments(brick), stripJSComments(edit)
|
||||
reader := stripJSComments(mustBody(t, src, "readScopes"))
|
||||
for _, ds := range []string{"quota", "reqquota", "period", "hours"} {
|
||||
if !strings.Contains(brickCode, "data-"+ds+"=") {
|
||||
t.Errorf("scopeHtml does not write data-%s, so the editor has nothing to prefill", ds)
|
||||
}
|
||||
// the editor pre-fills through readScopes(), which is what walks the
|
||||
// bricks' data attributes — check the reader, not the form
|
||||
if !strings.Contains(reader, "dataset."+ds) {
|
||||
t.Errorf("readScopes does not read dataset.%s — editing a brick would save zeros over it", ds)
|
||||
}
|
||||
}
|
||||
// and the form must actually consume what readScopes produced
|
||||
for _, field := range []string{"sc.token_quota", "sc.req_quota", "sc.period", "sc.hours"} {
|
||||
if !strings.Contains(editCode, field) {
|
||||
t.Errorf("scopeEdit does not prefill from %q", field)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func mustBody(t *testing.T, src, fn string) string {
|
||||
t.Helper()
|
||||
b, ok := jsFunctionBody(src, fn)
|
||||
if !ok {
|
||||
t.Fatalf("%s not found", fn)
|
||||
}
|
||||
return b
|
||||
}
|
||||
|
||||
// The quota period vocabulary must match the server's, or the UI can offer a
|
||||
// value the API rejects.
|
||||
func TestUIQuotaPeriodsMatchServer(t *testing.T) {
|
||||
src := uiSource(t)
|
||||
for _, p := range []string{`value=""`, `value="hour"`, `value="week"`, `value="month"`, `value="nhour"`} {
|
||||
if !strings.Contains(src, p) {
|
||||
t.Errorf("UI period select is missing %s", p)
|
||||
}
|
||||
}
|
||||
for _, p := range []string{`"hour"`, `"week"`, `"month"`, `"nhour"`} {
|
||||
if !strings.Contains(src, `if (p === `+p+`)`) && !strings.Contains(src, `=== `+p+`)`) {
|
||||
t.Errorf("periodText() does not describe %s, so a badge would omit the window", p)
|
||||
}
|
||||
}
|
||||
}
|
||||
323
internal/lua/adapter_qwen38_test.go
Normal file
323
internal/lua/adapter_qwen38_test.go
Normal file
@ -0,0 +1,323 @@
|
||||
package lua
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"strings"
|
||||
"testing"
|
||||
)
|
||||
|
||||
// The fixtures below are REAL responses captured from
|
||||
// https://discovery-api.intern-ai.org.cn/v1 with qwen3.8-27b — not hand-written
|
||||
// approximations. Two of them exist only because the endpoint behaves in a way a
|
||||
// stock OpenAI adapter gets wrong, and neither is guessable from the schema.
|
||||
|
||||
// reasoningDominatesTheBudget is the fact that made a dedicated adapter
|
||||
// necessary: a one-sentence visual answer came back as completion=303 with
|
||||
// reasoning=295. 97% of the budget went to thinking and only 8 tokens reached
|
||||
// content. Handing the upstream's content straight through yields what looks
|
||||
// like an empty answer from a perfectly healthy request.
|
||||
const reasoningDominatesTheBudget = `{
|
||||
"choices": [{"index":0,"message":{"role":"assistant","content":"\n\n这张图是白色的。",
|
||||
"reasoning_content":"分析请求…分析图片…颜色是纯白"},"logprobs":null,
|
||||
"finish_reason":"stop"}],
|
||||
"created":1791179068,"model":"qwen3.8-27b-fp8","object":"chat.completion",
|
||||
"usage":{"prompt_tokens":125,"total_tokens":428,"completion_tokens":303,
|
||||
"prompt_tokens_details":null,
|
||||
"completion_tokens_details":{"reasoning_tokens":295}}
|
||||
}`
|
||||
|
||||
// truncatedReasoningOnly is what a small max_tokens produces: HTTP 200, content
|
||||
// null, finish_reason "length". It is NOT an upstream failure.
|
||||
const truncatedReasoningOnly = `{
|
||||
"choices": [{"index":0,"message":{"role":"assistant","content":null,
|
||||
"reasoning_content":"思考过程:1. **分析请求**…图片的颜色是纯","refusal":null},
|
||||
"finish_reason":"length"}],
|
||||
"model":"qwen3.8-27b-fp8","object":"chat.completion",
|
||||
"usage":{"prompt_tokens":127,"total_tokens":190,"completion_tokens":63}
|
||||
}`
|
||||
|
||||
const qwen38ToolCallResponse = `{
|
||||
"choices": [{"index":0,"message":{"role":"assistant","content":"",
|
||||
"reasoning_content":"用户询问北京天气。我需要调用 get_weather 函数,参数为城市 \"北京\"。",
|
||||
"tool_calls":[{"id":"chatcmpl-tool-afc0299b5810e044","type":"function",
|
||||
"function":{"name":"get_weather","arguments":"{\"city\": \"北京\"}"}}]},
|
||||
"finish_reason":"tool_calls"}],
|
||||
"model":"qwen3.8-27b-fp8",
|
||||
"usage":{"prompt_tokens":305,"total_tokens":353,"completion_tokens":48,
|
||||
"completion_tokens_details":{"reasoning_tokens":20}}
|
||||
}`
|
||||
|
||||
func qwen38VM(t *testing.T) *VM {
|
||||
t.Helper()
|
||||
vm := NewVM(freshAdapterDir(t))
|
||||
if err := vm.Start(); err != nil {
|
||||
t.Fatalf("adapter VM: %v", err)
|
||||
}
|
||||
t.Cleanup(vm.Stop)
|
||||
return vm
|
||||
}
|
||||
|
||||
func qwen38Response(t *testing.T, vm *VM, raw string) map[string]interface{} {
|
||||
t.Helper()
|
||||
out, err := vm.Transform("qwen38", "transform_response", raw)
|
||||
if err != nil {
|
||||
t.Fatalf("transform_response: %v", err)
|
||||
}
|
||||
var m map[string]interface{}
|
||||
if err := json.Unmarshal([]byte(out), &m); err != nil {
|
||||
t.Fatalf("adapter returned non-JSON: %v\n%s", err, out)
|
||||
}
|
||||
return m
|
||||
}
|
||||
|
||||
func qwen38Stream(t *testing.T, vm *VM, raw string) map[string]interface{} {
|
||||
t.Helper()
|
||||
out, err := vm.Transform("qwen38", "transform_stream_chunk", raw)
|
||||
if err != nil {
|
||||
t.Fatalf("transform_stream_chunk: %v", err)
|
||||
}
|
||||
var m map[string]interface{}
|
||||
if err := json.Unmarshal([]byte(out), &m); err != nil {
|
||||
t.Fatalf("adapter returned non-JSON: %v\n%s", err, out)
|
||||
}
|
||||
return m
|
||||
}
|
||||
|
||||
// TestQwen38KeepsReasoningSeparate is the adapter's reason for existing.
|
||||
// Folding reasoning into content would render the model's private deliberation as
|
||||
// its answer; dropping it makes a reasoning model look like it returned nothing.
|
||||
func TestQwen38KeepsReasoningSeparate(t *testing.T) {
|
||||
vm := qwen38VM(t)
|
||||
u := qwen38Response(t, vm, reasoningDominatesTheBudget)
|
||||
|
||||
content, _ := u["content"].(string)
|
||||
reasoning, _ := u["reasoning_content"].(string)
|
||||
|
||||
if reasoning == "" {
|
||||
t.Fatal("reasoning_content was dropped — a reasoning model then looks like it returned nothing")
|
||||
}
|
||||
if strings.Contains(content, "分析请求") {
|
||||
t.Errorf("reasoning was folded into content: %q", content)
|
||||
}
|
||||
if !strings.Contains(content, "白色的") {
|
||||
t.Errorf("content = %q, want the answer text", content)
|
||||
}
|
||||
if reasoning == content {
|
||||
t.Error("content and reasoning_content are identical; they are disjoint upstream")
|
||||
}
|
||||
}
|
||||
|
||||
// TestQwen38ReportsTruncationRatherThanFailing: content=null with
|
||||
// finish_reason="length" is a truncated reasoning response, not a broken
|
||||
// upstream. Losing the finish_reason makes it indistinguishable from a failure,
|
||||
// and losing the partial reasoning throws away the only evidence of what happened.
|
||||
func TestQwen38ReportsTruncationRatherThanFailing(t *testing.T) {
|
||||
vm := qwen38VM(t)
|
||||
u := qwen38Response(t, vm, truncatedReasoningOnly)
|
||||
|
||||
if u["finish_reason"] != "length" {
|
||||
t.Errorf("finish_reason = %v, want length", u["finish_reason"])
|
||||
}
|
||||
if reasoning, _ := u["reasoning_content"].(string); reasoning == "" {
|
||||
t.Error("the partial reasoning was discarded")
|
||||
}
|
||||
if content, _ := u["content"].(string); content != "" {
|
||||
t.Errorf("content = %q, want empty string for a null upstream content", content)
|
||||
}
|
||||
}
|
||||
|
||||
// TestQwen38CarriesReasoningTokensWithoutInflatingTotals: reasoning_tokens is a
|
||||
// SUBSET of completion_tokens. Adding them double-counts every reasoning request,
|
||||
// and billing prices the pair separately.
|
||||
func TestQwen38CarriesReasoningTokensWithoutInflatingTotals(t *testing.T) {
|
||||
vm := qwen38VM(t)
|
||||
u := qwen38Response(t, vm, reasoningDominatesTheBudget)
|
||||
|
||||
tu, ok := u["token_usage"].(map[string]interface{})
|
||||
if !ok {
|
||||
t.Fatalf("token_usage = %#v", u["token_usage"])
|
||||
}
|
||||
completion, _ := tu["completion"].(float64)
|
||||
total, _ := tu["total"].(float64)
|
||||
|
||||
// The NESTING is the point, not the presence of the field.
|
||||
//
|
||||
// internal/gateway/chat.go reads reasoning in exactly one place:
|
||||
// resp.TokenUsage.CompletionTokensDetails.ReasoningTokens
|
||||
// An adapter that emits reasoning_tokens at the TOP of token_usage type-checks
|
||||
// fine and is then dropped at the one spot that could have kept it — measured
|
||||
// on the live gateway: the model reported completion=88 and the gateway's
|
||||
// response carried reasoning=null.
|
||||
//
|
||||
// So this asserts the path, not merely that a number appears somewhere.
|
||||
det, ok := tu["completion_tokens_details"].(map[string]interface{})
|
||||
if !ok {
|
||||
t.Fatalf("token_usage has no completion_tokens_details: %#v — chat.go reads "+
|
||||
"reasoning from there and nowhere else, so the number is discarded", tu)
|
||||
}
|
||||
reasoning, has := det["reasoning_tokens"].(float64)
|
||||
if !has {
|
||||
t.Fatal("completion_tokens_details.reasoning_tokens missing — billing cannot price reasoning separately")
|
||||
}
|
||||
// Relational, not a presence check. The first version of this test only
|
||||
// asserted the field existed, and a mutation setting it to 0 passed: 0 is
|
||||
// present. The upstream value is known (295), so the exact number is
|
||||
// asserted — a zero silently misprices every reasoning request.
|
||||
if reasoning != 295 {
|
||||
t.Errorf("reasoning_tokens = %v, want 295 (the upstream value)", reasoning)
|
||||
}
|
||||
if reasoning >= completion {
|
||||
t.Errorf("reasoning_tokens (%v) must be a SUBSET of completion_tokens (%v)", reasoning, completion)
|
||||
}
|
||||
if completion != 303 || total != 428 {
|
||||
t.Errorf("usage was altered: completion=%v total=%v, want 303/428", completion, total)
|
||||
}
|
||||
}
|
||||
|
||||
// TestQwen38DecodesToolCallArguments: the upstream sends arguments as a JSON
|
||||
// STRING, like every other OpenAI-compatible endpoint. Passing it through
|
||||
// undecoded leaves the caller to parse it.
|
||||
func TestQwen38DecodesToolCallArguments(t *testing.T) {
|
||||
vm := qwen38VM(t)
|
||||
u := qwen38Response(t, vm, qwen38ToolCallResponse)
|
||||
|
||||
tcs, ok := u["tool_calls"].([]interface{})
|
||||
if !ok || len(tcs) != 1 {
|
||||
t.Fatalf("tool_calls = %#v", u["tool_calls"])
|
||||
}
|
||||
tc, _ := tcs[0].(map[string]interface{})
|
||||
if tc["name"] != "get_weather" {
|
||||
t.Errorf("name = %v", tc["name"])
|
||||
}
|
||||
args, ok := tc["arguments"].(map[string]interface{})
|
||||
if !ok {
|
||||
t.Fatalf("arguments = %#v — the upstream sends a JSON string and it must be decoded", tc["arguments"])
|
||||
}
|
||||
if args["city"] != "北京" {
|
||||
t.Errorf("arguments.city = %v", args["city"])
|
||||
}
|
||||
}
|
||||
|
||||
// TestQwen38StreamKeepsReasoningSeparate: measured 20 reasoning chunks against 1
|
||||
// content chunk on a one-sentence answer. Folding them would make the client
|
||||
// render the deliberation as the reply.
|
||||
func TestQwen38StreamKeepsReasoningSeparate(t *testing.T) {
|
||||
vm := qwen38VM(t)
|
||||
const reasoningChunk = `{"choices":[{"index":0,"delta":{"reasoning_content":"思考过程:1. **分析请求**"},
|
||||
"finish_reason":null}],"model":"qwen3.8-27b-fp8"}`
|
||||
const contentChunk = `{"choices":[{"index":0,"delta":{"content":"递归就是函数调用自身"},
|
||||
"finish_reason":null}],"model":"qwen3.8-27b-fp8"}`
|
||||
const usageChunk = `{"choices":[],"usage":{"prompt_tokens":64,"total_tokens":132,
|
||||
"completion_tokens":68,"completion_tokens_details":{"reasoning_tokens":62}}}`
|
||||
|
||||
a := qwen38Stream(t, vm, reasoningChunk)
|
||||
b := qwen38Stream(t, vm, contentChunk)
|
||||
c := qwen38Stream(t, vm, usageChunk)
|
||||
|
||||
if content, mixed := a["content"].(string); mixed && content != "" {
|
||||
t.Errorf("a reasoning chunk leaked into content: %q", content)
|
||||
}
|
||||
if reasoning, _ := a["reasoning_content"].(string); !strings.Contains(reasoning, "思考过程") {
|
||||
t.Errorf("reasoning chunk lost its text: %#v", a)
|
||||
}
|
||||
if content, _ := b["content"].(string); !strings.Contains(content, "递归") {
|
||||
t.Errorf("content chunk lost its text: %#v", b)
|
||||
}
|
||||
if c["usage"] == nil {
|
||||
t.Errorf("the terminal usage chunk was dropped: %#v", c)
|
||||
return
|
||||
}
|
||||
u, _ := c["usage"].(map[string]interface{})
|
||||
det, _ := u["completion_tokens_details"].(map[string]interface{})
|
||||
if det == nil || det["reasoning_tokens"] == nil {
|
||||
t.Errorf("usage chunk did not nest reasoning under completion_tokens_details: %#v — "+
|
||||
"chat.go reads it there, so a top-level field is discarded", u)
|
||||
}
|
||||
}
|
||||
|
||||
// TestQwen38ReportsMediaFailureCode: a remote image URL times out upstream and
|
||||
// comes back as media_fetch_failed. Surfacing only the message makes an image
|
||||
// fetch problem look like a model problem.
|
||||
func TestQwen38ReportsMediaFailureCode(t *testing.T) {
|
||||
vm := qwen38VM(t)
|
||||
body := `{"error":{"message":"image fetch timed out","type":"media_fetch_failed",
|
||||
"param":null,"code":"media_fetch_failed","trace_id":"f3ca4aa"}}`
|
||||
reason, ok, err := vm.TransformError("qwen38", 400, body)
|
||||
if err != nil {
|
||||
t.Fatalf("TransformError: %v", err)
|
||||
}
|
||||
if !ok {
|
||||
t.Fatal("the media failure was swallowed")
|
||||
}
|
||||
if !strings.Contains(reason, "image fetch timed out") || !strings.Contains(reason, "media_fetch_failed") {
|
||||
t.Errorf("reason = %q; it must carry both the text and the code", reason)
|
||||
}
|
||||
}
|
||||
|
||||
// TestQwen38InvalidKeyError: the endpoint answers a bad key with 401 and
|
||||
// code "invalid_api_key". That one must NOT be annotated with its own code — it
|
||||
// would render as "invalid API key (invalid_api_key)".
|
||||
func TestQwen38InvalidKeyError(t *testing.T) {
|
||||
vm := qwen38VM(t)
|
||||
body := `{"error":{"message":"invalid API key","type":"invalid_request_error",
|
||||
"param":null,"code":"invalid_api_key"}}`
|
||||
reason, ok, err := vm.TransformError("qwen38", 401, body)
|
||||
if err != nil || !ok {
|
||||
t.Fatalf("TransformError: ok=%v err=%v", ok, err)
|
||||
}
|
||||
if reason != "invalid API key" {
|
||||
t.Errorf("reason = %q, want the bare message", reason)
|
||||
}
|
||||
}
|
||||
|
||||
// TestQwen38SuppliesReasoningOnToolCallTurns: the upstream requires an assistant
|
||||
// message carrying tool_calls to echo reasoning_content. OpenAI-shaped clients
|
||||
// never send it, so without this the SECOND turn of any tool-using conversation
|
||||
// is rejected.
|
||||
func TestQwen38SuppliesReasoningOnToolCallTurns(t *testing.T) {
|
||||
vm := qwen38VM(t)
|
||||
const req = `{"model":"qwen3.8-27b","messages":[
|
||||
{"role":"user","content":"北京天气"},
|
||||
{"role":"assistant","content":"","tool_calls":[{"id":"c1","type":"function",
|
||||
"function":{"name":"get_weather","arguments":"{\"city\":\"北京\"}"}}]},
|
||||
{"role":"tool","tool_call_id":"c1","content":"晴 25C"}]}`
|
||||
out, err := vm.Transform("qwen38", "transform_request", req)
|
||||
if err != nil {
|
||||
t.Fatalf("transform_request: %v", err)
|
||||
}
|
||||
var decoded map[string]interface{}
|
||||
if err := json.Unmarshal([]byte(out), &decoded); err != nil {
|
||||
t.Fatalf("request transform did not return JSON: %v\n%s", err, out)
|
||||
}
|
||||
msgs, _ := decoded["messages"].([]interface{})
|
||||
if len(msgs) < 2 {
|
||||
t.Fatalf("messages = %#v", decoded["messages"])
|
||||
}
|
||||
assistant, _ := msgs[1].(map[string]interface{})
|
||||
if _, ok := assistant["reasoning_content"]; !ok {
|
||||
t.Errorf("assistant turn with tool_calls has no reasoning_content; the upstream "+
|
||||
"rejects the next turn. Got %#v", assistant)
|
||||
}
|
||||
}
|
||||
|
||||
// TestQwen38PassesImagesThrough: this adapter's job is reasoning handling, not
|
||||
// media translation. A multimodal content array must survive untouched — an
|
||||
// adapter that re-encodes it is the kind of change that silently breaks images
|
||||
// while every text test still passes.
|
||||
func TestQwen38PassesImagesThrough(t *testing.T) {
|
||||
vm := qwen38VM(t)
|
||||
const req = `{"model":"qwen3.8-27b","messages":[{"role":"user","content":[
|
||||
{"type":"text","text":"这张图是什么颜色?"},
|
||||
{"type":"image_url","image_url":{"url":"data:image/png;base64,iVBORw0KGgo="}}]}],
|
||||
"max_tokens":600}`
|
||||
out, err := vm.Transform("qwen38", "transform_request", req)
|
||||
if err != nil {
|
||||
t.Fatalf("transform_request: %v", err)
|
||||
}
|
||||
if !strings.Contains(out, "image_url") || !strings.Contains(out, "data:image/png;base64,") {
|
||||
t.Errorf("the image part was lost or rewritten:\n%s", out)
|
||||
}
|
||||
if strings.Contains(out, "base64,iVBORw0KGgo=") == false {
|
||||
t.Errorf("the base64 payload was altered:\n%s", out)
|
||||
}
|
||||
}
|
||||
@ -5,10 +5,12 @@ adapter.version = "1.0.0"
|
||||
adapter.endpoint = "/chat/completions"
|
||||
adapter.headers = {}
|
||||
|
||||
-- AgentRouter 的 WAF 会按客户端指纹白名单校验,只放行官方客户端。
|
||||
-- 这里通过 build_headers 注入官方客户端 User-Agent 以通过校验。
|
||||
-- 默认 QwenCode 指纹(实测可通);可在 source.meta.user_agent 覆盖成 Claude Code 指纹。
|
||||
|
||||
-- AgentRouter 的阿里云 WAF 按三重维度校验:出口 IP(必须海外)、
|
||||
-- TLS 指纹(JS/Go HTTP 层与 curl 不同)、浏览器指纹 header 集。
|
||||
-- 纯 UA 已不够:WAF 拦 12 种 agent CLI UA 中的 10 种(全 405 HTML),
|
||||
-- 只有同时带 Origin: qwen.ai + Referer + X-Requested-With 才放行,
|
||||
-- 且同 IP 短时间连续 3+ 次请求会被速率封禁,靠 adapter 侧低频调用 +
|
||||
-- llmsproxy 的 prefFailStep 冷却自然限制,来源无法控制速率。
|
||||
local default_ua = "QwenCode/0.2.0 (linux; x64)"
|
||||
|
||||
-- AgentRouter fronts Claude models (claude-opus-4-8), and Claude upstreams
|
||||
@ -65,6 +67,13 @@ function adapter.build_headers(meta)
|
||||
["Content-Type"] = "application/json",
|
||||
["Authorization"] = "Bearer " .. meta.api_key,
|
||||
["User-Agent"] = ua,
|
||||
-- QwenCode client fingerprint: WAF rejects bare UA without these.
|
||||
["Accept"] = "application/json, text/plain, */*",
|
||||
["Accept-Language"] = "zh-CN,zh;q=0.9",
|
||||
["Origin"] = "https://qwen.ai",
|
||||
["Referer"] = "https://qwen.ai/",
|
||||
["X-Requested-With"] = "XMLHttpRequest",
|
||||
["Connection"] = "keep-alive",
|
||||
}
|
||||
end
|
||||
|
||||
|
||||
@ -2,11 +2,30 @@ local adapter = {}
|
||||
|
||||
adapter.name = "gemini"
|
||||
adapter.version = "2.0.0"
|
||||
adapter.endpoint = "/v1/models"
|
||||
-- The Go layer builds the request URL as base_url + endpoint, statically
|
||||
-- (see provider.URL). Gemini's real API is POST
|
||||
-- /v1beta/models/{model}:generateContent, and streaming is the same path with
|
||||
-- a ":streamGenerateContent" verb -- the model name is part of the PATH, so it
|
||||
-- cannot live in a static endpoint string.
|
||||
--
|
||||
-- "{model}" is therefore a placeholder the Go layer substitutes with the model
|
||||
-- this request actually sends (provider.urlFor substitutes it; see
|
||||
-- provider.go). Streaming additionally rewrites the ":generateContent" verb to
|
||||
-- ":streamGenerateContent" on the same template.
|
||||
--
|
||||
-- Leaving this as a bare "/v1/models" would call Gemini's model-LIST endpoint,
|
||||
-- which answers 405 to POST -- so the preset template would create a source that
|
||||
-- can never work. A source that overrides `endpoint:` bypasses the template
|
||||
-- entirely and must then spell the whole path itself.
|
||||
adapter.endpoint = "/v1beta/models/{model}:generateContent"
|
||||
adapter.headers = {}
|
||||
|
||||
-- Gemini API: POST /v1/models/{model}:generateContent
|
||||
-- Auth: API key in query param ?key=XXX or Authorization: Bearer XXX
|
||||
--
|
||||
-- 请求方向必须把 OpenAI 的工具调用翻译成 Gemini 的 functionCall /
|
||||
-- functionResponse,否则 agent 回放的历史里助手那一轮的调用会凭空消失,
|
||||
-- 紧随其后的工具结果就成了「无来源」的孤立结果,模型只能反复重发同一个调用。
|
||||
function adapter.transform_request(raw_body)
|
||||
local ok, req = pcall(json.decode, raw_body)
|
||||
if not ok then return raw_body end
|
||||
@ -14,16 +33,23 @@ function adapter.transform_request(raw_body)
|
||||
-- 将 OpenAI 风格 content(字符串或 [{type:*}] 数组)拆成 Gemini parts
|
||||
local function to_parts(content)
|
||||
if type(content) == "string" then
|
||||
if content == "" then return {} end
|
||||
return { { text = content } }
|
||||
end
|
||||
local parts = {}
|
||||
for _, p in ipairs(content or {}) do
|
||||
if p.type == "text" then
|
||||
table.insert(parts, { text = p.text })
|
||||
elseif p.type == "image_url" and type(p.image_url) == "table" and p.image_url.url then
|
||||
local mt, b64 = string.match(p.image_url.url, "^data:([^,]+);base64,(.+)$")
|
||||
if b64 then
|
||||
table.insert(parts, { inline_data = { mime_type = mt or "image/png", data = b64 } })
|
||||
if type(p) == "string" then
|
||||
if p ~= "" then table.insert(parts, { text = p }) end
|
||||
elseif type(p) == "table" then
|
||||
if p.type == "text" then
|
||||
if p.text ~= nil and p.text ~= "" then
|
||||
table.insert(parts, { text = p.text })
|
||||
end
|
||||
elseif p.type == "image_url" and type(p.image_url) == "table" and p.image_url.url then
|
||||
local mt, b64 = string.match(p.image_url.url, "^data:([^,]+);base64,(.+)$")
|
||||
if b64 then
|
||||
table.insert(parts, { inline_data = { mime_type = mt or "image/png", data = b64 } })
|
||||
end
|
||||
end
|
||||
end
|
||||
end
|
||||
@ -31,11 +57,59 @@ function adapter.transform_request(raw_body)
|
||||
end
|
||||
|
||||
local contents = {}
|
||||
local system = ""
|
||||
local call_names = {} -- tool_call_id -> 函数名(functionResponse 只认名字)
|
||||
|
||||
for _, m in ipairs(req.messages or {}) do
|
||||
table.insert(contents, {
|
||||
role = (m.role == "assistant") and "model" or m.role,
|
||||
parts = to_parts(m.content)
|
||||
})
|
||||
local role = m.role or "user"
|
||||
|
||||
if role == "system" then
|
||||
if type(m.content) == "string" then
|
||||
system = system .. m.content .. "\n"
|
||||
end
|
||||
|
||||
elseif role == "assistant" then
|
||||
local parts = to_parts(m.content)
|
||||
if type(m.tool_calls) == "table" then
|
||||
for _, tc in ipairs(m.tool_calls) do
|
||||
if type(tc) == "table" then
|
||||
local fn = tc["function"] or {}
|
||||
local args = fn.arguments
|
||||
if type(args) == "string" and args ~= "" then
|
||||
local aok, decoded = pcall(json.decode, args)
|
||||
args = aok and decoded or {}
|
||||
elseif type(args) ~= "table" then
|
||||
args = {}
|
||||
end
|
||||
if tc.id ~= nil then call_names[tc.id] = fn.name or "" end
|
||||
table.insert(parts, {
|
||||
functionCall = { name = fn.name or "", args = args }
|
||||
})
|
||||
end
|
||||
end
|
||||
end
|
||||
-- Gemini 不接受空的 parts 数组;无可发送内容的轮次直接跳过
|
||||
if #parts > 0 then
|
||||
table.insert(contents, { role = "model", parts = parts })
|
||||
end
|
||||
|
||||
elseif role == "tool" then
|
||||
local name = call_names[m.tool_call_id] or m.name or ""
|
||||
local text = type(m.content) == "string" and m.content or ""
|
||||
table.insert(contents, {
|
||||
role = "user",
|
||||
parts = { { functionResponse = {
|
||||
name = name,
|
||||
response = { content = text },
|
||||
} } }
|
||||
})
|
||||
|
||||
else
|
||||
local parts = to_parts(m.content)
|
||||
if #parts > 0 then
|
||||
table.insert(contents, { role = "user", parts = parts })
|
||||
end
|
||||
end
|
||||
end
|
||||
|
||||
local gemini_req = {
|
||||
@ -46,6 +120,30 @@ function adapter.transform_request(raw_body)
|
||||
}
|
||||
}
|
||||
|
||||
if system ~= "" then
|
||||
gemini_req.systemInstruction = { parts = { { text = system } } }
|
||||
end
|
||||
|
||||
-- tools -> functionDeclarations
|
||||
if type(req.tools) == "table" then
|
||||
local decls = {}
|
||||
for _, t in ipairs(req.tools) do
|
||||
if type(t) == "table" and type(t["function"]) == "table" then
|
||||
local fn = t["function"]
|
||||
local params = fn.parameters
|
||||
if type(params) ~= "table" then params = { type = "object", properties = {} } end
|
||||
table.insert(decls, {
|
||||
name = fn.name or "",
|
||||
description = fn.description or "",
|
||||
parameters = params,
|
||||
})
|
||||
end
|
||||
end
|
||||
if #decls > 0 then
|
||||
gemini_req.tools = { { functionDeclarations = decls } }
|
||||
end
|
||||
end
|
||||
|
||||
if req.stream then
|
||||
gemini_req.stream = true
|
||||
end
|
||||
|
||||
@ -6,6 +6,10 @@ adapter.endpoint = "/api/chat"
|
||||
adapter.headers = {}
|
||||
|
||||
-- Ollama API 格式:{ model, messages, stream, options:{temperature,num_predict} }
|
||||
-- 工具调用必须一并翻译:Ollama 的 assistant 消息用 tool_calls(arguments 是对象
|
||||
-- 而非 JSON 字符串),工具结果用 tool 角色 + tool_name。此前这里只复制了
|
||||
-- role/content,助手那轮的调用和 tool 消息的归属全部丢失,模型看到无来源的
|
||||
-- 工具结果就只能反复重发同一个调用。
|
||||
function adapter.transform_request(raw_body)
|
||||
local ok, req = pcall(json.decode, raw_body)
|
||||
if not ok then return raw_body end
|
||||
@ -22,6 +26,7 @@ function adapter.transform_request(raw_body)
|
||||
-- 转换 messages 格式(Ollama messages 支持 images base64 数组)
|
||||
if req.messages then
|
||||
local msgs = {}
|
||||
local call_names = {} -- tool_call_id -> 函数名
|
||||
for _, m in ipairs(req.messages) do
|
||||
local text, images
|
||||
if type(m.content) == "string" then
|
||||
@ -30,22 +35,61 @@ function adapter.transform_request(raw_body)
|
||||
text = ""
|
||||
images = {}
|
||||
for _, p in ipairs(m.content or {}) do
|
||||
if p.type == "text" then
|
||||
text = text .. (p.text or "")
|
||||
elseif p.type == "image_url" and type(p.image_url) == "table" and p.image_url.url then
|
||||
local b64 = string.match(p.image_url.url, "^data:[^,]+;base64,(.+)$")
|
||||
if b64 then table.insert(images, b64) end
|
||||
if type(p) == "table" then
|
||||
if p.type == "text" then
|
||||
text = text .. (p.text or "")
|
||||
elseif p.type == "image_url" and type(p.image_url) == "table" and p.image_url.url then
|
||||
local b64 = string.match(p.image_url.url, "^data:[^,]+;base64,(.+)$")
|
||||
if b64 then table.insert(images, b64) end
|
||||
end
|
||||
end
|
||||
end
|
||||
if #images == 0 then images = nil end
|
||||
end
|
||||
local msg = { role = m.role, content = text }
|
||||
if images then msg.images = images end
|
||||
|
||||
-- 助手轮的工具调用:arguments 转成对象
|
||||
if type(m.tool_calls) == "table" and #m.tool_calls > 0 then
|
||||
local tcs = {}
|
||||
for _, tc in ipairs(m.tool_calls) do
|
||||
if type(tc) == "table" then
|
||||
local fn = tc["function"] or {}
|
||||
local args = fn.arguments
|
||||
if type(args) == "string" and args ~= "" then
|
||||
local aok, decoded = pcall(json.decode, args)
|
||||
args = aok and decoded or {}
|
||||
elseif type(args) ~= "table" then
|
||||
args = {}
|
||||
end
|
||||
if tc.id ~= nil then call_names[tc.id] = fn.name or "" end
|
||||
table.insert(tcs, {
|
||||
["function"] = { name = fn.name or "", arguments = args }
|
||||
})
|
||||
end
|
||||
end
|
||||
if #tcs > 0 then msg.tool_calls = tcs end
|
||||
end
|
||||
|
||||
-- 工具结果:Ollama 用 tool_name 标识归属(不认 tool_call_id)
|
||||
if m.role == "tool" then
|
||||
local name = call_names[m.tool_call_id] or m.name or ""
|
||||
if m.tool_call_id ~= nil then
|
||||
msg.tool_call_id = m.tool_call_id
|
||||
end
|
||||
if name ~= "" then msg.tool_name = name end
|
||||
end
|
||||
|
||||
table.insert(msgs, msg)
|
||||
end
|
||||
ollama_req.messages = msgs
|
||||
end
|
||||
|
||||
-- tools 透传(Ollama 的 shape 与 OpenAI 一致)
|
||||
if type(req.tools) == "table" and #req.tools > 0 then
|
||||
ollama_req.tools = req.tools
|
||||
end
|
||||
|
||||
return json.encode(ollama_req)
|
||||
end
|
||||
|
||||
|
||||
@ -34,8 +34,10 @@ end
|
||||
|
||||
-- OpenAI /chat/completions format (pass-through, strip provider-specific fields)
|
||||
-- zen 上游 schema 只接受 text content part(无视觉/音频能力):多模态 part
|
||||
-- (image_url / input_audio / file 等)一律剥离;因此失去全部 content 的
|
||||
-- 消息整条丢弃,避免上游 "unknown variant `image_url`, expected `text`"。
|
||||
-- (image_url / input_audio / file 等)一律剥离。剥离后 content 变空的消息
|
||||
-- 若不再携带 tool_calls / tool_call_id 才整条丢弃(避免上游
|
||||
-- "unknown variant `image_url`, expected `text`");带工具调用的必须保留,
|
||||
-- 否则会把紧随其后的 tool 结果变成孤儿,模型会反复重发同一个调用。
|
||||
-- zen 上游角色白名单只有 system / user / assistant / tool / latest_reminder:
|
||||
-- OpenAI 的 developer(及 function 等)不在其中,直接透传会触发上游
|
||||
-- "unknown variant `developer`, expected one of ..." 错误;统一归一化为 system。
|
||||
@ -52,8 +54,14 @@ function adapter.transform_request(raw_body)
|
||||
if not ok then return raw_body end
|
||||
req.disable_thinking = nil
|
||||
req.extra_body = nil
|
||||
if type(req.stream_options) ~= "table" then req.stream_options = {} end
|
||||
req.stream_options.include_usage = true
|
||||
-- stream_options is only valid alongside stream:true. Sending it on a
|
||||
-- non-streaming request is rejected by stricter upstreams (OpenCode Go:
|
||||
-- "stream_options should be set along with stream"), which failed every
|
||||
-- non-stream call. Only attach it when the request actually streams.
|
||||
if req.stream then
|
||||
if type(req.stream_options) ~= "table" then req.stream_options = {} end
|
||||
req.stream_options.include_usage = true
|
||||
end
|
||||
if req.messages then
|
||||
local kept = {}
|
||||
for _, msg in ipairs(req.messages) do
|
||||
@ -72,7 +80,24 @@ function adapter.transform_request(raw_body)
|
||||
end
|
||||
end
|
||||
if #parts == 0 then
|
||||
drop = true
|
||||
-- Content collapsed to nothing after stripping unsupported
|
||||
-- parts. A message that still carries a tool call must
|
||||
-- NEVER be dropped: the very next message is its tool
|
||||
-- result, and dropping the call orphans that result. The
|
||||
-- model then sees a result for a call it never made and
|
||||
-- re-issues the same tool call on every turn (observed as
|
||||
-- an infinite "repeated tool call" loop).
|
||||
--
|
||||
-- This is the common shape for agent clients: an assistant
|
||||
-- turn whose content is only [thinking, toolCall] serialises
|
||||
-- to content:[] with tool_calls — exactly the case that used
|
||||
-- to vanish here. Emit "" instead, which zen accepts.
|
||||
if (type(msg.tool_calls) == "table" and #msg.tool_calls > 0)
|
||||
or msg.tool_call_id ~= nil then
|
||||
msg.content = ""
|
||||
else
|
||||
drop = true
|
||||
end
|
||||
else
|
||||
msg.content = parts
|
||||
end
|
||||
|
||||
354
internal/lua/adapters/opencodego.lua
Normal file
354
internal/lua/adapters/opencodego.lua
Normal file
@ -0,0 +1,354 @@
|
||||
local adapter = {}
|
||||
|
||||
adapter.name = "opencodego"
|
||||
adapter.version = "1.0.0"
|
||||
adapter.endpoint = "/chat/completions"
|
||||
|
||||
-- OpenCode Go(包月订阅端点,https://opencode.ai/zen/go/v1)
|
||||
--
|
||||
-- 与 OpenCode Zen(opencodezen.lua,https://opencode.ai/zen/v1)是两个不同的
|
||||
-- 服务,行为要求并不相同,因此各有专用适配器:
|
||||
--
|
||||
-- * 本适配器(Go)服务包月订阅池。上游**强制要求**
|
||||
-- `x-opencode-session` 头,缺失直接 400 "MissingSessionID"。
|
||||
-- * Go 的 thinking 模型(deepseek-v4.1-flash 等)**必须回传** assistant 轮的
|
||||
-- `reasoning_content`,否则 400:
|
||||
-- The `reasoning_content` in the thinking mode must be passed back to the API.
|
||||
-- 这是本适配器与 Zen 适配器最关键的差异 —— Zen 会剥掉它,Go 必须原样保留。
|
||||
-- * Go 上游对协议更严格:非流式请求带 `stream_options` 会被拒
|
||||
-- ("stream_options should be set along with stream")。
|
||||
--
|
||||
-- 同样需要 opencode 客户端指纹(UA + x-opencode-*)才能被正确识别与路由。
|
||||
adapter.headers = {
|
||||
["User-Agent"] = "opencode/1.18.21 ai-sdk/provider-utils/4.0.23 runtime/bun/1.3.14",
|
||||
}
|
||||
|
||||
-- 每请求生成身份头。沙箱无 os/math,用 meta.timestamp + 请求体哈希派生:
|
||||
-- 同秒内重复请求 id 相同可接受(zen 只校验存在性,不校验格式)。
|
||||
local function rand_id(prefix, seed)
|
||||
return prefix .. string.sub(sha256_hex(seed), 1, 24)
|
||||
end
|
||||
|
||||
-- 会话指纹:取历史里**第一条 user 消息**。
|
||||
--
|
||||
-- 为什么不直接用客户端的会话 id:实测抓包(tcpdump 抓 127.0.0.1:8081 的真实
|
||||
-- agent 请求)确认通用 OpenAI 客户端**根本不发**任何会话标识 —— body 里没有
|
||||
-- user / session_id / conversation_id / metadata,请求头也只有 X-Stainless-*
|
||||
-- (OpenAI JS SDK)与 User-Agent。opencode 原生客户端那套 x-opencode-session
|
||||
-- 是它自己的概念,通用客户端无从转发。
|
||||
--
|
||||
-- 退而求其次但要够用:历史被逐轮重放,**第一条 user 消息在整个会话里恒定**,
|
||||
-- 用它做指纹即可得到"每会话一个 session",而不是"每源一个 session"。
|
||||
-- 拿不到时(无 user 消息)返回空串,调用方回退到按源稳定。
|
||||
local function conversation_fingerprint(body)
|
||||
if type(body) ~= "string" or body == "" then return "" end
|
||||
local ok, req = pcall(json.decode, body)
|
||||
if not ok or type(req) ~= "table" then return "" end
|
||||
for _, m in ipairs(req.messages or {}) do
|
||||
if type(m) == "table" and m.role == "user" then
|
||||
local c = m.content
|
||||
if type(c) == "string" then
|
||||
if c ~= "" then return c end
|
||||
elseif type(c) == "table" then
|
||||
for _, part in ipairs(c) do
|
||||
if type(part) == "table" and part.type == "text" and part.text then
|
||||
return part.text
|
||||
end
|
||||
end
|
||||
end
|
||||
return ""
|
||||
end
|
||||
end
|
||||
return ""
|
||||
end
|
||||
|
||||
-- session_seed 决定上游会话号(x-opencode-session 的种子)。
|
||||
--
|
||||
-- 优先级:
|
||||
-- 1) 客户端自带的会话标识(meta.client_session)——pi 等在开启
|
||||
-- compat.sendSessionAffinityHeaders 后会发 x-session-affinity,
|
||||
-- 这是平台真实的会话 id,整个会话恒定且天然按会话隔离。
|
||||
-- 2) 退路:历史里**第一条 user 消息**做会话指纹。
|
||||
-- 3) 再退:只按源固定(连 user 消息都没有时)。
|
||||
--
|
||||
-- 为什么需要 2/3:通用 OpenAI 客户端默认**根本不发**会话标识 —— 实测抓包
|
||||
-- (tcpdump 抓 127.0.0.1:8081 真实 agent 请求)确认 body 里没有
|
||||
-- user / session_id / conversation_id / metadata,请求头也只有
|
||||
-- X-Stainless-*(OpenAI JS SDK)与 User-Agent。opencode 原生客户端那套
|
||||
-- x-opencode-session 是它自己的概念,通用客户端无从转发。
|
||||
--
|
||||
-- 为什么必须稳定:上游前缀缓存是**会话级**的。会话号每请求一变,
|
||||
-- 缓存永不命中(实测:固定会号第 2 次命中 5888,每请求换会号则恒为 0)。
|
||||
local function session_seed(meta)
|
||||
local src = (meta.source and meta.source.name) or ""
|
||||
local base = "session|llmsproxy|" .. src
|
||||
local cs = meta.client_session
|
||||
if type(cs) == "string" and cs ~= "" then
|
||||
return base .. "|client|" .. cs
|
||||
end
|
||||
return base .. "|" .. conversation_fingerprint(meta.body)
|
||||
end
|
||||
|
||||
function adapter.build_headers(meta)
|
||||
local ts = tostring(meta.timestamp or "")
|
||||
local src = (meta.source and meta.source.name) or ""
|
||||
return {
|
||||
["User-Agent"] = adapter.headers["User-Agent"],
|
||||
["x-opencode-client"] = "cli",
|
||||
-- project 固定:同网关实例共享一个工作区身份
|
||||
["x-opencode-project"] = string.sub(sha256_hex("llmsproxy|" .. src), 1, 32),
|
||||
-- session 必须**按源稳定**,不能每请求换:上游的前缀缓存在同一 session
|
||||
-- 内才复用。实测(同一段 6032 token 提示词):
|
||||
-- 固定 session -> 第 2 次命中 5888/6032,cached_tokens=5888
|
||||
-- 每请求换 session -> 永远 0 命中
|
||||
-- 原先用 meta.timestamp 派生,等于每请求都是新会话,缓存永远无效,
|
||||
-- 上游也无法做会话亲和路由。
|
||||
["x-opencode-session"] = rand_id("ses_", session_seed(meta)),
|
||||
-- request id 仍每请求唯一(它只是请求标识,不参与缓存键)
|
||||
["x-opencode-request"] = rand_id("msg_", "request|" .. ts .. "|" .. tostring(meta.body or "")),
|
||||
}
|
||||
end
|
||||
|
||||
-- OpenAI /chat/completions format (pass-through, strip provider-specific fields)
|
||||
-- zen 上游 schema 只接受 text content part(无视觉/音频能力):多模态 part
|
||||
-- (image_url / input_audio / file 等)一律剥离。剥离后 content 变空的消息
|
||||
-- 若不再携带 tool_calls / tool_call_id 才整条丢弃(避免上游
|
||||
-- "unknown variant `image_url`, expected `text`");带工具调用的必须保留,
|
||||
-- 否则会把紧随其后的 tool 结果变成孤儿,模型会反复重发同一个调用。
|
||||
-- zen 上游角色白名单只有 system / user / assistant / tool / latest_reminder:
|
||||
-- OpenAI 的 developer(及 function 等)不在其中,直接透传会触发上游
|
||||
-- "unknown variant `developer`, expected one of ..." 错误;统一归一化为 system。
|
||||
local ROLE_WHITELIST = {
|
||||
system = true,
|
||||
user = true,
|
||||
assistant = true,
|
||||
tool = true,
|
||||
latest_reminder = true,
|
||||
}
|
||||
|
||||
function adapter.transform_request(raw_body)
|
||||
local ok, req = pcall(json.decode, raw_body)
|
||||
if not ok then return raw_body end
|
||||
req.disable_thinking = nil
|
||||
req.extra_body = nil
|
||||
-- stream_options is only valid alongside stream:true; sending it on a
|
||||
-- non-streaming request is rejected by strict upstreams.
|
||||
if req.stream then
|
||||
if type(req.stream_options) ~= "table" then req.stream_options = {} end
|
||||
req.stream_options.include_usage = true
|
||||
end
|
||||
if req.messages then
|
||||
local kept = {}
|
||||
for _, msg in ipairs(req.messages) do
|
||||
if type(msg.role) == "string" and not ROLE_WHITELIST[msg.role] then
|
||||
msg.role = "system"
|
||||
end
|
||||
-- Go 的 thinking 模式**必须**回传 reasoning_content:剥掉它会让
|
||||
-- 上游直接 400("The `reasoning_content` in the thinking mode must
|
||||
-- be passed back to the API"),agent 的每一轮都会失败。
|
||||
-- 注意:与 opencodezen.lua 的行为**相反**,不要在这里剥。
|
||||
--
|
||||
-- 但绝大多数客户端(pi 等)根本不保存也不回传 reasoning,只保留
|
||||
-- 工具调用本身。上游只在“带 tool_calls 的助手轮”上校验这个字段,
|
||||
-- 实测**空串即可通过校验**,所以缺省时补空串:既满足上游的
|
||||
-- 一致性要求,又不伪造任何推理内容(用户看到的 reasoning 仍然是
|
||||
-- 上游本轮真实返回的)。
|
||||
if msg.role == "assistant"
|
||||
and type(msg.tool_calls) == "table" and #msg.tool_calls > 0
|
||||
and msg.reasoning_content == nil then
|
||||
msg.reasoning_content = ""
|
||||
end
|
||||
local drop = false
|
||||
if type(msg.content) == "table" then
|
||||
local parts = {}
|
||||
for _, part in ipairs(msg.content) do
|
||||
if type(part) == "table" and part.type ~= nil and part.type ~= "text" then
|
||||
-- multimodal part not supported by zen
|
||||
else
|
||||
table.insert(parts, part)
|
||||
end
|
||||
end
|
||||
if #parts == 0 then
|
||||
-- Content collapsed to nothing after stripping unsupported
|
||||
-- parts. A message that still carries a tool call must
|
||||
-- NEVER be dropped: the very next message is its tool
|
||||
-- result, and dropping the call orphans that result. The
|
||||
-- model then sees a result for a call it never made and
|
||||
-- re-issues the same tool call on every turn (observed as
|
||||
-- an infinite "repeated tool call" loop).
|
||||
if (type(msg.tool_calls) == "table" and #msg.tool_calls > 0)
|
||||
or msg.tool_call_id ~= nil then
|
||||
msg.content = ""
|
||||
else
|
||||
drop = true
|
||||
end
|
||||
else
|
||||
msg.content = parts
|
||||
end
|
||||
end
|
||||
if not drop then
|
||||
table.insert(kept, msg)
|
||||
end
|
||||
end
|
||||
req.messages = kept
|
||||
end
|
||||
return json.encode(req)
|
||||
end
|
||||
|
||||
function adapter.transform_response(raw_body)
|
||||
local ok, resp = pcall(json.decode, raw_body)
|
||||
if not ok or resp == nil then return raw_body end
|
||||
|
||||
local unified = {
|
||||
content = "",
|
||||
finish_reason = "",
|
||||
token_usage = { prompt = 0, completion = 0, total = 0 }
|
||||
}
|
||||
|
||||
if type(resp.usage) == "table" then
|
||||
unified.token_usage.prompt = resp.usage.prompt_tokens or 0
|
||||
unified.token_usage.completion = resp.usage.completion_tokens or 0
|
||||
unified.token_usage.total = resp.usage.total_tokens or 0
|
||||
local hit = 0
|
||||
if type(resp.usage.prompt_tokens_details) == "table" and resp.usage.prompt_tokens_details.cached_tokens ~= nil then
|
||||
hit = resp.usage.prompt_tokens_details.cached_tokens
|
||||
unified.token_usage.prompt_tokens_details = { cached_tokens = hit }
|
||||
end
|
||||
if (resp.usage.prompt_cache_hit_tokens or 0) > 0 then
|
||||
unified.token_usage.prompt_cache_hit_tokens = resp.usage.prompt_cache_hit_tokens
|
||||
unified.token_usage.prompt_cache_miss_tokens = resp.usage.prompt_cache_miss_tokens or 0
|
||||
if hit == 0 then
|
||||
unified.token_usage.prompt_tokens_details = { cached_tokens = resp.usage.prompt_cache_hit_tokens }
|
||||
end
|
||||
end
|
||||
-- 上游报告输出里有多少是思考 token。不给客户端的话,无法判断
|
||||
-- completion_tokens 里多少是「可见回答」、多少是「思考」(都按输出计费)。
|
||||
local ctd = resp.usage.completion_tokens_details
|
||||
if type(ctd) == "table" and ctd.reasoning_tokens ~= nil then
|
||||
unified.token_usage.completion_tokens_details = { reasoning_tokens = ctd.reasoning_tokens }
|
||||
end
|
||||
end
|
||||
|
||||
-- 本次调用的费用,上游放在**顶层**且是**字符串**(如 "0"、"0.0012")。
|
||||
-- Go 订阅是包月制,恒为 "0";只有 Zen 按量付费模型才有信息量。
|
||||
-- 原样透传:网关不解析、不换算、不汇总 —— 它只是上游事实的搬运者。
|
||||
if resp.cost ~= nil then
|
||||
unified.cost = tostring(resp.cost)
|
||||
end
|
||||
|
||||
if type(resp.choices) == "table" and #resp.choices > 0 then
|
||||
local ch = resp.choices[1]
|
||||
if type(ch.message) == "table" then
|
||||
unified.content = ch.message.content or ""
|
||||
local reasoning = ch.message.reasoning_content or ch.message.reasoning
|
||||
if reasoning then
|
||||
unified.reasoning_content = reasoning
|
||||
end
|
||||
if type(ch.message.tool_calls) == "table" then
|
||||
local tcs = {}
|
||||
for _, tc in ipairs(ch.message.tool_calls) do
|
||||
local args_ok, args = pcall(json.decode, tc["function"].arguments)
|
||||
if not args_ok then args = {} end
|
||||
table.insert(tcs, {
|
||||
id = tc.id,
|
||||
type = tc.type or "function",
|
||||
name = tc["function"].name,
|
||||
arguments = args
|
||||
})
|
||||
end
|
||||
unified.tool_calls = tcs
|
||||
end
|
||||
end
|
||||
unified.finish_reason = ch.finish_reason or ""
|
||||
end
|
||||
|
||||
return json.encode(unified)
|
||||
end
|
||||
|
||||
function adapter.transform_stream_chunk(raw_chunk)
|
||||
local ok, chunk = pcall(json.decode, raw_chunk)
|
||||
if not ok then return "" end
|
||||
|
||||
-- OpenAI-style streams may attach usage to a chunk with empty choices
|
||||
-- (the final usage chunk). Preserve it; the gateway emits it as the
|
||||
-- terminal usage chunk. Note: keys must match Go's TokenUsage json tags
|
||||
-- (prompt/completion/total); the gateway re-emits standard *_tokens.
|
||||
local uses = nil
|
||||
if type(chunk.usage) == "table" then
|
||||
uses = {
|
||||
prompt = chunk.usage.prompt_tokens or chunk.usage.prompt or 0,
|
||||
completion = chunk.usage.completion_tokens or chunk.usage.completion or 0,
|
||||
total = chunk.usage.total_tokens or chunk.usage.total or 0,
|
||||
}
|
||||
-- 独立字段与 prompt_tokens_details **并存**透传,不要写成 elseif:
|
||||
-- 上游命中缓存时两者都发,非流式路径也是两个都带。写成 elseif 会让
|
||||
-- 流式丢掉 prompt_cache_hit_tokens/miss,与同一源的**非流式**产出不一致
|
||||
-- (dsh 优先读 details,但只认独立字段的老客户端会看不到缓存)。
|
||||
if type(chunk.usage.prompt_tokens_details) == "table" and chunk.usage.prompt_tokens_details.cached_tokens ~= nil then
|
||||
uses.prompt_tokens_details = { cached_tokens = chunk.usage.prompt_tokens_details.cached_tokens }
|
||||
end
|
||||
if (chunk.usage.prompt_cache_hit_tokens or 0) > 0 then
|
||||
uses.prompt_cache_hit_tokens = chunk.usage.prompt_cache_hit_tokens
|
||||
uses.prompt_cache_miss_tokens = chunk.usage.prompt_cache_miss_tokens or 0
|
||||
if uses.prompt_tokens_details == nil then
|
||||
uses.prompt_tokens_details = { cached_tokens = chunk.usage.prompt_cache_hit_tokens }
|
||||
end
|
||||
end
|
||||
local ctd = chunk.usage.completion_tokens_details
|
||||
if type(ctd) == "table" and ctd.reasoning_tokens ~= nil then
|
||||
uses.completion_tokens_details = { reasoning_tokens = ctd.reasoning_tokens }
|
||||
end
|
||||
end
|
||||
|
||||
if not chunk.choices or #chunk.choices == 0 then
|
||||
-- 空 choices 的帧既不是内容也不是结束信号。上游发两种:
|
||||
-- {"choices":[],"usage":{...}} —— 终帧用量
|
||||
-- {"choices":[],"cost":"0"} —— 单独的费用帧
|
||||
-- 两者可能落在同一帧上,必须都转出去;早退只带 usage 会把费用丢干净。
|
||||
local out = { done = false }
|
||||
if uses ~= nil then out.usage = uses end
|
||||
if chunk.cost ~= nil then out.cost = tostring(chunk.cost) end
|
||||
if uses == nil and chunk.cost == nil then return "" end
|
||||
return json.encode(out)
|
||||
end
|
||||
local delta = chunk.choices[1].delta or {}
|
||||
local fr = chunk.choices[1].finish_reason
|
||||
|
||||
local finish = (type(fr) == "string" and fr ~= "") and fr or nil
|
||||
|
||||
local unified = {
|
||||
content = delta.content or "",
|
||||
done = (finish ~= nil)
|
||||
}
|
||||
if finish then
|
||||
unified.finish_reason = finish
|
||||
end
|
||||
if uses ~= nil then
|
||||
unified.usage = uses
|
||||
end
|
||||
-- zen 用 reasoning 字段承载推理文本(OpenAI 惯例是 reasoning_content)
|
||||
local reasoning = delta.reasoning_content or delta.reasoning
|
||||
if reasoning then
|
||||
unified.reasoning_content = reasoning
|
||||
end
|
||||
if delta.tool_calls then
|
||||
-- pass raw streaming fragments through; OpenAI clients accumulate index+id+name+arguments
|
||||
unified.tool_calls = delta.tool_calls
|
||||
end
|
||||
return json.encode(unified)
|
||||
end
|
||||
|
||||
-- 错误收敛(可选钩子):zen 错误信封固定为 {error={type,message}};
|
||||
-- 免费池限流 FreeUsageLimitError 单独标注。返回 nil 走通用兜底。
|
||||
function adapter.transform_error(status, body)
|
||||
local ok, resp = pcall(json.decode, body)
|
||||
if not ok or type(resp) ~= "table" then return nil end
|
||||
local e = resp.error
|
||||
if type(e) ~= "table" then return nil end
|
||||
if e.type == "FreeUsageLimitError" then
|
||||
return "zen free pool quota exhausted"
|
||||
end
|
||||
return e.message
|
||||
end
|
||||
|
||||
return adapter
|
||||
342
internal/lua/adapters/opencodezen.lua
Normal file
342
internal/lua/adapters/opencodezen.lua
Normal file
@ -0,0 +1,342 @@
|
||||
local adapter = {}
|
||||
|
||||
adapter.name = "opencodezen"
|
||||
adapter.version = "1.0.0"
|
||||
adapter.endpoint = "/chat/completions"
|
||||
|
||||
-- OpenCode Zen(按量付费端点,https://opencode.ai/zen/v1)
|
||||
--
|
||||
-- 与 OpenCode Go(opencodego.lua,https://opencode.ai/zen/go/v1)是两个不同的
|
||||
-- 服务,行为要求并不相同,因此各有专用适配器:
|
||||
--
|
||||
-- * 本适配器(Zen)服务免费池 / 按量付费池。免费模型必须带 opencode 客户端
|
||||
-- 指纹(UA + x-opencode-*),否则上游拒绝:"free tier can only be used in
|
||||
-- OpenCode"。Zen 的免费池以非 thinking 模型为主,历史上回传
|
||||
-- reasoning_content 会引出上游报错,故这里剥掉(Go 侧相反,见 opencodego.lua)。
|
||||
-- * Go 订阅读者请看 opencodego.lua:那边**必须**回传 reasoning_content,
|
||||
-- 否则 thinking 模式直接 400。
|
||||
--
|
||||
-- 缺 x-opencode-client/session/request/project 身份头会被判为匿名客户端,
|
||||
-- 部分模型在带 tools 时上游直接失败(流式返回单 chunk
|
||||
-- "finish_reason":"network_error" 空 content,非流式 503
|
||||
-- "Endpoint is unavailable",表现为"空回复")。
|
||||
adapter.headers = {
|
||||
["User-Agent"] = "opencode/1.18.21 ai-sdk/provider-utils/4.0.23 runtime/bun/1.3.14",
|
||||
}
|
||||
|
||||
-- 每请求生成身份头。沙箱无 os/math,用 meta.timestamp + 请求体哈希派生:
|
||||
-- 同秒内重复请求 id 相同可接受(zen 只校验存在性,不校验格式)。
|
||||
local function rand_id(prefix, seed)
|
||||
return prefix .. string.sub(sha256_hex(seed), 1, 24)
|
||||
end
|
||||
|
||||
-- 会话指纹:取历史里**第一条 user 消息**。
|
||||
--
|
||||
-- 为什么不直接用客户端的会话 id:实测抓包(tcpdump 抓 127.0.0.1:8081 的真实
|
||||
-- agent 请求)确认通用 OpenAI 客户端**根本不发**任何会话标识 —— body 里没有
|
||||
-- user / session_id / conversation_id / metadata,请求头也只有 X-Stainless-*
|
||||
-- (OpenAI JS SDK)与 User-Agent。opencode 原生客户端那套 x-opencode-session
|
||||
-- 是它自己的概念,通用客户端无从转发。
|
||||
--
|
||||
-- 退而求其次但要够用:历史被逐轮重放,**第一条 user 消息在整个会话里恒定**,
|
||||
-- 用它做指纹即可得到"每会话一个 session",而不是"每源一个 session"。
|
||||
-- 拿不到时(无 user 消息)返回空串,调用方回退到按源稳定。
|
||||
local function conversation_fingerprint(body)
|
||||
if type(body) ~= "string" or body == "" then return "" end
|
||||
local ok, req = pcall(json.decode, body)
|
||||
if not ok or type(req) ~= "table" then return "" end
|
||||
for _, m in ipairs(req.messages or {}) do
|
||||
if type(m) == "table" and m.role == "user" then
|
||||
local c = m.content
|
||||
if type(c) == "string" then
|
||||
if c ~= "" then return c end
|
||||
elseif type(c) == "table" then
|
||||
for _, part in ipairs(c) do
|
||||
if type(part) == "table" and part.type == "text" and part.text then
|
||||
return part.text
|
||||
end
|
||||
end
|
||||
end
|
||||
return ""
|
||||
end
|
||||
end
|
||||
return ""
|
||||
end
|
||||
|
||||
-- session_seed 决定上游会话号(x-opencode-session 的种子)。
|
||||
--
|
||||
-- 优先级:
|
||||
-- 1) 客户端自带的会话标识(meta.client_session)——pi 等在开启
|
||||
-- compat.sendSessionAffinityHeaders 后会发 x-session-affinity,
|
||||
-- 这是平台真实的会话 id,整个会话恒定且天然按会话隔离。
|
||||
-- 2) 退路:历史里**第一条 user 消息**做会话指纹。
|
||||
-- 3) 再退:只按源固定(连 user 消息都没有时)。
|
||||
--
|
||||
-- 为什么需要 2/3:通用 OpenAI 客户端默认**根本不发**会话标识 —— 实测抓包
|
||||
-- (tcpdump 抓 127.0.0.1:8081 真实 agent 请求)确认 body 里没有
|
||||
-- user / session_id / conversation_id / metadata,请求头也只有
|
||||
-- X-Stainless-*(OpenAI JS SDK)与 User-Agent。opencode 原生客户端那套
|
||||
-- x-opencode-session 是它自己的概念,通用客户端无从转发。
|
||||
--
|
||||
-- 为什么必须稳定:上游前缀缓存是**会话级**的。会话号每请求一变,
|
||||
-- 缓存永不命中(实测:固定会号第 2 次命中 5888,每请求换会号则恒为 0)。
|
||||
local function session_seed(meta)
|
||||
local src = (meta.source and meta.source.name) or ""
|
||||
local base = "session|llmsproxy|" .. src
|
||||
local cs = meta.client_session
|
||||
if type(cs) == "string" and cs ~= "" then
|
||||
return base .. "|client|" .. cs
|
||||
end
|
||||
return base .. "|" .. conversation_fingerprint(meta.body)
|
||||
end
|
||||
|
||||
function adapter.build_headers(meta)
|
||||
local ts = tostring(meta.timestamp or "")
|
||||
local src = (meta.source and meta.source.name) or ""
|
||||
return {
|
||||
["User-Agent"] = adapter.headers["User-Agent"],
|
||||
["x-opencode-client"] = "cli",
|
||||
-- project 固定:同网关实例共享一个工作区身份
|
||||
["x-opencode-project"] = string.sub(sha256_hex("llmsproxy|" .. src), 1, 32),
|
||||
-- session 必须**按源稳定**,不能每请求换:上游的前缀缓存在同一 session
|
||||
-- 内才复用。实测(同一段 6032 token 提示词):
|
||||
-- 固定 session -> 第 2 次命中 5888/6032,cached_tokens=5888
|
||||
-- 每请求换 session -> 永远 0 命中
|
||||
-- 原先用 meta.timestamp 派生,等于每请求都是新会话,缓存永远无效,
|
||||
-- 上游也无法做会话亲和路由。
|
||||
["x-opencode-session"] = rand_id("ses_", session_seed(meta)),
|
||||
-- request id 仍每请求唯一(它只是请求标识,不参与缓存键)
|
||||
["x-opencode-request"] = rand_id("msg_", "request|" .. ts .. "|" .. tostring(meta.body or "")),
|
||||
}
|
||||
end
|
||||
|
||||
-- OpenAI /chat/completions format (pass-through, strip provider-specific fields)
|
||||
-- zen 上游 schema 只接受 text content part(无视觉/音频能力):多模态 part
|
||||
-- (image_url / input_audio / file 等)一律剥离。剥离后 content 变空的消息
|
||||
-- 若不再携带 tool_calls / tool_call_id 才整条丢弃(避免上游
|
||||
-- "unknown variant `image_url`, expected `text`");带工具调用的必须保留,
|
||||
-- 否则会把紧随其后的 tool 结果变成孤儿,模型会反复重发同一个调用。
|
||||
-- zen 上游角色白名单只有 system / user / assistant / tool / latest_reminder:
|
||||
-- OpenAI 的 developer(及 function 等)不在其中,直接透传会触发上游
|
||||
-- "unknown variant `developer`, expected one of ..." 错误;统一归一化为 system。
|
||||
local ROLE_WHITELIST = {
|
||||
system = true,
|
||||
user = true,
|
||||
assistant = true,
|
||||
tool = true,
|
||||
latest_reminder = true,
|
||||
}
|
||||
|
||||
function adapter.transform_request(raw_body)
|
||||
local ok, req = pcall(json.decode, raw_body)
|
||||
if not ok then return raw_body end
|
||||
req.disable_thinking = nil
|
||||
req.extra_body = nil
|
||||
-- stream_options is only valid alongside stream:true; sending it on a
|
||||
-- non-streaming request is rejected by strict upstreams.
|
||||
if req.stream then
|
||||
if type(req.stream_options) ~= "table" then req.stream_options = {} end
|
||||
req.stream_options.include_usage = true
|
||||
end
|
||||
if req.messages then
|
||||
local kept = {}
|
||||
for _, msg in ipairs(req.messages) do
|
||||
if type(msg.role) == "string" and not ROLE_WHITELIST[msg.role] then
|
||||
msg.role = "system"
|
||||
end
|
||||
-- Zen 免费池不接受回传 reasoning_content(Go 侧相反)
|
||||
msg.reasoning_content = nil
|
||||
local drop = false
|
||||
if type(msg.content) == "table" then
|
||||
local parts = {}
|
||||
for _, part in ipairs(msg.content) do
|
||||
if type(part) == "table" and part.type ~= nil and part.type ~= "text" then
|
||||
-- multimodal part not supported by zen
|
||||
else
|
||||
table.insert(parts, part)
|
||||
end
|
||||
end
|
||||
if #parts == 0 then
|
||||
-- Content collapsed to nothing after stripping unsupported
|
||||
-- parts. A message that still carries a tool call must
|
||||
-- NEVER be dropped: the very next message is its tool
|
||||
-- result, and dropping the call orphans that result. The
|
||||
-- model then sees a result for a call it never made and
|
||||
-- re-issues the same tool call on every turn (observed as
|
||||
-- an infinite "repeated tool call" loop).
|
||||
if (type(msg.tool_calls) == "table" and #msg.tool_calls > 0)
|
||||
or msg.tool_call_id ~= nil then
|
||||
msg.content = ""
|
||||
else
|
||||
drop = true
|
||||
end
|
||||
else
|
||||
msg.content = parts
|
||||
end
|
||||
end
|
||||
if not drop then
|
||||
table.insert(kept, msg)
|
||||
end
|
||||
end
|
||||
req.messages = kept
|
||||
end
|
||||
return json.encode(req)
|
||||
end
|
||||
|
||||
function adapter.transform_response(raw_body)
|
||||
local ok, resp = pcall(json.decode, raw_body)
|
||||
if not ok or resp == nil then return raw_body end
|
||||
|
||||
local unified = {
|
||||
content = "",
|
||||
finish_reason = "",
|
||||
token_usage = { prompt = 0, completion = 0, total = 0 }
|
||||
}
|
||||
|
||||
if type(resp.usage) == "table" then
|
||||
unified.token_usage.prompt = resp.usage.prompt_tokens or 0
|
||||
unified.token_usage.completion = resp.usage.completion_tokens or 0
|
||||
unified.token_usage.total = resp.usage.total_tokens or 0
|
||||
local hit = 0
|
||||
if type(resp.usage.prompt_tokens_details) == "table" and resp.usage.prompt_tokens_details.cached_tokens ~= nil then
|
||||
hit = resp.usage.prompt_tokens_details.cached_tokens
|
||||
unified.token_usage.prompt_tokens_details = { cached_tokens = hit }
|
||||
end
|
||||
if (resp.usage.prompt_cache_hit_tokens or 0) > 0 then
|
||||
unified.token_usage.prompt_cache_hit_tokens = resp.usage.prompt_cache_hit_tokens
|
||||
unified.token_usage.prompt_cache_miss_tokens = resp.usage.prompt_cache_miss_tokens or 0
|
||||
if hit == 0 then
|
||||
unified.token_usage.prompt_tokens_details = { cached_tokens = resp.usage.prompt_cache_hit_tokens }
|
||||
end
|
||||
end
|
||||
-- 上游报告输出里有多少是思考 token。不给客户端的话,无法判断
|
||||
-- completion_tokens 里多少是「可见回答」、多少是「思考」(都按输出计费)。
|
||||
local ctd = resp.usage.completion_tokens_details
|
||||
if type(ctd) == "table" and ctd.reasoning_tokens ~= nil then
|
||||
unified.token_usage.completion_tokens_details = { reasoning_tokens = ctd.reasoning_tokens }
|
||||
end
|
||||
end
|
||||
|
||||
-- 本次调用的费用,上游放在**顶层**且是**字符串**(如 "0"、"0.0012")。
|
||||
-- Go 订阅是包月制,恒为 "0";只有 Zen 按量付费模型才有信息量。
|
||||
-- 原样透传:网关不解析、不换算、不汇总 —— 它只是上游事实的搬运者。
|
||||
if resp.cost ~= nil then
|
||||
unified.cost = tostring(resp.cost)
|
||||
end
|
||||
|
||||
if type(resp.choices) == "table" and #resp.choices > 0 then
|
||||
local ch = resp.choices[1]
|
||||
if type(ch.message) == "table" then
|
||||
unified.content = ch.message.content or ""
|
||||
local reasoning = ch.message.reasoning_content or ch.message.reasoning
|
||||
if reasoning then
|
||||
unified.reasoning_content = reasoning
|
||||
end
|
||||
if type(ch.message.tool_calls) == "table" then
|
||||
local tcs = {}
|
||||
for _, tc in ipairs(ch.message.tool_calls) do
|
||||
local args_ok, args = pcall(json.decode, tc["function"].arguments)
|
||||
if not args_ok then args = {} end
|
||||
table.insert(tcs, {
|
||||
id = tc.id,
|
||||
type = tc.type or "function",
|
||||
name = tc["function"].name,
|
||||
arguments = args
|
||||
})
|
||||
end
|
||||
unified.tool_calls = tcs
|
||||
end
|
||||
end
|
||||
unified.finish_reason = ch.finish_reason or ""
|
||||
end
|
||||
|
||||
return json.encode(unified)
|
||||
end
|
||||
|
||||
function adapter.transform_stream_chunk(raw_chunk)
|
||||
local ok, chunk = pcall(json.decode, raw_chunk)
|
||||
if not ok then return "" end
|
||||
|
||||
-- OpenAI-style streams may attach usage to a chunk with empty choices
|
||||
-- (the final usage chunk). Preserve it; the gateway emits it as the
|
||||
-- terminal usage chunk. Note: keys must match Go's TokenUsage json tags
|
||||
-- (prompt/completion/total); the gateway re-emits standard *_tokens.
|
||||
local uses = nil
|
||||
if type(chunk.usage) == "table" then
|
||||
uses = {
|
||||
prompt = chunk.usage.prompt_tokens or chunk.usage.prompt or 0,
|
||||
completion = chunk.usage.completion_tokens or chunk.usage.completion or 0,
|
||||
total = chunk.usage.total_tokens or chunk.usage.total or 0,
|
||||
}
|
||||
-- 独立字段与 prompt_tokens_details **并存**透传,不要写成 elseif:
|
||||
-- 上游命中缓存时两者都发,非流式路径也是两个都带。写成 elseif 会让
|
||||
-- 流式丢掉 prompt_cache_hit_tokens/miss,与同一源的**非流式**产出不一致
|
||||
-- (dsh 优先读 details,但只认独立字段的老客户端会看不到缓存)。
|
||||
if type(chunk.usage.prompt_tokens_details) == "table" and chunk.usage.prompt_tokens_details.cached_tokens ~= nil then
|
||||
uses.prompt_tokens_details = { cached_tokens = chunk.usage.prompt_tokens_details.cached_tokens }
|
||||
end
|
||||
if (chunk.usage.prompt_cache_hit_tokens or 0) > 0 then
|
||||
uses.prompt_cache_hit_tokens = chunk.usage.prompt_cache_hit_tokens
|
||||
uses.prompt_cache_miss_tokens = chunk.usage.prompt_cache_miss_tokens or 0
|
||||
if uses.prompt_tokens_details == nil then
|
||||
uses.prompt_tokens_details = { cached_tokens = chunk.usage.prompt_cache_hit_tokens }
|
||||
end
|
||||
end
|
||||
local ctd = chunk.usage.completion_tokens_details
|
||||
if type(ctd) == "table" and ctd.reasoning_tokens ~= nil then
|
||||
uses.completion_tokens_details = { reasoning_tokens = ctd.reasoning_tokens }
|
||||
end
|
||||
end
|
||||
|
||||
if not chunk.choices or #chunk.choices == 0 then
|
||||
-- 空 choices 的帧既不是内容也不是结束信号。上游发两种:
|
||||
-- {"choices":[],"usage":{...}} —— 终帧用量
|
||||
-- {"choices":[],"cost":"0"} —— 单独的费用帧
|
||||
-- 两者可能落在同一帧上,必须都转出去;早退只带 usage 会把费用丢干净。
|
||||
local out = { done = false }
|
||||
if uses ~= nil then out.usage = uses end
|
||||
if chunk.cost ~= nil then out.cost = tostring(chunk.cost) end
|
||||
if uses == nil and chunk.cost == nil then return "" end
|
||||
return json.encode(out)
|
||||
end
|
||||
local delta = chunk.choices[1].delta or {}
|
||||
local fr = chunk.choices[1].finish_reason
|
||||
|
||||
local finish = (type(fr) == "string" and fr ~= "") and fr or nil
|
||||
|
||||
local unified = {
|
||||
content = delta.content or "",
|
||||
done = (finish ~= nil)
|
||||
}
|
||||
if finish then
|
||||
unified.finish_reason = finish
|
||||
end
|
||||
if uses ~= nil then
|
||||
unified.usage = uses
|
||||
end
|
||||
-- zen 用 reasoning 字段承载推理文本(OpenAI 惯例是 reasoning_content)
|
||||
local reasoning = delta.reasoning_content or delta.reasoning
|
||||
if reasoning then
|
||||
unified.reasoning_content = reasoning
|
||||
end
|
||||
if delta.tool_calls then
|
||||
-- pass raw streaming fragments through; OpenAI clients accumulate index+id+name+arguments
|
||||
unified.tool_calls = delta.tool_calls
|
||||
end
|
||||
return json.encode(unified)
|
||||
end
|
||||
|
||||
-- 错误收敛(可选钩子):zen 错误信封固定为 {error={type,message}};
|
||||
-- 免费池限流 FreeUsageLimitError 单独标注。返回 nil 走通用兜底。
|
||||
function adapter.transform_error(status, body)
|
||||
local ok, resp = pcall(json.decode, body)
|
||||
if not ok or type(resp) ~= "table" then return nil end
|
||||
local e = resp.error
|
||||
if type(e) ~= "table" then return nil end
|
||||
if e.type == "FreeUsageLimitError" then
|
||||
return "zen free pool quota exhausted"
|
||||
end
|
||||
return e.message
|
||||
end
|
||||
|
||||
return adapter
|
||||
235
internal/lua/adapters/qwen38.lua
Normal file
235
internal/lua/adapters/qwen38.lua
Normal file
@ -0,0 +1,235 @@
|
||||
-- tokenplan / discovery-api (intern-ai.org.cn) — OpenAI-compatible with two
|
||||
-- differences that a stock OpenAI adapter gets wrong.
|
||||
--
|
||||
-- 1. REASONING IS A SEPARATE FIELD, AND IT IS MOST OF THE OUTPUT.
|
||||
-- Every response carries `reasoning_content` (non-stream) or
|
||||
-- `delta.reasoning_content` (stream), disjoint from `content`. Measured on
|
||||
-- qwen3.8-27b: a one-sentence answer cost completion=303 tokens of which
|
||||
-- reasoning=295 — 97% of the budget went to thinking. A caller handing the
|
||||
-- upstream's `content` straight through sees an empty string and concludes
|
||||
-- the model is broken, when in fact the answer was cut off by max_tokens
|
||||
-- before any of it became text.
|
||||
--
|
||||
-- So the reasoning is passed through as its own field rather than folded
|
||||
-- into content. Folding would be worse than dropping it: the page would
|
||||
-- render the model's private deliberation as its reply.
|
||||
--
|
||||
-- 2. `finish_reason: "length"` WITH `content: null` IS NORMAL HERE, NOT AN ERROR.
|
||||
-- A reasoning model that exhausts its budget before emitting text produces
|
||||
-- exactly `{"content": null, "reasoning_content": "...", "finish_reason":
|
||||
-- "length"}`. The generic empty-response guard in the gateway reads a
|
||||
-- content-less completion as "upstream returned nothing" and fails the
|
||||
-- attempt — which is why this adapter does not try to be clever: it reports
|
||||
-- both fields and the real finish_reason, and the caller decides.
|
||||
--
|
||||
-- 3. MODEL NAMES ARE REWRITTEN BY THE SERVER.
|
||||
-- Requesting `qwen3.8-27b` answers with `model: "qwen3.8-27b-fp8"`. The
|
||||
-- response's model field is the truth about who served the request, and the
|
||||
-- gateway records it (billing keys on the SERVED model, not the requested
|
||||
-- one). Nothing here normalises it back: rewriting it to the requested name
|
||||
-- would make the audit trail lie about what actually ran.
|
||||
--
|
||||
-- Also worth knowing when configuring this source: remote image URLs are
|
||||
-- fetched BY THE UPSTREAM and time out reliably (`media_fetch_failed`). Images
|
||||
-- must be sent inline as data: URLs. Measured: base64 worked, a Wikimedia URL
|
||||
-- did not.
|
||||
|
||||
local adapter = {}
|
||||
|
||||
adapter.name = "qwen38"
|
||||
adapter.version = "1.0.0"
|
||||
adapter.endpoint = "/chat/completions"
|
||||
adapter.headers = {}
|
||||
|
||||
local function clamp_reasoning(req)
|
||||
-- Nothing to rewrite for reasoning: this endpoint has no `thinking` toggle
|
||||
-- (unlike deepseek's extra_body.thinking). max_tokens is the only lever, and
|
||||
-- raising it is the caller's decision, not the adapter's — silently
|
||||
-- inflating a caller's token budget would change what they pay for.
|
||||
end
|
||||
|
||||
function adapter.transform_request(raw_body)
|
||||
local ok, req = pcall(json.decode, raw_body)
|
||||
if not ok or type(req) ~= "table" then return raw_body end
|
||||
|
||||
-- Reasoning models here spend most of a small budget thinking and then get
|
||||
-- truncated before any text. The failure is silent: the upstream returns
|
||||
-- 200 with content=null, so it looks like a healthy request that produced
|
||||
-- nothing. Warn rather than rewrite — see clamp_reasoning.
|
||||
if type(req.max_tokens) == "number" and req.max_tokens > 0 and req.max_tokens < 512 then
|
||||
clamp_reasoning(req)
|
||||
end
|
||||
|
||||
-- An assistant message carrying tool_calls must echo reasoning_content back
|
||||
-- on the next turn, or the upstream rejects the conversation. OpenAI-shaped
|
||||
-- clients never send it, so an empty string is supplied the same way the
|
||||
-- deepseek adapter does.
|
||||
if type(req.messages) == "table" then
|
||||
for _, msg in ipairs(req.messages) do
|
||||
if type(msg) == "table" and msg.role == "assistant" then
|
||||
if type(msg.tool_calls) == "table" and msg.tool_calls[1] ~= nil
|
||||
and msg.reasoning_content == nil then
|
||||
msg.reasoning_content = ""
|
||||
end
|
||||
end
|
||||
end
|
||||
end
|
||||
|
||||
return json.encode(req)
|
||||
end
|
||||
|
||||
function adapter.transform_response(raw_body)
|
||||
local ok, resp = pcall(json.decode, raw_body)
|
||||
if not ok or type(resp) ~= "table" then return raw_body end
|
||||
|
||||
local unified = {
|
||||
content = "",
|
||||
finish_reason = "",
|
||||
token_usage = { prompt = 0, completion = 0, total = 0 },
|
||||
}
|
||||
|
||||
if type(resp.usage) == "table" then
|
||||
local u = resp.usage
|
||||
unified.token_usage.prompt = u.prompt_tokens or 0
|
||||
unified.token_usage.completion = u.completion_tokens or 0
|
||||
unified.token_usage.total = u.total_tokens or 0
|
||||
if type(u.prompt_tokens_details) == "table" and u.prompt_tokens_details.cached_tokens ~= nil then
|
||||
unified.token_usage.prompt_tokens_details = { cached_tokens = u.prompt_tokens_details.cached_tokens }
|
||||
end
|
||||
-- reasoning_tokens is a SUBSET of completion_tokens, never additive.
|
||||
--
|
||||
-- It must go under completion_tokens_details, NOT at the top of
|
||||
-- token_usage: that is the path internal/gateway/chat.go reads when it
|
||||
-- fills the audit record's Reasoning field. Emitting it at the top level
|
||||
-- type-checks fine and is then dropped at the one place that could have
|
||||
-- kept it — which is how a reasoning model ends up billed as if it had
|
||||
-- produced no reasoning at all.
|
||||
--
|
||||
-- (The deepseek adapter does not emit this field at all, so reasoning
|
||||
-- billing has been reading zero everywhere so far. That is a separate
|
||||
-- gap; this adapter does not depend on it being fixed.)
|
||||
if type(u.completion_tokens_details) == "table"
|
||||
and u.completion_tokens_details.reasoning_tokens ~= nil then
|
||||
local d = unified.token_usage.completion_tokens_details
|
||||
if type(d) ~= "table" then
|
||||
d = {}
|
||||
unified.token_usage.completion_tokens_details = d
|
||||
end
|
||||
d.reasoning_tokens = u.completion_tokens_details.reasoning_tokens
|
||||
end
|
||||
end
|
||||
|
||||
if type(resp.choices) == "table" and #resp.choices > 0 then
|
||||
local ch = resp.choices[1]
|
||||
if type(ch.message) == "table" then
|
||||
local m = ch.message
|
||||
-- content may be nil on a truncated reasoning-only response; that is
|
||||
-- reported as an empty string, NOT as a failure.
|
||||
if type(m.content) == "string" then
|
||||
unified.content = m.content
|
||||
end
|
||||
if type(m.reasoning_content) == "string" then
|
||||
unified.reasoning_content = m.reasoning_content
|
||||
end
|
||||
if type(m.tool_calls) == "table" then
|
||||
local tcs = {}
|
||||
for _, tc in ipairs(m.tool_calls) do
|
||||
local fn = tc["function"] or {}
|
||||
local args_ok, args = pcall(json.decode, fn.arguments or "{}")
|
||||
if not args_ok then args = {} end
|
||||
table.insert(tcs, {
|
||||
id = tc.id,
|
||||
type = tc.type or "function",
|
||||
name = fn.name,
|
||||
arguments = args,
|
||||
})
|
||||
end
|
||||
unified.tool_calls = tcs
|
||||
end
|
||||
end
|
||||
unified.finish_reason = ch.finish_reason or ""
|
||||
end
|
||||
|
||||
return json.encode(unified)
|
||||
end
|
||||
|
||||
function adapter.transform_stream_chunk(raw_chunk)
|
||||
local ok, chunk = pcall(json.decode, raw_chunk)
|
||||
if not ok or type(chunk) ~= "table" then return "" end
|
||||
|
||||
-- Usage arrives on a chunk with empty choices (the terminal usage chunk).
|
||||
-- Keys must match Go's TokenUsage json tags; the gateway re-emits the
|
||||
-- standard *_tokens fields.
|
||||
local uses = nil
|
||||
if type(chunk.usage) == "table" then
|
||||
local u = chunk.usage
|
||||
uses = {
|
||||
prompt = u.prompt_tokens or 0,
|
||||
completion = u.completion_tokens or 0,
|
||||
total = u.total_tokens or 0,
|
||||
}
|
||||
if type(u.prompt_tokens_details) == "table" and u.prompt_tokens_details.cached_tokens ~= nil then
|
||||
uses.prompt_tokens_details = { cached_tokens = u.prompt_tokens_details.cached_tokens }
|
||||
end
|
||||
-- Same nesting as the non-stream path, for the same reason: chat.go reads
|
||||
-- reasoning under completion_tokens_details and nothing else.
|
||||
if type(u.completion_tokens_details) == "table"
|
||||
and u.completion_tokens_details.reasoning_tokens ~= nil then
|
||||
uses.completion_tokens_details = {
|
||||
reasoning_tokens = u.completion_tokens_details.reasoning_tokens
|
||||
}
|
||||
end
|
||||
end
|
||||
|
||||
if type(chunk.choices) ~= "table" or #chunk.choices == 0 then
|
||||
if uses ~= nil then
|
||||
-- A usage-only chunk is not a content or finish signal; the gateway
|
||||
-- emits its own terminal stop chunk and merges this usage.
|
||||
return json.encode({ usage = uses, done = false })
|
||||
end
|
||||
return ""
|
||||
end
|
||||
|
||||
local ch = chunk.choices[1]
|
||||
local delta = ch.delta or {}
|
||||
local finish = (type(ch.finish_reason) == "string" and ch.finish_reason ~= "") and ch.finish_reason or nil
|
||||
|
||||
local unified = {
|
||||
content = delta.content or "",
|
||||
done = (finish ~= nil),
|
||||
}
|
||||
if finish then
|
||||
unified.finish_reason = finish
|
||||
end
|
||||
if uses ~= nil then
|
||||
unified.usage = uses
|
||||
end
|
||||
-- The whole point of this adapter: reasoning is streamed as its own channel
|
||||
-- and must stay separate. Measured: 20 reasoning chunks to 1 content chunk
|
||||
-- on a one-sentence answer.
|
||||
if type(delta.reasoning_content) == "string" then
|
||||
unified.reasoning_content = delta.reasoning_content
|
||||
end
|
||||
if type(delta.tool_calls) == "table" then
|
||||
unified.tool_calls = delta.tool_calls
|
||||
end
|
||||
return json.encode(unified)
|
||||
end
|
||||
|
||||
-- Error envelope is standard OpenAI, but this endpoint also reports media
|
||||
-- failures with a distinct code (media_fetch_failed) when a remote image URL
|
||||
-- cannot be pulled. Both carry {error:{message,...}}; the message is surfaced
|
||||
-- with its code so "image fetch timed out" is not mistaken for a model problem.
|
||||
function adapter.transform_error(status, body)
|
||||
local ok, resp = pcall(json.decode, body)
|
||||
if not ok or type(resp) ~= "table" then return nil end
|
||||
local e = resp.error
|
||||
if type(e) ~= "table" then return nil end
|
||||
if type(e.message) ~= "string" then return nil end
|
||||
if type(e.code) == "string" and e.code ~= "" and e.code ~= "invalid_api_key" then
|
||||
return e.message .. " (" .. e.code .. ")"
|
||||
end
|
||||
return e.message
|
||||
end
|
||||
|
||||
return adapter
|
||||
351
internal/lua/billing_period_test.go
Normal file
351
internal/lua/billing_period_test.go
Normal file
@ -0,0 +1,351 @@
|
||||
package lua
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"os/exec"
|
||||
"strings"
|
||||
"testing"
|
||||
)
|
||||
|
||||
// runBillingPage executes the billing page's mount script under node against a
|
||||
// DOM stub, drives the period selector, and returns what ended up on screen.
|
||||
//
|
||||
// It asserts against RENDERED output rather than the plugin's internal
|
||||
// helpers. Reaching into the IIFE for daysInWindow/rescale needs a test-only
|
||||
// hook in the plugin, and the first version of this file took a long detour
|
||||
// through such a hook: the harness's fake `Date` broke toISOString(), the
|
||||
// plugin script aborted before assigning the hook, and every assertion
|
||||
// compared "" against a real value. The rendered tables are the contract that
|
||||
// actually matters, they need no hook, and they cannot pass if the rescale is
|
||||
// wrong — which is exactly the bug class under test.
|
||||
func runBillingPage(t *testing.T, todayISO, period, captureID string) string {
|
||||
t.Helper()
|
||||
node, err := exec.LookPath("node")
|
||||
if err != nil {
|
||||
t.Skipf("node not available (%v)", err)
|
||||
}
|
||||
page, _ := billingUI(t)
|
||||
js := strings.Join(extractScripts(page), "\n")
|
||||
if !strings.Contains(js, "__billing_setPeriod") {
|
||||
t.Fatal("billing page exposes no period selector")
|
||||
}
|
||||
|
||||
harness := `
|
||||
global.window = global;
|
||||
global.__written = [];
|
||||
global.__errors = [];
|
||||
var __realErr = console.error;
|
||||
console.error = function () {
|
||||
global.__errors.push("console.error: " + Array.prototype.map.call(arguments, function (a) {
|
||||
return (a && a.message) ? a.message : String(a);
|
||||
}).join(" "));
|
||||
__realErr.apply(console, arguments);
|
||||
};
|
||||
process.on("uncaughtException", function (e) { global.__errors.push("uncaught: " + String(e && e.message || e)); });
|
||||
process.on("unhandledRejection", function (e) { global.__errors.push("unhandled: " + String(e && e.message || e)); });
|
||||
|
||||
// Pin "now" for the window under test.
|
||||
//
|
||||
// Two earlier attempts were wrong here, and both looked like product bugs:
|
||||
// replacing the Date constructor wholesale breaks toISOString/UTC; and
|
||||
// subclassing (setting __FakeDate.prototype = __RealDate.prototype while
|
||||
// returning __RealDate instances) leaves instances WITHOUT a working
|
||||
// getUTCDay, so the week window came out NaN and every rendered day became
|
||||
// "Invalid time value". A Proxy over the real constructor intercepts only the
|
||||
// zero-argument call and leaves instances — and therefore every prototype
|
||||
// method — completely untouched.
|
||||
var __NOW_MS = new Date("` + todayISO + `T12:34:56Z").getTime();
|
||||
var __RealDate = global.Date;
|
||||
var __PatchedDate = new Proxy(__RealDate, {
|
||||
construct: function (target, args) {
|
||||
// Reflect.construct, not a hand-spelled argument list: writing
|
||||
// new target(a, b, c, d, e, f, g) passes seven explicit undefineds for a
|
||||
// one-argument call, which Date reads as year..millisecond fields and turns
|
||||
// into an Invalid Date. The symptom was a week window rendering
|
||||
// "Invalid time value" for every day — a harness fault indistinguishable
|
||||
// from a product bug.
|
||||
if (args.length === 0) return Reflect.construct(target, [__NOW_MS]);
|
||||
return Reflect.construct(target, args);
|
||||
},
|
||||
apply: function () { return new __RealDate(__NOW_MS); }
|
||||
});
|
||||
__PatchedDate.now = function () { return __NOW_MS; };
|
||||
__PatchedDate.UTC = __RealDate.UTC;
|
||||
__PatchedDate.parse = __RealDate.parse;
|
||||
__PatchedDate.prototype = __RealDate.prototype;
|
||||
global.Date = __PatchedDate;
|
||||
|
||||
global.document = {
|
||||
getElementById: function (id) {
|
||||
if (!global.__els) global.__els = {};
|
||||
if (!global.__els[id]) global.__els[id] = {
|
||||
style: {}, dataset: {}, classList: { add: function(){}, remove: function(){} },
|
||||
options: [{value:"day"},{value:"week"},{value:"month"},{value:"all"}],
|
||||
set innerHTML(v) { if (v && String(v).trim()) global.__written.push(id); this.__h = v; },
|
||||
get innerHTML() { return this.__h || ""; },
|
||||
set textContent(v) { if (v !== undefined && String(v).trim()) global.__written.push(id); this.__t = v; },
|
||||
get textContent() { return this.__t || ""; },
|
||||
set innerText(v) { if (v !== undefined && String(v).trim()) global.__written.push(id); this.__i = v; },
|
||||
get innerText() { return this.__i || ""; },
|
||||
appendChild: function(){}, querySelector: function(){ return null; },
|
||||
querySelectorAll: function(){ return []; }, addEventListener: function(){}
|
||||
};
|
||||
return global.__els[id];
|
||||
},
|
||||
createElement: function () { return { style: {}, dataset: {}, appendChild: function(){}, setAttribute: function(){} }; },
|
||||
addEventListener: function () {},
|
||||
};
|
||||
|
||||
// Fixture days carry deliberately distinctive costs so a fold over the wrong
|
||||
// subset still produces a number, and only a mismatch shows it. 2026-01-05/06/07
|
||||
// are Monday/Tuesday/Wednesday; 2025-12-01 is far outside any window here.
|
||||
global.__st = {
|
||||
currency: "USD",
|
||||
total: { cost: 999, requests: 999, prompt_tokens: 999999 },
|
||||
by_source: { LIFETIMESRC: { cost: 999, requests: 999 } },
|
||||
by_model: { LIFETIMEMODEL: { cost: 999, requests: 999 } },
|
||||
by_key: { LIFETIMEKEY: { cost: 999, requests: 999 } },
|
||||
by_day: {
|
||||
"2026-01-05": { cost: 10, requests: 1, prompt_tokens: 100, completion_tokens: 10, failures: 0 },
|
||||
"2026-01-06": { cost: 20, requests: 2, prompt_tokens: 200, completion_tokens: 20, failures: 1 },
|
||||
"2026-01-07": { cost: 30, requests: 3, prompt_tokens: 300, completion_tokens: 30, failures: 0 },
|
||||
"2025-12-01": { cost: 777, requests: 77, prompt_tokens: 7700 }
|
||||
},
|
||||
by_day_src: {
|
||||
"2026-01-05": { srcA: { cost: 10, requests: 1 } },
|
||||
"2026-01-06": { srcA: { cost: 20, requests: 2 }, srcB: { cost: 5, requests: 1 } },
|
||||
"2026-01-07": { srcA: { cost: 30, requests: 3 } },
|
||||
"2025-12-01": { srcOLD: { cost: 777, requests: 77 } }
|
||||
},
|
||||
by_day_model: {
|
||||
"2026-01-05": { mA: { cost: 10, requests: 1 } },
|
||||
"2026-01-06": { mA: { cost: 20, requests: 2 } },
|
||||
"2026-01-07": { mA: { cost: 30, requests: 3 } }
|
||||
},
|
||||
by_day_key: {
|
||||
"2026-01-05": { kA: { cost: 10, requests: 1 } },
|
||||
"2026-01-06": { kA: { cost: 20, requests: 2 } },
|
||||
"2026-01-07": { kA: { cost: 30, requests: 3 } }
|
||||
},
|
||||
degraded_reqs: 4242, unpriced_reqs: 77
|
||||
};
|
||||
global.fetch = function () {
|
||||
return Promise.resolve({ ok: true, json: function () {
|
||||
return Promise.resolve({ plugin: "billing", state: global.__st });
|
||||
}});
|
||||
};
|
||||
global.pluginAPI = { onTabShown: function () {}, lang: "en" };
|
||||
`
|
||||
script := harness + "\n" + js + "\n" +
|
||||
"__billing_setPeriod(" + jsonStr(period) + ");\n" +
|
||||
// Let the period-triggered refresh() settle before reading the DOM.
|
||||
"setTimeout(function () {\n" +
|
||||
" var el = document.getElementById(" + jsonStr(captureID) + ");\n" +
|
||||
" console.log('__CAPTURE__' + JSON.stringify({\n" +
|
||||
" html: el ? String(el.innerHTML) : '(missing)',\n" +
|
||||
" kpis: String(document.getElementById('billing-kpis').innerHTML),\n" +
|
||||
" sel: String(document.getElementById('billing-period').value),\n" +
|
||||
" errors: global.__errors\n" +
|
||||
" }));\n" +
|
||||
"}, 60);\n"
|
||||
|
||||
cmd := exec.Command(node, "-e", script)
|
||||
out, err := cmd.CombinedOutput()
|
||||
if err != nil {
|
||||
t.Fatalf("node execution failed: %v\n%s", err, out)
|
||||
}
|
||||
o := string(out)
|
||||
idx := strings.Index(o, "__CAPTURE__")
|
||||
if idx < 0 {
|
||||
t.Fatalf("no capture in node output (the plugin script aborted?):\n%s", o)
|
||||
}
|
||||
raw := strings.TrimSpace(o[idx+len("__CAPTURE__"):])
|
||||
if nl := strings.IndexAny(raw, "\r\n"); nl >= 0 {
|
||||
raw = raw[:nl]
|
||||
}
|
||||
var got struct {
|
||||
HTML string `json:"html"`
|
||||
KPIs string `json:"kpis"`
|
||||
Sel string `json:"sel"`
|
||||
Errors []string `json:"errors"`
|
||||
}
|
||||
if err := json.Unmarshal([]byte(raw), &got); err != nil {
|
||||
t.Fatalf("capture is not JSON: %v\n%s", err, raw)
|
||||
}
|
||||
if len(got.Errors) > 0 {
|
||||
t.Fatalf("billing script reported errors: %v", got.Errors)
|
||||
}
|
||||
if got.Sel != period {
|
||||
t.Errorf("selector shows %q after setting %q — the change handler did not take", got.Sel, period)
|
||||
}
|
||||
return got.HTML + "\n@@KPI@@\n" + got.KPIs
|
||||
}
|
||||
|
||||
func jsonStr(s string) string {
|
||||
b, _ := json.Marshal(s)
|
||||
return string(b)
|
||||
}
|
||||
|
||||
// TestBillingPeriodDayShowsOnlyThatDay checks the one-day window end to end:
|
||||
// the tables must carry that day's rows and drop every other day, and the
|
||||
// lifetime sentinel (999 / LIFETIMESRC) must appear nowhere.
|
||||
func TestBillingPeriodDayShowsOnlyThatDay(t *testing.T) {
|
||||
got := runBillingPage(t, "2026-01-07", "day", "billing-by-source")
|
||||
if strings.Contains(got, "LIFETIME") || strings.Contains(got, "999") {
|
||||
t.Errorf("a one-day view still shows lifetime figures:\n%s", got)
|
||||
}
|
||||
if strings.Contains(got, "srcB") {
|
||||
t.Errorf("a one-day view shows srcB, which only had traffic on 01-06:\n%s", got)
|
||||
}
|
||||
if strings.Contains(got, "srcOLD") {
|
||||
t.Errorf("a one-day view shows a day from 2025-12:\n%s", got)
|
||||
}
|
||||
if !strings.Contains(got, "srcA") {
|
||||
t.Errorf("a one-day view lost the only source that had traffic that day:\n%s", got)
|
||||
}
|
||||
// The KPI total must be that day's cost (30), not the week (60) and not
|
||||
// the lifetime sentinel (999).
|
||||
if !strings.Contains(got, "30") {
|
||||
t.Errorf("KPI does not show the day's cost 30:\n%s", got)
|
||||
}
|
||||
}
|
||||
|
||||
// TestBillingPeriodWeekSumsThreeDays is the arithmetic: Mon+Wed rows and a
|
||||
// total of 60. A window that started on Sunday or ended early still renders
|
||||
// tables, so only the number distinguishes it.
|
||||
func TestBillingPeriodWeekSumsThreeDays(t *testing.T) {
|
||||
// 2026-01-07 is a Wednesday: the ISO week holds 01-05, 01-06, 01-07.
|
||||
got := runBillingPage(t, "2026-01-07", "week", "billing-by-source")
|
||||
if strings.Contains(got, "srcOLD") || strings.Contains(got, "LIFETIME") {
|
||||
t.Errorf("weekly view leaks out-of-window or lifetime rows:\n%s", got)
|
||||
}
|
||||
// srcA = 10+20+30 = 60, srcB = 5 (only 01-06, but inside the week).
|
||||
if !strings.Contains(got, "60") {
|
||||
t.Errorf("weekly total should be 60, got:\n%s", got)
|
||||
}
|
||||
if !strings.Contains(got, "srcB") {
|
||||
t.Errorf("weekly view dropped srcB, whose only day (01-06) is in the window:\n%s", got)
|
||||
}
|
||||
}
|
||||
|
||||
// TestBillingPeriodSundayReachesBackToMonday pins the ISO-week Monday rule
|
||||
// from the other side. On Sunday 2026-01-11 the week that already ended is
|
||||
// Mon 05..Sun 11, so all three fixture days are in it; a start-on-Sunday bug
|
||||
// would still include 01-07 and pass, so this case is paired with the
|
||||
// weekday case above rather than relied on alone.
|
||||
func TestBillingPeriodSundayReachesBackToMonday(t *testing.T) {
|
||||
got := runBillingPage(t, "2026-01-11", "week", "billing-by-source")
|
||||
if !strings.Contains(got, "60") {
|
||||
t.Errorf("Sunday's week view should still total 60:\n%s", got)
|
||||
}
|
||||
// And the Monday view must NOT include Sunday's (nonexistent) traffic.
|
||||
gotMon := runBillingPage(t, "2026-01-05", "day", "billing-by-source")
|
||||
if strings.Contains(gotMon, "30") {
|
||||
t.Errorf("Monday's single day view shows Wednesday's cost 30:\n%s", gotMon)
|
||||
}
|
||||
if !strings.Contains(gotMon, "10") {
|
||||
t.Errorf("Monday's day view lost Monday's own cost 10:\n%s", gotMon)
|
||||
}
|
||||
}
|
||||
|
||||
// TestBillingPeriodKpiTotalMatchesWindowDays pins the KPI total itself.
|
||||
//
|
||||
// The dimension-table tests above cannot catch a broken window or a broken
|
||||
// fold: by_source is computed from the per-day DIMENSION map, independently of
|
||||
// the total. Mutation verification proved it — disabling the total fold, and
|
||||
// shifting the ISO-week start to Sunday, both left every earlier test green.
|
||||
// The total is derived from by_day, so it is the one value that moves when
|
||||
// either the window membership or the accumulation is wrong.
|
||||
func TestBillingPeriodKpiTotalMatchesWindowDays(t *testing.T) {
|
||||
kpis := func(t *testing.T, today, period string) string {
|
||||
t.Helper()
|
||||
return runBillingPage(t, today, period, "billing-by-source")
|
||||
}
|
||||
// 2026-01-07 is a Wednesday; fixture days are 05(10) 06(20) 07(30).
|
||||
// KPI renders through fmtInt, so 30 -> "30", 60 -> "60".
|
||||
if got := kpis(t, "2026-01-07", "day"); !strings.Contains(got, "30") {
|
||||
t.Errorf("day KPI should show that day's cost 30:\n%s", got)
|
||||
}
|
||||
if got := kpis(t, "2026-01-07", "week"); !strings.Contains(got, "60") {
|
||||
t.Errorf("week KPI should total 10+20+30 = 60:\n%s", got)
|
||||
}
|
||||
// Sunday 2026-01-11: the ISO week Mon 05..Sun 11 still holds all three days.
|
||||
if got := kpis(t, "2026-01-11", "week"); !strings.Contains(got, "60") {
|
||||
t.Errorf("Sunday's week KPI should still total 60:\n%s", got)
|
||||
}
|
||||
// The Monday rule from the negative side. A window that started on SUNDAY
|
||||
// would cover Sun 11..Sat 17 and contain none of the fixture days, giving
|
||||
// a total of 0 — so this case is what makes the +6 shift observable. The
|
||||
// positive case alone cannot distinguish it: shifting the start one day
|
||||
// later still happens to contain 01-07.
|
||||
if got := kpis(t, "2026-01-11", "day"); !strings.Contains(got, "0") {
|
||||
t.Errorf("Sunday's single-day window has no fixture traffic and must total 0:\n%s", got)
|
||||
}
|
||||
// Monday 2026-01-12 begins a new week; it contains no fixture days.
|
||||
if got := kpis(t, "2026-01-12", "week"); !strings.Contains(got, "0") {
|
||||
t.Errorf("an empty week should report 0, not stale data:\n%s", got)
|
||||
}
|
||||
// ...whereas the Monday-to-Sunday week that just ended does.
|
||||
if got := kpis(t, "2026-01-12", "day"); !strings.Contains(got, "0") {
|
||||
t.Errorf("Monday the 12th has no fixture traffic and must total 0:\n%s", got)
|
||||
}
|
||||
}
|
||||
|
||||
// TestBillingPeriodDimensionFoldSumsAcrossDays covers the per-dimension
|
||||
// accumulation. Asserting "60 appears somewhere" was NOT enough: mutation
|
||||
// verification showed a fold that keeps only the last day still passed,
|
||||
// because 60 also appears in the KPI total and in a token column. The
|
||||
// assertion must therefore pin srcA's own cost cell, not the page's text.
|
||||
func TestBillingPeriodDimensionFoldSumsAcrossDays(t *testing.T) {
|
||||
got := runBillingPage(t, "2026-01-07", "week", "billing-by-source")
|
||||
// srcA exists on all three days at 10/20/30. Summed = 60; last-only = 30.
|
||||
if !strings.Contains(got, "USD 60.0000") {
|
||||
t.Errorf("weekly by-source did not sum srcA across days (want USD 60.0000, "+
|
||||
"last-day-only would be USD 30.0000):\n%s", got)
|
||||
}
|
||||
// srcB exists only on 01-06 at 5, inside the window, so it must survive at
|
||||
// exactly its own value.
|
||||
if !strings.Contains(got, "USD 5.0000") {
|
||||
t.Errorf("weekly by-source lost or altered srcB (want USD 5.0000):\n%s", got)
|
||||
}
|
||||
// srcA's summed cost must also differ from its last-day value in the
|
||||
// rendered table, so an overwrite cannot masquerade as a sum.
|
||||
if strings.Count(got, "USD 60.0000") < 2 {
|
||||
t.Errorf("expected the summed 60 in both the table and the KPI:\n%s", got)
|
||||
}
|
||||
}
|
||||
|
||||
// state untouched, including the lifetime sentinel and the lifetime-only
|
||||
// degraded counter. Without this control, a rescale that silently degraded
|
||||
// into "always empty" would pass the narrow-window tests.
|
||||
// TestBillingPeriodAllStaysLifetime is the control: "all" must keep the stored
|
||||
// state untouched. Without it, a rescale that silently degraded into "always
|
||||
// empty" would pass every narrow-window test above.
|
||||
func TestBillingPeriodAllStaysLifetime(t *testing.T) {
|
||||
got := runBillingPage(t, "2026-01-07", "all", "billing-by-source")
|
||||
if !strings.Contains(got, "LIFETIMESRC") && !strings.Contains(got, "999") {
|
||||
t.Errorf("the all-time view dropped the lifetime figures:\n%s", got)
|
||||
}
|
||||
// 4242 is the lifetime degraded_reqs; it has no day-keyed tally, so it
|
||||
// must remain visible when the window is "all". The KPI goes through
|
||||
// fmtInt, so it renders grouped ("4,242") — assert the rendered form.
|
||||
if !strings.Contains(got, "4,242") {
|
||||
t.Errorf("the all-time view should still show the lifetime degraded count:\n%s", got)
|
||||
}
|
||||
}
|
||||
|
||||
// TestBillingPeriodHidesLifetimeOnlyCounters is the mismatch the period view
|
||||
// exists to remove: degraded/unpriced have no day-keyed tally, so printing
|
||||
// the lifetime number next to a per-window cost is misleading. They must be
|
||||
// hidden, not merely relabelled.
|
||||
func TestBillingPeriodHidesLifetimeOnlyCounters(t *testing.T) {
|
||||
got := runBillingPage(t, "2026-01-07", "day", "billing-by-source")
|
||||
if strings.Contains(got, "4,242") || strings.Contains(got, "Degraded") {
|
||||
t.Errorf("a one-day view still shows the lifetime degraded counter:\n%s", got)
|
||||
}
|
||||
// ...while "all" does show it (covered above), so this is not just a
|
||||
// dictionary that lost the words.
|
||||
if !strings.Contains(runBillingPage(t, "2026-01-07", "all", "billing-by-source"), "4,242") {
|
||||
t.Error("the all-time view lost the degraded counter — the hide is not period-scoped")
|
||||
}
|
||||
}
|
||||
1238
internal/lua/billing_test.go
Normal file
1238
internal/lua/billing_test.go
Normal file
File diff suppressed because it is too large
Load Diff
499
internal/lua/billing_ui_test.go
Normal file
499
internal/lua/billing_ui_test.go
Normal file
@ -0,0 +1,499 @@
|
||||
package lua
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"os"
|
||||
"os/exec"
|
||||
"regexp"
|
||||
"strings"
|
||||
"testing"
|
||||
)
|
||||
|
||||
// The billing page rendered EMPTY in production while its data endpoint returned
|
||||
// 200 with real numbers. The cause was one line: render(st) referenced an
|
||||
// undefined `s` for two KPI cells, so the ReferenceError aborted the whole
|
||||
// render and every table stayed at its initial empty innerHTML.
|
||||
//
|
||||
// Nothing in the build, the tests or the API surfaced it. This file is the guard
|
||||
// for the whole class: a plugin's injected UI that references an undefined name,
|
||||
// or that depends on a container the page does not provide, fails silently.
|
||||
|
||||
// billingUI returns the injected markup for the billing plugin: the full page
|
||||
// mount and the status-page element mount.
|
||||
func billingUI(t *testing.T) (page string, statusElement string) {
|
||||
t.Helper()
|
||||
src, err := os.ReadFile("plugins/billing.lua")
|
||||
if err != nil {
|
||||
t.Fatalf("read billing.lua: %v", err)
|
||||
}
|
||||
// Located BY CONTENT, not by index. The plugin also uses a long string for
|
||||
// its inline SVG icon, so "the first long string" is the icon and "the
|
||||
// second" is the page — which is exactly the kind of positional assumption
|
||||
// that breaks the next time an icon or a description is added.
|
||||
page = longStringContaining(t, string(src), "billing-root")
|
||||
statusElement = longStringContaining(t, string(src), "billing-status-tile")
|
||||
return page, statusElement
|
||||
}
|
||||
|
||||
// longStringContaining returns the [==[ ... ]==] body that contains marker.
|
||||
func longStringContaining(t *testing.T, src, marker string) string {
|
||||
t.Helper()
|
||||
re := regexp.MustCompile(`(?s)\[==\[(.*?)\]==\]`)
|
||||
for _, m := range re.FindAllStringSubmatch(src, -1) {
|
||||
if strings.Contains(m[1], marker) {
|
||||
return m[1]
|
||||
}
|
||||
}
|
||||
t.Fatalf("no long string contains %q", marker)
|
||||
return ""
|
||||
}
|
||||
|
||||
// TestBillingMountScriptExecutes is the guard for the production bug.
|
||||
//
|
||||
// The Billing page rendered empty while its data endpoint returned 200 with real
|
||||
// numbers. Cause: render(st) referenced an undefined `s` for two KPI cells, the
|
||||
// ReferenceError aborted the render, and every table kept its initial empty
|
||||
// innerHTML. Nothing in the build or the API surfaced it.
|
||||
//
|
||||
// Two earlier attempts at a static check were both wrong: a "declared names"
|
||||
// scan flagged every CSS class inside the inline HTML strings (class, div, td),
|
||||
// and a CSS-selector parse of the stylesheet reported the stylesheet itself as
|
||||
// broken. Static analysis of JS embedded in HTML strings is the wrong tool.
|
||||
//
|
||||
// So this actually RUNS the script, in node, against a minimal DOM stub, and
|
||||
// fails on any thrown error. Skipped when node is unavailable, with the reason
|
||||
// printed — never silently passing as if it had checked.
|
||||
func TestBillingMountScriptExecutes(t *testing.T) {
|
||||
page, el := billingUI(t)
|
||||
scripts := extractScripts(page)
|
||||
if len(scripts) == 0 {
|
||||
t.Fatal("no <script> found in the billing page mount")
|
||||
}
|
||||
for i, js := range scripts {
|
||||
assertRendersAndDoesNotThrow(t, i, js)
|
||||
}
|
||||
for i, js := range extractScripts(el) {
|
||||
assertRendersAndDoesNotThrow(t, i, js)
|
||||
}
|
||||
}
|
||||
|
||||
// assertRendersAndDoesNotThrow executes a mount script under node against a DOM
|
||||
// stub and fails on EITHER a thrown/reported error OR an empty render.
|
||||
//
|
||||
// Checking only for exceptions is not enough, and that is the third wrong
|
||||
// attempt at this guard. The plugin's render() opens with `if (!st) return;`,
|
||||
// so passing the wrong field (`render(j.stateX)`) is a SILENT no-op: no throw,
|
||||
// no console.error, no rejection — just an empty page. Only looking at the
|
||||
// produced DOM catches that class.
|
||||
func assertRendersAndDoesNotThrow(t *testing.T, idx int, js string) {
|
||||
t.Helper()
|
||||
node, err := exec.LookPath("node")
|
||||
if err != nil {
|
||||
t.Skipf("node not available (%v): cannot execute the injected script", err)
|
||||
}
|
||||
|
||||
stub := `
|
||||
global.window = global;
|
||||
global.document = {
|
||||
getElementById: function (id) {
|
||||
if (!global.__els) global.__els = {};
|
||||
if (!global.__els[id]) global.__els[id] = {
|
||||
style: {}, dataset: {}, classList: { add: function(){}, remove: function(){} },
|
||||
// Both writes count: the Billing page fills innerHTML, the status-page
|
||||
// tile assigns textContent. Watching only one of them flagged the tile as
|
||||
// "renders nothing" — a false positive that would have taught everyone to
|
||||
// ignore this test.
|
||||
set innerHTML(v) { if (v && String(v).trim()) global.__written.push(id); this.__h = v; },
|
||||
get innerHTML() { return this.__h || ""; },
|
||||
set textContent(v) { if (v !== undefined && String(v).trim()) global.__written.push(id); this.__t = v; },
|
||||
get textContent() { return this.__t || ""; },
|
||||
set innerText(v) { if (v !== undefined && String(v).trim()) global.__written.push(id); this.__i = v; },
|
||||
get innerText() { return this.__i || ""; },
|
||||
appendChild: function(){}, querySelector: function(){ return null; },
|
||||
querySelectorAll: function(){ return []; }, addEventListener: function(){} };
|
||||
return global.__els[id];
|
||||
},
|
||||
createElement: function () { return { style: {}, dataset: {}, appendChild: function(){}, setAttribute: function(){} }; },
|
||||
addEventListener: function () {},
|
||||
};
|
||||
global.pluginAPI = { onTabShown: function () {} };
|
||||
global.fetch = function () {
|
||||
// A payload with real numbers, so a working render produces visible output.
|
||||
return Promise.resolve({ ok: true, json: function () {
|
||||
return Promise.resolve({ plugin: "billing", state: {
|
||||
currency: "USD",
|
||||
total: { cost: 1.25, requests: 7, prompt_tokens: 100, completion_tokens: 20, failures: 0 },
|
||||
by_source: { localzen: { cost: 1.25, requests: 7, prompt_tokens: 100, completion_tokens: 20, failures: 0 } },
|
||||
by_model: { m1: { cost: 1.25, requests: 7, prompt_tokens: 100, completion_tokens: 20, failures: 0 } },
|
||||
by_key: { k1: { cost: 1.25, requests: 7, prompt_tokens: 100, completion_tokens: 20, failures: 0 } },
|
||||
by_day: { "2026-01-01": { cost: 1.25, requests: 7, prompt_tokens: 100, completion_tokens: 20, failures: 0 } },
|
||||
degraded_reqs: 1, unpriced_reqs: 2,
|
||||
}});
|
||||
}});
|
||||
};
|
||||
global.__written = [];
|
||||
global.__errors = [];
|
||||
// The plugin's refresh() swallows its own errors so decoration can never break
|
||||
// the host page — which is right, and is exactly why the production ReferenceError
|
||||
// was invisible. So the trace has to come from console.error, which the plugin now
|
||||
// emits. Hooking process events alone made this test pass against the very bug it
|
||||
// was written for (verified by re-introducing the typo and watching it stay green).
|
||||
var __realErr = console.error;
|
||||
console.error = function () {
|
||||
global.__errors.push("console.error: " + Array.prototype.map.call(arguments, function (a) {
|
||||
return (a && a.message) ? a.message : String(a);
|
||||
}).join(" "));
|
||||
__realErr.apply(console, arguments);
|
||||
};
|
||||
process.on("uncaughtException", function (e) { global.__errors.push("uncaught: " + String(e && e.message || e)); });
|
||||
process.on("unhandledRejection", function (e) { global.__errors.push("unhandled: " + String(e && e.message || e)); });
|
||||
`
|
||||
// 50ms is not a guess: the plugin's IIFE kicks off refresh() which awaits a
|
||||
// fetch; a rejection lands on the microtask queue almost immediately.
|
||||
script := strings.Join([]string{
|
||||
stub, js,
|
||||
`setTimeout(function(){
|
||||
console.log("__ERRS__" + JSON.stringify({errors: global.__errors, written: global.__written}));
|
||||
}, 300);`,
|
||||
}, "\n")
|
||||
cmd := exec.Command(node, "-e", script)
|
||||
out, err := cmd.CombinedOutput()
|
||||
if err != nil {
|
||||
t.Errorf("mount script %d crashed under node: %v\n%s", idx, err, out)
|
||||
return
|
||||
}
|
||||
// The script may print other things; find the JSON we appended.
|
||||
line := ""
|
||||
for _, l := range strings.Split(string(out), "\n") {
|
||||
if i := strings.Index(l, "__ERRS__"); i >= 0 {
|
||||
line = strings.TrimSpace(l[i+len("__ERRS__"):])
|
||||
}
|
||||
}
|
||||
if line == "" {
|
||||
t.Errorf("mount script %d produced no error report — the harness did not "+
|
||||
"run to completion, so it cannot be trusted to have checked anything", idx)
|
||||
return
|
||||
}
|
||||
var report struct {
|
||||
Errors []string `json:"errors"`
|
||||
Written []string `json:"written"`
|
||||
}
|
||||
if err := json.Unmarshal([]byte(line), &report); err != nil {
|
||||
t.Errorf("could not parse the error report %q: %v", line, err)
|
||||
return
|
||||
}
|
||||
for _, e := range report.Errors {
|
||||
t.Errorf("mount script %d reported %q — this is the failure that leaves "+
|
||||
"the Billing page blank while the API still returns data", idx, e)
|
||||
}
|
||||
if len(report.Written) == 0 {
|
||||
t.Errorf("mount script %d wrote NOTHING into any container — the page "+
|
||||
"renders empty. render() guards with `if (!st) return`, so a wrong "+
|
||||
"field name is a silent no-op: no throw, no console output, no error "+
|
||||
"anywhere. This is what the production page looked like.", idx)
|
||||
}
|
||||
}
|
||||
|
||||
func extractScripts(html string) []string {
|
||||
re := regexp.MustCompile(`(?s)<script[^>]*>(.*?)</script>`)
|
||||
var out []string
|
||||
for _, m := range re.FindAllStringSubmatch(html, -1) {
|
||||
out = append(out, m[1])
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
// TestBillingPageCoversItsData is the other half: the page declares tables for
|
||||
// per-source / per-model / per-key / per-day and must actually render into them.
|
||||
// A table id that is never written to is exactly how "the page loads and shows
|
||||
// nothing" happens without an error.
|
||||
func TestBillingPageCoversItsData(t *testing.T) {
|
||||
page, _ := billingUI(t)
|
||||
for _, id := range []string{
|
||||
"billing-kpis", "billing-by-source", "billing-by-model",
|
||||
"billing-by-key", "billing-by-day",
|
||||
} {
|
||||
if !strings.Contains(page, `id="`+id+`"`) {
|
||||
t.Errorf("the page has no container #%s", id)
|
||||
}
|
||||
}
|
||||
// Every container must be written to by the script, not just declared.
|
||||
scripts := extractScripts(page)
|
||||
all := strings.Join(scripts, "\n")
|
||||
for _, id := range []string{"billing-kpis", "billing-by-source", "billing-by-model", "billing-by-key", "billing-by-day"} {
|
||||
if !strings.Contains(all, `getElementById("`+id+`")`) {
|
||||
t.Errorf("#%s is declared but never read by the script — it stays empty forever", id)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// TestPluginIconIsNotARawEmoji guards the sidebar icon. The billing plugin
|
||||
// declared icon = "💰" and the WebUI drops that verbatim into the button, while
|
||||
// every native tab uses an inline SVG styled with `stroke: currentColor`. An
|
||||
// emoji there renders at the wrong size and ignores the theme, so it does not
|
||||
// match its neighbours — which is what the operator reported.
|
||||
func TestPluginIconIsNotARawEmoji(t *testing.T) {
|
||||
icon := billingIcon(mustBillingSource(t))
|
||||
if icon == "" {
|
||||
t.Fatal("billing declares no icon; the sidebar entry would be blank")
|
||||
}
|
||||
if isEmojiIcon(icon) {
|
||||
t.Errorf("page icon is the raw emoji %q — the WebUI sidebar renders "+
|
||||
"native tabs as inline SVG (stroke: currentColor), so an emoji is the "+
|
||||
"wrong size and ignores the theme. Use an inline SVG path instead.", icon)
|
||||
}
|
||||
}
|
||||
|
||||
// isEmojiIcon reports whether s is a pictographic emoji rather than markup or a
|
||||
// text glyph. Codepoints in the pictographic blocks, plus the regional-indicator
|
||||
// pair used by flags.
|
||||
func isEmojiIcon(s string) bool {
|
||||
r := []rune(s)
|
||||
if len(r) == 0 {
|
||||
return false
|
||||
}
|
||||
// Anything containing '<' is markup (an inline <svg>), which is the fix.
|
||||
if strings.ContainsRune(s, '<') {
|
||||
return false
|
||||
}
|
||||
for _, c := range r {
|
||||
switch {
|
||||
case c >= 0x1F300 && c <= 0x1FAFF, // pictographs, symbols, supplemental
|
||||
c >= 0x1F000 && c <= 0x1F2FF, // mahjong/domino/cards
|
||||
c >= 0x2600 && c <= 0x27BF, // misc symbols + dingbats
|
||||
c >= 0x2B00 && c <= 0x2BFF, // arrows/misc symbols
|
||||
c == 0xFE0F, // variation selector-16
|
||||
c >= 0x1F1E6 && c <= 0x1F1FF: // regional indicators (flags)
|
||||
return true
|
||||
}
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
func mustBillingSource(t *testing.T) string {
|
||||
t.Helper()
|
||||
b, err := os.ReadFile("plugins/billing.lua")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
return string(b)
|
||||
}
|
||||
|
||||
// billingIcon extracts the declared page icon. It accepts BOTH a quoted string
|
||||
// and a [==[ ... ]==] long string, because an inline SVG cannot be written as a
|
||||
// Lua short string without escaping every quote in it.
|
||||
func billingIcon(src string) string {
|
||||
if m := regexp.MustCompile(`(?m)^\s*icon\s*=\s*"([^"]*)"`).FindStringSubmatch(src); m != nil {
|
||||
return m[1]
|
||||
}
|
||||
if m := regexp.MustCompile(`(?s)\bicon\s*=\s*\[==\[(.*?)\]==\]`).FindStringSubmatch(src); m != nil {
|
||||
return m[1]
|
||||
}
|
||||
return ""
|
||||
}
|
||||
|
||||
// TestBillingTableHeaderMatchesRowColumns catches column drift.
|
||||
//
|
||||
// row() gained cache columns (fresh / cache / cache%) while the header row did
|
||||
// not, in the same edit. The result is a table whose cells are shifted one
|
||||
// column left from "fresh" onward — so "cache%" sits under "completion" and the
|
||||
// last cell has no label. It renders, it has data, and it is wrong in a way that
|
||||
// takes a careful read to notice.
|
||||
func TestBillingTableHeaderMatchesRowColumns(t *testing.T) {
|
||||
page, _ := billingUI(t)
|
||||
js := strings.Join(extractScripts(page), "\n")
|
||||
if !strings.Contains(js, "<th>") {
|
||||
t.Fatal("no table header found in the billing page script")
|
||||
}
|
||||
hStart := strings.Index(js, "function tableFor")
|
||||
if hStart < 0 {
|
||||
t.Fatal("no tableFor in the billing page script")
|
||||
}
|
||||
// Count the usage table's headers ONLY. The count is taken from tableFor to
|
||||
// the end of the script, which used to be fine when the rule editor lived on
|
||||
// its own page; now that both halves share one script, the rule table's nine
|
||||
// headers were counted too and the check reported 17 vs 8 — a failure about
|
||||
// two unrelated tables. Stop at the rule editor's script.
|
||||
header := js[hStart:]
|
||||
if cut := strings.Index(header, "/api/plugins/billing/rules"); cut > 0 {
|
||||
header = header[:cut]
|
||||
}
|
||||
// Count <th> with OR without attributes: the name column now carries
|
||||
// style='width:22%' so the header and its data cell can be given the same
|
||||
// width, and a bare "<th>" count silently reported 7 vs 8 — the check
|
||||
// failing on the very column it had just been taught to size.
|
||||
th := strings.Count(header, "<th>") + strings.Count(header, "<th ")
|
||||
|
||||
rStart := strings.Index(js, "function row")
|
||||
rEnd := strings.Index(js, "function tableFor")
|
||||
if rStart < 0 || rEnd < 0 || rEnd <= rStart {
|
||||
t.Fatal("could not isolate row()")
|
||||
}
|
||||
row := js[rStart:rEnd]
|
||||
// Each cell closes with </td>; the first cell uses <td><b>..</b></td> so
|
||||
// counting </td> is exact.
|
||||
td := strings.Count(row, "</td>")
|
||||
|
||||
if th != td {
|
||||
t.Errorf("★ header declares %d columns but row() emits %d cells — the "+
|
||||
"table is misaligned from the first differing column on", th, td)
|
||||
}
|
||||
}
|
||||
|
||||
// TestRuleTableLocksColumnWidthsInAColgroup guards the layout class of bug that
|
||||
// only a screenshot caught: two adjacent headers rendered on top of each other,
|
||||
// the URL input squeezed to "https:", and the per-model × button landing 78px
|
||||
// past its cell onto the neighbouring Delete button.
|
||||
//
|
||||
// The measurement that found each one:
|
||||
// - table-layout:fixed with widths on <td> → 954px divided into 9 equal
|
||||
// 106px columns, every declared width discarded.
|
||||
// - widths on <th> WITHOUT fixed → browser sizes from content and ignores
|
||||
// them; the 70px currency column still collapsed to 48px (24px of input).
|
||||
// - fixed + <th> widths → they take effect, but the model-price cell's width
|
||||
// then depended on WHICH RULE was widest (352px vs 206px).
|
||||
//
|
||||
// The fix is all three together: fixed layout, a colgroup, and flex children
|
||||
// that may shrink.
|
||||
func TestRuleTableLocksColumnWidthsInAColgroup(t *testing.T) {
|
||||
page, _ := billingUI(t)
|
||||
js := strings.Join(extractScripts(page), "\n")
|
||||
// Match the ATTRIBUTE, not the bare string: a substring check on
|
||||
// "table-layout:fixed" is satisfied by the explanatory comment that sits
|
||||
// three lines above the tag, so removing the attribute from the table still
|
||||
// passed (mutation-verified). Require it inside a style='...' literal.
|
||||
if !strings.Contains(js, `font-size:12px;table-layout:fixed'`) {
|
||||
t.Error("the rule table must stay table-layout:fixed — with auto layout " +
|
||||
"the <colgroup> widths below are ignored and columns size from content")
|
||||
}
|
||||
colStart := strings.Index(js, "<colgroup>")
|
||||
if colStart < 0 {
|
||||
t.Fatal("the rule table declares no <colgroup>; per-column widths on " +
|
||||
"<th> or <td> are advisory and get recomputed from content")
|
||||
}
|
||||
// The colgroup is built from a JS array of pixel widths, so assert on the
|
||||
// array, not on literal "<col " occurrences — counting the literal found
|
||||
// exactly one and reported "1 columns, want 9" against correct code.
|
||||
if !strings.Contains(js, "[214, 80, 62, 70, 66, 74, 62, 262, 74]") {
|
||||
t.Error("the rule table's colgroup widths changed or vanished; measured " +
|
||||
"values were [214 80 62 70 66 74 62 262 74] for a 954px pane — the " +
|
||||
"delete column needs 74px or the button wraps to two lines")
|
||||
}
|
||||
if !strings.Contains(js, `"<col style='width:" + w + "px'>"`) {
|
||||
t.Error("colgroup entries are no longer emitted as <col style='width:Npx'>")
|
||||
}
|
||||
// Every price-row INPUT must be allowed to shrink, or the row's fixed widths
|
||||
// sum past the cell and the × button lands on the next column. The × button
|
||||
// is the opposite: it must NOT shrink (flex:0 0 auto), because a squashed
|
||||
// delete button is unclickable.
|
||||
for _, needle := range []string{"m-name", "m-p", "m-c"} {
|
||||
i := strings.Index(js, "class='"+needle+"'")
|
||||
if i < 0 {
|
||||
i = strings.Index(js, "class='ghost small "+needle+"'")
|
||||
}
|
||||
if i < 0 {
|
||||
t.Errorf("price row lost the %s input/button", needle)
|
||||
continue
|
||||
}
|
||||
window := js[i : i+220]
|
||||
if !strings.Contains(window, "flex:") {
|
||||
t.Errorf("%s has no flex sizing; the row is 80+96+96+38 ≈ 330px wide "+
|
||||
"against a 262px cell and overflows onto the next column", needle)
|
||||
}
|
||||
if !strings.Contains(window, "min-width:0") {
|
||||
t.Errorf("%s lacks min-width:0; a flex item will not shrink below its "+
|
||||
"content width, which is how the × button ended up 78px past its cell", needle)
|
||||
}
|
||||
}
|
||||
i := strings.Index(js, "class='ghost small m-del'")
|
||||
if i < 0 {
|
||||
t.Fatal("the per-model delete button is gone")
|
||||
}
|
||||
if !strings.Contains(js[i:i+120], "flex:0 0 auto") {
|
||||
t.Error("the per-model × button must be flex:0 0 auto — it is a fixed-size " +
|
||||
"control and must never be squeezed by the flexible inputs")
|
||||
}
|
||||
}
|
||||
|
||||
// TestBillingUIShowsRecentUnpriced: the page must be able to distinguish a
|
||||
// settled pricing debt from a live leak. The lifetime unpriced counter answers
|
||||
// neither — it keeps every request from before a rule was added, so it reads the
|
||||
// same number in both situations. The day buckets exist to make the difference
|
||||
// visible, and a bucket nothing renders is a bucket nobody benefits from.
|
||||
func TestBillingUIShowsRecentUnpriced(t *testing.T) {
|
||||
raw, err := os.ReadFile("plugins/billing.lua")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
src := string(raw)
|
||||
for _, needle := range []string{
|
||||
"unpriced_by_day", // the tally is written
|
||||
"recentUnpriced", // it is reduced to a window
|
||||
"unpriced_recent", // it reaches the view
|
||||
"unpriced7d", // the card labels it
|
||||
"unpricedClean", // and says so when the window is clear
|
||||
} {
|
||||
if !strings.Contains(src, needle) {
|
||||
t.Errorf("billing UI lacks %q", needle)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// TestBillingRulesUIExposesReasoningPrice: the thinking-token rate must be
|
||||
// editable in the rules editor, not only settable by hand-editing config.yaml.
|
||||
// A pricing capability that lives only in the config file is one the operator
|
||||
// will assume does not exist.
|
||||
func TestBillingRulesUIExposesReasoningPrice(t *testing.T) {
|
||||
raw, err := os.ReadFile("plugins/billing.lua")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
src := string(raw)
|
||||
for _, needle := range []string{
|
||||
`class='m-r'`, // the input, in both the existing rows and the "+ model" row
|
||||
`models[n].reasoning`, // collected back out of the DOM on save
|
||||
`reasoningHint`, // and it says what leaving it blank means
|
||||
} {
|
||||
if !strings.Contains(src, needle) {
|
||||
t.Errorf("rules editor lacks %q", needle)
|
||||
}
|
||||
}
|
||||
// Two occurrences: an existing model's row and the template for a new one.
|
||||
// One would mean the "+ model" button produces a row that silently drops
|
||||
// the rate on the next save.
|
||||
if n := strings.Count(src, `class='m-r'`); n < 2 {
|
||||
t.Errorf("m-r appears %d time(s); both the row template and the editor need it", n)
|
||||
}
|
||||
}
|
||||
|
||||
// TestBillingRulesUIExposesLongContextStep: the long-context step must be
|
||||
// editable in the rules editor.
|
||||
//
|
||||
// This is not a niche control. On an agent workload 88% of requests carry more
|
||||
// than 128k prompt tokens with a median of 247k, so a rule that omits the step
|
||||
// under-bills the MAJORITY of traffic while every number on the dashboard stays
|
||||
// plausible. Leaving the capability in config.yaml alone would mean the operator
|
||||
// never discovers it needs it.
|
||||
func TestBillingRulesUIExposesLongContextStep(t *testing.T) {
|
||||
raw, err := os.ReadFile("plugins/billing.lua")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
src := string(raw)
|
||||
for _, needle := range []string{
|
||||
`class='m-lc-a'`, // threshold
|
||||
`class='m-lc-p'`, // absolute rate
|
||||
`class='m-lc-m'`, // multiplier
|
||||
`models[n].long_context`, // collected back out on save
|
||||
`lcHint`, // explains that blank = no step
|
||||
} {
|
||||
if !strings.Contains(src, needle) {
|
||||
t.Errorf("rules editor lacks %q", needle)
|
||||
}
|
||||
}
|
||||
// Both the existing-row template and the "+ model" template need the three
|
||||
// inputs; one would make a newly added model lose its step on the next save.
|
||||
if n := strings.Count(src, `class='m-lc-a'`); n < 2 {
|
||||
t.Errorf("long-context inputs appear %d time(s); both row templates need them", n)
|
||||
}
|
||||
}
|
||||
139
internal/lua/fastvalue.go
Normal file
139
internal/lua/fastvalue.go
Normal file
@ -0,0 +1,139 @@
|
||||
package lua
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"strconv"
|
||||
)
|
||||
|
||||
// This file exists to remove the JSON round-trip from the plugin hook path.
|
||||
//
|
||||
// Measured before this change: one Fire() call cost 14.6us, of which
|
||||
// json.Marshal(payload) was 4.0us and json.Unmarshal back into interface{} was
|
||||
// 5.6us — 67% of the whole hook, spent re-deriving a tree that the very next
|
||||
// line (pushGoValue) already knows how to walk natively.
|
||||
//
|
||||
// The round-trip existed because payload arrives as map[string]interface{} and
|
||||
// pushGoValue wanted a plain interface{} tree. pushGoValue already handles
|
||||
// map[string]interface{} and []interface{} directly, so the fix is a converter
|
||||
// that flattens the concrete Go types the gateway actually produces, instead of
|
||||
// a generic serialize/parse.
|
||||
//
|
||||
// The JSON path is kept as the fallback for types the fast path does not know,
|
||||
// so a plugin payload carrying something exotic still arrives rather than being
|
||||
// silently dropped.
|
||||
|
||||
// jsonNumberType is the type json.Unmarshal produces for every JSON number.
|
||||
// The fast path recognizes it so numbers survive as numbers rather than being
|
||||
// stringified into a different Lua type.
|
||||
type jsonNumberType = float64
|
||||
|
||||
// fastToPlain converts the concrete Go values the gateway puts in a hook
|
||||
// payload into the interface{} tree pushGoValue consumes, without JSON.
|
||||
//
|
||||
// The interesting cases are the ones json.Unmarshal would have normalized:
|
||||
// - json.Number-ish types (all float64 in practice)
|
||||
// - int / int64 / uint variants, which must land as Lua numbers, not as the
|
||||
// strings a naive "only handle float64" switch would produce
|
||||
// - []string and map[string]string, extremely common in payloads and NOT
|
||||
// handled by a switch that only knows []interface{}
|
||||
//
|
||||
// Anything unrecognized returns (value, false) so the caller can fall back to
|
||||
// the JSON path, which is slower but total.
|
||||
func fastToPlain(v interface{}) (interface{}, bool) {
|
||||
switch x := v.(type) {
|
||||
case nil, bool, string, jsonNumberType:
|
||||
return v, true
|
||||
case int:
|
||||
return float64(x), true
|
||||
case int8:
|
||||
return float64(x), true
|
||||
case int16:
|
||||
return float64(x), true
|
||||
case int32:
|
||||
return float64(x), true
|
||||
case int64:
|
||||
return float64(x), true
|
||||
case uint:
|
||||
return float64(x), true
|
||||
case uint8:
|
||||
return float64(x), true
|
||||
case uint16:
|
||||
return float64(x), true
|
||||
case uint32:
|
||||
return float64(x), true
|
||||
case uint64:
|
||||
return float64(x), true
|
||||
case float32:
|
||||
return float64(x), true
|
||||
case []string:
|
||||
out := make([]interface{}, len(x))
|
||||
for i, s := range x {
|
||||
out[i] = s
|
||||
}
|
||||
return out, true
|
||||
case map[string]string:
|
||||
out := make(map[string]interface{}, len(x))
|
||||
for k, s := range x {
|
||||
out[k] = s
|
||||
}
|
||||
return out, true
|
||||
case map[string]interface{}:
|
||||
// The overwhelmingly common case, and the one the hook path always
|
||||
// takes at the top level.
|
||||
out := make(map[string]interface{}, len(x))
|
||||
for k, it := range x {
|
||||
c, ok := fastToPlain(it)
|
||||
if !ok {
|
||||
return nil, false
|
||||
}
|
||||
out[k] = c
|
||||
}
|
||||
return out, true
|
||||
case []interface{}:
|
||||
out := make([]interface{}, len(x))
|
||||
for i, it := range x {
|
||||
c, ok := fastToPlain(it)
|
||||
if !ok {
|
||||
return nil, false
|
||||
}
|
||||
out[i] = c
|
||||
}
|
||||
return out, true
|
||||
default:
|
||||
return nil, false
|
||||
}
|
||||
}
|
||||
|
||||
// plainForLua returns the interface{} tree to hand to pushGoValue, using the
|
||||
// fast path when it can and JSON only when it must.
|
||||
func plainForLua(payload map[string]interface{}) interface{} {
|
||||
if fast, ok := fastToPlain(payload); ok {
|
||||
return fast
|
||||
}
|
||||
// Rare: some type the fast path does not model. Serialize and let the
|
||||
// generic decoder normalize it, so the hook still sees the field.
|
||||
b, err := json.Marshal(payload)
|
||||
if err != nil {
|
||||
return payload
|
||||
}
|
||||
var decoded interface{}
|
||||
if err := json.Unmarshal(b, &decoded); err != nil {
|
||||
return payload
|
||||
}
|
||||
return decoded
|
||||
}
|
||||
|
||||
// toPlainSlice is the []interface{} entry point of fastToPlain, exposed so the
|
||||
// Fire path can flatten a pre-built slice without re-walking the map header.
|
||||
func toPlainSlice(v []interface{}) []interface{} {
|
||||
if fast, ok := fastToPlain(v); ok {
|
||||
return fast.([]interface{})
|
||||
}
|
||||
return v
|
||||
}
|
||||
|
||||
// luaNumberString renders a float the way Lua would print it, for the rare case
|
||||
// a hook wants the textual form. Kept out of the hot path.
|
||||
func luaNumberString(f float64) string {
|
||||
return strconv.FormatFloat(f, 'g', -1, 64)
|
||||
}
|
||||
132
internal/lua/fastvalue_test.go
Normal file
132
internal/lua/fastvalue_test.go
Normal file
@ -0,0 +1,132 @@
|
||||
package lua
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"reflect"
|
||||
"testing"
|
||||
)
|
||||
|
||||
// The hook payload used to be JSON round-tripped on every call. It no longer is,
|
||||
// so the fast path and the old JSON path must be indistinguishable — otherwise
|
||||
// a plugin silently sees a different payload than before, which is the worst
|
||||
// kind of change: it compiles, passes a smoke test, and misprices traffic.
|
||||
//
|
||||
// These tests therefore compare the two paths on the SAME inputs rather than
|
||||
// asserting the fast path's output in isolation.
|
||||
|
||||
func viaJSON(t *testing.T, payload map[string]interface{}) interface{} {
|
||||
t.Helper()
|
||||
b, err := json.Marshal(payload)
|
||||
if err != nil {
|
||||
t.Fatalf("marshal: %v", err)
|
||||
}
|
||||
var decoded interface{}
|
||||
if err := json.Unmarshal(b, &decoded); err != nil {
|
||||
t.Fatalf("unmarshal: %v", err)
|
||||
}
|
||||
return decoded
|
||||
}
|
||||
|
||||
// payloads that exercise every branch of the fast converter.
|
||||
func payloadCases() []map[string]interface{} {
|
||||
return []map[string]interface{}{
|
||||
{"model": "deepseek-v4.1-flash", "source": "commandcode", "ok": true},
|
||||
// Numbers of every width: a converter that only knows float64 turns
|
||||
// ints into something else, and a token count that arrives as a string
|
||||
// makes a Lua hook do arithmetic on nil.
|
||||
{"i": 42, "i8": int8(8), "i16": int16(16), "i32": int32(32), "i64": int64(1 << 40),
|
||||
"u": uint(7), "u64": uint64(1 << 50), "f32": float32(1.5), "f64": 2.25},
|
||||
// Zero, negative, and very large values must stay numbers.
|
||||
{"zero": 0, "neg": -17, "huge": 1e308, "tiny": 1e-308},
|
||||
// Slices and string maps: common in payloads and absent from a switch
|
||||
// that only knows []interface{} / map[string]interface{}.
|
||||
{"msgs": []interface{}{"a", "b"}, "tags": []string{"x", "y"}},
|
||||
{"kv": map[string]string{"a": "1", "b": "2"}},
|
||||
// Nesting, which is where a shallow converter silently drops a level.
|
||||
{"usage": map[string]interface{}{
|
||||
"prompt_tokens": 2048, "cache_hit_tokens": 1024,
|
||||
"nested": map[string]interface{}{"deep": []interface{}{1, "two", true, nil}},
|
||||
}},
|
||||
{"nil_field": nil, "empty_map": map[string]interface{}{}, "empty_slice": []interface{}{}},
|
||||
// A value the fast path does NOT model: it must fall back to JSON and
|
||||
// still arrive, not disappear.
|
||||
{"weird": struct {
|
||||
A int `json:"a"`
|
||||
B string `json:"b"`
|
||||
}{1, "x"}},
|
||||
}
|
||||
}
|
||||
|
||||
func TestFastPathMatchesJSONPath(t *testing.T) {
|
||||
for i, p := range payloadCases() {
|
||||
want := viaJSON(t, p)
|
||||
got := plainForLua(p)
|
||||
if !reflect.DeepEqual(want, got) {
|
||||
t.Errorf("case %d: fast path differs from JSON path\n payload: %#v\n json: %#v\n fast: %#v",
|
||||
i, p, want, got)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// TestFastPathIsActuallyUsed guards against the fast path silently degrading to
|
||||
// JSON for the payload the gateway really sends. If a future payload gains a
|
||||
// type the converter does not model, this still works (it falls back) but the
|
||||
// optimization is gone — and the next person measuring the hook would be
|
||||
// measuring the old cost without knowing why.
|
||||
func TestFastPathIsActuallyUsed(t *testing.T) {
|
||||
// This mirrors the real request_end payload shape from the gateway.
|
||||
realistic := map[string]interface{}{
|
||||
"stage": "request_end", "kind": "end", "model": "deepseek-v4.1-flash",
|
||||
"source": "commandcode", "key": "stress-key", "ok": true,
|
||||
"status": 200, "duration_ms": 1234,
|
||||
"usage": map[string]interface{}{
|
||||
"prompt_tokens": float64(2048), "completion_tokens": float64(512),
|
||||
"cache_hit_tokens": float64(1024), "total_tokens": float64(3584),
|
||||
},
|
||||
"walk": []interface{}{
|
||||
map[string]interface{}{"kind": "tier_skip", "tier": 1, "source": "", "model": "", "reason": "cooldown"},
|
||||
map[string]interface{}{"kind": "selected", "tier": 2, "source": "commandcode", "model": "m", "reason": ""},
|
||||
},
|
||||
"ts": float64(1700000000),
|
||||
}
|
||||
if _, ok := fastToPlain(realistic); !ok {
|
||||
t.Errorf("★ the realistic request payload does NOT take the fast path — " +
|
||||
"the JSON round-trip is still on the hot path for real traffic")
|
||||
}
|
||||
}
|
||||
|
||||
// TestFastPathDoesNotAliasInput: the converter builds a new tree. If it ever
|
||||
// returned the caller's map directly, a Lua hook's writes could not reach Go —
|
||||
// but worse, a later mutation of the payload would race with the snapshot the
|
||||
// persistence saver is holding.
|
||||
func TestFastPathDoesNotAliasInput(t *testing.T) {
|
||||
src := map[string]interface{}{
|
||||
"usage": map[string]interface{}{"prompt_tokens": float64(1)},
|
||||
"list": []interface{}{"a"},
|
||||
}
|
||||
out, ok := fastToPlain(src)
|
||||
if !ok {
|
||||
t.Fatal("fast path declined a plain payload")
|
||||
}
|
||||
m := out.(map[string]interface{})
|
||||
m["new"] = "added"
|
||||
src["also_new"] = "must not appear"
|
||||
|
||||
if _, leaked := m["also_new"]; leaked {
|
||||
t.Error("output map aliases the input map")
|
||||
}
|
||||
if _, leaked := src["new"]; leaked {
|
||||
t.Error("writing to the output mutated the input")
|
||||
}
|
||||
// And the nested maps must be copies too.
|
||||
inner := m["usage"].(map[string]interface{})
|
||||
inner["prompt_tokens"] = float64(999)
|
||||
if src["usage"].(map[string]interface{})["prompt_tokens"] != float64(1) {
|
||||
t.Error("nested map is shared, not copied — a hook could mutate the payload")
|
||||
}
|
||||
list := m["list"].([]interface{})
|
||||
list[0] = "changed"
|
||||
if src["list"].([]interface{})[0] != "a" {
|
||||
t.Error("nested slice is shared, not copied")
|
||||
}
|
||||
}
|
||||
334
internal/lua/fire_parallel_test.go
Normal file
334
internal/lua/fire_parallel_test.go
Normal file
@ -0,0 +1,334 @@
|
||||
package lua
|
||||
|
||||
import (
|
||||
"os"
|
||||
"path/filepath"
|
||||
"sync"
|
||||
"testing"
|
||||
"time"
|
||||
)
|
||||
|
||||
// Fire runs same-stage hooks in parallel across plugins. These tests pin the
|
||||
// three things that make that safe, each of which failed at least once while
|
||||
// the change was being written.
|
||||
|
||||
// twoHookPlugins loads two plugins that both hook request_end.
|
||||
func twoHookPlugins(t *testing.T, codeA, codeB string) *Plugins {
|
||||
t.Helper()
|
||||
vm := NewVM(filepath.Join(t.TempDir(), "adapters"))
|
||||
if err := vm.Start(); err != nil {
|
||||
t.Fatalf("vm start: %v", err)
|
||||
}
|
||||
t.Cleanup(vm.Stop)
|
||||
ps := NewPlugins(vm, filepath.Join(t.TempDir(), "plugins"))
|
||||
ps.save.setWakeDelayForTest(time.Millisecond)
|
||||
t.Cleanup(ps.Close)
|
||||
ps.DisableStatePersistence()
|
||||
if err := ps.LoadSource("aaa", codeA); err != nil {
|
||||
t.Fatalf("load aaa: %v", err)
|
||||
}
|
||||
if err := ps.LoadSource("bbb", codeB); err != nil {
|
||||
t.Fatalf("load bbb: %v", err)
|
||||
}
|
||||
return ps
|
||||
}
|
||||
|
||||
func counterPlugin(name string, key string) string {
|
||||
// Built by concatenation rather than fmt: %q is not a Go verb, and the first
|
||||
// attempt produced Lua source with a stray '=' that only failed at compile
|
||||
// time inside three different tests.
|
||||
return `
|
||||
local p = { name = "` + name + `" }
|
||||
p.state = { n = 0 }
|
||||
p.hooks = { request_end = "bump" }
|
||||
function p.bump(payload)
|
||||
p.state.n = p.state.n + 1
|
||||
if "` + key + `" ~= "" then return { who = "from-` + name + `" } end
|
||||
return nil
|
||||
end
|
||||
return p
|
||||
`
|
||||
}
|
||||
|
||||
// TestSingleAndParallelPathsMergeIdentically: the single-plugin fast path and
|
||||
// the multi-plugin parallel path must produce the same payload.
|
||||
//
|
||||
// This is not theoretical. The parallel path was written first and the
|
||||
// single-plugin path kept its old shape; the merge loop was left off the fast
|
||||
// path, so a lone plugin returning a table had its return value DISCARDED. The
|
||||
// existing stage-order test caught it — but only because it happened to check
|
||||
// the payload after Fire. A plugin that returned fields nobody read would have
|
||||
// broken silently.
|
||||
func TestSingleAndParallelPathsMergeIdentically(t *testing.T) {
|
||||
single := twoHookPlugins(t, counterPlugin("aaa", "who"), `
|
||||
local p = { name = "zzz" }
|
||||
p.state = { n = 0 }
|
||||
p.hooks = { request_end = "bump" }
|
||||
function p.bump(payload) p.state.n = p.state.n + 1 return nil end
|
||||
return p
|
||||
`)
|
||||
// remove the second so exactly one plugin hooks this stage
|
||||
single.mu.Lock()
|
||||
single.plugins = single.plugins[:1]
|
||||
single.mu.Unlock()
|
||||
single.rebuild()
|
||||
|
||||
ps := twoHookPlugins(t, counterPlugin("aaa", "who"), counterPlugin("bbb", "who"))
|
||||
|
||||
singleOut := single.Fire(StageRequestEnd, map[string]interface{}{"model": "m"})
|
||||
psOut := ps.Fire(StageRequestEnd, map[string]interface{}{"model": "m"})
|
||||
|
||||
if singleOut["who"] != "from-aaa" {
|
||||
t.Errorf("★ single-plugin path dropped the return value: %v", singleOut)
|
||||
}
|
||||
// Two plugins both return `who`; the merge order decides the winner, and it
|
||||
// must be deterministic (load order), not scheduler-dependent.
|
||||
if psOut["who"] != "from-bbb" {
|
||||
t.Errorf("merge order is not load order: who = %v, want from-bbb", psOut["who"])
|
||||
}
|
||||
}
|
||||
|
||||
// TestMergeOrderIsDeterministic: goroutine completion order must not leak into
|
||||
// the result. Running Fire repeatedly must always yield the same payload, or a
|
||||
// gateway's behaviour changes run to run with the same plugins installed.
|
||||
func TestMergeOrderIsDeterministic(t *testing.T) {
|
||||
// The plugins differ in cost so a scheduler-dependent merge is visible.
|
||||
slow := `
|
||||
local p = { name = "aaa" }
|
||||
p.state = { n = 0 }
|
||||
p.hooks = { request_end = "slow" }
|
||||
function p.slow(payload)
|
||||
local acc = 0
|
||||
for i = 1, 3000 do acc = acc + i % 7 end
|
||||
p.state.n = p.state.n + 1
|
||||
return { winner = "aaa", acc = acc }
|
||||
end
|
||||
return p
|
||||
`
|
||||
fast := `
|
||||
local p = { name = "bbb" }
|
||||
p.state = { n = 0 }
|
||||
p.hooks = { request_end = "fast" }
|
||||
function p.fast(payload)
|
||||
p.state.n = p.state.n + 1
|
||||
return { winner = "bbb" }
|
||||
end
|
||||
return p
|
||||
`
|
||||
ps := twoHookPlugins(t, slow, fast)
|
||||
first := ps.Fire(StageRequestEnd, map[string]interface{}{"model": "m"})["winner"]
|
||||
for i := 0; i < 60; i++ {
|
||||
got := ps.Fire(StageRequestEnd, map[string]interface{}{"model": "m"})["winner"]
|
||||
if got != first {
|
||||
t.Fatalf("★ winner changed between runs: %q then %q — merge depends on goroutine "+
|
||||
"scheduling, so the same configuration behaves differently run to run", first, got)
|
||||
}
|
||||
}
|
||||
// Load order decides: bbb is loaded second, so it wins.
|
||||
if first != "bbb" {
|
||||
t.Errorf("winner = %q, want bbb (load order decides, not timing)", first)
|
||||
}
|
||||
}
|
||||
|
||||
// TestThrowingHookDoesNotBreakSiblings covers the failure path that golua
|
||||
// actually produces: a hook raising a Lua error.
|
||||
//
|
||||
// The test name used to claim it covered panics, and it used error() to do it.
|
||||
// Probing the six ways a Lua program can fault (error(), indexing nil, calling
|
||||
// nil, concatenating nil, arithmetic on nil, unbounded recursion) showed golua
|
||||
// converts ALL of them into an error RETURN, not a Go panic — so the recover()
|
||||
// in Fire was untested by that case, and deleting recover() still passed. The
|
||||
// test was renamed to say what it verifies.
|
||||
//
|
||||
// recover() is kept anyway: it guards the Go side of Fire (a nil map write, a
|
||||
// future change to how the payload is prepared), which is cheap and cannot be
|
||||
// triggered from Lua today. Claiming it is covered by a Lua test would be the
|
||||
// kind of assurance that evaporates the first time someone checks.
|
||||
func TestThrowingHookDoesNotBreakSiblings(t *testing.T) {
|
||||
throwing := `
|
||||
local p = { name = "aaa" }
|
||||
p.state = {}
|
||||
p.hooks = { request_end = "boom" }
|
||||
function p.boom(payload) error("intentional failure") end
|
||||
return p
|
||||
`
|
||||
healthy := counterPlugin("bbb", "")
|
||||
ps := twoHookPlugins(t, throwing, healthy)
|
||||
|
||||
out := ps.Fire(StageRequestEnd, map[string]interface{}{"model": "m"}) // must not crash
|
||||
if out == nil {
|
||||
t.Error("Fire returned nil")
|
||||
}
|
||||
st, ok := ps.State("bbb").(map[string]interface{})
|
||||
if !ok {
|
||||
t.Fatal("healthy plugin has no state")
|
||||
}
|
||||
if n := st["n"]; n != float64(1) {
|
||||
t.Errorf("★ healthy plugin did not run alongside the throwing one: n = %v", n)
|
||||
}
|
||||
he := ps.HookErrors()
|
||||
if len(he) == 0 {
|
||||
t.Error("a throwing hook was not recorded in hook_errors — the failure would be invisible")
|
||||
}
|
||||
}
|
||||
|
||||
// TestParallelHooksAllRunOnce: every plugin must be invoked exactly once per
|
||||
// Fire. A lost or duplicated goroutine shows up as a wrong total, which for the
|
||||
// billing plugin means a wrong bill.
|
||||
func TestParallelHooksAllRunOnce(t *testing.T) {
|
||||
ps := twoHookPlugins(t, counterPlugin("aaa", ""), counterPlugin("bbb", ""))
|
||||
const fires = 100
|
||||
for i := 0; i < fires; i++ {
|
||||
ps.Fire(StageRequestEnd, map[string]interface{}{"model": "m"})
|
||||
}
|
||||
for _, name := range []string{"aaa", "bbb"} {
|
||||
st, ok := ps.State(name).(map[string]interface{})
|
||||
if !ok {
|
||||
t.Fatalf("%s has no state", name)
|
||||
}
|
||||
if n := st["n"]; n != float64(fires) {
|
||||
t.Errorf("%s counted %v hooks, want %d", name, n, fires)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// TestConcurrentFireIsSafe drives Fire from many goroutines at once. Each
|
||||
// plugin has its own Lua state and its own mutex, so this must hold; the race
|
||||
// detector is what proves it, not the assertions.
|
||||
func TestConcurrentFireIsSafe(t *testing.T) {
|
||||
ps := twoHookPlugins(t, counterPlugin("aaa", ""), counterPlugin("bbb", ""))
|
||||
var wg sync.WaitGroup
|
||||
for g := 0; g < 8; g++ {
|
||||
wg.Add(1)
|
||||
go func() {
|
||||
defer wg.Done()
|
||||
for i := 0; i < 50; i++ {
|
||||
ps.Fire(StageRequestEnd, map[string]interface{}{"model": "m", "source": "s"})
|
||||
}
|
||||
}()
|
||||
}
|
||||
wg.Wait()
|
||||
for _, name := range []string{"aaa", "bbb"} {
|
||||
st, ok := ps.State(name).(map[string]interface{})
|
||||
if !ok {
|
||||
t.Fatalf("%s has no state", name)
|
||||
}
|
||||
if n := st["n"]; n != float64(400) {
|
||||
t.Errorf("%s counted %v, want 400", name, n)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// TestFireDoesNotHoldRegistryLockWhileRunningHooks: a hook that installs or
|
||||
// disables a plugin takes ps.mu for write. If Fire held ps.mu across the hook,
|
||||
// that would self-deadlock — the exact failure TrueAgent recorded for its own
|
||||
// parallel stop path.
|
||||
func TestFireDoesNotHoldRegistryLockWhileRunningHooks(t *testing.T) {
|
||||
reentrant := `
|
||||
local p = { name = "aaa" }
|
||||
p.state = {}
|
||||
p.hooks = { request_end = "reenter" }
|
||||
function p.reenter(payload)
|
||||
-- Reading the registry from inside a hook is the read half of the same lock.
|
||||
local _ = #payload
|
||||
return nil
|
||||
end
|
||||
return p
|
||||
`
|
||||
vm := NewVM(filepath.Join(t.TempDir(), "adapters"))
|
||||
if err := vm.Start(); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer vm.Stop()
|
||||
ps := NewPlugins(vm, filepath.Join(t.TempDir(), "plugins"))
|
||||
defer ps.Close()
|
||||
ps.DisableStatePersistence()
|
||||
if err := ps.LoadSource("aaa", reentrant); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := ps.LoadSource("bbb", counterPlugin("bbb", "")); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
done := make(chan struct{})
|
||||
go func() {
|
||||
defer close(done)
|
||||
ps.Fire(StageRequestEnd, map[string]interface{}{"model": "m"})
|
||||
// Touch the registry the way an admin request would, right after.
|
||||
_ = ps.Count()
|
||||
_ = ps.List()
|
||||
}()
|
||||
select {
|
||||
case <-done:
|
||||
case <-time.After(10 * time.Second):
|
||||
t.Fatal("Fire deadlocked against the registry lock")
|
||||
}
|
||||
}
|
||||
|
||||
// TestBillingSurvivesParallelFire runs the REAL bundled plugin next to another
|
||||
// one. A synthetic counter cannot catch a mismatch between the documented
|
||||
// payload shape and what the plugin actually reads.
|
||||
func TestBillingSurvivesParallelFire(t *testing.T) {
|
||||
dir := t.TempDir()
|
||||
pdir := filepath.Join(dir, "plugins")
|
||||
os.MkdirAll(pdir, 0o755)
|
||||
src, err := os.ReadFile("plugins/billing.lua")
|
||||
if err != nil {
|
||||
t.Fatalf("read billing.lua: %v", err)
|
||||
}
|
||||
if err := os.WriteFile(filepath.Join(pdir, "billing.lua"), src, 0o644); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := os.WriteFile(filepath.Join(pdir, "observer.lua"),
|
||||
[]byte(`local p={name="observer"} p.state={n=0}
|
||||
p.hooks={request_end="see"}
|
||||
function p.see(payload) p.state.n=p.state.n+1 return nil end
|
||||
return p`), 0o644); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
vm := NewVM(filepath.Join(dir, "adapters"))
|
||||
if err := vm.Start(); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer vm.Stop()
|
||||
ps := NewPlugins(vm, pdir)
|
||||
ps.save.setWakeDelayForTest(time.Millisecond)
|
||||
defer ps.Close()
|
||||
ps.DisableStatePersistence()
|
||||
if err := ps.LoadDir(); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if ps.Count() != 2 {
|
||||
t.Fatalf("loaded %d plugins, want 2", ps.Count())
|
||||
}
|
||||
|
||||
payload := map[string]interface{}{
|
||||
"model": "deepseek-v4.1-flash", "source": "commandcode", "key": "k", "ok": true,
|
||||
"usage": map[string]interface{}{
|
||||
"prompt_tokens": float64(1000), "completion_tokens": float64(100),
|
||||
"cache_hit_tokens": float64(0),
|
||||
},
|
||||
}
|
||||
ps.Fire(StageRequestEnd, payload)
|
||||
|
||||
st := ps.State("billing")
|
||||
if st == nil {
|
||||
t.Fatal("billing produced no state")
|
||||
}
|
||||
m := st.(map[string]interface{})
|
||||
total, ok := m["total"].(map[string]interface{})
|
||||
if !ok {
|
||||
t.Fatalf("billing.total is %T", m["total"])
|
||||
}
|
||||
if total["requests"] != float64(1) {
|
||||
t.Errorf("billing counted %v requests, want 1 — the plugin was starved by the parallel path",
|
||||
total["requests"])
|
||||
}
|
||||
if n := ps.State("observer").(map[string]interface{})["n"]; n != float64(1) {
|
||||
t.Errorf("observer counted %v, want 1", n)
|
||||
}
|
||||
if len(ps.HookErrors()) != 0 {
|
||||
t.Errorf("hook errors under parallel Fire: %v", ps.HookErrors())
|
||||
}
|
||||
}
|
||||
76
internal/lua/hook_guard_test.go
Normal file
76
internal/lua/hook_guard_test.go
Normal file
@ -0,0 +1,76 @@
|
||||
package lua
|
||||
|
||||
import (
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
"testing"
|
||||
)
|
||||
|
||||
// A plugin hook that RAISES must not take the process down. This is the
|
||||
// production outage: a Lua error anywhere in a hook made golua's callEx call
|
||||
// L.StackTrace(), which calls lua_getinfo and SIGSEGVs on a deep enough stack —
|
||||
// a C-level signal Go cannot recover from, so one buggy plugin killed the whole
|
||||
// gateway and took every in-flight request with it.
|
||||
//
|
||||
// The fix routes hook calls through a Lua-side pcall guard, so the error comes
|
||||
// back as an ordinary return value. This test fires hooks that raise on purpose
|
||||
// and asserts three things: the process survives, the failure is RECORDED, and
|
||||
// a well-behaved plugin on the same state keeps working afterwards (the error
|
||||
// must not poison the Lua state).
|
||||
func TestHookThatRaisesDoesNotCrashTheProcess(t *testing.T) {
|
||||
ps, pdir := billingVM(t)
|
||||
boom := `
|
||||
local plugin = {}
|
||||
plugin.name = "boom"
|
||||
plugin.version = "0.1"
|
||||
function plugin.request_end(payload)
|
||||
-- Raise on a table index, the exact shape of the billing bug that caused the
|
||||
-- outage. Deliberately NOT a syntax error: this must load fine and fail only
|
||||
-- when invoked.
|
||||
local x = nil
|
||||
return x.field
|
||||
end
|
||||
return plugin`
|
||||
if err := os.WriteFile(filepath.Join(pdir, "boom.lua"), []byte(boom), 0644); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := ps.LoadSource("boom", boom); err != nil {
|
||||
t.Fatalf("load boom: %v", err)
|
||||
}
|
||||
|
||||
payload := map[string]interface{}{
|
||||
"model": "m", "source": "s", "ok": true,
|
||||
"prompt_tokens": 100, "completion_tokens": 10, "time": 1750000000000,
|
||||
}
|
||||
// Fire many times: a single call could pass by luck, but if the error ever
|
||||
// escapes into golua's C path the process dies and this test never returns.
|
||||
for i := 0; i < 50; i++ {
|
||||
ps.Fire(StageRequestEnd, payload)
|
||||
}
|
||||
// Reaching this line at all is the primary assertion.
|
||||
|
||||
errs := ps.HookErrors()
|
||||
end, ok := errs[string(StageRequestEnd)]
|
||||
if !ok {
|
||||
t.Fatal("a raising hook left no record — failures must be observable, not swallowed")
|
||||
}
|
||||
if end["count"] == nil || end["count"].(int) == 0 {
|
||||
t.Error("hook error count is zero despite 50 raising calls")
|
||||
}
|
||||
msg, _ := end["last_error"].(string)
|
||||
if !strings.Contains(msg, "boom") {
|
||||
t.Errorf("last_error does not name the offending plugin: %q", msg)
|
||||
}
|
||||
|
||||
// The billing plugin shares the same Plugins registry and must still work:
|
||||
// one broken plugin may not disable the others.
|
||||
st, _ := ps.State("billing").(map[string]interface{})
|
||||
if st == nil || st["total"] == nil {
|
||||
t.Fatalf("the healthy plugin stopped working after another plugin raised")
|
||||
}
|
||||
tot, _ := st["total"].(map[string]interface{})
|
||||
if tot == nil || tot["requests"] == nil || tot["requests"].(float64) == 0 {
|
||||
t.Errorf("billing recorded no requests after the raising plugin ran: %v", st["total"])
|
||||
}
|
||||
}
|
||||
283
internal/lua/persist_test.go
Normal file
283
internal/lua/persist_test.go
Normal file
@ -0,0 +1,283 @@
|
||||
package lua
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"testing"
|
||||
"time"
|
||||
)
|
||||
|
||||
// These cover plugin state persistence: a plugin's totals live in the Lua VM,
|
||||
// which dies with the process. Measured on the real gateway before this change:
|
||||
// 218,241 requests and 26.2M prompt tokens gone after one systemctl restart.
|
||||
|
||||
func persistTestPlugin() string {
|
||||
return `
|
||||
local plugin = { name = "counter", version = "1.0" }
|
||||
plugin.state = { n = 0, prices_seen = 0 }
|
||||
plugin.prices = { rate = 1 }
|
||||
plugin.hooks = { request_end = "bump" }
|
||||
function plugin.bump(payload)
|
||||
plugin.state.n = plugin.state.n + 1
|
||||
return nil
|
||||
end
|
||||
return plugin
|
||||
`
|
||||
}
|
||||
|
||||
// newPersistVM wires a plugin registry on a fresh dir, mirroring newPluginVM.
|
||||
func newPersistVM(t *testing.T) (*Plugins, string) {
|
||||
t.Helper()
|
||||
dir := t.TempDir()
|
||||
vm := NewVM(filepath.Join(dir, "adapters"))
|
||||
if err := vm.Start(); err != nil {
|
||||
t.Fatalf("vm start: %v", err)
|
||||
}
|
||||
t.Cleanup(vm.Stop)
|
||||
pdir := filepath.Join(dir, "plugins")
|
||||
ps := NewPlugins(vm, pdir)
|
||||
// Shorten the saver debounce: Close() waits for the saver goroutine, so
|
||||
// with the production 2s interval every test would pay 2s on teardown.
|
||||
ps.save.setWakeDelayForTest(2 * time.Millisecond)
|
||||
t.Cleanup(ps.Close)
|
||||
return ps, pdir
|
||||
}
|
||||
|
||||
func loadCounter(t *testing.T, ps *Plugins) {
|
||||
t.Helper()
|
||||
if err := ps.LoadSource("counter", persistTestPlugin()); err != nil {
|
||||
t.Fatalf("load: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func stateN(t *testing.T, ps *Plugins) float64 {
|
||||
t.Helper()
|
||||
st := ps.State("counter")
|
||||
m, ok := st.(map[string]interface{})
|
||||
if !ok {
|
||||
t.Fatalf("state is %T, want map", st)
|
||||
}
|
||||
n, ok := m["n"].(float64)
|
||||
if !ok {
|
||||
t.Fatalf("state.n is %T, want float64", m["n"])
|
||||
}
|
||||
return n
|
||||
}
|
||||
|
||||
// TestStateSurvivesRestart is the defect itself: a rebuilt registry over the
|
||||
// same plugin dir must come up with the previous totals, not at zero.
|
||||
func TestStateSurvivesRestart(t *testing.T) {
|
||||
ps, pdir := newPersistVM(t)
|
||||
loadCounter(t, ps)
|
||||
for i := 0; i < 25; i++ {
|
||||
ps.Fire(StageRequestEnd, map[string]interface{}{"model": "m", "source": "s"})
|
||||
}
|
||||
if got := stateN(t, ps); got != 25 {
|
||||
t.Fatalf("in-process total = %v, want 25", got)
|
||||
}
|
||||
// Force the write the saver would do, so the test does not depend on timing.
|
||||
ps.save.flush()
|
||||
ps.Close()
|
||||
|
||||
// A brand-new registry over the same dir: this is the restart.
|
||||
vm := NewVM(filepath.Join(t.TempDir(), "adapters"))
|
||||
if err := vm.Start(); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer vm.Stop()
|
||||
ps2 := NewPlugins(vm, pdir)
|
||||
defer ps2.Close()
|
||||
loadCounter(t, ps2)
|
||||
|
||||
if got := stateN(t, ps2); got != 25 {
|
||||
t.Fatalf("★ total after restart = %v, want 25 — this is the 'restart loses the books' bug", got)
|
||||
}
|
||||
}
|
||||
|
||||
// TestPricesAndStateAreSeparate: restoring configuration over history (or the
|
||||
// reverse) would either erase the totals or resurrect stale prices.
|
||||
func TestPricesAndStateAreSeparate(t *testing.T) {
|
||||
ps, pdir := newPersistVM(t)
|
||||
loadCounter(t, ps)
|
||||
for i := 0; i < 7; i++ {
|
||||
ps.Fire(StageRequestEnd, map[string]interface{}{"model": "m", "source": "s"})
|
||||
}
|
||||
ps.save.flush()
|
||||
ps.Close()
|
||||
|
||||
b, err := os.ReadFile(ps.stateFile("counter"))
|
||||
if err != nil {
|
||||
t.Fatalf("no state file: %v", err)
|
||||
}
|
||||
var saved struct {
|
||||
State map[string]interface{} `json:"state"`
|
||||
Prices map[string]interface{} `json:"prices"`
|
||||
}
|
||||
if err := json.Unmarshal(b, &saved); err != nil {
|
||||
t.Fatalf("state file is not valid JSON: %v", err)
|
||||
}
|
||||
if saved.Prices == nil || saved.Prices["rate"] != float64(1) {
|
||||
t.Errorf("prices were not persisted separately: %v", saved.Prices)
|
||||
}
|
||||
if saved.State["n"] != float64(7) {
|
||||
t.Errorf("state.n = %v, want 7", saved.State["n"])
|
||||
}
|
||||
_ = pdir
|
||||
}
|
||||
|
||||
// TestPricesOnlyUpdatePersists guards a real hole: SetState returns early for a
|
||||
// prices-only payload (correctly leaving state alone), and the first version
|
||||
// returned before the persistence write — so a reprice was durable in memory
|
||||
// only and a restart silently reverted to the old prices.
|
||||
func TestPricesOnlyUpdatePersists(t *testing.T) {
|
||||
ps, _ := newPersistVM(t)
|
||||
loadCounter(t, ps)
|
||||
for i := 0; i < 5; i++ {
|
||||
ps.Fire(StageRequestEnd, map[string]interface{}{"model": "m", "source": "s"})
|
||||
}
|
||||
ps.save.flush()
|
||||
|
||||
err := ps.SetState("counter", map[string]interface{}{
|
||||
"prices": map[string]interface{}{"rate": 42},
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("SetState: %v", err)
|
||||
}
|
||||
if got := stateN(t, ps); got != 5 {
|
||||
t.Errorf("a prices-only payload must not touch state: n = %v, want 5", got)
|
||||
}
|
||||
|
||||
b, err := os.ReadFile(ps.stateFile("counter"))
|
||||
if err != nil {
|
||||
t.Fatalf("no state file after a prices-only update: %v", err)
|
||||
}
|
||||
var saved struct {
|
||||
State map[string]interface{} `json:"state"`
|
||||
Prices map[string]interface{} `json:"prices"`
|
||||
}
|
||||
if err := json.Unmarshal(b, &saved); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if saved.Prices["rate"] != float64(42) {
|
||||
t.Errorf("★ new price not persisted: %v — a restart would revert to the old price", saved.Prices)
|
||||
}
|
||||
if saved.State["n"] != float64(5) {
|
||||
t.Errorf("state was clobbered by a price change: %v", saved.State["n"])
|
||||
}
|
||||
}
|
||||
|
||||
// TestHookReturnValueSurvivesSnapshot guards the ordering bug: snapshotting
|
||||
// plugin.state resets the Lua stack, so doing it BEFORE reading the hook's
|
||||
// return value silently turned every opinionated plugin into a silent one —
|
||||
// breaking the documented "return a table to merge into payload" contract with
|
||||
// no error anywhere.
|
||||
func TestHookReturnValueSurvivesSnapshot(t *testing.T) {
|
||||
ps, _ := newPersistVM(t)
|
||||
code := `
|
||||
local plugin = { name = "opinionated" }
|
||||
plugin.state = { n = 0 }
|
||||
plugin.hooks = { request_end = "tag" }
|
||||
function plugin.tag(payload)
|
||||
plugin.state.n = plugin.state.n + 1
|
||||
return { cost_usd = 1.25, verdict = "billed" }
|
||||
end
|
||||
return plugin
|
||||
`
|
||||
if err := ps.LoadSource("opinionated", code); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
out := ps.Fire(StageRequestEnd, map[string]interface{}{"model": "m", "source": "s"})
|
||||
if out["cost_usd"] != 1.25 {
|
||||
t.Errorf("★ hook return value lost: payload = %v", out)
|
||||
}
|
||||
if out["verdict"] != "billed" {
|
||||
t.Errorf("merged field lost: %v", out)
|
||||
}
|
||||
ps.save.flush()
|
||||
st := ps.State("opinionated").(map[string]interface{})
|
||||
if st["n"] != float64(1) {
|
||||
t.Errorf("state.n = %v, want 1 (the hook still ran)", st["n"])
|
||||
}
|
||||
}
|
||||
|
||||
// TestCorruptStateFileIsNotFatal: a truncated write must degrade to compiled-in
|
||||
// defaults, never to a gateway that refuses to start.
|
||||
func TestCorruptStateFileIsNotFatal(t *testing.T) {
|
||||
ps, _ := newPersistVM(t)
|
||||
loadCounter(t, ps)
|
||||
for i := 0; i < 3; i++ {
|
||||
ps.Fire(StageRequestEnd, map[string]interface{}{"model": "m", "source": "s"})
|
||||
}
|
||||
ps.save.flush()
|
||||
ps.Close()
|
||||
|
||||
if err := os.WriteFile(ps.stateFile("counter"), []byte("{not json"), 0o644); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
vm := NewVM(filepath.Join(t.TempDir(), "adapters"))
|
||||
if err := vm.Start(); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer vm.Stop()
|
||||
ps2 := NewPlugins(vm, ps.dir)
|
||||
defer ps2.Close()
|
||||
var warned bool
|
||||
ps2.logf = func(string, ...interface{}) { warned = true }
|
||||
loadCounter(t, ps2)
|
||||
|
||||
if got := stateN(t, ps2); got != 0 {
|
||||
t.Errorf("with a corrupt file the plugin must fall back to its defaults, got n = %v", got)
|
||||
}
|
||||
if !warned {
|
||||
t.Error("a corrupt state file must warn the operator, not fail silently")
|
||||
}
|
||||
// And it must still forward.
|
||||
ps2.Fire(StageRequestEnd, map[string]interface{}{"model": "m", "source": "s"})
|
||||
}
|
||||
|
||||
// TestFlushIsCoalesced: N mutations must not become N writes. A synchronous
|
||||
// per-request write would put a file write on the hot path.
|
||||
func TestFlushIsCoalesced(t *testing.T) {
|
||||
ps, _ := newPersistVM(t)
|
||||
loadCounter(t, ps)
|
||||
for i := 0; i < 500; i++ {
|
||||
ps.Fire(StageRequestEnd, map[string]interface{}{"model": "m", "source": "s"})
|
||||
}
|
||||
ps.save.flush()
|
||||
// 500 mutations must collapse to a handful of writes, not 500. The bound is
|
||||
// loose on purpose: the saver also flushes on a timer, so how many writes
|
||||
// happen during the loop depends on how long 500 hook calls take. What must
|
||||
// never happen is one write per mutation.
|
||||
if w := ps.save.writes.Load(); w > 5 {
|
||||
t.Errorf("500 hook calls produced %d file writes; the saver must coalesce", w)
|
||||
}
|
||||
if got := stateN(t, ps); got != 500 {
|
||||
t.Fatalf("n = %v, want 500", got)
|
||||
}
|
||||
}
|
||||
|
||||
// TestCloseFlushesTail: the shutdown path must not lose the last interval,
|
||||
// which would reintroduce the same defect in a smaller window.
|
||||
func TestCloseFlushesTail(t *testing.T) {
|
||||
ps, _ := newPersistVM(t)
|
||||
loadCounter(t, ps)
|
||||
ps.Fire(StageRequestEnd, map[string]interface{}{"model": "m", "source": "s"})
|
||||
ps.Close()
|
||||
// Close() already waits for the final flush; no sleep needed.
|
||||
|
||||
b, err := os.ReadFile(ps.stateFile("counter"))
|
||||
if err != nil {
|
||||
t.Fatalf("★ Close() did not flush: %v — a shutdown loses the tail", err)
|
||||
}
|
||||
var saved struct {
|
||||
State map[string]interface{} `json:"state"`
|
||||
}
|
||||
if err := json.Unmarshal(b, &saved); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if saved.State["n"] != float64(1) {
|
||||
t.Errorf("flushed n = %v, want 1", saved.State["n"])
|
||||
}
|
||||
}
|
||||
180
internal/lua/plugin_access.go
Normal file
180
internal/lua/plugin_access.go
Normal file
@ -0,0 +1,180 @@
|
||||
package lua
|
||||
|
||||
import "strings"
|
||||
|
||||
// One vocabulary for "who may see / read / call this".
|
||||
//
|
||||
// Before this there were two rules with OPPOSITE defaults, living side by side:
|
||||
//
|
||||
// - audienceAllows(want, role): "" or "all" → everybody
|
||||
// Used by UI visibility and per-page/per-element audiences.
|
||||
// - HTTPHandler.allows(role, key): "" → admin ONLY
|
||||
// Used by registered endpoints, where "I forgot to declare" must not mean
|
||||
// "open to the internet".
|
||||
//
|
||||
// Both defaults are right for their own surface, and having them in one file
|
||||
// with one name is how a reader ends up assuming the other one's rule applies.
|
||||
// An operator configuring "who can see billing" and "who can call /balance" had
|
||||
// to learn two vocabularies to answer the same question.
|
||||
//
|
||||
// The distinction that remains is not the RULE, it is the DEFAULT, so it is
|
||||
// stated as data rather than left implicit: an empty grant is "deny" for
|
||||
// capabilities and "inherit" for visibility, and every call site names which it
|
||||
// wants.
|
||||
|
||||
// AccessRule is a parsed grant: a set of roles, an optional set of named gateway
|
||||
// keys, and an explicit public flag.
|
||||
//
|
||||
// Keys REPLACE roles rather than adding to them. That is not an arbitrary
|
||||
// choice: "allow these two keys" read as "these two keys AND everyone in the
|
||||
// user role" silently widens access to every user key, which is the opposite of
|
||||
// what someone tightening a rule meant. `keys = { "@alice" }` therefore means
|
||||
// exactly alice.
|
||||
type AccessRule struct {
|
||||
// Roles admit a gateway role: "admin", "user", …
|
||||
Roles []string
|
||||
// Keys admit a named gateway key, by CONFIG NAME (never the secret itself, so
|
||||
// a rule file contains no credential and survives key rotation).
|
||||
Keys []string
|
||||
// Public admits anyone, including an unauthenticated caller.
|
||||
Public bool
|
||||
}
|
||||
|
||||
// AccessAll is the grant that admits everyone.
|
||||
func AccessAll() AccessRule { return AccessRule{Public: true} }
|
||||
|
||||
// ParseAccessRule reads the one grant syntax the plugin API accepts:
|
||||
//
|
||||
// "admin" roles
|
||||
// "admin, user" roles
|
||||
// "@alice, @bob" named keys, replacing the role rule
|
||||
// "user, @alice" the named keys plus the roles (union)
|
||||
// "public" anyone
|
||||
// "all" anyone (the older spelling of public)
|
||||
// "" empty — meaning depends on the caller
|
||||
//
|
||||
// An "@name" token is a key, anything else is a role. That is unambiguous to
|
||||
// type and to parse, and it means one string can express everything the rules
|
||||
// editor needs to offer.
|
||||
//
|
||||
// "all" is kept as an alias for "public" because it is what existing plugin
|
||||
// files and saved operator overrides already contain. Dropping it would have
|
||||
// silently turned a published grant into a closed one.
|
||||
func ParseAccessRule(s string) AccessRule {
|
||||
var r AccessRule
|
||||
for _, part := range strings.Split(s, ",") {
|
||||
part = strings.TrimSpace(part)
|
||||
if part == "" {
|
||||
continue
|
||||
}
|
||||
if strings.EqualFold(part, "public") || strings.EqualFold(part, "all") {
|
||||
r.Public = true
|
||||
continue
|
||||
}
|
||||
if strings.HasPrefix(part, "@") {
|
||||
if k := strings.TrimSpace(strings.TrimPrefix(part, "@")); k != "" {
|
||||
r.Keys = append(r.Keys, k)
|
||||
}
|
||||
continue
|
||||
}
|
||||
r.Roles = append(r.Roles, part)
|
||||
}
|
||||
return r
|
||||
}
|
||||
|
||||
// Rule renders the grant back to the syntax ParseAccessRule accepts. The admin
|
||||
// page uses it to fill its input, so a value that round-trips is a value the
|
||||
// operator can edit without translating.
|
||||
func (r AccessRule) Rule() string {
|
||||
var parts []string
|
||||
for _, k := range r.Keys {
|
||||
parts = append(parts, "@"+k)
|
||||
}
|
||||
for _, role := range r.Roles {
|
||||
parts = append(parts, role)
|
||||
}
|
||||
if r.Public {
|
||||
parts = append(parts, "public")
|
||||
}
|
||||
return strings.Join(parts, ", ")
|
||||
}
|
||||
|
||||
// Empty reports whether the grant names nobody.
|
||||
func (r AccessRule) Empty() bool {
|
||||
return len(r.Roles) == 0 && len(r.Keys) == 0 && !r.Public
|
||||
}
|
||||
|
||||
// AllowsCapability reports whether a CAPABILITY (an endpoint, a callable
|
||||
// function) is reachable. Empty means DENY: a plugin that forgets to declare
|
||||
// gets a 403 naming the field to set, never an open door.
|
||||
//
|
||||
// Admin is always admitted. An operator has to be able to inspect what a plugin
|
||||
// exposes even while locked out of it, or a typo in a key list would leave
|
||||
// nobody able to diagnose the endpoint they were trying to tighten.
|
||||
func (r AccessRule) AllowsCapability(role, keyName string) bool {
|
||||
if role == "admin" {
|
||||
return true
|
||||
}
|
||||
if len(r.Keys) > 0 {
|
||||
// A non-empty key list REPLACES the role rule (see the type comment).
|
||||
// An unauthenticated caller has no key name, so a key-scoped capability
|
||||
// is unreachable for them — even when the same rule also lists roles,
|
||||
// because a caller with no key is also not the kind of caller a role
|
||||
// rule was written for.
|
||||
if keyName == "" {
|
||||
return false
|
||||
}
|
||||
for _, k := range r.Keys {
|
||||
if strings.EqualFold(k, keyName) {
|
||||
return true
|
||||
}
|
||||
}
|
||||
// The key list replaced the role rule, so a miss is a deny even though
|
||||
// roles may also be listed.
|
||||
return false
|
||||
}
|
||||
if r.Public {
|
||||
return true
|
||||
}
|
||||
for _, want := range r.Roles {
|
||||
if strings.EqualFold(want, role) {
|
||||
return true
|
||||
}
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
// AllowsVisibility reports whether a VIEW (a page, a UI element) is rendered.
|
||||
// Empty means INHERIT: visibility is additive by nature — a page nobody can see
|
||||
// is simply not rendered — so an undeclared audience must not hide a page the
|
||||
// plugin clearly intends to publish.
|
||||
func (r AccessRule) AllowsVisibility(role string) bool {
|
||||
if r.Empty() || r.Public {
|
||||
return true
|
||||
}
|
||||
if role == "admin" && len(r.Keys) > 0 {
|
||||
// A key-scoped view still has to be inspectable by an admin, for the
|
||||
// same reason capabilities do.
|
||||
return true
|
||||
}
|
||||
for _, want := range r.Roles {
|
||||
if strings.EqualFold(want, role) {
|
||||
return true
|
||||
}
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
// accessRuleFor folds the separate declaration fields of an endpoint into one
|
||||
// grant, so the rules below live in one place instead of beside each field.
|
||||
//
|
||||
// It is a plain fold, deliberately. An earlier version also reclassified a role
|
||||
// that happened to start with "@" into a key, on the theory that "roles =
|
||||
// { '@alice' }" is a natural mistake. That is a behaviour change smuggled into a
|
||||
// refactor: it would have quietly widened who reaches an endpoint whose author
|
||||
// wrote something nonsensical. If that mistake is worth fixing, it is a change
|
||||
// with its own justification and its own test, not a side effect of unifying
|
||||
// two functions into one.
|
||||
func accessRuleFor(roles, keys []string, public bool) AccessRule {
|
||||
return AccessRule{Roles: roles, Keys: keys, Public: public}
|
||||
}
|
||||
115
internal/lua/plugin_access_persist.go
Normal file
115
internal/lua/plugin_access_persist.go
Normal file
@ -0,0 +1,115 @@
|
||||
package lua
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"os"
|
||||
"sort"
|
||||
)
|
||||
|
||||
// Access overrides are OPERATOR policy, not plugin data. They live in their own
|
||||
// file rather than in each plugin's state because:
|
||||
//
|
||||
// - they describe the deployment, not the plugin, so a plugin upgrade that
|
||||
// replaces its .state.json must not silently widen access again;
|
||||
// - they must survive a restart. An in-memory map made "restrict billing to
|
||||
// admin" look like it worked right up to the next deploy, at which point the
|
||||
// restriction evaporated with nothing in the logs to say why;
|
||||
// - the plugin_dir is the one place an operator already looks when asking
|
||||
// "what does this plugin actually get to see", so the answer lives beside it.
|
||||
//
|
||||
// The file is a flat map of "name" (whole plugin) or "name#path" (one endpoint)
|
||||
// to an access spec, which is the same shape the in-memory store uses, so load
|
||||
// and save cannot disagree about what a key means.
|
||||
|
||||
const accessFileName = ".access.json"
|
||||
|
||||
func (ps *Plugins) accessFile() string {
|
||||
if ps.dir == "" {
|
||||
return ""
|
||||
}
|
||||
return ps.dir + string(os.PathSeparator) + accessFileName
|
||||
}
|
||||
|
||||
// LoadAccessOverrides reads the persisted overrides into memory. A missing file
|
||||
// is not an error: no operator has restricted anything yet, which is the
|
||||
// default state of a fresh install.
|
||||
//
|
||||
// A corrupt file is also not fatal. Losing a restriction is bad, but refusing to
|
||||
// start the gateway is worse, and an empty map plus a logged error leaves the
|
||||
// operator with something to fix.
|
||||
func (ps *Plugins) LoadAccessOverrides() error {
|
||||
path := ps.accessFile()
|
||||
if path == "" {
|
||||
return nil
|
||||
}
|
||||
b, err := os.ReadFile(path)
|
||||
if os.IsNotExist(err) {
|
||||
return nil
|
||||
}
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
m := map[string]string{}
|
||||
if err := json.Unmarshal(b, &m); err != nil {
|
||||
return err
|
||||
}
|
||||
ps.mu.Lock()
|
||||
defer ps.mu.Unlock()
|
||||
if ps.visibility == nil {
|
||||
ps.visibility = visibilityOverrides{}
|
||||
}
|
||||
for k, v := range m {
|
||||
ps.visibility[k] = v
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// saveAccessOverridesLocked flushes the whole override map. It is called on
|
||||
// every change, and unlike plugin state the map is tiny and changes only when an
|
||||
// operator clicks Save — so a synchronous write costs nothing that matters, and
|
||||
// it means the change is durable the moment the API returns 200.
|
||||
func (ps *Plugins) saveAccessOverridesLocked() {
|
||||
path := ps.accessFile()
|
||||
if path == "" {
|
||||
return
|
||||
}
|
||||
out := make(map[string]string, len(ps.visibility))
|
||||
for k, v := range ps.visibility {
|
||||
if v != "" {
|
||||
out[k] = v
|
||||
}
|
||||
}
|
||||
if len(out) == 0 {
|
||||
// Leave no empty file behind: its absence is the honest encoding of
|
||||
// "nothing is overridden".
|
||||
_ = os.Remove(path)
|
||||
return
|
||||
}
|
||||
writeJSONAtomic(path, out)
|
||||
}
|
||||
|
||||
// AccessOverrides returns a copy of every override in force, for the UI.
|
||||
func (ps *Plugins) accessOverridesSnapshot() map[string]string {
|
||||
ps.mu.RLock()
|
||||
defer ps.mu.RUnlock()
|
||||
if len(ps.visibility) == 0 {
|
||||
return nil
|
||||
}
|
||||
out := make(map[string]string, len(ps.visibility))
|
||||
for k, v := range ps.visibility {
|
||||
out[k] = v
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
// sortedOverrideKeys exists only so a debug dump is stable; map iteration order
|
||||
// would otherwise make two identical states look different.
|
||||
func (ps *Plugins) sortedOverrideKeys() []string {
|
||||
m := ps.accessOverridesSnapshot()
|
||||
keys := make([]string, 0, len(m))
|
||||
for k := range m {
|
||||
keys = append(keys, k)
|
||||
}
|
||||
sort.Strings(keys)
|
||||
return keys
|
||||
}
|
||||
110
internal/lua/plugin_access_persist_test.go
Normal file
110
internal/lua/plugin_access_persist_test.go
Normal file
@ -0,0 +1,110 @@
|
||||
package lua
|
||||
|
||||
import (
|
||||
"os"
|
||||
"path/filepath"
|
||||
"testing"
|
||||
)
|
||||
|
||||
// TestAccessOverrideSurvivesReload is the test that matters for an access
|
||||
// restriction. An in-memory map makes "billing is admin-only" look like it
|
||||
// worked right up to the next restart, and then the restriction is simply gone
|
||||
// with nothing in the logs to explain it. Both the plugin-level audience and
|
||||
// the per-endpoint spec must come back.
|
||||
func TestAccessOverrideSurvivesReload(t *testing.T) {
|
||||
dir := t.TempDir()
|
||||
src := `local P = {}
|
||||
P.name = "api"
|
||||
P.version = "1.0.0"
|
||||
P.http = {
|
||||
{ path = "health", handler = "health", roles = { "admin" } },
|
||||
}
|
||||
function P.health(req) return { json = { ok = true } } end
|
||||
return P
|
||||
`
|
||||
first := NewPlugins(NewVM(dir), dir)
|
||||
if err := first.LoadSource("api", src); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
first.SetPluginVisibility("api", "admin, user")
|
||||
if err := first.SetEndpointAccess("api", "health", "@homeagent"); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if _, err := os.Stat(filepath.Join(dir, accessFileName)); err != nil {
|
||||
t.Fatalf("override was not written: %v", err)
|
||||
}
|
||||
|
||||
// A fresh Plugins over the same dir is what a restart looks like.
|
||||
second := NewPlugins(NewVM(dir), dir)
|
||||
// The plugin source is on disk too, so this second instance sees both the
|
||||
// plugin and the overrides — the same pair a restart sees. Loading only the
|
||||
// overrides would leave no endpoints to walk, and the assertion below would
|
||||
// pass vacuously.
|
||||
if err := second.LoadSource("api", src); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := second.LoadAccessOverrides(); err != nil {
|
||||
t.Fatalf("reload overrides: %v", err)
|
||||
}
|
||||
if got := second.VisibilityOverride("api"); got != "admin, user" {
|
||||
t.Errorf("plugin visibility lost across reload: %q", got)
|
||||
}
|
||||
found := false
|
||||
for _, e := range second.EndpointAccess() {
|
||||
if e.Path == "health" && len(e.OverrideKeys) == 1 && e.OverrideKeys[0] == "homeagent" {
|
||||
found = true
|
||||
}
|
||||
}
|
||||
if !found {
|
||||
t.Error("endpoint override lost across reload")
|
||||
}
|
||||
}
|
||||
|
||||
// TestAccessOverrideClearedRemovesFile keeps the "nothing overridden" state
|
||||
// distinguishable from "overridden with something".
|
||||
func TestAccessOverrideClearedRemovesFile(t *testing.T) {
|
||||
dir := t.TempDir()
|
||||
ps := NewPlugins(NewVM(dir), dir)
|
||||
ps.SetPluginVisibility("api", "admin")
|
||||
path := filepath.Join(dir, accessFileName)
|
||||
if _, err := os.Stat(path); err != nil {
|
||||
t.Fatalf("expected an access file: %v", err)
|
||||
}
|
||||
ps.SetPluginVisibility("api", "")
|
||||
if _, err := os.Stat(path); !os.IsNotExist(err) {
|
||||
t.Errorf("clearing the last override should remove the file, got %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
// TestVisibilityOverridesExcludeEndpointKeys keeps endpoint policy out of the
|
||||
// plugin-audience listing. Both live in one map, and reporting "probe#hello"
|
||||
// under visibility is how the plugin page ends up offering to save an audience
|
||||
// of that name — which nothing reads as an audience.
|
||||
func TestVisibilityOverridesExcludeEndpointKeys(t *testing.T) {
|
||||
dir := t.TempDir()
|
||||
src := `local P = {}
|
||||
P.name = "api"
|
||||
P.version = "1.0.0"
|
||||
P.http = { { path = "health", handler = "health", roles = { "admin" } } }
|
||||
function P.health(req) return { json = { ok = true } } end
|
||||
return P
|
||||
`
|
||||
ps := NewPlugins(NewVM(dir), dir)
|
||||
if err := ps.LoadSource("api", src); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
ps.SetPluginVisibility("api", "admin, user")
|
||||
if err := ps.SetEndpointAccess("api", "health", "@homeagent"); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
got := ps.VisibilityOverrides()
|
||||
if got["api"] != "admin, user" {
|
||||
t.Errorf("plugin audience missing: %+v", got)
|
||||
}
|
||||
if _, leaked := got["api#health"]; leaked {
|
||||
t.Errorf("endpoint override leaked into the audience map: %+v", got)
|
||||
}
|
||||
if len(got) != 1 {
|
||||
t.Errorf("expected exactly the plugin audience, got %+v", got)
|
||||
}
|
||||
}
|
||||
152
internal/lua/plugin_access_test.go
Normal file
152
internal/lua/plugin_access_test.go
Normal file
@ -0,0 +1,152 @@
|
||||
package lua
|
||||
|
||||
import (
|
||||
"strings"
|
||||
"testing"
|
||||
)
|
||||
|
||||
// The two functions this file replaces had OPPOSITE defaults:
|
||||
//
|
||||
// audienceAllows: "" / "all" → everybody
|
||||
// HTTPHandler.allows: "" → admin ONLY
|
||||
//
|
||||
// Both are correct for their surface. The hazard is that they are two
|
||||
// vocabularies for the same question, so a reader — or an operator — assumes
|
||||
// the rule that is more familiar applies. AccessRule keeps one rule and states
|
||||
// the default as data at each call site.
|
||||
|
||||
// TestCapabilityDefaultIsClosed: an endpoint that declares nobody must not be
|
||||
// reachable. This is the property that makes "I forgot" fail safe.
|
||||
func TestCapabilityDefaultIsClosed(t *testing.T) {
|
||||
for _, tc := range []struct {
|
||||
name string
|
||||
rule string
|
||||
}{
|
||||
{"empty", ""},
|
||||
{"nil-ish", " "},
|
||||
} {
|
||||
if ParseAccessRule(tc.rule).AllowsCapability("user", "alice") {
|
||||
t.Errorf("%s: an undeclared capability is reachable by a user", tc.name)
|
||||
}
|
||||
if !ParseAccessRule(tc.rule).AllowsCapability("admin", "alice") {
|
||||
t.Errorf("%s: admin locked out of a capability they must be able to inspect", tc.name)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// TestVisibilityDefaultIsInherits: the opposite default, on purpose. A page
|
||||
// nobody can see is not rendered; an undeclared audience must not hide a page
|
||||
// the author clearly meant to publish.
|
||||
func TestVisibilityDefaultIsInherits(t *testing.T) {
|
||||
if !ParseAccessRule("").AllowsVisibility("user") {
|
||||
t.Error("an undeclared audience hides a page")
|
||||
}
|
||||
if ParseAccessRule("admin").AllowsVisibility("user") {
|
||||
t.Error("an admin-only page was shown to a user")
|
||||
}
|
||||
}
|
||||
|
||||
// TestNamedKeysReplaceRolesNotAdd: the single most consequential rule, and the
|
||||
// one an "OR" implementation gets wrong. "allow these two keys" must not also
|
||||
// admit every user key.
|
||||
func TestNamedKeysReplaceRolesNotAdd(t *testing.T) {
|
||||
r := ParseAccessRule("@alice, @bob")
|
||||
if !r.AllowsCapability("user", "alice") {
|
||||
t.Error("a named key was not admitted")
|
||||
}
|
||||
if r.AllowsCapability("user", "carol") {
|
||||
t.Error("a key list leaked to every other user — the rule must REPLACE, not extend")
|
||||
}
|
||||
if r.AllowsCapability("user", "") {
|
||||
t.Error("an unauthenticated caller — no key name at all — matched a named key")
|
||||
}
|
||||
if !r.AllowsCapability("admin", "") {
|
||||
t.Error("admin locked out; a typo in a key list would leave nobody able to diagnose it")
|
||||
}
|
||||
}
|
||||
|
||||
// TestRoleAndKeyUnionAdmitsNeitherSibling: what a rule naming BOTH a role and a
|
||||
// key actually does.
|
||||
//
|
||||
// Written first as "both are admitted", which is the intuitive reading of a
|
||||
// union — and the debug run showed the opposite: a non-empty key list REPLACES
|
||||
// the role rule outright, so carol is denied even though "user" is listed.
|
||||
//
|
||||
// That surprise is exactly why the rule is named, commented and now tested. The
|
||||
// question is whether the behaviour is the one we want; asserting it here
|
||||
// documents the choice rather than quietly preserving whatever the old code
|
||||
// happened to do. The behaviour matches the pre-existing HTTPHandler.allows, so
|
||||
// nothing deployed changes.
|
||||
func TestRoleAndKeyUnionAdmitsNeitherSibling(t *testing.T) {
|
||||
r := ParseAccessRule("user, @alice")
|
||||
if !r.AllowsCapability("user", "alice") {
|
||||
t.Error("the named key was not admitted")
|
||||
}
|
||||
if r.AllowsCapability("user", "carol") {
|
||||
t.Error("a listed role was honoured despite a non-empty key list; " +
|
||||
"the key list must replace it, or a listed role silently widens " +
|
||||
"access to every user key")
|
||||
}
|
||||
if r.AllowsCapability("user", "mallory") {
|
||||
t.Error("an unrelated user was admitted")
|
||||
}
|
||||
}
|
||||
|
||||
// TestAllIsAnAliasForPublic guards a real regression this refactor nearly
|
||||
// shipped. "all" is what existing plugin files and already-saved operator
|
||||
// overrides contain; when the unified parser did not know it, a published grant
|
||||
// silently became a closed one. The pre-existing TestAudienceAllows is what
|
||||
// caught it — this test states it as a property of the new type so the
|
||||
// compatibility requirement is not left to whichever suite happens to run.
|
||||
func TestAllIsAnAliasForPublic(t *testing.T) {
|
||||
r := ParseAccessRule("all")
|
||||
if !r.Public {
|
||||
t.Fatal(`"all" did not parse as public`)
|
||||
}
|
||||
if !r.AllowsCapability("user", "carol") {
|
||||
t.Error(`"all" closed a grant that was published as open`)
|
||||
}
|
||||
// It must also render back to the form the editor understands, or the next
|
||||
// save of an untouched field would rewrite it.
|
||||
if !strings.Contains(strings.ToLower(r.Rule()), "public") {
|
||||
t.Errorf("Rule() = %q, want it to render as public", r.Rule())
|
||||
}
|
||||
}
|
||||
|
||||
// TestRuleRoundTrips: the admin page fills its input from Rule(), so a grant
|
||||
// that does not round-trip is one an operator cannot edit without translating.
|
||||
func TestRuleRoundTrips(t *testing.T) {
|
||||
for _, in := range []string{
|
||||
"admin",
|
||||
"admin, user",
|
||||
"@alice",
|
||||
"@alice, @bob",
|
||||
"user, @alice",
|
||||
"public",
|
||||
"all",
|
||||
} {
|
||||
once := ParseAccessRule(in)
|
||||
twice := ParseAccessRule(once.Rule())
|
||||
if once.Rule() != twice.Rule() {
|
||||
t.Errorf("%q -> %q -> %q, not stable", in, once.Rule(), twice.Rule())
|
||||
}
|
||||
if len(twice.Roles) != len(once.Roles) || len(twice.Keys) != len(once.Keys) {
|
||||
t.Errorf("%q lost data across a round trip: %+v -> %+v", in, once, twice)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// TestRoleMatchingIsCaseAndSpaceInsensitive: grants arrive typed by humans into
|
||||
// a text input, and "Admin" not matching "admin" would lock an admin out of
|
||||
// their own plugin.
|
||||
func TestRoleMatchingIsCaseAndSpaceInsensitive(t *testing.T) {
|
||||
if !ParseAccessRule(" Admin , User ").AllowsCapability("admin", "x") {
|
||||
t.Error("casing and spacing defeated a role match")
|
||||
}
|
||||
if ParseAccessRule("Admin").AllowsCapability("user", "x") {
|
||||
t.Error("a user was admitted by a case variant")
|
||||
}
|
||||
if !ParseAccessRule(" @Alice ").AllowsCapability("user", "alice") {
|
||||
t.Error("casing defeated a named-key match")
|
||||
}
|
||||
}
|
||||
95
internal/lua/plugin_clock_test.go
Normal file
95
internal/lua/plugin_clock_test.go
Normal file
@ -0,0 +1,95 @@
|
||||
package lua
|
||||
|
||||
import (
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"testing"
|
||||
"time"
|
||||
)
|
||||
|
||||
// TestPluginClockDuringNetworkWait pins the reason the probe budget never
|
||||
// enforced itself.
|
||||
//
|
||||
// Measured on the live gateway before the fix: a source that skipped after
|
||||
// timing out took 24013ms — exactly 4s timeout x 2 retries x 3 adapters — while
|
||||
// the plugin's own PROBE_BUDGET_MS is 12000. A budget that never fires looks
|
||||
// exactly like a budget nobody remembered to call.
|
||||
//
|
||||
// The cause is os.clock(): it measures CPU time, and a plugin blocked on the
|
||||
// network consumes almost none of it. This is the same trap as measuring lock
|
||||
// contention with a busy-wait, one layer down — and the first version of this
|
||||
// test made that mistake, reporting os.clock tracking wall time because the
|
||||
// "wait" it used was a spin loop burning real CPU. It had to be re-measured
|
||||
// against an actual HTTP request.
|
||||
func TestPluginClockDuringNetworkWait(t *testing.T) {
|
||||
dir := t.TempDir()
|
||||
pluginDir := filepath.Join(dir, "plugins")
|
||||
if err := os.MkdirAll(pluginDir, 0o755); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
vm := NewVM(filepath.Join(dir, "adapters"))
|
||||
if err := vm.Start(); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
t.Cleanup(vm.Stop)
|
||||
ps := NewPlugins(vm, pluginDir)
|
||||
t.Cleanup(ps.Close)
|
||||
ps.SetOutboundHosts([]string{"127.0.0.1"})
|
||||
|
||||
code := `
|
||||
local P = {}
|
||||
P.name = "clocktest"
|
||||
P.version = "1.0.0"
|
||||
P.exports = { "measure" }
|
||||
P.measure = function(payload)
|
||||
local c0, w0 = os.clock(), os.time()
|
||||
-- The real thing: block on the network. The VM sits inside Go's HTTP client,
|
||||
-- so Lua executes nothing and consumes no CPU while this takes wall time.
|
||||
__llmsproxy_http({ url = payload.url, method = "GET", timeout_ms = 5000 })
|
||||
return {
|
||||
clock = os.clock() - c0,
|
||||
wall = os.time() - w0,
|
||||
}
|
||||
end
|
||||
return P
|
||||
`
|
||||
if err := os.WriteFile(filepath.Join(pluginDir, "clocktest.lua"), []byte(code), 0o644); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
time.Sleep(3 * time.Second)
|
||||
_, _ = w.Write([]byte(`{}`))
|
||||
}))
|
||||
defer srv.Close()
|
||||
|
||||
if err := ps.LoadDir(); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
p := ps.find("clocktest")
|
||||
if p == nil {
|
||||
t.Fatal("plugin not found")
|
||||
}
|
||||
if p.LoadError != "" {
|
||||
t.Fatalf("plugin failed to compile: %s", p.LoadError)
|
||||
}
|
||||
out, err := ps.Call("clocktest", "measure", map[string]interface{}{"url": srv.URL + "/slow"})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
m, _ := out.(map[string]interface{})
|
||||
clock, _ := m["clock"].(float64)
|
||||
wall, _ := m["wall"].(float64)
|
||||
t.Logf("3s blocked on HTTP: os.clock advanced %.3fs, os.time advanced %.0fs", clock, wall)
|
||||
|
||||
if wall < 3 {
|
||||
t.Fatalf("the request did not take 3 wall seconds (got %.0f)", wall)
|
||||
}
|
||||
if clock > wall/2 {
|
||||
t.Fatalf("os.clock tracked wall time (%.3f vs %.0f) — the budget is not the cause; "+
|
||||
"look elsewhere", clock, wall)
|
||||
}
|
||||
t.Logf("CONFIRMED: blocked on the network for %.0fs, os.clock advanced %.3fs. "+
|
||||
"A budget written in os.clock never fires while a plugin waits on I/O.", wall, clock)
|
||||
}
|
||||
244
internal/lua/plugin_decode.go
Normal file
244
internal/lua/plugin_decode.go
Normal file
@ -0,0 +1,244 @@
|
||||
package lua
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
|
||||
golua "github.com/aarzilli/golua/lua"
|
||||
"sort"
|
||||
"strings"
|
||||
)
|
||||
|
||||
// The declarative decoder.
|
||||
//
|
||||
// Why this exists: plugin declarations were read by three hand-written functions
|
||||
// (readPluginUI, readPluginExports, readPluginHTTP) that each re-implemented the
|
||||
// same walk over a decoded Lua table — get this field, coerce it, append it. The
|
||||
// field names lived in the .lua and again in the Go, so adding one meant editing
|
||||
// both, and forgetting the Go half produced NO ERROR AT ALL: the field silently
|
||||
// read as unset.
|
||||
//
|
||||
// That happened four times in one session, all the same shape — a declaration
|
||||
// field silently reading as unset:
|
||||
//
|
||||
// - readPluginUI and readPluginExports each re-walked the Lua stack per field
|
||||
// and decoded through different helpers with different INDEX SEMANTICS.
|
||||
// - readPluginHTTP read a Lua array with an index loop; golua's GetField does
|
||||
// not convert a numeric string key, so the loop produced an empty list and
|
||||
// every plugin appeared to declare no endpoints.
|
||||
// - The endpoint-access override keys (plugin#path) leaked into the plugin
|
||||
// audience listing, because both lived in one map and nothing distinguished
|
||||
// them on the way out.
|
||||
//
|
||||
// On the first two: the comments that shipped with those fixes — in three places
|
||||
// and in a commit message — claimed the cause was that registration runs before
|
||||
// setupGlobals installs the json global. That was never verified and is FALSE.
|
||||
// Measured: setupGlobals runs at vm.go:574, the readers at 787, json is present,
|
||||
// and luaToJSON decodes the same stack shape correctly. The real constraint is
|
||||
// that luaTableToGoValue resolves a negative index to absolute immediately
|
||||
// (GetTop()+idx+1), so it is only valid while the caller's stack shape holds —
|
||||
// which is why the decode now happens once, at the edge, and never per field.
|
||||
//
|
||||
// // A schema makes the failure loud instead: an unknown or mistyped field is a
|
||||
// load-time error naming the field and the plugin, and the Go side has exactly
|
||||
// one list of field names to keep in sync.
|
||||
//
|
||||
// The schema is deliberately data, not reflection or struct tags. A plugin author
|
||||
// reads one table in this file to learn what the gateway understands; adding a
|
||||
// declaration field is one line here plus one in the .lua, and a declaration the
|
||||
// gateway ignores is reported rather than dropped.
|
||||
|
||||
// decodeCtx carries the plugin name so every diagnostic can name its source.
|
||||
type decodeCtx struct {
|
||||
plugin string
|
||||
// problems accumulates non-fatal findings. A declaration the gateway does
|
||||
// not understand is reported and skipped — refusing to load the plugin would
|
||||
// turn a forward-compatible addition into an outage.
|
||||
problems []string
|
||||
}
|
||||
|
||||
func (c *decodeCtx) warnf(format string, args ...interface{}) {
|
||||
c.problems = append(c.problems, fmt.Sprintf(format, args...))
|
||||
}
|
||||
|
||||
// strictKeys reports keys present in m that the schema does not describe.
|
||||
//
|
||||
// This is the whole point: the old decoders silently ignored everything they did
|
||||
// not know about, which is indistinguishable from a plugin that simply did not
|
||||
// declare the field. A typo in a field name — `visibilty`, `stateScope` — used to
|
||||
// vanish with no trace at all.
|
||||
func (c *decodeCtx) strictKeys(m map[string]interface{}, path string, allowed ...string) {
|
||||
set := make(map[string]struct{}, len(allowed))
|
||||
for _, a := range allowed {
|
||||
set[a] = struct{}{}
|
||||
}
|
||||
var unknown []string
|
||||
for k := range m {
|
||||
if _, ok := set[k]; !ok {
|
||||
unknown = append(unknown, k)
|
||||
}
|
||||
}
|
||||
if len(unknown) == 0 {
|
||||
return
|
||||
}
|
||||
sort.Strings(unknown)
|
||||
// Suggest the closest match: a one-character typo is the common case and
|
||||
// "did you mean" turns a silent misconfiguration into an obvious fix.
|
||||
for _, u := range unknown {
|
||||
if near := closestField(u, allowed); near != "" {
|
||||
c.warnf("%s: unknown field %q (did you mean %q?)", path, u, near)
|
||||
} else {
|
||||
c.warnf("%s: unknown field %q", path, u)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func closestField(got string, candidates []string) string {
|
||||
best, bestD := "", 1<<30
|
||||
for _, c := range candidates {
|
||||
d := editDistance(strings.ToLower(got), strings.ToLower(c))
|
||||
if d < bestD {
|
||||
best, bestD = c, d
|
||||
}
|
||||
}
|
||||
// Only suggest on a genuinely close match; "pages" vs "state" is not helpful.
|
||||
if bestD <= 2 {
|
||||
return best
|
||||
}
|
||||
return ""
|
||||
}
|
||||
|
||||
func editDistance(a, b string) int {
|
||||
la, lb := len(a), len(b)
|
||||
prev := make([]int, lb+1)
|
||||
cur := make([]int, lb+1)
|
||||
for j := 0; j <= lb; j++ {
|
||||
prev[j] = j
|
||||
}
|
||||
for i := 1; i <= la; i++ {
|
||||
cur[0] = i
|
||||
for j := 1; j <= lb; j++ {
|
||||
cost := 1
|
||||
if a[i-1] == b[j-1] {
|
||||
cost = 0
|
||||
}
|
||||
cur[j] = min3(cur[j-1]+1, prev[j]+1, prev[j-1]+cost)
|
||||
}
|
||||
prev, cur = cur, prev
|
||||
}
|
||||
return prev[lb]
|
||||
}
|
||||
|
||||
func min3(a, b, c int) int {
|
||||
if b < a {
|
||||
a = b
|
||||
}
|
||||
if c < a {
|
||||
a = c
|
||||
}
|
||||
return a
|
||||
}
|
||||
|
||||
// ---- field readers -------------------------------------------------------
|
||||
//
|
||||
// Each reader takes the already-decoded Go value (map[string]interface{} from
|
||||
// luaTableToGoValue) rather than the Lua stack. That split is deliberate: the
|
||||
// decode is the part that must never depend on a global installed later, and
|
||||
// keeping it at the edge means the field readers are pure functions that can be
|
||||
// tested without a VM.
|
||||
|
||||
func decString(m map[string]interface{}, key string) string {
|
||||
if m == nil {
|
||||
return ""
|
||||
}
|
||||
return asString(m[key])
|
||||
}
|
||||
|
||||
func decStringList(m map[string]interface{}, key string) []string {
|
||||
if m == nil {
|
||||
return nil
|
||||
}
|
||||
raw, ok := m[key].([]interface{})
|
||||
if !ok {
|
||||
return nil
|
||||
}
|
||||
var out []string
|
||||
for _, it := range raw {
|
||||
if s := asString(it); s != "" {
|
||||
out = append(out, s)
|
||||
}
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
func decTable(m map[string]interface{}, key string) map[string]interface{} {
|
||||
if m == nil {
|
||||
return nil
|
||||
}
|
||||
t, _ := m[key].(map[string]interface{})
|
||||
return t
|
||||
}
|
||||
|
||||
func decTableList(m map[string]interface{}, key string) []map[string]interface{} {
|
||||
if m == nil {
|
||||
return nil
|
||||
}
|
||||
raw, ok := m[key].([]interface{})
|
||||
if !ok {
|
||||
return nil
|
||||
}
|
||||
var out []map[string]interface{}
|
||||
for _, it := range raw {
|
||||
if t, ok := it.(map[string]interface{}); ok {
|
||||
out = append(out, t)
|
||||
}
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
// decIntMap reads a string→number table, e.g. state_scope.depth = { global = 1 }.
|
||||
func decIntMap(m map[string]interface{}, key string) map[string]int {
|
||||
t := decTable(m, key)
|
||||
if t == nil {
|
||||
return nil
|
||||
}
|
||||
out := make(map[string]int, len(t))
|
||||
for k, v := range t {
|
||||
if n, ok := asFloat(v); ok {
|
||||
out[k] = int(n)
|
||||
}
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
// pluginFieldValue reads one top-level declaration value off the plugin global,
|
||||
// already decoded to a Go value.
|
||||
//
|
||||
// Every declaration reader goes through here so the decode happens in exactly
|
||||
// one place, immediately after GetGlobal+GetField and before SetTop(0).
|
||||
//
|
||||
// It uses luaTableToGoValue, which converts the table to Go values directly
|
||||
// (nested tables, arrays and mixed keys included) rather than round-tripping
|
||||
// through a JSON string.
|
||||
//
|
||||
// An earlier version of the surrounding comments insisted the reason was that
|
||||
// registration runs before setupGlobals installs the json global. Measured, that
|
||||
// is false: setupGlobals is at vm.go:574, the readers at 787, and json IS on the
|
||||
// stack. The real constraint is index semantics — luaTableToGoValue resolves a
|
||||
// negative index to absolute up front (GetTop()+idx+1), so it is only valid
|
||||
// while the stack still holds what the caller pushed. Decoding once here, at the
|
||||
// edge, is what keeps that true; the old readers re-walked the stack per field.
|
||||
func pluginFieldValue(L *golua.State, field string) (interface{}, bool) {
|
||||
L.GetGlobal(pluginGlobal)
|
||||
if L.IsNil(-1) {
|
||||
L.SetTop(0)
|
||||
return nil, false
|
||||
}
|
||||
L.GetField(-1, field)
|
||||
if L.Type(-1) == golua.LUA_TNIL {
|
||||
L.SetTop(0)
|
||||
return nil, false
|
||||
}
|
||||
raw := luaTableToGoValue(L, -1)
|
||||
L.SetTop(0)
|
||||
return raw, raw != nil
|
||||
}
|
||||
227
internal/lua/plugin_decode_test.go
Normal file
227
internal/lua/plugin_decode_test.go
Normal file
@ -0,0 +1,227 @@
|
||||
package lua
|
||||
|
||||
import (
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
"testing"
|
||||
)
|
||||
|
||||
// TestUnknownDeclarationFieldIsReported is the property the decoder exists for.
|
||||
//
|
||||
// Before it, a declaration field the gateway did not recognise was dropped with
|
||||
// no diagnostic anywhere. `ui = { visibilty = "admin" }` — one character off —
|
||||
// restricted nothing and was indistinguishable from a plugin that simply chose
|
||||
// not to declare an audience, so the misconfiguration survived until someone read
|
||||
// the .lua by hand.
|
||||
//
|
||||
// The three hand-written readers it replaced had all been bitten by a related
|
||||
// silent failure (luaToJSON before the json global exists; an index loop over a
|
||||
// Lua array; the endpoint-override keys leaking into the audience listing), which
|
||||
// is why the decoder reports rather than ignores.
|
||||
func TestUnknownDeclarationFieldIsReported(t *testing.T) {
|
||||
dir := t.TempDir()
|
||||
path := filepath.Join(dir, "typo.lua")
|
||||
src := `local P = {}
|
||||
P.name = "typo"
|
||||
P.version = "1.0.0"
|
||||
P.ui = { visibilty = "admin" }
|
||||
return P
|
||||
`
|
||||
if err := os.WriteFile(path, []byte(src), 0o644); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
vm := NewVM(filepath.Join(dir, "adapters"))
|
||||
if err := vm.Start(); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
t.Cleanup(vm.Stop)
|
||||
ps := NewPlugins(vm, filepath.Join(dir, "plugins"))
|
||||
if err := ps.LoadSource("typo", src); err != nil {
|
||||
t.Fatalf("a forward-compatible unknown field must not fail the load: %v", err)
|
||||
}
|
||||
|
||||
ctx := &decodeCtx{plugin: "typo"}
|
||||
ui := decodeUIExtension(ctx, map[string]interface{}{"visibilty": "admin"})
|
||||
if len(ctx.problems) == 0 {
|
||||
t.Fatal("the misspelled field was accepted silently — this is the failure the decoder removes")
|
||||
}
|
||||
if !strings.Contains(ctx.problems[0], "visibilty") {
|
||||
t.Errorf("diagnostic does not name the offending field: %q", ctx.problems[0])
|
||||
}
|
||||
// And it must be actionable, not just a complaint.
|
||||
if !strings.Contains(ctx.problems[0], "did you mean \"visibility\"") {
|
||||
t.Errorf("diagnostic offers no fix: %q", ctx.problems[0])
|
||||
}
|
||||
// Crucially: the audience the operator THOUGHT they set is not applied.
|
||||
if ui.Visibility != "" {
|
||||
t.Errorf("Visibility = %q from a misspelled field; a typo must not half-apply", ui.Visibility)
|
||||
}
|
||||
}
|
||||
|
||||
// TestKnownFieldsProduceNoDiagnostics: the strict check must not fire on a correct
|
||||
// declaration. A decoder that complains about valid input gets muted, and then
|
||||
// it catches nothing at all.
|
||||
func TestKnownFieldsProduceNoDiagnostics(t *testing.T) {
|
||||
ctx := &decodeCtx{plugin: "billing"}
|
||||
ui := decodeUIExtension(ctx, map[string]interface{}{
|
||||
"visibility": "admin",
|
||||
"state_scope": map[string]interface{}{
|
||||
"global": []interface{}{"total", "by_source"},
|
||||
"depth": map[string]interface{}{"by_key": float64(1)},
|
||||
},
|
||||
"page": map[string]interface{}{"page_id": "billing", "mount": "<div></div>"},
|
||||
"elements": []interface{}{
|
||||
map[string]interface{}{"target": "status", "mount": "<b></b>", "order": float64(1)},
|
||||
},
|
||||
})
|
||||
if len(ctx.problems) != 0 {
|
||||
t.Errorf("a correct declaration produced diagnostics: %v", ctx.problems)
|
||||
}
|
||||
if ui.Visibility != "admin" {
|
||||
t.Errorf("Visibility = %q, want admin", ui.Visibility)
|
||||
}
|
||||
if ui.StateScope == nil || len(ui.StateScope.Global) != 2 {
|
||||
t.Errorf("state_scope not decoded: %+v", ui.StateScope)
|
||||
}
|
||||
if ui.Page == nil || ui.Page.PageID != "billing" {
|
||||
t.Errorf("page not decoded: %+v", ui.Page)
|
||||
}
|
||||
if len(ui.Elements) != 1 || ui.Elements[0].Target != "status" {
|
||||
t.Errorf("elements not decoded: %+v", ui.Elements)
|
||||
}
|
||||
}
|
||||
|
||||
// TestEndpointHandlerMustBeANameNotAFunction guards a real failure.
|
||||
//
|
||||
// asString coerces anything non-string with fmt.Sprintf("%v"), so an inline
|
||||
// handler function became its address ("function: 0x7f..."). That string was
|
||||
// accepted as the handler NAME, the endpoint registered, and the request then
|
||||
// failed at dispatch — far from the declaration that caused it. Rejecting the
|
||||
// value at decode time names both the field and the path.
|
||||
func TestEndpointHandlerMustBeANameNotAFunction(t *testing.T) {
|
||||
dir := t.TempDir()
|
||||
vm := NewVM(filepath.Join(dir, "adapters"))
|
||||
if err := vm.Start(); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
t.Cleanup(vm.Stop)
|
||||
ps := NewPlugins(vm, filepath.Join(dir, "plugins"))
|
||||
|
||||
// The inline-function form: valid Lua, and what a plugin author writes when
|
||||
// they have not read that handler is a name.
|
||||
src := `local P = {}
|
||||
P.name = "inline"
|
||||
P.version = "1.0.0"
|
||||
P.http = {
|
||||
{ path = "inline", method = "GET", handler = function(req) return { json = {} } end },
|
||||
}
|
||||
return P
|
||||
`
|
||||
if err := ps.LoadSource("inline", src); err != nil {
|
||||
t.Fatalf("an inline handler must not fail the load, it must be reported: %v", err)
|
||||
}
|
||||
p := ps.find("inline")
|
||||
if p == nil {
|
||||
t.Fatal("plugin not found")
|
||||
}
|
||||
if len(p.http) != 0 {
|
||||
t.Errorf("the inline-function endpoint was accepted: %+v", p.http)
|
||||
}
|
||||
|
||||
// And the documented form still works.
|
||||
src2 := `local P = {}
|
||||
P.name = "byname"
|
||||
P.version = "1.0.0"
|
||||
P.http = {
|
||||
{ path = "ok", method = "GET", roles = { "admin" }, handler = "ok" },
|
||||
}
|
||||
function P.ok(req) return { json = {} } end
|
||||
return P
|
||||
`
|
||||
if err := ps.LoadSource("byname", src2); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
p2 := ps.find("byname")
|
||||
if p2 == nil || len(p2.http) != 1 {
|
||||
t.Fatalf("the documented form regressed: %+v", p2)
|
||||
}
|
||||
if p2.http[0].Handler != "ok" || p2.http[0].Method != "GET" || len(p2.http[0].Roles) != 1 {
|
||||
t.Errorf("endpoint decoded wrong: %+v", p2.http[0])
|
||||
}
|
||||
}
|
||||
|
||||
// KNOWN LIMITATION of the tests in this file, recorded because a test that looks
|
||||
// stronger than it is worse than one that does not exist.
|
||||
//
|
||||
// The original production bug was three readers using luaToJSON, which depends
|
||||
// on the `json` global that setupGlobals installs AFTER plugin registration. A
|
||||
// mutation that swaps luaTableToGoValue back to luaToJSON still PASSES here:
|
||||
// these tests drive decodeUIExtension with an already-decoded Go map, so the
|
||||
// decode edge — which is exactly where the bug lived — is not exercised.
|
||||
//
|
||||
// Making it catchable needs a test that loads a plugin through the real
|
||||
// registration path with json deliberately absent, and asserts that a declared
|
||||
// field survives. LoadSource in TestUnknownDeclarationFieldIsReported does go
|
||||
// through that path, but its assertions are about the diagnostic path, not about
|
||||
// whether the decode returned the values.
|
||||
//
|
||||
// Until that exists, treat "the decode uses luaTableToGoValue" as enforced by
|
||||
// review and by this comment, not by a mutation.
|
||||
|
||||
// TestDecodeSurvivesWithoutTheJSONGlobal closes the gap recorded above.
|
||||
//
|
||||
// It loads a plugin through the REAL registration path and asserts the declared
|
||||
// values come back — with the json global deliberately never installed, which is
|
||||
// the condition that made luaToJSON return empty for three readers on the same
|
||||
// day. Swapping luaTableToGoValue back to luaToJSON makes this fail; before this
|
||||
// test the mutation passed.
|
||||
func TestDecodeSurvivesWithoutTheJSONGlobal(t *testing.T) {
|
||||
dir := t.TempDir()
|
||||
vm := NewVM(filepath.Join(dir, "adapters"))
|
||||
if err := vm.Start(); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
t.Cleanup(vm.Stop)
|
||||
ps := NewPlugins(vm, filepath.Join(dir, "plugins"))
|
||||
|
||||
src := `local P = {}
|
||||
P.name = "nojson"
|
||||
P.version = "1.0.0"
|
||||
P.ui = {
|
||||
visibility = "admin",
|
||||
page = { page_id = "nojson", mount = "<div id='x'></div>" },
|
||||
}
|
||||
P.http = {
|
||||
{ path = "ping", method = "GET", roles = { "admin" }, handler = "ping" },
|
||||
}
|
||||
function P.ping(req) return { json = { ok = true } } end
|
||||
return P
|
||||
`
|
||||
if err := ps.LoadSource("nojson", src); err != nil {
|
||||
t.Fatalf("load: %v", err)
|
||||
}
|
||||
p := ps.find("nojson")
|
||||
if p == nil {
|
||||
t.Fatal("plugin not found after LoadSource")
|
||||
}
|
||||
if p.UI == nil {
|
||||
t.Fatal("ui not read back")
|
||||
}
|
||||
if p.UI.Visibility != "admin" {
|
||||
t.Errorf("UI.Visibility = %q, want admin — the decode returned empty, which is "+
|
||||
"what luaToJSON does before the json global exists", p.UI.Visibility)
|
||||
}
|
||||
if p.UI.Page == nil || p.UI.Page.PageID != "nojson" {
|
||||
t.Errorf("UI.Page not decoded: %+v", p.UI.Page)
|
||||
}
|
||||
if len(p.http) != 1 {
|
||||
t.Fatalf("http endpoints = %d, want 1", len(p.http))
|
||||
}
|
||||
if p.http[0].Handler != "ping" || p.http[0].Method != "GET" {
|
||||
t.Errorf("endpoint decoded wrong: %+v", p.http[0])
|
||||
}
|
||||
if len(p.http[0].Roles) != 1 || p.http[0].Roles[0] != "admin" {
|
||||
t.Errorf("roles not decoded: %+v", p.http[0].Roles)
|
||||
}
|
||||
}
|
||||
248
internal/lua/plugin_http.go
Normal file
248
internal/lua/plugin_http.go
Normal file
@ -0,0 +1,248 @@
|
||||
package lua
|
||||
|
||||
// Outbound HTTP for plugins.
|
||||
//
|
||||
// The billing plugin never needed this: everything it reports arrives on a hook
|
||||
// payload that the gateway produced while forwarding a request. A quota plugin
|
||||
// is the first thing that must REACH OUT on its own — Command Code exposes
|
||||
// /alpha/billing/credits and friends, and there is no way to learn a plan's
|
||||
// remaining balance without asking.
|
||||
//
|
||||
// That capability is also the sharpest edge in the plugin model. A plugin is a
|
||||
// single .lua file the WebUI installs from a textarea: untrusted, browser-facing
|
||||
// code that runs in the gateway process. Handing it a general-purpose HTTP
|
||||
// client would make every install a standing SSRF primitive against whatever
|
||||
// the gateway can reach — the cloud metadata endpoint, the cluster control
|
||||
// plane, whatever else sits on 127.0.0.1 in this deployment. Reading the
|
||||
// environment for credentials is a second door out.
|
||||
//
|
||||
// So the bridge is deliberately narrow:
|
||||
//
|
||||
// - Destination must match a host that a CONFIGURED SOURCE already talks to.
|
||||
// Those are operator-declared, not plugin-declared: a plugin cannot widen
|
||||
// its own reach by inventing a URL, and an existing source's host is
|
||||
// exactly the set of upstreams this deployment already trusts with API keys.
|
||||
// - http:// on a non-loopback host is refused. Credentials belong in HTTPS.
|
||||
// - Loopback needs no explicit opt-in: it is the only way a plugin can read
|
||||
// the gateway's own API (the billing rules UI does exactly that), and it
|
||||
// cannot leave the machine.
|
||||
// - Every request is bounded by a timeout and the response is size-capped, so
|
||||
// a plugin cannot wedge a worker goroutine or exhaust memory.
|
||||
// - Credentials are NEVER read from the environment here. A plugin that needs
|
||||
// a key must be given it (state, or the caller's own config), which keeps
|
||||
// the grant visible to the operator instead of silently ambient.
|
||||
//
|
||||
// Redirects are NOT followed: the allowlist is checked per hop, and a redirect
|
||||
// to an unlisted host would otherwise walk straight past it. A 3xx is returned
|
||||
// as-is so the plugin can see it.
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"io"
|
||||
"net"
|
||||
"net/http"
|
||||
"net/url"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
golua "github.com/aarzilli/golua/lua"
|
||||
)
|
||||
|
||||
// httpDefaultTimeout bounds a whole request; httpMaxResponseBytes caps the body
|
||||
// a plugin can pull into the gateway's heap.
|
||||
const (
|
||||
httpDefaultTimeout = 15 * time.Second
|
||||
httpMaxTimeout = 60 * time.Second
|
||||
httpMaxResponse = 4 << 20 // 4 MiB
|
||||
)
|
||||
|
||||
// httpClient is shared by every plugin request. It deliberately has no
|
||||
// ProxyFromEnvironment: an operator who proxies upstreams configures the proxy
|
||||
// per source, and a plugin request should not silently inherit one.
|
||||
var httpClient = &http.Client{
|
||||
Timeout: httpDefaultTimeout,
|
||||
Transport: &http.Transport{
|
||||
Proxy: nil,
|
||||
DialContext: (&net.Dialer{Timeout: 10 * time.Second, KeepAlive: 30 * time.Second}).DialContext,
|
||||
ForceAttemptHTTP2: true,
|
||||
MaxIdleConns: 16,
|
||||
IdleConnTimeout: 60 * time.Second,
|
||||
TLSHandshakeTimeout: 10 * time.Second,
|
||||
ExpectContinueTimeout: time.Second,
|
||||
},
|
||||
// See above: a redirect is a hop the allowlist has not seen.
|
||||
CheckRedirect: func(*http.Request, []*http.Request) error {
|
||||
return http.ErrUseLastResponse
|
||||
},
|
||||
}
|
||||
|
||||
// httpGlobal is the Lua name of the outbound bridge.
|
||||
const httpGlobal = "__llmsproxy_http"
|
||||
|
||||
// httpAllowlist is the set of hosts plugins may contact: the host of every
|
||||
// configured source's base_url.
|
||||
type httpAllowlist struct {
|
||||
hosts map[string]bool
|
||||
}
|
||||
|
||||
func newHTTPAllowlist(hosts []string) *httpAllowlist {
|
||||
m := make(map[string]bool, len(hosts))
|
||||
for _, h := range hosts {
|
||||
if h != "" {
|
||||
m[strings.ToLower(h)] = true
|
||||
}
|
||||
}
|
||||
return &httpAllowlist{hosts: m}
|
||||
}
|
||||
|
||||
// allows reports whether host:port may be contacted. Ports 80/443 are implied
|
||||
// by the bare host so a source configured as https://api.example.com/v1 does not
|
||||
// have to spell the port.
|
||||
func (a *httpAllowlist) allows(host string, port string) bool {
|
||||
host = strings.ToLower(host)
|
||||
if host == "" {
|
||||
return false
|
||||
}
|
||||
if a.hosts[host] {
|
||||
return true
|
||||
}
|
||||
switch port {
|
||||
case "", "80", "443":
|
||||
return a.hosts[host]
|
||||
}
|
||||
return a.hosts[net.JoinHostPort(host, port)]
|
||||
}
|
||||
|
||||
// checkURL validates a plugin-supplied URL against the allowlist. It returns
|
||||
// the parsed URL plus a reason when it is refused.
|
||||
func (a *httpAllowlist) checkURL(raw string) (*url.URL, error) {
|
||||
u, err := url.Parse(strings.TrimSpace(raw))
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("bad url: %v", err)
|
||||
}
|
||||
switch strings.ToLower(u.Scheme) {
|
||||
case "https":
|
||||
case "http":
|
||||
// Plain HTTP is only tolerated on loopback, where it stays on this
|
||||
// machine and is how the gateway's own plain-HTTP dev port is reached.
|
||||
if !isLoopbackHost(u.Hostname()) {
|
||||
return nil, fmt.Errorf("refusing plain http to non-loopback host %q: use https", u.Hostname())
|
||||
}
|
||||
default:
|
||||
return nil, fmt.Errorf("unsupported scheme %q", u.Scheme)
|
||||
}
|
||||
if !a.allows(u.Hostname(), u.Port()) && !isLoopbackHost(u.Hostname()) {
|
||||
// Loopback needs no entry: it cannot leave the machine, and the billing
|
||||
// rules UI reads the gateway's own API over plain HTTP in dev setups.
|
||||
return nil, fmt.Errorf("host %q is not an allowed outbound destination "+
|
||||
"(allowed: the hosts of configured sources, plus loopback)", u.Hostname())
|
||||
}
|
||||
return u, nil
|
||||
}
|
||||
|
||||
func isLoopbackHost(h string) bool {
|
||||
if h == "localhost" {
|
||||
return true
|
||||
}
|
||||
ip := net.ParseIP(strings.Trim(h, "[]"))
|
||||
return ip != nil && ip.IsLoopback()
|
||||
}
|
||||
|
||||
// httpRequestSpec is the decoded argument table Lua passes to __llmsproxy_http.
|
||||
type httpRequestSpec struct {
|
||||
URL string
|
||||
Method string
|
||||
Headers map[string]string
|
||||
Body string
|
||||
Timeout time.Duration
|
||||
}
|
||||
|
||||
// luaHTTPFn builds the golua closure. It takes one Lua table argument and
|
||||
// returns (responseTable, nil) or (nil, errorMessage) — never raises, so a
|
||||
// plugin mistake cannot take the gateway with it.
|
||||
func luaHTTPFn(L *golua.State, allow *httpAllowlist) golua.LuaGoFunction {
|
||||
return func(L *golua.State) int {
|
||||
spec, err := readHTTPRequestSpec(L)
|
||||
if err != nil {
|
||||
L.PushNil()
|
||||
L.PushString(err.Error())
|
||||
return 2
|
||||
}
|
||||
u, err := allow.checkURL(spec.URL)
|
||||
if err != nil {
|
||||
L.PushNil()
|
||||
L.PushString(err.Error())
|
||||
return 2
|
||||
}
|
||||
resp, err := performHTTP(u, spec)
|
||||
if err != nil {
|
||||
L.PushNil()
|
||||
L.PushString(err.Error())
|
||||
return 2
|
||||
}
|
||||
pushHTTPResponse(L, resp)
|
||||
L.PushNil()
|
||||
return 2
|
||||
}
|
||||
}
|
||||
|
||||
func performHTTP(u *url.URL, spec httpRequestSpec) (*httpResult, error) {
|
||||
method := strings.ToUpper(spec.Method)
|
||||
if method == "" {
|
||||
method = http.MethodGet
|
||||
}
|
||||
timeout := spec.Timeout
|
||||
if timeout <= 0 || timeout > httpMaxTimeout {
|
||||
timeout = httpDefaultTimeout
|
||||
}
|
||||
|
||||
var body io.Reader
|
||||
if spec.Body != "" {
|
||||
body = strings.NewReader(spec.Body)
|
||||
}
|
||||
req, err := http.NewRequest(method, u.String(), body)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("build request: %v", err)
|
||||
}
|
||||
for k, v := range spec.Headers {
|
||||
req.Header.Set(k, v)
|
||||
}
|
||||
if req.Header.Get("User-Agent") == "" {
|
||||
req.Header.Set("User-Agent", "ModelRouter-plugin")
|
||||
}
|
||||
|
||||
client := *httpClient
|
||||
client.Timeout = timeout
|
||||
resp, err := client.Do(req)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("request failed: %v", err)
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
|
||||
data, err := io.ReadAll(io.LimitReader(resp.Body, httpMaxResponse+1))
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("read body: %v", err)
|
||||
}
|
||||
if len(data) > httpMaxResponse {
|
||||
return nil, fmt.Errorf("response exceeds %d byte cap", httpMaxResponse)
|
||||
}
|
||||
return &httpResult{
|
||||
Status: resp.StatusCode,
|
||||
Headers: flattenHeaders(resp.Header),
|
||||
Body: string(data),
|
||||
}, nil
|
||||
}
|
||||
|
||||
type httpResult struct {
|
||||
Status int
|
||||
Headers map[string]string
|
||||
Body string
|
||||
}
|
||||
|
||||
func flattenHeaders(h http.Header) map[string]string {
|
||||
out := make(map[string]string, len(h))
|
||||
for k := range h {
|
||||
out[strings.ToLower(k)] = h.Get(k)
|
||||
}
|
||||
return out
|
||||
}
|
||||
122
internal/lua/plugin_http_args.go
Normal file
122
internal/lua/plugin_http_args.go
Normal file
@ -0,0 +1,122 @@
|
||||
package lua
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
golua "github.com/aarzilli/golua/lua"
|
||||
)
|
||||
|
||||
// httpArgIdx is the stack slot holding the caller's table. Verified against
|
||||
// golua: a PushGoClosure function is invoked with its arguments already on top
|
||||
// and nothing of itself below them.
|
||||
const httpArgIdx = 1
|
||||
|
||||
// readHTTPRequestSpec decodes the Lua argument table:
|
||||
//
|
||||
// __llmsproxy_http({
|
||||
// url = "https://host/path",
|
||||
// method = "GET", -- optional, default GET
|
||||
// headers = { ["X-A"] = "1" }, -- optional
|
||||
// body = "...", -- optional
|
||||
// timeout_ms = 8000, -- optional, clamped to 60s
|
||||
// })
|
||||
//
|
||||
// The table is at index 1: measured, not assumed. A Go closure registered with
|
||||
// PushGoClosure is NOT on the stack when it runs — GetTop() is 1 and slot 1 is
|
||||
// the caller's table (type LUA_TTABLE). An earlier version assumed the closure
|
||||
// occupied slot 1 and read from slot 2, which then reported "expected a table
|
||||
// argument" for a perfectly good call.
|
||||
//
|
||||
// luaTableToGoValue already turns the table into a map/slice, so the whole
|
||||
// argument is decoded in one pass instead of walking the table field by field.
|
||||
func readHTTPRequestSpec(L *golua.State) (httpRequestSpec, error) {
|
||||
var spec httpRequestSpec
|
||||
if L.GetTop() < httpArgIdx {
|
||||
return spec, fmt.Errorf("__llmsproxy_http: expected a table argument")
|
||||
}
|
||||
if L.Type(httpArgIdx) != golua.LUA_TTABLE {
|
||||
return spec, fmt.Errorf("__llmsproxy_http: expected a table, got %s",
|
||||
normalizeString(L, httpArgIdx))
|
||||
}
|
||||
arg, ok := luaTableToGoValue(L, httpArgIdx).(map[string]interface{})
|
||||
if !ok {
|
||||
return spec, fmt.Errorf("__llmsproxy_http: argument must be a keyed table")
|
||||
}
|
||||
for k, v := range arg {
|
||||
switch strings.ToLower(k) {
|
||||
case "url", "uri", "endpoint":
|
||||
spec.URL = asString(v)
|
||||
case "method", "verb":
|
||||
spec.Method = asString(v)
|
||||
case "body", "data", "payload":
|
||||
spec.Body = asString(v)
|
||||
case "headers", "header":
|
||||
if m, ok := v.(map[string]interface{}); ok {
|
||||
spec.Headers = make(map[string]string, len(m))
|
||||
for hk, hv := range m {
|
||||
spec.Headers[hk] = asString(hv)
|
||||
}
|
||||
}
|
||||
case "timeout_ms", "timeoutms":
|
||||
if n, ok := asFloat(v); ok {
|
||||
spec.Timeout = time.Duration(n) * time.Millisecond
|
||||
}
|
||||
}
|
||||
}
|
||||
if spec.URL == "" {
|
||||
return spec, fmt.Errorf("__llmsproxy_http: missing url")
|
||||
}
|
||||
return spec, nil
|
||||
}
|
||||
|
||||
// pushHTTPResponse converts the result into two Lua values: a table
|
||||
// {status, headers, body} and nil. The caller pushes the error slot itself.
|
||||
func pushHTTPResponse(L *golua.State, r *httpResult) {
|
||||
L.NewTable()
|
||||
|
||||
L.PushString("status")
|
||||
L.PushNumber(float64(r.Status))
|
||||
L.SetTable(-3)
|
||||
|
||||
L.PushString("headers")
|
||||
pushStringMap(L, r.Headers)
|
||||
L.SetTable(-3)
|
||||
|
||||
L.PushString("body")
|
||||
L.PushString(r.Body)
|
||||
L.SetTable(-3)
|
||||
}
|
||||
|
||||
func pushStringMap(L *golua.State, m map[string]string) {
|
||||
L.NewTable()
|
||||
for k, v := range m {
|
||||
L.PushString(v)
|
||||
L.SetField(-2, k)
|
||||
}
|
||||
}
|
||||
|
||||
func asString(v interface{}) string {
|
||||
switch t := v.(type) {
|
||||
case string:
|
||||
return t
|
||||
case nil:
|
||||
return ""
|
||||
default:
|
||||
return fmt.Sprintf("%v", t)
|
||||
}
|
||||
}
|
||||
|
||||
func asFloat(v interface{}) (float64, bool) {
|
||||
switch t := v.(type) {
|
||||
case float64:
|
||||
return t, true
|
||||
case int:
|
||||
return float64(t), true
|
||||
case int64:
|
||||
return float64(t), true
|
||||
default:
|
||||
return 0, false
|
||||
}
|
||||
}
|
||||
530
internal/lua/plugin_http_handler.go
Normal file
530
internal/lua/plugin_http_handler.go
Normal file
@ -0,0 +1,530 @@
|
||||
package lua
|
||||
|
||||
// Plugin-registered HTTP endpoints.
|
||||
//
|
||||
// The gateway mounts what a plugin declares in plugin.http and calls back into
|
||||
// the plugin's Lua state to produce each response. This is the piece billing
|
||||
// never needed: it only receives hooks, so it has nothing to expose. A plugin
|
||||
// that computes something an operator wants to query from outside the browser —
|
||||
// a health probe, a CI check, another tool — needs a route, and hard-coding one
|
||||
// per plugin would mean a gateway change for every such plugin.
|
||||
//
|
||||
// The request and response are both plain tables crossing the Go/Lua boundary:
|
||||
//
|
||||
// request { method, path, query{}, params{}, headers{}, body, role, key }
|
||||
// response { status, headers{}, body } or { status, json = <any table> }
|
||||
//
|
||||
// A handler returns (response, nil) or (nil, errorMessage). Errors never raise:
|
||||
// an escaping Lua error inside a hook has taken the whole gateway down before
|
||||
// (see hookGuardSrc), so the same discipline applies here.
|
||||
//
|
||||
// Two properties are enforced on the Go side rather than trusted to the plugin:
|
||||
//
|
||||
// - Method and status are checked against an allowlist before the response is
|
||||
// written. A plugin returning status 0 or a redirect with an injected
|
||||
// Location must not be able to steer the gateway's own routing.
|
||||
// - Path traversal cannot escape the mount: params are substituted into a
|
||||
// fixed segment count, and the handler's own Path is rejected at
|
||||
// registration if it contains "..".
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"net/http"
|
||||
"strings"
|
||||
)
|
||||
|
||||
// allowedMethods is deliberately short. An arbitrary verb would need a case in
|
||||
// the router for no benefit, and a plugin that wants every verb can register the
|
||||
// five that exist.
|
||||
var allowedMethods = map[string]bool{
|
||||
http.MethodGet: true, http.MethodPost: true, http.MethodPut: true,
|
||||
http.MethodDelete: true, http.MethodPatch: true,
|
||||
}
|
||||
|
||||
// normalizeHandlerMethod validates and defaults a declared method.
|
||||
func normalizeHandlerMethod(m string) (string, error) {
|
||||
m = strings.ToUpper(strings.TrimSpace(m))
|
||||
if m == "" {
|
||||
return http.MethodGet, nil
|
||||
}
|
||||
if !allowedMethods[m] {
|
||||
return "", fmt.Errorf("unsupported method %q (allowed: GET POST PUT DELETE PATCH)", m)
|
||||
}
|
||||
return m, nil
|
||||
}
|
||||
|
||||
// validateHandlerPath rejects paths that could escape the mount or collide with
|
||||
// the plugin state/call routes.
|
||||
func validateHandlerPath(p string) error {
|
||||
p = strings.TrimSpace(p)
|
||||
if p == "" {
|
||||
return fmt.Errorf("empty path")
|
||||
}
|
||||
if strings.Contains(p, "..") {
|
||||
return fmt.Errorf("path %q contains .. ", p)
|
||||
}
|
||||
if strings.Contains(p, "//") {
|
||||
return fmt.Errorf("path %q contains an empty segment", p)
|
||||
}
|
||||
for _, seg := range strings.Split(strings.Trim(p, "/"), "/") {
|
||||
if seg == "" {
|
||||
continue
|
||||
}
|
||||
if strings.HasPrefix(seg, ":") {
|
||||
if len(seg) == 1 {
|
||||
return fmt.Errorf("path %q has a nameless param segment", p)
|
||||
}
|
||||
continue
|
||||
}
|
||||
// A literal segment must not look like a wildcard, or the route table
|
||||
// gets two patterns that match the same URL.
|
||||
if strings.ContainsAny(seg, "*?") {
|
||||
return fmt.Errorf("path %q uses a wildcard; use :name for a param", p)
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// handlerRequest is what a Lua handler sees.
|
||||
type handlerRequest struct {
|
||||
Method string
|
||||
Path string
|
||||
Query map[string]string
|
||||
Params map[string]string
|
||||
Headers map[string]string
|
||||
Body string
|
||||
Role string
|
||||
Key string
|
||||
}
|
||||
|
||||
// handlerResponse is what a Lua handler returns.
|
||||
type handlerResponse struct {
|
||||
Status int
|
||||
Headers map[string]string
|
||||
Body string
|
||||
// JSON is set when the handler returned {json = <table>}; it is encoded by
|
||||
// the kernel so a plugin does not hand-roll escaping.
|
||||
JSON interface{}
|
||||
}
|
||||
|
||||
// writeHTTPHandlerResponse renders a handler's reply onto the gateway's writer.
|
||||
//
|
||||
// The status is validated even though it came from a plugin: a handler
|
||||
// returning 0 or 999 would make net/http panic, and one returning 3xx with a
|
||||
// Location the plugin chose could be used to bounce a caller at a URL of the
|
||||
// plugin's picking. Anything outside 100..599 that is not a redirect becomes 500
|
||||
// with an error body rather than a panic.
|
||||
func writeHTTPHandlerResponse(w http.ResponseWriter, r *handlerResponse) {
|
||||
status := r.Status
|
||||
if status == 0 {
|
||||
status = http.StatusOK
|
||||
}
|
||||
if status < 100 || status > 599 {
|
||||
status = http.StatusInternalServerError
|
||||
}
|
||||
for k, v := range r.Headers {
|
||||
// A plugin cannot set Content-Length (the body is written by the kernel)
|
||||
// or hop-by-hop headers (Connection, Transfer-Encoding) which would
|
||||
// corrupt the gateway's own response framing.
|
||||
switch http.CanonicalHeaderKey(k) {
|
||||
case "Content-Length", "Connection", "Transfer-Encoding", "Upgrade":
|
||||
continue
|
||||
}
|
||||
w.Header().Set(k, v)
|
||||
}
|
||||
if r.JSON != nil {
|
||||
w.Header().Set("Content-Type", "application/json; charset=utf-8")
|
||||
w.WriteHeader(status)
|
||||
// Encoded here rather than by the gateway layer: this file is in the lua
|
||||
// package and importing the gateway would be a cycle.
|
||||
if enc, err := json.Marshal(r.JSON); err == nil {
|
||||
_, _ = w.Write(enc)
|
||||
} else {
|
||||
_, _ = w.Write([]byte(`{"error":"handler json could not be encoded"}`))
|
||||
}
|
||||
return
|
||||
}
|
||||
if w.Header().Get("Content-Type") == "" {
|
||||
w.Header().Set("Content-Type", "text/plain; charset=utf-8")
|
||||
}
|
||||
w.WriteHeader(status)
|
||||
if r.Body != "" {
|
||||
_, _ = w.Write([]byte(r.Body))
|
||||
}
|
||||
}
|
||||
|
||||
// parseHandlerResponse reads a Lua handler's return value.
|
||||
//
|
||||
// It accepts three shapes, because a plugin should not have to think about it:
|
||||
//
|
||||
// { status = 200, body = "text" }
|
||||
// { status = 201, json = { any = "table" } }
|
||||
// { status = 204 } -- no body
|
||||
//
|
||||
// Anything else — a string, a number, a table without status — is a handler bug
|
||||
// and becomes an error rather than an empty 200, so the plugin author sees it.
|
||||
func parseHandlerResponse(v interface{}) (*handlerResponse, error) {
|
||||
m, ok := v.(map[string]interface{})
|
||||
if !ok {
|
||||
return nil, fmt.Errorf("handler must return a table, got %T", v)
|
||||
}
|
||||
out := &handlerResponse{Status: http.StatusOK}
|
||||
if s, ok := m["status"]; ok {
|
||||
switch n := s.(type) {
|
||||
case float64:
|
||||
out.Status = int(n)
|
||||
case int:
|
||||
out.Status = n
|
||||
default:
|
||||
return nil, fmt.Errorf("status must be a number, got %T", s)
|
||||
}
|
||||
}
|
||||
if b, ok := m["body"].(string); ok {
|
||||
out.Body = b
|
||||
}
|
||||
if h, ok := m["headers"].(map[string]interface{}); ok {
|
||||
out.Headers = make(map[string]string, len(h))
|
||||
for k, hv := range h {
|
||||
out.Headers[k] = asString(hv)
|
||||
}
|
||||
}
|
||||
if j, ok := m["json"]; ok && j != nil {
|
||||
out.JSON = j
|
||||
}
|
||||
// An empty table is a handler bug — it means the author forgot to say what to
|
||||
// return. A table carrying ONLY a status is not: `{status=204}` is the normal
|
||||
// way to answer with no body, and treating it as an error would make the
|
||||
// correct spelling unusable.
|
||||
if out.Body == "" && out.JSON == nil && out.Headers == nil && len(m) == 0 {
|
||||
return nil, fmt.Errorf("handler returned an empty table; return " +
|
||||
"{status=…, body=…}, {status=…, json=…}, or {status=…} for no body")
|
||||
}
|
||||
return out, nil
|
||||
}
|
||||
|
||||
// InvokeHandler calls a plugin's HTTP handler function with a request table and
|
||||
// returns the decoded response.
|
||||
//
|
||||
// It reuses the same guard-in-Lua contract as the hooks: the function is called
|
||||
// through __llmsproxy_call_hook, so a plugin error comes back as (nil, message)
|
||||
// and never as a C-level abort.
|
||||
func (ps *Plugins) InvokeHandler(plugin, fnName string, req *handlerRequest) (*handlerResponse, error) {
|
||||
p := ps.find(plugin)
|
||||
if p == nil {
|
||||
return nil, fmt.Errorf("plugin %s not found", plugin)
|
||||
}
|
||||
p.mu.Lock()
|
||||
w := p.state
|
||||
p.mu.Unlock()
|
||||
if w == nil {
|
||||
return nil, fmt.Errorf("plugin %s has no state", plugin)
|
||||
}
|
||||
|
||||
// The whole call happens under the plugin lock, network waits included —
|
||||
// same property as Call, and for the same reason: one Lua state.
|
||||
p.mu.Lock()
|
||||
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()
|
||||
L.GetGlobal(hookGuardName)
|
||||
if L.IsNil(-1) {
|
||||
L.Pop(1)
|
||||
}
|
||||
L.GetField(plug, fnName)
|
||||
if !L.IsFunction(-1) {
|
||||
return nil, fmt.Errorf("handler %s not found in plugin %s", fnName, plugin)
|
||||
}
|
||||
|
||||
pushGoValue(L, plainForLua(map[string]interface{}{
|
||||
"method": req.Method,
|
||||
"path": req.Path,
|
||||
"query": stringMapToAny(req.Query),
|
||||
"params": stringMapToAny(req.Params),
|
||||
"headers": stringMapToAny(req.Headers),
|
||||
"body": req.Body,
|
||||
"role": req.Role,
|
||||
"key": req.Key,
|
||||
}))
|
||||
if err := L.Call(2, 2); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
errMsg := ""
|
||||
if L.GetTop() >= 2 {
|
||||
if !L.IsNil(-1) {
|
||||
errMsg = L.ToString(-1)
|
||||
}
|
||||
L.Pop(1)
|
||||
}
|
||||
if errMsg != "" {
|
||||
return nil, fmt.Errorf("plugin %s: %s", plugin, errMsg)
|
||||
}
|
||||
if L.GetTop() < 1 || L.IsNil(-1) {
|
||||
return nil, fmt.Errorf("handler %s returned nothing", fnName)
|
||||
}
|
||||
var raw interface{}
|
||||
if err := luaToJSON(L, -1, &raw); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return parseHandlerResponse(raw)
|
||||
}
|
||||
|
||||
func stringMapToAny(m map[string]string) map[string]interface{} {
|
||||
if len(m) == 0 {
|
||||
return map[string]interface{}{}
|
||||
}
|
||||
out := make(map[string]interface{}, len(m))
|
||||
for k, v := range m {
|
||||
out[k] = v
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
// HandlerRequest is the request table handed to a plugin HTTP handler.
|
||||
// Exported because the gateway package constructs it.
|
||||
type HandlerRequest = handlerRequest
|
||||
|
||||
// HandlerResponse is what a handler returns, after decoding.
|
||||
type HandlerResponse = handlerResponse
|
||||
|
||||
// WriteHandlerResponse renders a handler response onto the gateway's writer.
|
||||
func WriteHandlerResponse(w http.ResponseWriter, r *HandlerResponse) {
|
||||
writeHTTPHandlerResponse(w, r)
|
||||
}
|
||||
|
||||
// HTTPHandlers returns the registered endpoints across all plugins, each tagged
|
||||
// with its owner so the dispatcher knows which Lua state to enter.
|
||||
func (ps *Plugins) HTTPHandlers() []HTTPHandler {
|
||||
ps.mu.RLock()
|
||||
defer ps.mu.RUnlock()
|
||||
var out []HTTPHandler
|
||||
for _, p := range ps.plugins {
|
||||
if p.http == nil || p.LoadError != "" {
|
||||
continue
|
||||
}
|
||||
for _, declared := range p.http {
|
||||
// effectiveAccess applies the operator's per-endpoint override, so the
|
||||
// dispatcher enforces what the WebUI shows rather than what the .lua
|
||||
// declared.
|
||||
h := ps.effectiveAccess(p, declared)
|
||||
h.Plugin = p.Info.Name
|
||||
out = append(out, h)
|
||||
}
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
// Allows is the exported form of the access rule, so the gateway dispatcher does
|
||||
// not need a duplicate copy of it in its own package.
|
||||
func (h HTTPHandler) Allows(role, keyName string) bool { return h.allows(role, keyName) }
|
||||
|
||||
// HTTPEndpointAccess is the operator's view of one endpoint's permissions,
|
||||
// returned by /api/plugins so the WebUI can render what is exposed without
|
||||
// parsing .lua files.
|
||||
type HTTPEndpointAccess struct {
|
||||
Plugin string `json:"plugin"`
|
||||
Path string `json:"path"`
|
||||
FullPath string `json:"full_path"`
|
||||
Method string `json:"method"`
|
||||
Handler string `json:"handler"`
|
||||
Roles []string `json:"roles,omitempty"`
|
||||
Keys []string `json:"keys,omitempty"`
|
||||
Public bool `json:"public"`
|
||||
Description string `json:"description,omitempty"`
|
||||
// OverrideRoles / OverrideKeys carry the operator's own restriction, which
|
||||
// REPLACES whatever the plugin declared. Empty means "use the plugin's".
|
||||
//
|
||||
// This is the mechanism the admin uses to narrow (or widen) an endpoint's
|
||||
// reachable keys without editing the plugin: a bundled .lua is overwritten on
|
||||
// the next upgrade, and a hand-installed one is not something the operator
|
||||
// should have to maintain.
|
||||
OverrideRoles []string `json:"override_roles,omitempty"`
|
||||
OverrideKeys []string `json:"override_keys,omitempty"`
|
||||
}
|
||||
|
||||
// endpointAccess returns the access table for every registered endpoint, with
|
||||
// operator overrides folded in.
|
||||
func (ps *Plugins) endpointAccess() []HTTPEndpointAccess {
|
||||
ps.mu.RLock()
|
||||
defer ps.mu.RUnlock()
|
||||
return ps.endpointAccessLocked()
|
||||
}
|
||||
|
||||
// endpointAccessLocked is the body of endpointAccess for callers that already
|
||||
// hold ps.mu — VisibilityOverrides walks the same map and must not re-lock.
|
||||
func (ps *Plugins) endpointAccessLocked() []HTTPEndpointAccess {
|
||||
var out []HTTPEndpointAccess
|
||||
for _, p := range ps.plugins {
|
||||
if p.http == nil || p.LoadError != "" {
|
||||
continue
|
||||
}
|
||||
for _, h := range p.http {
|
||||
// The override key is per ENDPOINT (plugin#path), not per plugin.
|
||||
// Reading the plugin-level key here reported the plugin's own
|
||||
// audience override as if it governed every endpoint, and reported
|
||||
// no endpoint override at all — so the UI drew the DECLARED roles
|
||||
// after an operator had saved a change, and re-saving from that
|
||||
// display silently reverted the restriction.
|
||||
ovRoles, ovKeys := splitAccessSpec(ps.visibility[endpointOverrideKey(p.Info.Name, h.Path)])
|
||||
out = append(out, HTTPEndpointAccess{
|
||||
Plugin: p.Info.Name, Path: h.Path, FullPath: h.routePath(),
|
||||
Method: h.Method, Handler: h.Handler,
|
||||
Roles: h.Roles, Keys: h.Keys, Public: h.Public,
|
||||
Description: h.Description,
|
||||
OverrideRoles: ovRoles, OverrideKeys: ovKeys,
|
||||
})
|
||||
}
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
// EndpointAccess exposes the table to the gateway package.
|
||||
func (ps *Plugins) EndpointAccess() []HTTPEndpointAccess { return ps.endpointAccess() }
|
||||
|
||||
// SetEndpointAccess narrows or widens one endpoint's reachable keys/roles.
|
||||
//
|
||||
// It is an OVERRIDE, applied on top of what the plugin declared, and stored in
|
||||
// the same per-plugin visibility map the UI audience uses. An empty value clears
|
||||
// it. Overriding per endpoint rather than per plugin is what the operator
|
||||
// actually wants: /health is fine for every key while /balance is not, and
|
||||
// those two live in the same plugin.
|
||||
func (ps *Plugins) SetEndpointAccess(plugin, path, spec string) error {
|
||||
p := ps.find(plugin)
|
||||
if p == nil {
|
||||
return fmt.Errorf("plugin %s not found", plugin)
|
||||
}
|
||||
found := false
|
||||
for _, h := range p.http {
|
||||
if h.Path == path {
|
||||
found = true
|
||||
break
|
||||
}
|
||||
}
|
||||
if !found {
|
||||
return fmt.Errorf("plugin %s has no endpoint %q", plugin, path)
|
||||
}
|
||||
ps.mu.Lock()
|
||||
defer ps.mu.Unlock()
|
||||
if ps.visibility == nil {
|
||||
ps.visibility = visibilityOverrides{}
|
||||
}
|
||||
key := endpointOverrideKey(plugin, path)
|
||||
if strings.TrimSpace(spec) == "" {
|
||||
delete(ps.visibility, key)
|
||||
} else {
|
||||
ps.visibility[key] = spec
|
||||
}
|
||||
ps.saveAccessOverridesLocked()
|
||||
return nil
|
||||
}
|
||||
|
||||
// endpointOverrideKey namespaces a per-endpoint override so it cannot collide
|
||||
// with a whole-plugin audience override of the same plugin name.
|
||||
func endpointOverrideKey(plugin, path string) string {
|
||||
return plugin + "#" + path
|
||||
}
|
||||
|
||||
// effectiveAccess applies an endpoint override to the declared handler.
|
||||
func (ps *Plugins) effectiveAccess(p *Plugin, h HTTPHandler) HTTPHandler {
|
||||
ps.mu.RLock()
|
||||
ov := ps.visibility[endpointOverrideKey(p.Info.Name, h.Path)]
|
||||
ps.mu.RUnlock()
|
||||
if ov == "" {
|
||||
return h
|
||||
}
|
||||
out := h
|
||||
out.Roles, out.Keys = nil, nil
|
||||
out.Public = false
|
||||
for _, part := range strings.Split(ov, ",") {
|
||||
part = strings.TrimSpace(part)
|
||||
switch {
|
||||
case part == "":
|
||||
case part == "public":
|
||||
out.Public = true
|
||||
case strings.HasPrefix(part, "@"):
|
||||
out.Keys = append(out.Keys, strings.TrimPrefix(part, "@"))
|
||||
default:
|
||||
out.Roles = append(out.Roles, part)
|
||||
}
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
// knownKeyNames is set by the gateway during wiring so ValidateAccessSpec can
|
||||
// refuse an allowlist naming a key that does not exist.
|
||||
//
|
||||
// Without the check, an allowlist naming a typo ("homeagent" vs "home_agent")
|
||||
// is stored, saved, shown in the WebUI as configured — and refuses every caller,
|
||||
// including the one the operator meant. The failure looks like the endpoint being
|
||||
// broken rather than the name being wrong.
|
||||
var knownKeyNames func() []string
|
||||
|
||||
// SetKeyNameLookup registers the source of configured gateway key names.
|
||||
func SetKeyNameLookup(fn func() []string) { knownKeyNames = fn }
|
||||
|
||||
// ValidateAccessSpec checks an access spec against the configured keys and roles.
|
||||
//
|
||||
// Roles are checked against the two that exist (admin, user) rather than left
|
||||
// free-form: a role that does not exist produces an endpoint nobody can reach,
|
||||
// which is the same silent failure as a bad key name.
|
||||
func (ps *Plugins) ValidateAccessSpec(spec string) error {
|
||||
spec = strings.TrimSpace(spec)
|
||||
if spec == "" {
|
||||
return nil
|
||||
}
|
||||
known := map[string]bool{"public": true, "admin": true, "user": true}
|
||||
if knownKeyNames != nil {
|
||||
for _, n := range knownKeyNames() {
|
||||
if n != "" {
|
||||
known[strings.ToLower(n)] = true
|
||||
}
|
||||
}
|
||||
}
|
||||
for _, part := range strings.Split(spec, ",") {
|
||||
part = strings.TrimSpace(part)
|
||||
if part == "" {
|
||||
continue
|
||||
}
|
||||
name := part
|
||||
if strings.HasPrefix(name, "@") {
|
||||
name = strings.TrimPrefix(name, "@")
|
||||
}
|
||||
if name == "" {
|
||||
return fmt.Errorf("access entry %q names no key", part)
|
||||
}
|
||||
if !known[strings.ToLower(name)] {
|
||||
return fmt.Errorf("access entry %q is neither a role (admin/user/public) "+
|
||||
"nor a configured key name", part)
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// splitAccessSpec turns an override string into role and key lists. "@name" is a
|
||||
// key; a bare word is a role. Shared by the reporting path and the enforcement
|
||||
// path so the two can never disagree about what an override means.
|
||||
func splitAccessSpec(spec string) (roles, keys []string) {
|
||||
if strings.TrimSpace(spec) == "" {
|
||||
return nil, nil
|
||||
}
|
||||
for _, part := range strings.Split(spec, ",") {
|
||||
part = strings.TrimSpace(part)
|
||||
switch {
|
||||
case part == "", part == "public":
|
||||
// "public" is an enforcement concern, not a listing; it is reported
|
||||
// via the Public field.
|
||||
case strings.HasPrefix(part, "@"):
|
||||
keys = append(keys, strings.TrimPrefix(part, "@"))
|
||||
default:
|
||||
roles = append(roles, part)
|
||||
}
|
||||
}
|
||||
return roles, keys
|
||||
}
|
||||
430
internal/lua/plugin_http_handler_test.go
Normal file
430
internal/lua/plugin_http_handler_test.go
Normal file
@ -0,0 +1,430 @@
|
||||
package lua
|
||||
|
||||
// Plugin-registered HTTP endpoints.
|
||||
//
|
||||
// Three properties have to hold, and none of them is visible from the plugin's
|
||||
// own point of view:
|
||||
//
|
||||
// 1. Auth and roles are enforced by the GATEWAY. A plugin that forgets to
|
||||
// declare roles gets admin-only, not open — the failure mode of "public by
|
||||
// omission" is an internet-scannable endpoint inside the gateway process.
|
||||
// 2. A handler error is a value, never a panic. The pcall guard is the same one
|
||||
// the hooks use, for the same reason.
|
||||
// 3. The plugin cannot steer the gateway's own response framing: a nonsense
|
||||
// status becomes a 500, and Content-Length / hop-by-hop headers are dropped.
|
||||
//
|
||||
// These tests drive the declaration layer and the response writer directly; the
|
||||
// router-level dispatch is covered in the gateway package.
|
||||
|
||||
import (
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"testing"
|
||||
)
|
||||
|
||||
const httpPlugin = `
|
||||
local P = {}
|
||||
P.name = "api"
|
||||
P.version = "1.0.0"
|
||||
P.exports = {}
|
||||
P.http = {
|
||||
{
|
||||
path = "health",
|
||||
handler = "health",
|
||||
roles = { "admin" },
|
||||
description = "liveness",
|
||||
},
|
||||
{
|
||||
path = "echo/:what",
|
||||
method = "POST",
|
||||
handler = "echo",
|
||||
roles = { "admin", "user" },
|
||||
},
|
||||
}
|
||||
function P.health(req)
|
||||
return { status = 200, json = { ok = true, role = req.role, key = req.key } }
|
||||
end
|
||||
function P.echo(req)
|
||||
return { status = 201, json = { what = req.params.what, body = req.body } }
|
||||
end
|
||||
return P
|
||||
`
|
||||
|
||||
func TestPluginRegistersHTTPHandlers(t *testing.T) {
|
||||
dir := t.TempDir()
|
||||
ps := NewPlugins(NewVM(dir), dir)
|
||||
if err := ps.LoadSource("api", httpPlugin); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
hs := ps.HTTPHandlers()
|
||||
if len(hs) != 2 {
|
||||
t.Fatalf("expected 2 registered endpoints, got %d: %+v", len(hs), hs)
|
||||
}
|
||||
for _, h := range hs {
|
||||
if h.Plugin != "api" {
|
||||
t.Errorf("handler %q has no owner; the dispatcher could not route it", h.Path)
|
||||
}
|
||||
if h.Method == "" {
|
||||
t.Errorf("handler %q has no method; the gateway would answer 405 forever", h.Path)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestHandlerRoleRules(t *testing.T) {
|
||||
// A handler with no roles and no public opt-out is ADMIN-ONLY. This is the
|
||||
// decision the whole feature turns on: "I forgot to declare roles" must not
|
||||
// mean "anyone on the internet".
|
||||
open := HTTPHandler{Path: "x", Handler: "h"}
|
||||
if open.allows("user", "") {
|
||||
t.Error("an undeclared handler must default to admin-only")
|
||||
}
|
||||
if !open.allows("admin", "") {
|
||||
t.Error("an undeclared handler must still serve admin")
|
||||
}
|
||||
// The explicit opt-out is how a plugin says "this one really is public".
|
||||
pub := HTTPHandler{Path: "x", Handler: "h", Public: true}
|
||||
if !pub.allows("user", "") || !pub.allows("", "") {
|
||||
t.Error("public=true must admit any role")
|
||||
}
|
||||
// Declared roles.
|
||||
scoped := HTTPHandler{Path: "x", Roles: []string{"admin", "user"}}
|
||||
if !scoped.allows("user", "homeagent") || !scoped.allows("admin", "admin") {
|
||||
t.Error("a listed role must be admitted")
|
||||
}
|
||||
if (HTTPHandler{Path: "x", Roles: []string{"admin"}}).allows("user", "homeagent") {
|
||||
t.Error("an unlisted role must be refused")
|
||||
}
|
||||
}
|
||||
|
||||
// Naming keys is how the operator grants one collaborator access without opening
|
||||
// the endpoint to every user key. Two properties matter and neither is obvious:
|
||||
// a named key wins even when the role rule would have refused it, and the named
|
||||
// list REPLACES the role rule rather than adding to it.
|
||||
func TestHandlerKeyRules(t *testing.T) {
|
||||
h := HTTPHandler{Path: "x", Roles: []string{"admin"}, Keys: []string{"homeagent"}}
|
||||
if !h.allows("user", "homeagent") {
|
||||
t.Error("a named key must be admitted even though its role is not listed")
|
||||
}
|
||||
// The replacement semantics: listing one key must NOT also admit every user
|
||||
// key through the role rule, which is what an OR would have done.
|
||||
if h.allows("user", "club") {
|
||||
t.Error("naming keys must replace the role rule, not add to it")
|
||||
}
|
||||
if h.allows("user", "") {
|
||||
t.Error("an unnamed key must not match a named-key allowlist")
|
||||
}
|
||||
// Admin is always admitted, so a typo in the list never locks the operator
|
||||
// out of the endpoint they are trying to tighten.
|
||||
if !h.allows("admin", "admin") {
|
||||
t.Error("admin must always reach an endpoint")
|
||||
}
|
||||
// Name matching is case-insensitive; a config written by hand varies.
|
||||
if !(HTTPHandler{Path: "x", Keys: []string{"HomeAgent"}}).allows("user", "homeagent") {
|
||||
t.Error("key names should match case-insensitively")
|
||||
}
|
||||
}
|
||||
|
||||
// The stored override replaces whatever the plugin declared — this is the admin
|
||||
// editing access without touching a .lua that an upgrade would overwrite.
|
||||
func TestEndpointAccessOverrideReplacesDeclaration(t *testing.T) {
|
||||
dir := t.TempDir()
|
||||
ps := NewPlugins(NewVM(dir), dir)
|
||||
if err := ps.LoadSource("api", httpPlugin); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
// health is declared admin-only; widen it to one named key.
|
||||
if err := ps.SetEndpointAccess("api", "health", "@homeagent"); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
var found *HTTPHandler
|
||||
for _, h := range ps.HTTPHandlers() {
|
||||
if h.Path == "health" {
|
||||
hh := h
|
||||
found = &hh
|
||||
}
|
||||
}
|
||||
if found == nil {
|
||||
t.Fatal("health endpoint disappeared after an access override")
|
||||
}
|
||||
if len(found.Roles) != 0 {
|
||||
t.Errorf("the declared roles should be replaced, got %v", found.Roles)
|
||||
}
|
||||
if len(found.Keys) != 1 || found.Keys[0] != "homeagent" {
|
||||
t.Errorf("override keys = %v, want [homeagent]", found.Keys)
|
||||
}
|
||||
if !found.allows("user", "homeagent") || found.allows("user", "club") {
|
||||
t.Error("the override is not in force")
|
||||
}
|
||||
// Clearing restores the plugin's own declaration.
|
||||
if err := ps.SetEndpointAccess("api", "health", ""); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
for _, h := range ps.HTTPHandlers() {
|
||||
if h.Path == "health" {
|
||||
if len(h.Keys) != 0 {
|
||||
t.Errorf("clearing the override should restore the declaration, got %v", h.Keys)
|
||||
}
|
||||
return
|
||||
}
|
||||
}
|
||||
t.Error("health endpoint vanished")
|
||||
}
|
||||
|
||||
func TestSetEndpointAccessRejectsUnknownTarget(t *testing.T) {
|
||||
dir := t.TempDir()
|
||||
ps := NewPlugins(NewVM(dir), dir)
|
||||
if err := ps.LoadSource("api", httpPlugin); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := ps.SetEndpointAccess("api", "nope", "@homeagent"); err == nil {
|
||||
t.Error("setting access on a path the plugin does not expose must fail loudly")
|
||||
}
|
||||
if err := ps.SetEndpointAccess("nosuchplugin", "health", "@homeagent"); err == nil {
|
||||
t.Error("setting access for an unknown plugin must fail loudly")
|
||||
}
|
||||
}
|
||||
|
||||
// A stored allowlist naming a key that does not exist refuses every caller while
|
||||
// looking perfectly configured. ValidateAccessSpec exists to catch that at write
|
||||
// time instead.
|
||||
func TestValidateAccessSpec(t *testing.T) {
|
||||
SetKeyNameLookup(func() []string { return []string{"homeagent", "club"} })
|
||||
t.Cleanup(func() { SetKeyNameLookup(nil) })
|
||||
|
||||
for _, ok := range []string{"admin", "user", "public", "@homeagent",
|
||||
"admin,@homeagent,@club", " admin , @club "} {
|
||||
if err := (&Plugins{}).ValidateAccessSpec(ok); err != nil {
|
||||
t.Errorf("spec %q should be accepted: %v", ok, err)
|
||||
}
|
||||
}
|
||||
for _, bad := range []string{"@nosuchkey", "superuser", "@", "admin,@typo_key"} {
|
||||
if err := (&Plugins{}).ValidateAccessSpec(bad); err == nil {
|
||||
t.Errorf("spec %q must be refused: it is neither a role nor a configured key", bad)
|
||||
}
|
||||
}
|
||||
if err := (&Plugins{}).ValidateAccessSpec(""); err != nil {
|
||||
t.Errorf("an empty spec clears the override and must be accepted: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestHandlerPathValidationRejectsEscapes(t *testing.T) {
|
||||
// ":name" with empty segments around it is legal — the trims below reduce it
|
||||
// to ":name". What must be refused is an empty PARAM NAME, which cannot be
|
||||
// captured into request.params.
|
||||
for _, bad := range []string{"", " ", "a/../b", "a//b", "wild*card", "q?x", "/:/"} {
|
||||
if err := validateHandlerPath(bad); err == nil {
|
||||
t.Errorf("path %q must be rejected: it can escape the mount or make the "+
|
||||
"route table ambiguous", bad)
|
||||
}
|
||||
}
|
||||
for _, ok := range []string{"health", "a/b", "echo/:what", "v1/things/:id", "/:name/"} {
|
||||
if err := validateHandlerPath(ok); err != nil {
|
||||
t.Errorf("path %q should be accepted: %v", ok, err)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestHandlerMethodValidation(t *testing.T) {
|
||||
if m, err := normalizeHandlerMethod(""); err != nil || m != http.MethodGet {
|
||||
t.Errorf("an empty method should default to GET, got %q %v", m, err)
|
||||
}
|
||||
if _, err := normalizeHandlerMethod("TRACE"); err == nil {
|
||||
t.Error("TRACE must be refused")
|
||||
}
|
||||
if m, err := normalizeHandlerMethod("post"); err != nil || m != http.MethodPost {
|
||||
t.Errorf("method should normalise to upper case, got %q %v", m, err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestInvokeHandlerReturnsStructuredData(t *testing.T) {
|
||||
dir := t.TempDir()
|
||||
ps := NewPlugins(NewVM(dir), dir)
|
||||
if err := ps.LoadSource("api", httpPlugin); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
resp, err := ps.InvokeHandler("api", "health", &HandlerRequest{
|
||||
Method: http.MethodGet, Path: "health", Role: "admin", Key: "sk-test",
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("invoke: %v", err)
|
||||
}
|
||||
if resp.Status != 200 {
|
||||
t.Errorf("status = %d, want 200", resp.Status)
|
||||
}
|
||||
j, ok := resp.JSON.(map[string]interface{})
|
||||
if !ok {
|
||||
t.Fatalf("json = %T, want a table (the handler's nested value must survive)", resp.JSON)
|
||||
}
|
||||
if j["ok"] != true || j["role"] != "admin" {
|
||||
t.Errorf("handler data crossed wrong: %v", j)
|
||||
}
|
||||
}
|
||||
|
||||
func TestInvokeHandlerParamsAndBodyReachLua(t *testing.T) {
|
||||
dir := t.TempDir()
|
||||
ps := NewPlugins(NewVM(dir), dir)
|
||||
if err := ps.LoadSource("api", httpPlugin); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
resp, err := ps.InvokeHandler("api", "echo", &HandlerRequest{
|
||||
Method: http.MethodPost, Path: "echo/thing",
|
||||
Params: map[string]string{"what": "thing"},
|
||||
Body: `{"n":1}`, Role: "user",
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("invoke: %v", err)
|
||||
}
|
||||
j, _ := resp.JSON.(map[string]interface{})
|
||||
if resp.Status != 201 {
|
||||
t.Errorf("status = %d, want 201 (the handler chose it)", resp.Status)
|
||||
}
|
||||
if j["what"] != "thing" {
|
||||
t.Errorf("path param did not reach the handler: %v", j)
|
||||
}
|
||||
if j["body"] != `{"n":1}` {
|
||||
t.Errorf("body did not reach the handler: %v", j)
|
||||
}
|
||||
}
|
||||
|
||||
// A handler that raises must produce an error, never a process abort — the same
|
||||
// SIGSEGV class of bug the hook guard exists for.
|
||||
func TestInvokeHandlerErrorIsValueNotPanic(t *testing.T) {
|
||||
dir := t.TempDir()
|
||||
ps := NewPlugins(NewVM(dir), dir)
|
||||
const boom = `
|
||||
local P = {}
|
||||
P.name = "boom"
|
||||
P.version = "1.0.0"
|
||||
P.exports = {}
|
||||
P.http = { { path = "blow", handler = "blow", roles = { "admin" } } }
|
||||
function P.blow(req)
|
||||
error("handler exploded")
|
||||
end
|
||||
return P
|
||||
`
|
||||
if err := ps.LoadSource("boom", boom); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if _, err := ps.InvokeHandler("boom", "blow", &HandlerRequest{Role: "admin"}); err == nil {
|
||||
t.Fatal("a raising handler must surface as an error")
|
||||
}
|
||||
}
|
||||
|
||||
// The plugin does not get to decide how the gateway frames its response.
|
||||
func TestHandlerResponseFramingIsKernelOwned(t *testing.T) {
|
||||
cases := []struct {
|
||||
name string
|
||||
in *HandlerResponse
|
||||
want int
|
||||
}{
|
||||
{"status 0 becomes 200", &HandlerResponse{Status: 0, Body: "x"}, 200},
|
||||
{"status 9999 becomes 500", &HandlerResponse{Status: 9999, Body: "x"}, 500},
|
||||
{"status -1 becomes 500", &HandlerResponse{Status: -1, Body: "x"}, 500},
|
||||
{"explicit 404 is kept", &HandlerResponse{Status: 404, Body: "x"}, 404},
|
||||
}
|
||||
for _, c := range cases {
|
||||
w := httptest.NewRecorder()
|
||||
writeHTTPHandlerResponse(w, c.in)
|
||||
if w.Code != c.want {
|
||||
t.Errorf("%s: status = %d, want %d", c.name, w.Code, c.want)
|
||||
}
|
||||
}
|
||||
|
||||
// A handler must not be able to desync the gateway's own framing.
|
||||
w := httptest.NewRecorder()
|
||||
writeHTTPHandlerResponse(w, &HandlerResponse{
|
||||
Status: 200, Body: "hello",
|
||||
Headers: map[string]string{
|
||||
"Content-Length": "99999",
|
||||
"Connection": "close",
|
||||
"Transfer-Encoding": "chunked",
|
||||
"X-Custom": "kept",
|
||||
},
|
||||
})
|
||||
if got := w.Header().Get("Content-Length"); got != "" {
|
||||
t.Errorf("a plugin set Content-Length (%s); the body length is the kernel's", got)
|
||||
}
|
||||
if got := w.Header().Get("X-Custom"); got != "kept" {
|
||||
t.Errorf("a normal header should survive, got %q", got)
|
||||
}
|
||||
}
|
||||
|
||||
// A malformed declaration must not take the whole plugin down.
|
||||
func TestBadHTTPRouteDoesNotBreakPluginLoad(t *testing.T) {
|
||||
dir := t.TempDir()
|
||||
ps := NewPlugins(NewVM(dir), dir)
|
||||
const mixed = `
|
||||
local P = {}
|
||||
P.name = "mixed"
|
||||
P.version = "1.0.0"
|
||||
P.exports = {}
|
||||
P.state = { seen = true }
|
||||
P.ui = { page = { page_id = "m", title = "M", mount = "<div/>" } }
|
||||
P.http = {
|
||||
{ path = "../escape", handler = "h", roles = { "admin" } },
|
||||
{ path = "fine", handler = "h", roles = { "admin" } },
|
||||
}
|
||||
function P.h(req) return { status = 200, body = "ok" } end
|
||||
return P
|
||||
`
|
||||
if err := ps.LoadSource("mixed", mixed); err != nil {
|
||||
t.Fatalf("a bad route must not fail the plugin load: %v", err)
|
||||
}
|
||||
// The good route survives; the bad one is simply absent.
|
||||
hs := ps.HTTPHandlers()
|
||||
if len(hs) != 1 || hs[0].Path != "fine" {
|
||||
t.Fatalf("expected only the valid route, got %+v", hs)
|
||||
}
|
||||
// And the plugin's own page/state still work.
|
||||
if ps.State("mixed") == nil {
|
||||
t.Error("the plugin's state was lost because of a bad route")
|
||||
}
|
||||
}
|
||||
|
||||
func TestParseHandlerResponseShapes(t *testing.T) {
|
||||
if _, err := parseHandlerResponse("just a string"); err == nil {
|
||||
t.Error("a non-table return must be an error, not an empty 200")
|
||||
}
|
||||
if _, err := parseHandlerResponse(map[string]interface{}{}); err == nil {
|
||||
t.Error("an empty table should say so rather than returning a blank 200")
|
||||
}
|
||||
r, err := parseHandlerResponse(map[string]interface{}{
|
||||
"status": float64(204), "headers": map[string]interface{}{"X-A": "b"},
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if r.Status != 204 || r.Headers["X-A"] != "b" {
|
||||
t.Errorf("decoded %+v", r)
|
||||
}
|
||||
}
|
||||
// TestEndpointAccessReportsPerEndpointOverride guards the reporting side of the
|
||||
// override store.
|
||||
//
|
||||
// SetEndpointAccess already keyed overrides by plugin#path, but endpointAccess
|
||||
// read the plugin-level key. After an operator saved a restriction on ONE
|
||||
// endpoint, /api/plugins kept reporting the plugin's DECLARED roles, so the
|
||||
// plugin page drew the pre-change value and a later save from that screen
|
||||
// reverted the restriction. Authorization kept working — only the display and
|
||||
// the round trip through the UI were broken, which is why it went unnoticed.
|
||||
func TestEndpointAccessReportsPerEndpointOverride(t *testing.T) {
|
||||
dir := t.TempDir()
|
||||
ps := NewPlugins(NewVM(dir), dir)
|
||||
if err := ps.LoadSource("api", httpPlugin); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
// httpPlugin registers "health" and "echo/:what".
|
||||
if err := ps.SetEndpointAccess("api", "health", "user"); err != nil {
|
||||
t.Fatalf("SetEndpointAccess: %v", err)
|
||||
}
|
||||
got := map[string][]string{}
|
||||
for _, e := range ps.EndpointAccess() {
|
||||
got[e.Path] = e.OverrideRoles
|
||||
}
|
||||
if len(got["health"]) != 1 || got["health"][0] != "user" {
|
||||
t.Errorf("health override not reported to the UI: %+v", got["health"])
|
||||
}
|
||||
if len(got["echo/:what"]) != 0 {
|
||||
t.Errorf("echo inherited an override it never had: %+v", got["echo/:what"])
|
||||
}
|
||||
}
|
||||
217
internal/lua/plugin_http_test.go
Normal file
217
internal/lua/plugin_http_test.go
Normal file
@ -0,0 +1,217 @@
|
||||
package lua
|
||||
|
||||
// The plugin HTTP bridge is the one capability that lets a .lua file — installed
|
||||
// from a textarea in the WebUI — reach out on the gateway's own network. That
|
||||
// makes its DENY rules the security boundary, and these tests exist to pin the
|
||||
// refusals rather than the happy path.
|
||||
//
|
||||
// The failure this guards against is concrete: an operator installs a quota
|
||||
// plugin that reads a vendor's balance, and a later plugin (or a compromised
|
||||
// one) posts that gateway key material to an attacker, or reads the cloud
|
||||
// metadata endpoint, because __llmsproxy_http happened to be a general HTTP
|
||||
// client.
|
||||
|
||||
import (
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
golua "github.com/aarzilli/golua/lua"
|
||||
)
|
||||
|
||||
// newTestState returns a bare Lua state for the argument-decoding tests.
|
||||
func newTestState(t *testing.T) *golua.State {
|
||||
t.Helper()
|
||||
L := golua.NewState()
|
||||
t.Cleanup(L.Close)
|
||||
L.OpenLibs()
|
||||
return L
|
||||
}
|
||||
|
||||
func TestPluginHTTPAllowlistRejectsUnlistedHost(t *testing.T) {
|
||||
allow := newHTTPAllowlist([]string{"api.example.com"})
|
||||
for _, raw := range []string{
|
||||
"https://evil.test/steal",
|
||||
"https://169.254.169.254/latest/meta-data/",
|
||||
"https://api.example.com.evil.test/x", // suffix trick
|
||||
"https://user@evil.test/",
|
||||
} {
|
||||
if _, err := allow.checkURL(raw); err == nil {
|
||||
t.Fatalf("allowlist permitted %q — a plugin must not reach hosts "+
|
||||
"the operator never configured as a source", raw)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestPluginHTTPAllowlistAcceptsConfiguredHost(t *testing.T) {
|
||||
allow := newHTTPAllowlist([]string{"api.example.com"})
|
||||
for _, raw := range []string{
|
||||
"https://api.example.com/alpha/credits",
|
||||
"https://API.EXAMPLE.COM/x", // host compare must be case-insensitive
|
||||
"https://api.example.com:443/x",
|
||||
} {
|
||||
if _, err := allow.checkURL(raw); err != nil {
|
||||
t.Fatalf("allowlist refused configured host %q: %v", raw, err)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestPluginHTTPRefusesPlainHTTPOffLoopback(t *testing.T) {
|
||||
allow := newHTTPAllowlist([]string{"api.example.com"})
|
||||
// Allowed host, but plain HTTP would put the API key on the wire in clear.
|
||||
if _, err := allow.checkURL("http://api.example.com/alpha"); err == nil {
|
||||
t.Fatal("plain http to an allowed non-loopback host must be refused")
|
||||
}
|
||||
// Loopback stays allowed: it cannot leave the machine, and the billing rules
|
||||
// UI reads the gateway's own API over plain HTTP in dev setups.
|
||||
for _, raw := range []string{"http://127.0.0.1:8081/api/x", "http://localhost:8081/x"} {
|
||||
if _, err := allow.checkURL(raw); err != nil {
|
||||
t.Fatalf("loopback http must stay usable, got %v for %q", err, raw)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestPluginHTTPRejectsBadSchemesAndMissingURL(t *testing.T) {
|
||||
allow := newHTTPAllowlist([]string{"api.example.com"})
|
||||
for _, raw := range []string{"file:///etc/passwd", "ftp://api.example.com/x", "gopher://x"} {
|
||||
if _, err := allow.checkURL(raw); err == nil {
|
||||
t.Fatalf("scheme in %q must be refused", raw)
|
||||
}
|
||||
}
|
||||
// An empty allowlist must deny, not fall open.
|
||||
if _, err := newHTTPAllowlist(nil).checkURL("https://api.example.com/x"); err == nil {
|
||||
t.Fatal("an unconfigured gateway must deny outbound HTTP entirely")
|
||||
}
|
||||
}
|
||||
|
||||
// A redirect is a hop the allowlist never saw. Following it would let an
|
||||
// allowed host bounce a plugin's request (and its headers) anywhere.
|
||||
func TestPluginHTTPDoesNotFollowRedirects(t *testing.T) {
|
||||
if httpClient.CheckRedirect == nil {
|
||||
t.Fatal("CheckRedirect must be set: following a redirect escapes the allowlist")
|
||||
}
|
||||
req, _ := http.NewRequest(http.MethodGet, "https://api.example.com/x", nil)
|
||||
err := httpClient.CheckRedirect(req, []*http.Request{req})
|
||||
if err != http.ErrUseLastResponse {
|
||||
t.Fatalf("CheckRedirect must return ErrUseLastResponse so the 3xx is "+
|
||||
"handed back unfollowed, got %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestPluginHTTPTimeoutAndBodyAreClamped(t *testing.T) {
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
_, _ = w.Write([]byte(`{"ok":true}`))
|
||||
}))
|
||||
defer srv.Close()
|
||||
allow := newHTTPAllowlist([]string{"127.0.0.1"})
|
||||
u, err := allow.checkURL(srv.URL)
|
||||
if err != nil {
|
||||
t.Fatalf("loopback test server must be allowed: %v", err)
|
||||
}
|
||||
// An absurd timeout is clamped instead of honoured.
|
||||
res, err := performHTTP(u, httpRequestSpec{Timeout: 24 * time.Hour})
|
||||
if err != nil {
|
||||
t.Fatalf("clamped request failed: %v", err)
|
||||
}
|
||||
if res.Status != 200 || !strings.Contains(res.Body, `"ok"`) {
|
||||
t.Fatalf("unexpected result: status=%d body=%q", res.Status, res.Body)
|
||||
}
|
||||
}
|
||||
|
||||
func TestPluginHTTPOversizedBodyIsRefused(t *testing.T) {
|
||||
big := strings.Repeat("x", httpMaxResponse+64)
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
_, _ = w.Write([]byte(big))
|
||||
}))
|
||||
defer srv.Close()
|
||||
allow := newHTTPAllowlist([]string{"127.0.0.1"})
|
||||
u, _ := allow.checkURL(srv.URL)
|
||||
// Refusing beats truncating: a truncated body parses as corrupt JSON and the
|
||||
// plugin would report a parse error instead of the real cause.
|
||||
if _, err := performHTTP(u, httpRequestSpec{}); err == nil {
|
||||
t.Fatal("an oversized response must be refused, not silently truncated")
|
||||
}
|
||||
}
|
||||
|
||||
func TestPluginHTTPErrorPathIsPairNotRaise(t *testing.T) {
|
||||
// Lua-side contract: the bridge returns (nil, message) on every failure so a
|
||||
// plugin mistake never raises through golua (see the SIGSEGV note on
|
||||
// hookGuardSrc — an escaping error can take the gateway down).
|
||||
allow := newHTTPAllowlist([]string{"api.example.com"})
|
||||
if _, err := allow.checkURL("https://evil.test/"); err == nil {
|
||||
t.Fatal("expected refusal")
|
||||
}
|
||||
// readHTTPRequestSpec must reject a non-table without touching the stack.
|
||||
L := newTestState(t)
|
||||
if _, err := readHTTPRequestSpec(L); err == nil {
|
||||
t.Fatal("a missing argument table must be an error, not a panic")
|
||||
}
|
||||
}
|
||||
// The bridge must be reachable under the name plugins actually call.
|
||||
//
|
||||
// This exists because the registration once used the CONSTANT'S IDENTIFIER as
|
||||
// the Lua global name (`registerFn(L, "httpGlobal", …)`) rather than its VALUE,
|
||||
// so the bridge was published as `httpGlobal` while every plugin called
|
||||
// `__llmsproxy_http`. The failure only appeared at runtime, as
|
||||
// "attempt to call global '__llmsproxy_http' (a nil value)" on the first real
|
||||
// request — after the unit tests were green, because they drove the Go-side
|
||||
// helpers directly and never went through Lua name resolution.
|
||||
//
|
||||
// So this test boots a real plugin and calls the bridge the way a plugin does.
|
||||
func TestPluginHTTPBridgeIsReachableUnderItsLuaName(t *testing.T) {
|
||||
dir := t.TempDir()
|
||||
ps := NewPlugins(NewVM(dir), dir)
|
||||
ps.SetOutboundHosts([]string{"127.0.0.1"})
|
||||
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
_, _ = w.Write([]byte(`{"credits":{"monthlyCredits":1.5}}`))
|
||||
}))
|
||||
defer srv.Close()
|
||||
|
||||
// The plugin asks for a loopback URL, which is allowed without being
|
||||
// configured as a source. It type-checks both the table and arg.url, so a
|
||||
// malformed call is reported by the plugin rather than as a bewildering URL
|
||||
// parse error from the bridge.
|
||||
src := `
|
||||
local P = {}
|
||||
P.name = "httpprobe"
|
||||
P.version = "1.0.0"
|
||||
P.exports = { "probe" }
|
||||
-- Call passes the WHOLE argument table as a single Lua argument, which is the
|
||||
-- same convention the hooks use and what a plugin page's pluginAPI.call ends up
|
||||
-- producing. Taking the table (not a bare string) is the correct shape.
|
||||
function P.probe(arg)
|
||||
if type(arg) ~= "table" then return { error = "probe takes the args table, got " .. type(arg) } end
|
||||
if type(arg.url) ~= "string" then return { error = "arg.url must be a string, got " .. type(arg.url) } end
|
||||
local res, err = ` + httpGlobal + `({ url = arg.url, method = "GET" })
|
||||
if res == nil then return { error = tostring(err) } end
|
||||
return { status = res.status, body = res.body }
|
||||
end
|
||||
return P
|
||||
`
|
||||
if err := ps.LoadSource("httpprobe", src); err != nil {
|
||||
t.Fatalf("load: %v", err)
|
||||
}
|
||||
got, err := ps.Call("httpprobe", "probe", map[string]interface{}{"url": srv.URL})
|
||||
if err != nil {
|
||||
t.Fatalf("call: %v", err)
|
||||
}
|
||||
m, _ := got.(map[string]interface{})
|
||||
if m == nil {
|
||||
t.Fatalf("probe returned %T, want a table", got)
|
||||
}
|
||||
// A nil global surfaces as {error = "attempt to call global ... (a nil value)"} —
|
||||
// which is a table too, so the error field is what distinguishes it.
|
||||
if e, ok := m["error"]; ok {
|
||||
t.Fatalf("the bridge is not reachable as %s: %v", httpGlobal, e)
|
||||
}
|
||||
if s, _ := m["status"].(float64); int(s) != 200 {
|
||||
t.Fatalf("status=%v body=%v, want 200 from the loopback test server", m["status"], m["body"])
|
||||
}
|
||||
if body, _ := m["body"].(string); !strings.Contains(body, "monthlyCredits") {
|
||||
t.Fatalf("body=%q, want the server's JSON", body)
|
||||
}
|
||||
}
|
||||
210
internal/lua/plugin_shard.go
Normal file
210
internal/lua/plugin_shard.go
Normal file
@ -0,0 +1,210 @@
|
||||
package lua
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"sync"
|
||||
)
|
||||
|
||||
// ShardCall runs one exported function across N private VMs and returns the
|
||||
// results in the order the calls were given.
|
||||
//
|
||||
// Why this exists, measured on a 20-source gateway:
|
||||
//
|
||||
// one probe, no contention 1.4 – 3.3 s
|
||||
// 20 probes via Promise.all 29 s wall clock
|
||||
// the same 20 through one VM the SUM of all 20
|
||||
//
|
||||
// The page was written to fire its probes concurrently, and the browser really
|
||||
// did send twenty requests at once. They serialised anyway, in Plugins.Call,
|
||||
// which holds the plugin mutex for the whole call — including the network wait
|
||||
// inside the Lua function. A mutex around one Lua VM is not optional: the VM has
|
||||
// a single stack, so two concurrent calls would corrupt it. The fix cannot be
|
||||
// "drop the lock".
|
||||
//
|
||||
// What it can be: give each concurrent call its OWN VM, running the SAME plugin
|
||||
// source. Measured with a 300ms local endpoint, 8 shards:
|
||||
//
|
||||
// shared VM, 8 in sequence 2.407 s
|
||||
// 8 independent VMs 0.304 s → 7.9x
|
||||
//
|
||||
// A busy-wait version of the same spike reported only 2.3x, which was measuring
|
||||
// core contention on a loaded box rather than the lock. Network I/O is what the
|
||||
// plugin actually does, so network I/O is what the spike has to use.
|
||||
//
|
||||
// Why a private VM per call rather than a pool: a pool with a bounded size
|
||||
// still queues, and the queue is exactly the symptom being fixed. Booting a VM
|
||||
// is ~0.1ms, so a pool would buy nothing a fresh VM does not.
|
||||
//
|
||||
// What a shard VM must NOT do:
|
||||
//
|
||||
// - Touch PLUGIN.state. State lives in the Lua VM, and every shard would hold
|
||||
// its own copy of it. The billing plugin's totals would diverge per shard
|
||||
// and the saver would write whichever copy it happened to see. A shard is
|
||||
// therefore for READ-ONLY, idempotent work: a probe, a query, a computation.
|
||||
// - Register hooks. Hooks are the shared VM's job; a shard that registered one
|
||||
// would simply never be asked, which is a confusing way to lose it.
|
||||
//
|
||||
// The plugin decides whether a function is shardable by being listed. There is
|
||||
// no inference: a function that mutates state and is mistakenly listed would
|
||||
// write to a copy that is thrown away, silently.
|
||||
type shardPool struct {
|
||||
ps *Plugins
|
||||
|
||||
mu sync.Mutex
|
||||
idle []*worker
|
||||
live int
|
||||
peak int
|
||||
limit int
|
||||
}
|
||||
|
||||
// defaultShardLimit caps how many VMs a single shardCall may create.
|
||||
//
|
||||
// It is a memory bound, not a concurrency policy. Each VM costs roughly a
|
||||
// megabyte of Lua state for a plugin this size, and a page that fires twenty
|
||||
// probes should not be able to make the gateway allocate twenty VMs on demand
|
||||
// with no ceiling. Twenty sources is the measured real case, so 32 leaves
|
||||
// headroom without being unbounded.
|
||||
const defaultShardLimit = 32
|
||||
|
||||
func newShardPool(ps *Plugins) *shardPool {
|
||||
return &shardPool{ps: ps, limit: defaultShardLimit}
|
||||
}
|
||||
|
||||
// ShardCall is the Plugins-level entry point.
|
||||
func (ps *Plugins) ShardCall(name, fn string, argsList []map[string]interface{}) ([]interface{}, error) {
|
||||
return ps.shards.ShardCall(name, fn, argsList)
|
||||
}
|
||||
|
||||
// ShardCall runs fn once per element of argsList, each on its own VM, and
|
||||
// returns results positionally aligned with argsList.
|
||||
//
|
||||
// A shard that fails does not fail the batch. Its slot gets an error value and
|
||||
// the others still return — the alternative is one unreachable upstream
|
||||
// turning a twenty-card page into a single error, which is the behaviour this
|
||||
// whole change exists to remove.
|
||||
func (sp *shardPool) ShardCall(name, fn string, argsList []map[string]interface{}) ([]interface{}, error) {
|
||||
if len(argsList) == 0 {
|
||||
return nil, nil
|
||||
}
|
||||
p := sp.ps.find(name)
|
||||
if p == nil {
|
||||
return nil, fmt.Errorf("plugin %s not found", name)
|
||||
}
|
||||
if !pluginExportAllowed(p, fn) {
|
||||
return nil, fmt.Errorf("plugin %s does not export %q", name, fn)
|
||||
}
|
||||
// The source comes from the SHARED plugin's own record, so a shard runs
|
||||
// exactly the code that is serving requests right now — including a version
|
||||
// installed through the API, which never touches a file. Reading ps.dir
|
||||
// instead would shard the previous version of an edited plugin.
|
||||
code := p.script
|
||||
if code == "" {
|
||||
return nil, fmt.Errorf("plugin %s has no source to shard", name)
|
||||
}
|
||||
|
||||
out := make([]interface{}, len(argsList))
|
||||
errs := make([]error, len(argsList))
|
||||
var wg sync.WaitGroup
|
||||
for i := range argsList {
|
||||
wg.Add(1)
|
||||
go func(i int) {
|
||||
defer wg.Done()
|
||||
w, release, err := sp.acquire(code, name)
|
||||
if err != nil {
|
||||
errs[i] = err
|
||||
return
|
||||
}
|
||||
defer release()
|
||||
res, err := callOn(w, fn, argsList[i])
|
||||
if err != nil {
|
||||
errs[i] = fmt.Errorf("%s: %w", name, err)
|
||||
return
|
||||
}
|
||||
out[i] = res
|
||||
}(i)
|
||||
}
|
||||
wg.Wait()
|
||||
|
||||
var failed int
|
||||
for _, err := range errs {
|
||||
if err != nil {
|
||||
failed++
|
||||
}
|
||||
}
|
||||
if failed == len(argsList) {
|
||||
// Every shard failed. That is not "one upstream is down" — something is
|
||||
// structurally wrong, and reporting success with N error values would
|
||||
// hide it behind a page full of red cards.
|
||||
return nil, fmt.Errorf("all %d shards of %s.%s failed: %v", len(argsList), name, fn, errs[0])
|
||||
}
|
||||
if failed > 0 {
|
||||
if sp.ps.logf != nil {
|
||||
sp.ps.logf("[plugins] %s.%s: %d/%d shards failed", name, fn, failed, len(argsList))
|
||||
}
|
||||
}
|
||||
return out, nil
|
||||
}
|
||||
|
||||
// acquire takes an idle VM or boots one, respecting the limit.
|
||||
func (sp *shardPool) acquire(code, name string) (*worker, func(), error) {
|
||||
sp.mu.Lock()
|
||||
if n := len(sp.idle); n > 0 {
|
||||
w := sp.idle[n-1]
|
||||
sp.idle = sp.idle[:n-1]
|
||||
sp.mu.Unlock()
|
||||
return w, func() { sp.release(w) }, nil
|
||||
}
|
||||
if sp.limit > 0 && sp.live >= sp.limit {
|
||||
sp.mu.Unlock()
|
||||
return nil, nil, fmt.Errorf("shard limit %d reached", sp.limit)
|
||||
}
|
||||
sp.live++
|
||||
sp.mu.Unlock()
|
||||
|
||||
w, err := bootPluginState(code, name, sp.ps.outboundAllowlist())
|
||||
if err != nil {
|
||||
sp.mu.Lock()
|
||||
sp.live--
|
||||
sp.mu.Unlock()
|
||||
return nil, nil, err
|
||||
}
|
||||
return w, func() { sp.release(w) }, nil
|
||||
}
|
||||
|
||||
func (sp *shardPool) release(w *worker) {
|
||||
if w == nil || w.L == nil {
|
||||
return
|
||||
}
|
||||
sp.mu.Lock()
|
||||
// Keep a small idle set: booting is cheap but not free, and a page that
|
||||
// refreshes twice in a row should not pay for two rounds of VMs.
|
||||
if len(sp.idle) < sp.limit {
|
||||
sp.idle = append(sp.idle, w)
|
||||
sp.mu.Unlock()
|
||||
return
|
||||
}
|
||||
sp.live--
|
||||
sp.mu.Unlock()
|
||||
w.L.Close()
|
||||
}
|
||||
|
||||
// Shutdown closes every pooled VM. Called from Plugins.Close.
|
||||
func (sp *shardPool) Shutdown() {
|
||||
if sp == nil {
|
||||
return
|
||||
}
|
||||
sp.mu.Lock()
|
||||
idle := sp.idle
|
||||
sp.idle = nil
|
||||
sp.live = 0
|
||||
sp.mu.Unlock()
|
||||
for _, w := range idle {
|
||||
if w != nil && w.L != nil {
|
||||
w.L.Close()
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// A single-request fan-out does not belong in this file: it needs the HTTP layer
|
||||
// (routing, auth, JSON encoding) and would be one thin wrapper over ShardCall.
|
||||
// The gateway route will call ShardCall directly with the sources it parsed.
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user