mirror of
https://gitcode.com/JianFeeeee/ModelRouter.git
synced 2026-11-07 17:57:23 +00:00
Compare commits
104 Commits
| Author | SHA1 | Date | |
|---|---|---|---|
| 980f4a0e40 | |||
| 429afce67e | |||
| 5639bb662c | |||
| bf0657bb84 | |||
| 504c5ac9a0 | |||
| 70f1c879bd | |||
| 04e544c823 | |||
| 7e33d11d15 | |||
| cd82835f25 | |||
| 26ea782350 | |||
| 5c58244781 | |||
| 0121d23f91 | |||
| de7c372ad2 | |||
| c51066f0b6 | |||
| 5530912d32 | |||
| cc5e725226 | |||
| 652842783f | |||
| a21ae84cbe | |||
| d072a03c9a | |||
| f814ff7468 | |||
| ef631b43dd | |||
| 9c3aabb7f9 | |||
| a7355debed | |||
| f6baa13583 | |||
| fd03ef6e4a | |||
| 97bb9c6ef4 | |||
| d81074621a | |||
| 4dd3431c26 | |||
| 3d1407e6fa | |||
| db320d2f6c | |||
| b544671f5e | |||
| 0881a49286 | |||
| 4e2b5bfda9 | |||
| 61b847da86 | |||
| 6432b61eb6 | |||
| dd708c9680 | |||
| f0986751de | |||
| c6c3e0dcd7 | |||
| b9944332ff | |||
| 6bdb9fcc44 | |||
| 25eb30c15b | |||
| 131a42a169 | |||
| 3cdb16c906 | |||
| ef92ce4d82 | |||
| c524831616 | |||
| 868ac59692 | |||
| 4d197e4bd3 | |||
| 3806aaee03 | |||
| 5b628b4d6f | |||
| 2e3d5b79ad | |||
| 2ebc01e03b | |||
| 11776024e9 | |||
| bf1932c6c5 | |||
| 22ddf5a411 | |||
| 3698546d14 | |||
| ed05cccb72 | |||
| c309448414 | |||
| 5b6d6fc90d | |||
| 3ddae41f0c | |||
| 2378bc00ba | |||
| 882288f67f | |||
| 813de19bd0 | |||
| 25d8bd8632 | |||
| 3e27f4db24 | |||
| 23705042ee | |||
| 02cb8c0e33 | |||
| 6575058556 | |||
| d42c02b15d | |||
| df9aeed5fb | |||
| 21cd0429a2 | |||
| 2ed1f0ecde | |||
| 624fd74b45 | |||
| 94cbcb6771 | |||
| a5384d9fb6 | |||
| ab20f1be48 | |||
| 8334cffbc9 | |||
| a42ff62d06 | |||
| 66585549f1 | |||
| 747dff5b76 | |||
| 5d045f97a8 | |||
| e48baa1bb3 | |||
| 519b18518a | |||
| 77b00bfe3f | |||
| 21ec8f59d8 | |||
| ec89daad62 | |||
| 24609289e8 | |||
| 18c2385c61 | |||
| 2a5c7bbbc7 | |||
| 045ecf47bc | |||
| 334b984c25 | |||
| f3d5ba6cea | |||
| 98847adfd6 | |||
| 5bb94db08c | |||
| ef396f9b47 | |||
| c8a300c998 | |||
| 210243f724 | |||
| b2183df1e8 | |||
| 6eac80bc6c | |||
| e3a36bc7ca | |||
| 31b1ed2c96 | |||
| e19cff248f | |||
| bb3af3bdb3 | |||
| 9c99ded8c0 | |||
| 690f55c6f1 |
6
.gitignore
vendored
6
.gitignore
vendored
@ -12,11 +12,17 @@ cmd/gui/bin/
|
||||
cmd/gui/build/gui-dist/
|
||||
cmd/build/
|
||||
cmd/gui/*.log
|
||||
cmd/gui/.cache/
|
||||
# workspace (AI-assisted coding local state)
|
||||
.pi/
|
||||
.pi-glla/
|
||||
.codegraph/
|
||||
.omo/
|
||||
/build/
|
||||
|
||||
# local ops scripts (machine-specific, not for sharing)
|
||||
scripts/
|
||||
cmd/gui/dist/
|
||||
|
||||
# local design/working notes (not part of the shipped repo)
|
||||
/plan.md
|
||||
|
||||
18
Makefile
18
Makefile
@ -23,6 +23,14 @@ build:
|
||||
CGO_ENABLED=1 go build -tags luajit -trimpath -o $(BUILD_DIR)/$(BINARY) ./cmd/llmsproxy
|
||||
@echo "Built: $(BUILD_DIR)/$(BINARY) ($(VERSION))"
|
||||
|
||||
# ---- core distribution packages (deb + rpm + tar.gz) -----
|
||||
# Server admins install these with dpkg/rpm; unlike the GUI they are the bare
|
||||
# gateway binary + systemd unit (with memory tuning) + Lua adapters.
|
||||
core-dist:
|
||||
./packaging/core-dist.sh $(VERSION)
|
||||
|
||||
.PHONY: core-dist
|
||||
|
||||
# ---- GUI (Electron) ----
|
||||
gui:
|
||||
@cd $(GUI_DIR) && npm start
|
||||
@ -31,11 +39,12 @@ gui-dev:
|
||||
@cd $(GUI_DIR) && npm run dev
|
||||
|
||||
# build the embedded core into the GUI, then run electron-builder
|
||||
# targets: deb + AppImage (linux), rpm needs system rpmbuild
|
||||
# targets: deb + rpm + AppImage (linux); rpm needs system rpmbuild
|
||||
# NOTE: run in a clean env (no LD_LIBRARY_PATH pollution) so the packaged
|
||||
# binary does not link against host-only libs and crash on user machines.
|
||||
gui-dist:
|
||||
@cd $(GUI_DIR) && $(CLEAN_PACK_ENV) npm run dist:linux
|
||||
./packaging/verify-dist.sh
|
||||
|
||||
# deb only (for local sharing), rpm needs rpmbuild
|
||||
gui-deb:
|
||||
@ -48,6 +57,7 @@ gui-win:
|
||||
# windows nsis installer: full cross-build in docker (self-sufficient, no host mingw)
|
||||
gui-win-docker:
|
||||
@cd $(GUI_DIR) && $(CLEAN_PACK_ENV) ./scripts/dist-win-docker.sh
|
||||
./packaging/verify-dist.sh
|
||||
|
||||
gui-dist-dir:
|
||||
@cd $(GUI_DIR) && $(CLEAN_PACK_ENV) npm run dist:dir
|
||||
@ -56,15 +66,17 @@ test:
|
||||
go test ./...
|
||||
|
||||
clean:
|
||||
rm -rf $(BUILD_DIR) cmd/gui/build/gui-dist cmd/gui/bin
|
||||
rm -rf $(BUILD_DIR) cmd/build cmd/gui/dist cmd/gui/build/gui-dist cmd/gui/bin
|
||||
|
||||
help:
|
||||
@echo "Targets:"
|
||||
@echo " make build build core binary (luajit)"
|
||||
@echo " make core-dist build core packages: deb + rpm + tar.gz (needs nfpm)"
|
||||
@echo " make gui run desktop GUI (dev)"
|
||||
@echo " make gui-dev run GUI with devtools"
|
||||
@echo " make gui-dist build deb + AppImage (linux)"
|
||||
@echo " make gui-dist build deb + rpm + AppImage (linux)"
|
||||
@echo " make gui-deb build deb only"
|
||||
@echo " make gui-win build Windows nsis (needs host mingw + wine)"
|
||||
@echo " make gui-win-docker build Windows nsis all-in-docker (no host mingw needed)"
|
||||
@echo " make test run go tests"
|
||||
@echo " make clean remove build dirs and packaged installers (build/, cmd/build/, cmd/gui/dist/)"
|
||||
|
||||
266
README.md
266
README.md
@ -15,15 +15,18 @@
|
||||
|
||||
### 极致轻量
|
||||
|
||||
- **单二进制**:编译后约 10MB,零运行时依赖(仅依赖系统 libc),部署即用
|
||||
- **极低内存占用**:空闲状态仅 ~15MB RSS,满载并发 100+ 请求时峰值 < 100MB
|
||||
- **单二进制**:编译后 8~12MB(`-s -w` strip 后约 8MB),零运行时依赖(仅依赖系统 libc),部署即用
|
||||
- **低内存占用**:与源数量相关,非与运行时长相关——单源 ~10MB、16 源生产实例 ~32–35MB([实测分解与调优](#内存占用实测与调优))
|
||||
- **日志按需加载**:统计数字来自**全量**审计日志(启动时流式扫一遍即释放,29MB/22 万行约 260ms),但原始记录不常驻内存——默认只加载首屏,下滚自动分页,CSV 导出流式写出(O(1) 内存),离页即释放
|
||||
- **零运行时依赖**:纯 Go + LuaJIT 静态链接,无需安装 Python/Node/Java 等运行时
|
||||
- **启动极快**:冷启动 < 200ms,热重载配置 < 10ms
|
||||
|
||||
### 强大的多租户调度能力
|
||||
|
||||
- **多密钥多租户**:支持无限密钥,每个密钥独立角色、模型范围、Token 配额、重置周期
|
||||
- **按模型配额**:每把 key 的每个模型单独配 token + 请求数配额与重置周期(小时/周/月/自定义 N 小时);一个模型用满只影响该模型,同 key 其它模型照常;耗尽返 429 + `Retry-After` 并点名模型,admin key 永不受限
|
||||
- **AUTO 智能调度**:基于优先级档位的分级调度,同优先级源自动轮询负载均衡,故障自动毫秒级故障转移
|
||||
- **自愈冷却**:冷却上限 5 分钟,过半后放行 1 个探测请求,上游/额度恢复即刻回归轮询,无需等满冷却窗口
|
||||
- **Token 配额管理**:精确到模型级别的 Token 配额控制,支持小时/周/月/自定义小时周期自动重置
|
||||
- **同优先级源负载均衡**:同一优先级档位的多个源,请求自动 Round-Robin 均匀分发,故障毫秒级故障转移
|
||||
- **Source-Model 前缀路由**:支持 `source-model`/`source:model`/`source/model` 精确指定上游源
|
||||
@ -37,7 +40,7 @@
|
||||
|
||||
### 灵活的协议适配
|
||||
|
||||
- **LuaJIT VM**:每个适配器独立 VM + worker 池,安全并发,Lua 脚本热加载无需重启
|
||||
- **LuaJIT VM**:每个适配器独立 VM + **弹性 worker 池**(按需扩缩容,空闲时自动回收),安全并发,Lua 脚本热加载无需重启
|
||||
- **Lua 适配器协议**:`transform_request` / `transform_response` / `transform_stream_chunk` 双向转换
|
||||
- **动态请求头钩子**:`build_headers(meta)` 支持 HMAC 签名、动态 Header 注入
|
||||
- **多模态透传**:`image_url` 等多模态内容在多源间无损透传,Anthropic/Gemini/Ollama 自动转换
|
||||
@ -60,6 +63,12 @@ Groq、Mistral、Ollama、KimiCode…)通过 **Lua 适配器** 做协议转换
|
||||
- **生图**:`POST /v1/images/generations`,`kind: image` 的模型独立路由。
|
||||
- **多模态**:`content` 数组(`image_url` 等)在多源间无损透传,Anthropic/Gemini/Ollama 自动转换。
|
||||
- **LuaJIT VM**:基于 golua 绑定的 LuaJIT,每个适配器独立 VM + worker 池,安全并发。
|
||||
池是 **弹性** 的:`Σ max_concurrent` 只是上限而非预分配,状态按需创建、空闲时自动回收,
|
||||
因此空载网关几乎不持有 Lua 状态(启动 0 个)。
|
||||
- 扩容步长由适配器最大并发决定:`clamp(ceil(max/8), 1, 8)`,且仅在**真实并发争用**
|
||||
(所有现有状态均已占用)时批量预热,顺序流量始终只用 1 个状态。
|
||||
- 缩容步长由当前连接数决定:`clamp(ceil(冗余/(1+占用)), 1, 冗余)`——无连接时一轮收到
|
||||
常驻下限,忙时每轮只释放 1 个,保护热路径。WebUI「适配器」页展示实时池指标。
|
||||
- **disable_thinking**:请求 `disable_thinking:true`(或上游对应字段)动态开关推理。
|
||||
- **Lua 适配协议**:每个源挂一个 `.lua` 适配器,完成 `transform_request` /
|
||||
`transform_response` / `transform_stream_chunk` 双向转换,协议差异全在 Lua 层。
|
||||
@ -170,6 +179,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` 解析顺序:
|
||||
@ -184,16 +268,51 @@ sources:
|
||||
|
||||
### AUTO 链(优先级页)
|
||||
|
||||
AUTO 调度**只**由优先级页保存的规则(持久化到 `runtime_file` 的 `auto` 字段)决定,
|
||||
源配置里的 `models[].priority` 数字不再参与调度、也不再显示。
|
||||
AUTO 调度**只**由优先级页保存的规则决定,源配置里的 `models[].priority` 数字不再
|
||||
参与调度、也不再显示。**聊天与生图是两套独立的链**,在优先级页面通过「聊天 / 生图」
|
||||
按钮切换编辑:
|
||||
|
||||
- **聊天链**(`runtime_file` 的 `auto` 字段):服务 `POST /v1/chat/completions`
|
||||
的 `model: AUTO` 请求;只接受 `kind: chat` 模型,配置在链里的生图槽位会被忽略。
|
||||
- **生图链**(`runtime_file` 的 `auto_image` 字段):服务 `POST /v1/images/generations`
|
||||
的 `model: AUTO` 请求;只接受 `kind: image` 模型。未配置生图链时,AUTO 生图回落
|
||||
到「所有含生图模型的源、按注册顺序」的旧行为,兼容旧配置。
|
||||
|
||||
每条规则 = 一个「槽位」:`{ model, source, tier, token_quota, period, hours }`。
|
||||
|
||||
- 每条规则 = 一个「槽位」:`{ model, source, tier, token_quota, period, hours }`。
|
||||
- `tier` 表示优先级档位:同一档的模型并排、共享该优先级;档位从上到下递减。
|
||||
- 同一模型可配置多个槽位(如 A 源低配 → B 源低配 → A 源高配 → B 源高配),按档位顺延。
|
||||
- `token_quota` > 0 时该槽位在重置周期内用满即顺延到下一槽位;`period` 支持
|
||||
`hour` / `week` / `month` / `nhour`(配合 `hours`),空 = 不限。
|
||||
- 生图模型(`kind: image`)不参与 AUTO 链;生图走 `POST /v1/images/generations`
|
||||
的独立路径。
|
||||
- 同 tier 内按「偏好分」排序,失败槽位冷却后自动跳过;冷却 / 配额耗尽 / hard
|
||||
错误会顺延到下一 tier,全部失败时返回 503 并附上每档失败摘要。
|
||||
|
||||
### 冷却与半冷却探测(自愈调度)
|
||||
|
||||
槽位失败后进入指数退避冷却(5s → 10s → 20s …,上限 **5 分钟**)。冷却不是全时段
|
||||
硬闸:
|
||||
|
||||
- **前半段完全静默**:刚失败的槽位不接任何流量,避免对故障上游持续施压。
|
||||
- **后半段放行 1 个探测**:窗口过半后,同一 `(源, 模型)` 最多允许 **一个** 请求作为
|
||||
探测通过。探测槽数 = `clamp(max_concurrent / 10, 1, 2)`——`max_concurrent: 10`
|
||||
的源恰好放行 1 个探测请求。
|
||||
- **探测排在最后**:同 tier 内探测候选永远排在健康候选之后,只有没有正常槽可用时
|
||||
才承接真实流量,因此探测不会抢走可用容量。
|
||||
- **探测即真实请求**:成功则立即清零失败计数与冷却,槽位当场回到正常轮询——
|
||||
上游恢复/额度恢复后不必等满整个冷却窗口。失败则重开一个新窗口,下次探测顺延到
|
||||
新窗口的中点,不会连续重试。
|
||||
|
||||
分类冷却策略:
|
||||
|
||||
| 情况 | 冷却 | 说明 |
|
||||
|---|---|---|
|
||||
| 传输错误 / 5xx | 5s→10s→20s…,上限 5min | 指数退避 |
|
||||
| 401 / 403 凭据错误 | 10min | 换 key 后由下一次探测自动接回 |
|
||||
| 429 限流 | 30s 固定 | 「太快」不等于「坏了」,不进指数阶梯 |
|
||||
| 配额耗尽(`insufficient_quota` / `余额不足` 等) | 对齐配额窗口,上限 30min | **不进指数阶梯、不涨失败计数**,额度恢复即可用 |
|
||||
|
||||
WebUI「优先级」页的槽位健康标签会区分 `冷却` / `待探测` / `探测中`,悬停可看窗口
|
||||
起点、探测放行时刻与完全恢复时刻。
|
||||
|
||||
### 源(Source)与适配器(Adapter)的关系
|
||||
|
||||
@ -212,7 +331,7 @@ AUTO 调度**只**由优先级页保存的规则(持久化到 `runtime_file`
|
||||
2. `config.Load` 读取 `config.yaml`,`config.NewStore(runtime_file)` 读 WebUI 改动的
|
||||
运行时源,二者按名称合并成完整源列表。
|
||||
3. `rebuildRegistry` 为每个源创建 `provider.Provider`(持有目标适配器),并按源的
|
||||
并发上限配置适配器 worker 池大小。
|
||||
并发上限配置适配器 worker 池**上限**(不预分配,见下)。
|
||||
4. 请求进来时 `Registry` 按 `model` 路由到 Provider,Provider 调适配器
|
||||
`transform_request` → HTTP 发送 → `transform_response` / `transform_stream_chunk`。
|
||||
|
||||
@ -224,6 +343,70 @@ WebUI 上的"新增/编辑源"、"上传 Lua 适配器"、"改 AUTO 优先级链
|
||||
`config.yaml` 中定义的基源无法改写配置文件,采用删除标记隐藏(重启后仍隐藏),
|
||||
在 UI 里重新添加同名源即可恢复。
|
||||
|
||||
### 源模板(多 key 负载均衡)
|
||||
|
||||
多个上游 key 共享同一套配置(URL / 适配器 / 模型列表 / 并发 / rpm)时,逐一复制
|
||||
整段 source 太冗余。**源模板**保存除 `name` 和 `api_key` 以外的全部源字段,
|
||||
从模板一键创建多个仅 key 不同的源,天然支持多 key 负载均衡(每个展开源独立调度、
|
||||
独立冷却、独立健康状态)。
|
||||
|
||||
- WebUI「源」页右上角「**模板管理**」按钮:列出全部模板,可编辑 / 删除 / 新建。
|
||||
- 添加源弹窗右上角两个按钮:
|
||||
- **从模板创建** → 选模板 → 表单自动填充(name + key 仍需你填);
|
||||
- **存为模板** → 把当前表单非 key/name 字段存为新模板。
|
||||
- 模板持久化在 `runtime_file` 的 `source_templates` 字段,与运行时源一样热管理,
|
||||
改完即时生效。模板本身不含密钥,可以安全共享 / 纳管。
|
||||
|
||||
模板不会自己变成源——它只是「配方」。真正承担流量的是从模板创建出来的、带 key
|
||||
的具体源。
|
||||
|
||||
### 内存占用(实测与调优)
|
||||
|
||||
内存不是一个定数:它随**配置的源数量**增长(每个源一个 `http.Transport` 与连接池、
|
||||
一组健康状态),与运行时长无关。本机实测(Linux x86_64,12 核):
|
||||
|
||||
| 部署形态 | 启动 RSS | 稳态 RSS |
|
||||
|---|---|---|
|
||||
| 1 源 / 1 适配器(最小配置) | ~4 MB | ~10 MB |
|
||||
| 1 源 + 29 MB 历史审计日志 | ~19 MB | ~20 MB |
|
||||
| **16 源 / 13 适配器 / 59 模型(本机生产)** | ~28 MB | **~32–35 MB** |
|
||||
|
||||
> 历史参考:本项优化前同一生产配置为 **~105 MB**。降幅来自三处:审计日志的**原始记录**
|
||||
> 不再常驻内存(约 25 MB;统计聚合仍扫全量,但扫完即释放)、Lua 状态池不再单调增长、
|
||||
> 以及下面两个运行时开关。
|
||||
|
||||
内存构成(生产实例分段测量,`/proc/<pid>/smaps`):
|
||||
|
||||
| 区域 | RSS | 说明 |
|
||||
|---|---|---|
|
||||
| Go 堆 | ~14 MB | provider/registry/scheduler 结构 + 连接池缓冲 |
|
||||
| 其他匹名(线程栈 / LuaJIT chunk / runtime) | ~12 MB | 与线程数、已加载适配器数相关 |
|
||||
| 二进制 text+rodata | ~8 MB | 映射的可执行文件页(只读、可被内核回收) |
|
||||
| Go runtime 预留 | ~4 MB | VSZ 上看到的 2 GB+ 是地址空间预留,不占物理内存 |
|
||||
| 共享库 | ~3 MB | libc / libluajit / libm |
|
||||
|
||||
**两个推荐的部署开关**(只适用于环境变量,无需改代码):
|
||||
|
||||
```ini
|
||||
# /etc/systemd/system/llmsproxy.service
|
||||
Environment=GOGC=50
|
||||
Environment=MALLOC_ARENA_MAX=2
|
||||
```
|
||||
|
||||
- `MALLOC_ARENA_MAX=2`:LuaJIT 的分配走 cgo → glibc malloc,glibc 默认允许 `8×nproc`
|
||||
个 per-thread arena,每个碰到 malloc 的 OS 线程会占用一个(各约 1 MB,且**不归还给系统**)。
|
||||
实测从 8–15 个 arena(约 7–12 MB)降到 **0**。
|
||||
- `GOGC=50`:把 Go 堆增长目标减半。**单独使用无效**(省下的堆会立即被更多 glibc
|
||||
arena 吃掉,实测 20.3 → 21.5 MB 反而变大),必须与 `MALLOC_ARENA_MAX` 配合,
|
||||
两者同时开启才降 **~19%**。网关是 I/O 密集型(本机 9 小时仅消耗 1min10s CPU),
|
||||
多出的 GC 周期与它的空闲 CPU 相比可忽略。
|
||||
|
||||
> 桌面版(Electron)已在 `cmd/gui/main.js` 里默认为内嵌核心注入这两个开关;手动 export 同名
|
||||
> 环境变量可覆盖。如果你用自己的 systemd unit / 容器,建议照上面加上。
|
||||
|
||||
可选:`GOMEMLIMIT=48MiB` 作软上限,实测再省 ~1 MB;代价是逐近上限时 GC 转激进,
|
||||
温和场景不必开。
|
||||
|
||||
### disable_thinking
|
||||
|
||||
请求体带 `"disable_thinking": true` 时,网关透传给各适配器;DeepSeek 适配器将其
|
||||
@ -237,7 +420,10 @@ WebUI 上的"新增/编辑源"、"上传 Lua 适配器"、"改 AUTO 优先级链
|
||||
支持按时间范围导出 CSV;点击模型可生成 pin 到该模型的连接配置。
|
||||
- **对话页**:流式/非流式调试。
|
||||
- **密钥页**:创建/编辑网关 key,为每个 key 配模型范围(模型 + 源 + token 配额 + 周期),
|
||||
管理员管理全部 key,用户只看到自己的 key。
|
||||
管理员管理全部 key,用户只看到自己的 key。每个模型砖块显示自己的配额徽标
|
||||
(如 `1.0K 77×·1h`,未设配额显示 `∞`),点开可编辑该模型的 token 配额、
|
||||
请求数配额与重置周期 —— 配额按模型独立生效,一个用满不影响同一 key 的其它模型。
|
||||
「我的密钥」页对用户展示本 key 的模型范围与各自配额。
|
||||
- **优先级页**:拖拽积木配置 AUTO 链档位。
|
||||
- **源页**:在线增删改上游源(API key 等敏感字段加密落盘)。
|
||||
- **适配器页**:上传 / 删除 Lua 适配器脚本。
|
||||
@ -296,14 +482,24 @@ return {
|
||||
一键启停、系统托盘常驻、开机自启、静默启动,内容区一比一内嵌完整 WebUI
|
||||
(状态/对话/密钥/优先级/源/适配器 六个页面,免登录)。服务器用户继续使用纯 Go 二进制。
|
||||
|
||||
### 两种安装包:Headless 与 Desktop
|
||||
|
||||
发布到 Release 的安装包分为两类,面向不同用户:
|
||||
|
||||
- **Headless(服务器版)**:纯核心二进制(无 Electron),配置驱动,适合服务器、
|
||||
容器、systemd 等无人值守场景。发布包含 `deb` + `rpm`(含 systemd unit、内存调优
|
||||
与示例配置)与 `tar.gz`(二进制 + `config.example.yaml` + `adapters/` + README),
|
||||
Windows 为 `.zip`。
|
||||
- **Desktop(桌面版)**:Electron GUI 安装包,内嵌核心,适合个人桌面日常使用。
|
||||
|
||||
### 选择 GUI 还是纯后端
|
||||
|
||||
| 场景 | 推荐 | 理由 |
|
||||
| ------ | ------ | ------ |
|
||||
| 服务器 / 内网网关 / 无人值守常驻 | **纯后端**(单二进制) | 极轻量(~10MB、~15MB RSS),零依赖单进程,直接跑在 systemd/任意容器里,远程管理 |
|
||||
| 个人桌面日常使用 / 多设备内网共享 | **GUI**(Electron) | 免登录内嵌 WebUI、系统托盘一键启停、开机自启、静默后台,适合不懂命令行的使用者 |
|
||||
| Windows 桌面 | **GUI** | 纯后端在 Windows 上需自行注册服务,GUI 提供原生托盘/自启体验 |
|
||||
| CI 一键出三平台安装包 | **GUI 打包脚本** | `make gui-*` 系列出 deb/AppImage/NSIS,可进发行流水线 |
|
||||
| 服务器 / 内网网关 / 无人值守常驻 | **Headless**(单二进制) | 极轻量(二进制 ~8MB;RSS 随源数量,单源 ~10MB、16 源 ~32–35MB,见[内存占用](#内存占用实测与调优)),零依赖单进程,直接跑在 systemd/任意容器里,远程管理 |
|
||||
| 个人桌面日常使用 / 多设备内网共享 | **Desktop**(GUI) | 免登录内嵌 WebUI、系统托盘一键启停、开机自启、静默后台,适合不懂命令行的使用者 |
|
||||
| Windows 桌面 | **Desktop** | 纯后端在 Windows 上需自行注册服务,GUI 提供原生托盘/自启体验 |
|
||||
| CI 一键出三平台安装包 | **Desktop 打包脚本** | `make gui-*` 系列出 deb/AppImage/NSIS,可进发行流水线 |
|
||||
|
||||
两者完全同源:GUI 内嵌的就是纯后端的同一份 `llmsproxy` 核心(LuaJIT 版),
|
||||
配置/适配器格式完全一致,可随时互换。
|
||||
@ -315,27 +511,36 @@ cd cmd/gui
|
||||
npm install
|
||||
make gui # 或 npm run dev —— 开发运行(需先 make build 生成 bin/llmsproxy)
|
||||
make gui-deb # 本机构建 deb(分享给其它 Linux 用户)
|
||||
make gui-dist # 构建 deb + AppImage(linux)
|
||||
make gui-dist # 构建 deb + rpm + AppImage(linux,rpm 需系统 rpmbuild)
|
||||
make core-dist # 构建核心 deb + rpm + tar.gz(服务器版,需 nfpm)
|
||||
make gui-win # 构建 Windows nsis 安装包(本机 mingw + wine)
|
||||
make gui-win-docker # 构建 Windows nsis 安装包(全 docker 自足,无需本机 mingw)
|
||||
make gui-win-docker # 构建 Windows nsis 安装包(全 docker 自足,含 wine)
|
||||
make gui-dist-dir # 仅解包产物(调试用,不产安装包)
|
||||
```
|
||||
|
||||
产物在 `cmd/build/gui-dist/`:`ModelRouter-1.0.0.AppImage`、
|
||||
`modelrouter-gui_1.0.0_amd64.deb`、`ModelRouter Setup 1.0.0.exe`(win)。rpm 需系统 `rpmbuild`。
|
||||
产物在 `cmd/build/gui-dist/`:`ModelRouter-<ver>.AppImage`、
|
||||
`modelrouter-gui_<ver>_amd64.deb`、`ModelRouter Setup <ver>.exe`(win);
|
||||
核心包在 `cmd/build/dist/`:`llmsproxy_<ver>_amd64.deb`、`llmsproxy-<ver>.x86_64.rpm`。
|
||||
rpm 需系统 `rpmbuild`。每次打包后都会跑 `packaging/verify-dist.sh` 产物体检
|
||||
(大小下限门禁——1.3.0 曾因宿主 wine 损坏产出 264KB 的残缺 exe 却没有被拦住)。
|
||||
发布时会把产物重命名为 `ModelRouter-Desktop-*`(GUI)与 `ModelRouter-Headless-*`(纯后端)
|
||||
两类上传到 GitCode Release。
|
||||
|
||||
#### Windows 打包(docker 固化方案,推荐)
|
||||
|
||||
`make gui-win-docker` 一步完成:交叉编译核心 + NSIS 打包,宿主机**无需安装 mingw**。
|
||||
`make gui-win-docker` 一步完成:交叉编译核心 + NSIS 打包**全在 docker 内**,宿主机
|
||||
**无需 mingw、无需 wine、无需 node**。镜像内置 wine(含 i386 运行时),修复了
|
||||
此前宿主 wine 损坏(缺 syswow64 → `c0000135`)导致产出 264KB 残缺 exe 的问题。
|
||||
|
||||
- 构建镜像:`cmd/gui/docker/win-builder/Dockerfile`(`golang:1.25` + mingw-w64 +
|
||||
预编译 LuaJIT windows 版 `lua51.dll` 与 import lib)。LuaJIT 源码取自 gitcode
|
||||
镜像(国内可达),github 兜底;Go 依赖走 `https://goproxy.cn`。
|
||||
- 核心编译:`cmd/gui/scripts/win-core-docker.sh` —— 在容器内产出
|
||||
`cmd/gui/bin/{llmsproxy.exe,lua51.dll}`。
|
||||
- 一键发布:`cmd/gui/scripts/dist-win-docker.sh` —— 核心编译后调用
|
||||
electron-builder(宿主需有 `wine` 用于 NSIS)。
|
||||
- 首次运行自动构建 `modelrouter/win-builder` 镜像;已存在则直接复用。
|
||||
node + wine32/wine64 + 预编译 LuaJIT windows 版 `lua51.dll` 与 import lib)。
|
||||
LuaJIT 源码取自 gitcode 镜像(国内可达),github 兜底;Go 依赖走 `https://goproxy.cn`。
|
||||
- 一键发布:`cmd/gui/scripts/dist-win-docker.sh` —— 核心编译 + electron-builder
|
||||
NSIS 全部在容器内执行(`USE_SYSTEM_WINE=true` 走镜像内的 wine);electron-builder
|
||||
的 ~400MB 工具链缓存挂到 docker volume,二次构建增量复用。
|
||||
- 首次运行自动构建 `modelrouter/win-builder` 镜像(约 5–10 分钟);已存在则直接复用。
|
||||
- 打包完成后自动跑 `packaging/verify-dist.sh` 门禁:exe 低于 5MB / deb·rpm·7z 低于
|
||||
10MB 即失败——残缺安装包不可能再溜进 Release。
|
||||
|
||||
```bash
|
||||
# 全 Windows 发布(核心重编 + NSIS):
|
||||
@ -364,6 +569,15 @@ internal/gateway # OpenAI 兼容 HTTP 服务 + 鉴权 + SDK/流式 + Web
|
||||
internal/types # 统一格式 & OpenAI wire 类型
|
||||
```
|
||||
|
||||
## 开发协作
|
||||
|
||||
- **Git 分支工作流**:主分支 / 特性分支 / 发布分支的规范见
|
||||
**[docs/git-workflow.md](docs/git-workflow.md)**([English](docs/git-workflow-en.md))。
|
||||
主线:`main` 永可部署 → 特性合入 `main` → 切 `release/vX.Y.Z` 打 tag 发布 →
|
||||
hotfix 提交发布分支并 cherry-pick 回 `main`。
|
||||
- **Lua 适配器协议**:见 **[docs/lua-adapters.md](docs/lua-adapters.md)**
|
||||
([English](docs/lua-adapters-en.md))。
|
||||
|
||||
## 测试
|
||||
|
||||
```bash
|
||||
|
||||
292
README_EN.md
292
README_EN.md
@ -11,7 +11,7 @@ KimiCode…) behind a single endpoint. Pick a specific model or use **AUTO** mod
|
||||
which routes to the best healthy upstream by configured priority.
|
||||
|
||||
> **Lightweight, no recompile to add sources**: the gateway is a single Go
|
||||
> binary (~10 MB, zero runtime dependencies). Adding or switching an upstream is
|
||||
> binary (8-12 MB, ~8 MB stripped; zero runtime dependencies). Adding or switching an upstream is
|
||||
> just a `sources` entry in `config.yaml` (or via the WebUI) or a `.lua`
|
||||
> adapter — **no Go changes, no recompile**. Changes submitted through the
|
||||
> WebUI take effect immediately (hot reload); editing `config.yaml` or a `.lua`
|
||||
@ -38,10 +38,31 @@ 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
|
||||
pool for safe concurrency.
|
||||
pool for safe concurrency. Pools are **elastic**: the sum of `max_concurrent`
|
||||
over an adapter's sources is a ceiling, not a preallocation — states are booted
|
||||
on demand and reclaimed when demand drops, so an idle gateway holds close to
|
||||
zero Lua states. The grow step follows the adapter's max concurrency
|
||||
(`clamp(ceil(max/8), 1, 8)`, and never exceeds the number of queued callers);
|
||||
the shrink step follows the live connection count
|
||||
(`clamp(ceil(slack/(1+in_use)), 1, slack)`), so an idle pool collapses in one
|
||||
round while a busy one gives up a single state at a time.
|
||||
- **On-demand request logs**: dashboard totals are computed from the **full**
|
||||
audit history (streamed once at startup and released — 29 MB / 221k lines in
|
||||
~260 ms), while the raw records are never held in memory: the dashboard loads
|
||||
one screen, scrolling pages the rest straight off disk, CSV export streams in
|
||||
O(1) memory, and leaving the page releases everything.
|
||||
- **Self-healing cooldown**: cooldown is capped at 5 minutes and, past the
|
||||
window's midpoint, exactly one probe request is allowed through; a recovered
|
||||
upstream (or a reset quota) returns to full rotation on that probe instead of
|
||||
waiting out the window.
|
||||
- **disable_thinking**: `disable_thinking: true` toggles reasoning per-request.
|
||||
- **Lua adapter protocol**: each source mounts a `.lua` adapter with
|
||||
`transform_request` / `transform_response` / `transform_stream_chunk` — all
|
||||
@ -155,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:
|
||||
@ -176,12 +289,20 @@ tiered AUTO chain.
|
||||
|
||||
### AUTO chain (Priority page)
|
||||
|
||||
AUTO scheduling is driven **only** by the rules saved on the Priority page
|
||||
(persisted as the `auto` field of the runtime file). The numeric
|
||||
`models[].priority` in source configs no longer participates in scheduling and
|
||||
is no longer shown.
|
||||
AUTO scheduling is driven **only** by the rules saved on the Priority page. The
|
||||
numeric `models[].priority` in source configs no longer participates in
|
||||
scheduling and is no longer shown. **Chat and image generation are two
|
||||
independent chains**, switched in the Priority page via the `Chat / Image` toggle:
|
||||
|
||||
- **Chat chain** (`auto` field): serves `POST /v1/chat/completions` with
|
||||
`model: AUTO`; only `kind: chat` models are honored (image slots are skipped).
|
||||
- **Image chain** (`auto_image` field): serves `POST /v1/images/generations` with
|
||||
`model: AUTO`; only `kind: image` models are honored. With no image chain
|
||||
configured, AUTO image generation falls back to legacy discovery (all sources
|
||||
exposing an image model, in registry order) for backward compatibility.
|
||||
|
||||
Each rule is one "slot": `{ model, source, tier, token_quota, period, hours }`.
|
||||
|
||||
- Each rule is one "slot": `{ model, source, tier, token_quota, period, hours }`.
|
||||
- `tier` is the priority tier: models in the same tier sit side by side and
|
||||
share it; tiers run high → low.
|
||||
- The same model may appear in several slots (e.g. A low → B low → A high →
|
||||
@ -189,8 +310,9 @@ is no longer shown.
|
||||
- `token_quota` > 0 means the slot is skipped once its tokens within the reset
|
||||
window are exhausted; `period` supports `hour` / `week` / `month` / `nhour`
|
||||
(with `hours`); empty = unlimited.
|
||||
- Image models (`kind: image`) are kept out of the chain and are served by the
|
||||
separate `POST /v1/images/generations` path.
|
||||
- Within a tier, slots run by preference score; cooling / quota-exhausted /
|
||||
hard-failed slots fall through to the next tier, and total failure returns
|
||||
503 with a per-tier summary.
|
||||
|
||||
### Sources vs. adapters
|
||||
|
||||
@ -227,6 +349,82 @@ sources and `.lua` files; base sources defined in `config.yaml` cannot rewrite
|
||||
that file, so they are hidden via a deletion tombstone (still hidden after
|
||||
restart) and can be restored by re-adding the same name in the UI.
|
||||
|
||||
### Source templates (multi-key balancing)
|
||||
|
||||
When several upstream keys share one config (URL / adapter / model list /
|
||||
concurrency / rpm), copying the whole source block N times is wasteful. A
|
||||
**source template** stores every source field except `name` and `api_key`, so you
|
||||
create N key-only-different sources from one shared skeleton — natural multi-key
|
||||
load balancing (each expanded source schedules, cools down, and reports health
|
||||
independently).
|
||||
|
||||
- WebUI Sources page, top-right **Templates** button: list all templates with
|
||||
edit / delete / create.
|
||||
- Add-source dialog has two header buttons:
|
||||
- **From template** → pick a template → form auto-fills (name + key still
|
||||
yours to enter);
|
||||
- **As template** → save the current form's non-key/name fields as a template.
|
||||
- Templates persist in the runtime file's `source_templates` field alongside
|
||||
runtime sources (hot-managed, apply immediately). Templates hold no secret
|
||||
keys, so they're safe to share / version-control.
|
||||
|
||||
A template is just a recipe — it doesn't become a source by itself. Only the
|
||||
key-bearing sources created from it carry real traffic.
|
||||
|
||||
### Memory footprint (measured and tuned)
|
||||
|
||||
Memory is not a constant: it scales with the **number of configured sources**
|
||||
(each source owns an `http.Transport` with its connection pool plus per-model
|
||||
health state), not with uptime. Measured on this machine (Linux x86_64, 12 cores):
|
||||
|
||||
| Deployment shape | Startup RSS | Settled RSS |
|
||||
|---|---|---|
|
||||
| 1 source / 1 adapter (minimal) | ~4 MB | ~10 MB |
|
||||
| 1 source + a 29 MB audit history | ~19 MB | ~20 MB |
|
||||
| **16 sources / 13 adapters / 59 models (this host)** | ~28 MB | **~32-35 MB** |
|
||||
|
||||
> For reference, the same production config used **~105 MB** before this round of
|
||||
> work. The reduction comes from three places: raw audit records are no longer kept
|
||||
> resident (~25 MB — the aggregates still scan the full history, but the scan
|
||||
> releases as it goes), Lua state pools no longer grow monotonically, and the two
|
||||
> runtime knobs below.
|
||||
|
||||
Breakdown of the production instance (per-region, from `/proc/<pid>/smaps`):
|
||||
|
||||
| Region | RSS | Notes |
|
||||
|---|---|---|
|
||||
| Go heap | ~14 MB | provider/registry/scheduler state + connection-pool buffers |
|
||||
| Other anonymous (thread stacks / LuaJIT chunks / runtime) | ~12 MB | scales with thread count and loaded adapters |
|
||||
| Binary text+rodata | ~8 MB | mapped executable pages (read-only, reclaimable by the kernel) |
|
||||
| Go runtime reservations | ~4 MB | the 2 GB+ you see in VSZ is address space, not physical memory |
|
||||
| Shared libraries | ~3 MB | libc / libluajit / libm |
|
||||
|
||||
**Two recommended deployment knobs** (environment only, no code change):
|
||||
|
||||
```ini
|
||||
# /etc/systemd/system/llmsproxy.service
|
||||
Environment=GOGC=50
|
||||
Environment=MALLOC_ARENA_MAX=2
|
||||
```
|
||||
|
||||
- `MALLOC_ARENA_MAX=2`: LuaJIT allocates through cgo into glibc malloc, and glibc
|
||||
allows up to `8 x nproc` per-thread arenas. Every OS thread that touches malloc
|
||||
claims one (~1 MB each) and **never returns it to the OS**. Measured: 8-15
|
||||
arenas (7-12 MB) down to **0**.
|
||||
- `GOGC=50`: halves the Go heap growth target. **It does nothing on its own** —
|
||||
the heap it saves is immediately consumed by additional glibc arenas (measured
|
||||
20.3 -> 21.5 MB, i.e. slightly worse) — so it must be paired with
|
||||
`MALLOC_ARENA_MAX`. Together they cut settled RSS by **~19%**. The gateway is
|
||||
I/O bound (1min10s of CPU per 9 hours here), so the extra GC cycles are free.
|
||||
|
||||
> The desktop build (Electron) already injects both when spawning the embedded
|
||||
> core (`cmd/gui/main.js`); exporting either variable yourself overrides it. If you
|
||||
> write your own systemd unit or container spec, add them there.
|
||||
|
||||
Optional: `GOMEMLIMIT=48MiB` as a soft ceiling saves roughly 1 MB more, at the
|
||||
cost of GC turning aggressive as the limit approaches. Not worth it for a relaxed
|
||||
deployment.
|
||||
|
||||
### disable_thinking
|
||||
|
||||
With `"disable_thinking": true` in the request body, the gateway passes it to
|
||||
@ -308,14 +506,25 @@ tray, autostart, silent launch, and a pixel-for-pixel embedded full WebUI
|
||||
(Status/Chat/Keys/Priority/Sources/Adapters — no login needed). Server users
|
||||
keep using the plain Go binary.
|
||||
|
||||
### Two installer kinds: Headless and Desktop
|
||||
|
||||
Release installers ship in two flavors for different audiences:
|
||||
|
||||
- **Headless (server)**: the plain core binary (no Electron), config-driven,
|
||||
for servers / containers / systemd / unattended runs. Ships as `deb` + `rpm`
|
||||
(systemd unit with memory tuning + example config) and `tar.gz` (binary +
|
||||
`config.example.yaml` + `adapters/` + README), Windows as `.zip`.
|
||||
- **Desktop (desktop app)**: the Electron GUI installer with an embedded core,
|
||||
for personal daily desktop use.
|
||||
|
||||
### GUI vs. plain backend
|
||||
|
||||
| Scenario | Pick | Why |
|
||||
| ---- | ---- | ---- |
|
||||
| Server / intranet gateway / unattended long-running | **plain backend** (single binary) | ~10 MB, ~15 MB RSS, zero-dep single process — drop it into systemd or any container, remote admin |
|
||||
| Personal desktop daily use / multi-device intranet sharing | **GUI** (Electron) | no-login embedded WebUI, tray one-click, autostart, silent background — for non-CLI users |
|
||||
| Windows desktop | **GUI** | plain backend needs manual service registration; GUI ships native tray/autostart |
|
||||
| CI one-shot 3-platform installers | **GUI packaging scripts** | `make gui-*` emits deb / AppImage / NSIS, drops straight into a release pipeline |
|
||||
| Server / intranet gateway / unattended long-running | **Headless** (single binary) | ~8 MB binary; RSS scales with the number of sources (~10 MB for one, ~32-35 MB for 16 — see [Memory footprint](#memory-footprint-measured-and-tuned)), zero-dep single process — drop it into systemd or any container, remote admin |
|
||||
| Personal desktop daily use / multi-device intranet sharing | **Desktop** (GUI) | no-login embedded WebUI, tray one-click, autostart, silent background — for non-CLI users |
|
||||
| Windows desktop | **Desktop** | plain backend needs manual service registration; GUI ships native tray/autostart |
|
||||
| CI one-shot 3-platform installers | **Desktop packaging scripts** | `make gui-*` emits deb / AppImage / NSIS, drops straight into a release pipeline |
|
||||
|
||||
Both share the exact same `llmsproxy` core (LuaJIT build) — configs and
|
||||
adapters are fully compatible, freely interchangeable.
|
||||
@ -327,30 +536,42 @@ cd cmd/gui
|
||||
npm install
|
||||
make gui # or npm run dev — dev run (needs make build first for bin/llmsproxy)
|
||||
make gui-deb # local deb (share with other Linux users)
|
||||
make gui-dist # deb + AppImage (linux)
|
||||
make gui-dist # deb + rpm + AppImage (linux; rpm needs system rpmbuild)
|
||||
make core-dist # core deb + rpm + tar.gz (server edition; needs nfpm)
|
||||
make gui-win # Windows NSIS (host mingw + wine)
|
||||
make gui-win-docker # Windows NSIS, fully dockerized (no host mingw needed)
|
||||
make gui-win-docker # Windows NSIS, fully dockerized (wine in image; no host mingw/node/wine)
|
||||
```
|
||||
|
||||
Artifacts land in `cmd/build/gui-dist/`: `ModelRouter-1.0.0.AppImage`,
|
||||
`modelrouter-gui_1.0.0_amd64.deb`, `ModelRouter Setup 1.0.0.exe` (win). rpm
|
||||
needs system `rpmbuild`.
|
||||
Artifacts land in `cmd/build/gui-dist/`: `ModelRouter-<ver>.AppImage`,
|
||||
`modelrouter-gui_<ver>_amd64.deb`, `ModelRouter Setup <ver>.exe` (win); core
|
||||
packages land in `cmd/build/dist/`: `llmsproxy_<ver>_amd64.deb`,
|
||||
`llmsproxy-<ver>.x86_64.rpm`. Both target families run
|
||||
`packaging/verify-dist.sh` afterwards — a size floor that rejects degenerate
|
||||
installers (1.3.0 shipped a 264 KB NSIS stub because the host wine was broken
|
||||
and nothing caught it). At release time artifacts are renamed into the
|
||||
`ModelRouter-Desktop-*` (GUI) and `ModelRouter-Headless-*` (plain backend)
|
||||
families before upload to GitCode Releases.
|
||||
|
||||
#### Windows packaging (dockerized, recommended)
|
||||
|
||||
`make gui-win-docker` does cross-compile + NSIS in one step; the host needs
|
||||
**no mingw**.
|
||||
`make gui-win-docker` does cross-compile + NSIS **entirely in docker**; the host
|
||||
needs **no mingw, no wine, no node**. Wine (with the i386 runtime) lives inside
|
||||
the image, fixing the broken-host-wine failure that produced the 264 KB stub
|
||||
installer.
|
||||
|
||||
- Image: `cmd/gui/docker/win-builder/Dockerfile` (`golang:1.25` + mingw-w64 +
|
||||
prebuilt LuaJIT for Windows: `lua51.dll` + import lib). LuaJIT source comes
|
||||
from the gitcode mirror (reachable in CN) with github fallback; Go modules go
|
||||
through `https://goproxy.cn`.
|
||||
- Core build: `cmd/gui/scripts/win-core-docker.sh` — emits
|
||||
`cmd/gui/bin/{llmsproxy.exe,lua51.dll}` inside the container.
|
||||
- One-shot release: `cmd/gui/scripts/dist-win-docker.sh` — after the core,
|
||||
runs electron-builder (host needs `wine` for NSIS).
|
||||
- The `modelrouter/win-builder` image is auto-built on first run and reused
|
||||
afterwards.
|
||||
node + wine32/wine64 + prebuilt LuaJIT for Windows: `lua51.dll` + import lib).
|
||||
LuaJIT source comes from the gitcode mirror (reachable in CN) with github
|
||||
fallback; Go modules go through `https://goproxy.cn`.
|
||||
- One-shot release: `cmd/gui/scripts/dist-win-docker.sh` — core build + NSIS
|
||||
packaging run inside the container (`USE_SYSTEM_WINE=true` uses the image's
|
||||
wine); electron-builder's ~400 MB toolchain cache lives in a docker volume so
|
||||
later builds reuse it.
|
||||
- The `modelrouter/win-builder` image is auto-built on first run (~5-10 min)
|
||||
and reused afterwards.
|
||||
- After packaging, `packaging/verify-dist.sh` runs as a gate: an exe under
|
||||
5 MB (or a deb/rpm/7z under 10 MB) fails the build — a broken installer can
|
||||
never reach a Release again.
|
||||
|
||||
```bash
|
||||
# Full Windows release (rebuild core + NSIS):
|
||||
@ -383,6 +604,17 @@ internal/gateway # OpenAI-compatible HTTP + auth + SDK/streaming + WebUI
|
||||
internal/types # unified format & OpenAI wire types
|
||||
```
|
||||
|
||||
## Development & collaboration
|
||||
|
||||
- **Git branching workflow**: main / feature / release-branch conventions in
|
||||
**[docs/git-workflow-en.md](docs/git-workflow-en.md)**
|
||||
([中文](docs/git-workflow.md)). The line: `main` is always deployable →
|
||||
features merge into `main` → cut `release/vX.Y.Z` and tag → release → hotfixes
|
||||
land on the release branch and are cherry-picked back into `main`.
|
||||
- **Lua adapter protocol**: see
|
||||
**[docs/lua-adapters-en.md](docs/lua-adapters-en.md)**
|
||||
([中文](docs/lua-adapters.md)).
|
||||
|
||||
## Tests
|
||||
|
||||
```bash
|
||||
|
||||
@ -1,11 +1,20 @@
|
||||
# ModelRouter GUI Windows builder
|
||||
# Self-sufficient cross-compilation environment for the embedded llmsproxy core.
|
||||
# Produces llmsproxy.exe (Go, windows/amd64, luajit tag) + lua51.dll (LuaJIT mingw build)
|
||||
# inside a mounted workspace: cmd/gui/bin/{llmsproxy.exe,lua51.dll}
|
||||
# Self-sufficient cross-compilation environment for the embedded llmsproxy core
|
||||
# AND for packaging the Windows NSIS installer.
|
||||
#
|
||||
# Produces inside a mounted workspace:
|
||||
# - llmsproxy.exe (Go, windows/amd64, luajit tag) + lua51.dll
|
||||
# - ModelRouter Setup <ver>.exe (electron-builder --win nsis; requires wine
|
||||
# so the uninstaller is generated — set USE_SYSTEM_WINE=true)
|
||||
#
|
||||
# Build: docker build -t modelrouter/win-builder cmd/gui/docker/win-builder
|
||||
# Usage: make gui-win-docker (or cmd/gui/scripts/dist-win-docker.sh)
|
||||
#
|
||||
# This image is the SINGLE wine provider for the NSIS step, so a host that has
|
||||
# no / broken wine (e.g. amd64-only Debian: no syswow64 -> c0000135) can still
|
||||
# produce a real installer. The NSIS step MUST run inside this image; running it
|
||||
# on the host was how the 1.3.0 264 KB degenerate exe got shipped.
|
||||
#
|
||||
# LuaJIT source is fetched from gitcode.com (CN-reachable mirror) with github
|
||||
# fallback; pass --build-arg LUAJIT_REPO to override.
|
||||
|
||||
@ -13,11 +22,26 @@ FROM golang:1.25
|
||||
|
||||
ARG LUAJIT_REPO=https://gitcode.com/openresty/luajit2.git
|
||||
|
||||
RUN apt-get update -qq \
|
||||
# i386 arch is required for wine32 (the NSIS uninstaller step runs a win32 exe
|
||||
# through wine's wow64 path, which needs the 32-bit runtime under /usr/lib/i386).
|
||||
RUN dpkg --add-architecture i386 \
|
||||
&& apt-get update -qq \
|
||||
&& DEBIAN_FRONTEND=noninteractive apt-get install -y -qq \
|
||||
gcc-mingw-w64-x86-64 make ca-certificates git \
|
||||
# NSIS step needs node + the electron-builder toolchain
|
||||
nodejs npm \
|
||||
# wine (32+64) for generating the NSIS uninstaller
|
||||
wine64 wine32:i386 \
|
||||
&& rm -rf /var/lib/apt/lists/*
|
||||
|
||||
# electron-builder resolves the uninstaller via the system wine when
|
||||
# USE_SYSTEM_WINE=true (or by default on Linux without a toolset). Pin it.
|
||||
# The toolchain cache lives in $HOME/.cache (mounted as a volume by
|
||||
# dist-win-docker.sh), so it is reused across builds and never pollutes the
|
||||
# mounted workspace.
|
||||
ENV USE_SYSTEM_WINE=true
|
||||
ENV WINEDEBUG=-all
|
||||
|
||||
# LuaJIT cross-compiled for Windows; keeps lua51.dll and the cgo import lib
|
||||
# (libluajit-5.1.dll.a -> libluajit-5.1.a in the mingw lib dir).
|
||||
# NOTE: $LUAJIT_REPO deliberately unquoted in "sh -c" — quoting the ARG would
|
||||
|
||||
119
cmd/gui/main.js
119
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 = () =>
|
||||
@ -198,6 +273,17 @@ function startCore() {
|
||||
cwd: PROFILE_DIR,
|
||||
env: Object.assign({}, process.env, {
|
||||
LLMS_PROXY_GUI: "1",
|
||||
// Memory tuning, measured on a 16-source deployment (see README 内存占用):
|
||||
// MALLOC_ARENA_MAX caps glibc per-thread malloc arenas. LuaJIT allocates
|
||||
// through cgo -> glibc malloc, and glibc defaults to 8*nproc arenas, so
|
||||
// every OS thread touching malloc reserves its own ~1 MB arena that is
|
||||
// never returned to the OS.
|
||||
// GOGC halves the Go heap growth target. On its own it does not help (the
|
||||
// saved heap is immediately eaten by extra glibc arenas); only the pair
|
||||
// lowers settled RSS. The core is I/O bound, so the extra GC is free.
|
||||
// Both are overridable: a user who exports them keeps their own values.
|
||||
GOGC: process.env.GOGC || "50",
|
||||
MALLOC_ARENA_MAX: process.env.MALLOC_ARENA_MAX || "2",
|
||||
}),
|
||||
stdio: ["ignore", "pipe", "pipe"],
|
||||
});
|
||||
@ -263,7 +349,6 @@ function waitForCoreReady(timeoutMs = 10000) {
|
||||
let coreLastExit = 0;
|
||||
|
||||
function stopCore() {
|
||||
app.isQuitting = true;
|
||||
coreReady = false;
|
||||
if (coreProc && coreProc.exitCode === null) {
|
||||
try {
|
||||
|
||||
10
cmd/gui/package-lock.json
generated
10
cmd/gui/package-lock.json
generated
@ -1,12 +1,12 @@
|
||||
{
|
||||
"name": "modelrouter-gui",
|
||||
"version": "1.0.0",
|
||||
"version": "1.7.6",
|
||||
"lockfileVersion": 3,
|
||||
"requires": true,
|
||||
"packages": {
|
||||
"": {
|
||||
"name": "modelrouter-gui",
|
||||
"version": "1.0.0",
|
||||
"version": "1.7.6",
|
||||
"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,
|
||||
@ -4045,4 +4045,4 @@
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
@ -1,6 +1,6 @@
|
||||
{
|
||||
"name": "modelrouter-gui",
|
||||
"version": "1.0.0",
|
||||
"version": "1.7.6",
|
||||
"description": "ModelRouter Desktop — embedded ModelRouter core with tray",
|
||||
"author": "ModelRouter",
|
||||
"main": "main.js",
|
||||
@ -10,7 +10,7 @@
|
||||
"dev": "electron . --no-sandbox --dev",
|
||||
"core:linux": "node scripts/prepare-core.js linux",
|
||||
"core:win": "node scripts/prepare-core.js win",
|
||||
"dist:linux": "npm run core:linux && electron-builder --linux deb AppImage",
|
||||
"dist:linux": "npm run core:linux && electron-builder --linux deb rpm AppImage",
|
||||
"dist:debian": "npm run core:linux && electron-builder --linux deb",
|
||||
"dist:win": "npm run core:win && electron-builder --win nsis",
|
||||
"dist:dir": "npm run core:linux && electron-builder --dir"
|
||||
@ -22,7 +22,7 @@
|
||||
"build": {
|
||||
"appId": "com.modelrouter.gui",
|
||||
"productName": "ModelRouter",
|
||||
"linux": {
|
||||
"linux": {
|
||||
"target": [
|
||||
"deb",
|
||||
"rpm",
|
||||
@ -32,9 +32,9 @@
|
||||
"icon": "build/icon.png",
|
||||
"maintainer": "ModelRouter",
|
||||
"synopsis": "ModelRouter Desktop",
|
||||
"desktop": {
|
||||
"desktop": {
|
||||
"entry": {
|
||||
"StartupWMClass": "modelrouter-gui"
|
||||
"StartupWMClass": "modelrouter-gui"
|
||||
}
|
||||
},
|
||||
"extraResources": [
|
||||
|
||||
@ -26,16 +26,6 @@ const state = {
|
||||
|
||||
const $ = (s) => document.querySelector(s);
|
||||
|
||||
function esc(s) {
|
||||
return String(s == null ? "" : s).replace(
|
||||
/[&<>"']/g,
|
||||
(c) =>
|
||||
({ "&": "&", "<": "<", ">": ">", '"': """, "'": "'" })[
|
||||
c
|
||||
],
|
||||
);
|
||||
}
|
||||
|
||||
function toast(msg, isErr, ms) {
|
||||
const t = $("#toast");
|
||||
t.textContent = msg;
|
||||
|
||||
@ -1,23 +1,64 @@
|
||||
#!/usr/bin/env bash
|
||||
# One-shot Windows release build for the GUI:
|
||||
# One-shot Windows release build for the GUI — ENTIRELY inside docker.
|
||||
#
|
||||
# 1) cross-compile the embedded core in docker (llmsproxy.exe + lua51.dll)
|
||||
# 2) host-side electron-builder NSIS packaging (wine required on host)
|
||||
# 2) electron-builder NSIS packaging ALSO inside docker (wine for the
|
||||
# uninstaller step lives in the image — see Dockerfile). This is the fix
|
||||
# for the 1.3.0/1.4.0 broken installers: the host had amd64-only wine
|
||||
# (no syswow64 -> c0000135), which made electron-builder silently emit a
|
||||
# 264 KB degenerate exe without its 86 MB payload.
|
||||
#
|
||||
# Outputs: cmd/build/gui-dist/"ModelRouter Setup <version>.exe"
|
||||
# cmd/build/gui-dist/modelrouter-gui-<version>-x64.nsis.7z
|
||||
#
|
||||
# Needs: docker (any machine; no host mingw/wine/node required)
|
||||
set -euo pipefail
|
||||
|
||||
ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")/../../.." && pwd)"
|
||||
cd "$ROOT/cmd/gui"
|
||||
IMAGE="${WIN_BUILDER_IMAGE:-modelrouter/win-builder}"
|
||||
PROXY="${GOPROXY:-https://goproxy.cn,direct}"
|
||||
CACHE_VOL="modelrouter-eb-cache"
|
||||
|
||||
# 1) core (docker)
|
||||
"$ROOT/cmd/gui/scripts/win-core-docker.sh"
|
||||
command -v docker >/dev/null || { echo "docker required"; exit 1; }
|
||||
|
||||
# 2) packaging (host, needs wine for electron-builder's signtool/nsis)
|
||||
command -v wine >/dev/null || { echo "wine required on host for NSIS packaging"; exit 1; }
|
||||
if ! command -v npx >/dev/null; then echo "node/npm required"; exit 1; fi
|
||||
if ! docker image inspect "$IMAGE" >/dev/null 2>&1; then
|
||||
echo "[dist-win-docker] building image $IMAGE (first run, ~5-10 min)..."
|
||||
docker build -t "$IMAGE" "$ROOT/cmd/gui/docker/win-builder"
|
||||
fi
|
||||
|
||||
echo "[dist-win] electron-builder --win nsis ..."
|
||||
# bin artifacts already refreshed by win-core-docker.sh; skip the prepare step
|
||||
npx electron-builder --win nsis
|
||||
docker volume create "$CACHE_VOL" >/dev/null 2>&1 || true
|
||||
|
||||
ls -la "$ROOT/cmd/build/gui-dist/"*.exe 2>/dev/null
|
||||
echo "[dist-win] done"
|
||||
# 1) embedded core + 2) NSIS packaging — both inside the image so wine and node
|
||||
# are the image's, never the host's. electron-builder downloads electron + nsis
|
||||
# toolchains into ~/.cache on first run; the mounted volume keeps that warm so
|
||||
# subsequent builds are incremental instead of re-fetching ~400 MB.
|
||||
echo "[dist-win-docker] cross-compiling core + electron-builder NSIS in docker ..."
|
||||
docker run --rm \
|
||||
-v "$ROOT":/workspace \
|
||||
-v "$CACHE_VOL":/root/.cache \
|
||||
-e GOPROXY="$PROXY" \
|
||||
-w /workspace \
|
||||
"$IMAGE" bash -c '
|
||||
set -euo pipefail
|
||||
cp /opt/luajit2/src/lua51.dll /workspace/cmd/gui/bin/lua51.dll
|
||||
GOOS=windows GOARCH=amd64 CGO_ENABLED=1 CC=x86_64-w64-mingw32-gcc \
|
||||
go build -buildvcs=false -tags luajit -trimpath \
|
||||
-o /workspace/cmd/gui/bin/llmsproxy.exe ./cmd/llmsproxy
|
||||
echo " -> core built"
|
||||
# wine needs one prefix initialization so its wow64 runtime (syswow64/ntdll)
|
||||
# exists — the exact thing that was missing on the broken host. Use a shared
|
||||
# prefix in the cache volume so it is initialized only on the first run.
|
||||
export WINEPREFIX=/root/.cache/wineprefix
|
||||
if [ ! -f "$WINEPREFIX/drive_c/windows/syswow64/ntdll.dll" ]; then
|
||||
echo " -> initializing wine prefix (first run)"
|
||||
WINEDEBUG=-all wineboot -i || true
|
||||
fi
|
||||
cd /workspace/cmd/gui
|
||||
# node_modules ships in the repo (host npm install already ran); reuse it.
|
||||
npx electron-builder --win nsis
|
||||
'
|
||||
|
||||
echo "[dist-win-docker] verifying artifacts ..."
|
||||
"$ROOT/packaging/verify-dist.sh" "$ROOT/cmd/build/gui-dist"
|
||||
|
||||
echo "[dist-win-docker] done"
|
||||
@ -23,8 +23,39 @@ 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
|
||||
// without starting the Lua VM, touching runtime.json, or binding a port.
|
||||
// It deliberately does NOT call EnsureDefault, so a missing file is an
|
||||
// error here instead of being silently created.
|
||||
if *checkOnly {
|
||||
if _, err := os.Stat(*cfgPath); err != nil {
|
||||
log.Fatalf("[llmsproxy] check: %v", err)
|
||||
}
|
||||
if _, err := config.Load(*cfgPath); err != nil {
|
||||
log.Fatalf("[llmsproxy] check: %v", err)
|
||||
}
|
||||
log.Printf("[llmsproxy] check: %s is valid", *cfgPath)
|
||||
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)
|
||||
@ -42,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)
|
||||
}
|
||||
}
|
||||
}
|
||||
@ -88,4 +126,4 @@ func main() {
|
||||
signal.Notify(stop, syscall.SIGINT, syscall.SIGTERM)
|
||||
<-stop
|
||||
log.Printf("[llmsproxy] shutting down")
|
||||
}
|
||||
}
|
||||
|
||||
479
deploy.sh
Executable file
479
deploy.sh
Executable file
@ -0,0 +1,479 @@
|
||||
#!/usr/bin/env bash
|
||||
# ============================================================
|
||||
# llmsproxy 原子化滚动部署脚本
|
||||
#
|
||||
# 策略:
|
||||
# 1. /tmp (tmpfs) 构建,/usr/local/bin (ext4) 部署
|
||||
# 2. 构建后先 copy 到目标文件系统(.llmsproxy.tmp),
|
||||
# 再 rename(2) 做原子替换(同一文件系统内 mv 是原子的)
|
||||
# 3. runtime.json 不替换(保留运行时密钥、审计状态)
|
||||
# 4. systemctl restart → 服务中断仅在 Go 进程重启瞬间 (<500ms)
|
||||
# 5. 健康检查确认新进程正常
|
||||
#
|
||||
# 配置文件(2026-09-05 起):
|
||||
# 本脚本不再假设 config.yaml 与部署无关。实测教训:给源加 headers 时
|
||||
# 漏删原有的 `headers: {}` → YAML 重复键 → 进程启动即退出 → 健康检查
|
||||
# 失败 → 旧的 rollback 只恢复二进制,配置错误仍在,于是旧二进制也起不来,
|
||||
# 服务陷入 systemd 重启循环(当时累计 446 次)。
|
||||
# 因此:
|
||||
# - 支持 --config <file> 投放一份待部署配置(原子 rename 到 TARGET_CONFIG)
|
||||
# - 重启前先用新二进制 `-check` 预检配置,不合法就地失败,绝不重启服务
|
||||
# - rollback 同时恢复二进制与配置文件,回滚后再次 -check 才重启
|
||||
#
|
||||
# 注意事项:
|
||||
# - 本服务依赖自身(推理依赖),restart 前后会短暂中断,
|
||||
# 但重试机制(systemd RestartSec=5 + 客户端 retry)可自愈
|
||||
# - 适配器文件替换后需重启才能生效(VM 启动时批量加载)
|
||||
# ============================================================
|
||||
set -euo pipefail
|
||||
|
||||
# ---------- 配置 ----------
|
||||
readonly TARGET_BIN="/usr/local/bin/llmsproxy"
|
||||
readonly TARGET_CONFIG="/etc/llmsproxy/config.yaml"
|
||||
readonly TARGET_ADAPTERS="/etc/llmsproxy/adapters"
|
||||
readonly REPO_DIR="/home/program/llmsproxy"
|
||||
readonly HEALTH_URL="http://127.0.0.1:8081/v1/models"
|
||||
readonly BIN_NAME="llmsproxy"
|
||||
|
||||
# 构建目录在 /tmp(tmpfs 更快)
|
||||
readonly TMP_DIR="/tmp/llmsproxy-deploy"
|
||||
# 同文件系统暂存文件(用于原子 rename)
|
||||
readonly STAGING_BIN="${TARGET_BIN}.staging"
|
||||
readonly STAGING_CONFIG="${TARGET_CONFIG}.staging"
|
||||
# 回滚备份:健康检查失败时从这里恢复旧二进制 / 旧配置
|
||||
readonly BACKUP_BIN="${TARGET_BIN}.bak"
|
||||
readonly BACKUP_CONFIG="${TARGET_CONFIG}.rollback.bak"
|
||||
|
||||
# 待部署配置(--config <file> 指定;为空 = 不改动线上配置)
|
||||
NEW_CONFIG_SRC=""
|
||||
|
||||
usage() {
|
||||
cat <<'EOF'
|
||||
用法: ./deploy.sh [--config <file>]
|
||||
|
||||
--config <file> 投放这份配置到 /etc/llmsproxy/config.yaml(原子替换)。
|
||||
省略时不改动线上配置,只部署二进制与适配器。
|
||||
无论是否指定,重启前都会用新二进制 -check 预检线上配置。
|
||||
|
||||
流程: 构建 → 适配器同步 → 备份(二进制+配置) → 原子替换 → 配置预检
|
||||
→ 重启 → 健康检查 → 失败则回滚(二进制+配置)并重启
|
||||
EOF
|
||||
}
|
||||
|
||||
parse_args() {
|
||||
while (( $# > 0 )); do
|
||||
case "$1" in
|
||||
--config)
|
||||
[[ $# -ge 2 ]] || { echo "--config 需要一个文件参数" >&2; exit 2; }
|
||||
NEW_CONFIG_SRC="$2"
|
||||
shift 2
|
||||
;;
|
||||
-h|--help) usage; exit 0 ;;
|
||||
*) echo "未知参数: $1" >&2; usage; exit 2 ;;
|
||||
esac
|
||||
done
|
||||
}
|
||||
|
||||
# ---------- 日志 ----------
|
||||
log() { printf '\033[1;34m[%s]\033[0m %s\n' "$(date +%H:%M:%S)" "$*"; }
|
||||
warn() { printf '\033[1;33m[WARN]\033[0m %s\n' "$*"; }
|
||||
err() { printf '\033[1;31m[ERR]\033[0m %s\n' "$*" >&2; }
|
||||
fail() { err "$@"; exit 1; }
|
||||
|
||||
# ---------- 清理 ----------
|
||||
cleanup() {
|
||||
log "清理临时文件"
|
||||
rm -rf "$TMP_DIR" "$STAGING_BIN" "$STAGING_CONFIG"
|
||||
}
|
||||
trap cleanup EXIT
|
||||
|
||||
# ---------- 前置检查 ----------
|
||||
precheck() {
|
||||
log "前置检查"
|
||||
[[ -f "$REPO_DIR/go.mod" ]] || fail "仓库目录 $REPO_DIR 不存在"
|
||||
command -v go &>/dev/null || fail "未找到 go"
|
||||
command -v systemctl &>/dev/null || fail "未找到 systemctl"
|
||||
[[ $(id -u) -eq 0 ]] || fail "需要 root 权限"
|
||||
|
||||
if [[ -n "$NEW_CONFIG_SRC" ]]; then
|
||||
[[ -f "$NEW_CONFIG_SRC" ]] || fail "待部署配置不存在: $NEW_CONFIG_SRC"
|
||||
# 同一文件比对:--config 指向线上文件本身时没有"投放"动作可做,
|
||||
# 但仍然会走重启前的 -check 预检。
|
||||
if [[ "$(readlink -f "$NEW_CONFIG_SRC")" == "$(readlink -f "$TARGET_CONFIG" 2>/dev/null)" ]]; then
|
||||
warn "--config 指向线上配置本身,跳过投放(仍会预检)"
|
||||
NEW_CONFIG_SRC=""
|
||||
else
|
||||
log "待部署配置: $NEW_CONFIG_SRC → $TARGET_CONFIG"
|
||||
fi
|
||||
fi
|
||||
|
||||
OLD_SIZE=$(stat -c%s "$TARGET_BIN" 2>/dev/null || echo 0)
|
||||
OLD_SHA256=$(sha256sum "$TARGET_BIN" 2>/dev/null | awk '{print $1}')
|
||||
|
||||
log "旧二进制: $TARGET_BIN (${OLD_SIZE} bytes, sha256=${OLD_SHA256:0:16})"
|
||||
log "go: $(go version)"
|
||||
}
|
||||
|
||||
# ---------- 备份旧二进制 + 旧配置(用于回滚)----------
|
||||
# 配置也必须备份:启动失败常常是配置问题,只回滚二进制会让旧二进制
|
||||
# 继续吃坏配置,服务照样起不来。
|
||||
backup_old_state() {
|
||||
log "备份旧二进制到 $BACKUP_BIN"
|
||||
if [[ -f "$TARGET_BIN" ]]; then
|
||||
cp -f "$TARGET_BIN" "$BACKUP_BIN"
|
||||
log "✓ 旧二进制已备份 (sha256=$(sha256sum "$BACKUP_BIN" | awk '{print $1}' | cut -c1-16))"
|
||||
else
|
||||
warn "旧二进制不存在,跳过备份(无回滚点)"
|
||||
fi
|
||||
|
||||
if [[ -f "$TARGET_CONFIG" ]]; then
|
||||
cp -f "$TARGET_CONFIG" "$BACKUP_CONFIG"
|
||||
log "✓ 旧配置已备份到 $BACKUP_CONFIG (sha256=$(sha256sum "$BACKUP_CONFIG" | awk '{print $1}' | cut -c1-16))"
|
||||
else
|
||||
warn "线上配置不存在: $TARGET_CONFIG(首次部署?无配置回滚点)"
|
||||
fi
|
||||
}
|
||||
|
||||
# ---------- 构建 ----------
|
||||
build() {
|
||||
log "构建新二进制"
|
||||
rm -rf "$TMP_DIR"
|
||||
mkdir -p "$TMP_DIR/bin"
|
||||
|
||||
cd "$REPO_DIR"
|
||||
CGO_ENABLED=1 go build -tags luajit \
|
||||
-trimpath \
|
||||
-ldflags="-s -w" \
|
||||
-o "$TMP_DIR/bin/$BIN_NAME" \
|
||||
./cmd/llmsproxy
|
||||
|
||||
NEW_BIN="$TMP_DIR/bin/$BIN_NAME"
|
||||
NEW_SIZE=$(stat -c%s "$NEW_BIN")
|
||||
|
||||
if [[ "$NEW_SIZE" -eq "$OLD_SIZE" ]]; then
|
||||
NEW_SHA256=$(sha256sum "$NEW_BIN" | awk '{print $1}')
|
||||
if [[ "$NEW_SHA256" == "$OLD_SHA256" ]]; then
|
||||
# 二进制无变化。只有在同时也没有待部署配置时才能直接退出;
|
||||
# 否则会造成“--config 指定了新配置却被静默丢弃”。
|
||||
if [[ -z "$NEW_CONFIG_SRC" ]]; then
|
||||
warn "新二进制与旧版本完全一致(sha256 相同)且无待部署配置,跳过部署"
|
||||
exit 0
|
||||
fi
|
||||
log "二进制无变化(sha256 相同),但有待部署配置,继续"
|
||||
fi
|
||||
fi
|
||||
|
||||
log "新二进制: $NEW_BIN (${NEW_SIZE} bytes)"
|
||||
|
||||
# 验证:确认包含关键修复
|
||||
local checks=(
|
||||
"mergeUsage"
|
||||
"prompt_tokens"
|
||||
"completion_tokens"
|
||||
"message_start"
|
||||
)
|
||||
for c in "${checks[@]}"; do
|
||||
if strings "$NEW_BIN" | grep -qF "$c"; then
|
||||
log " ✓ 新二进制包含: $c"
|
||||
else
|
||||
warn " ⚠ 新二进制不包含: $c(非预期但非致命)"
|
||||
fi
|
||||
done
|
||||
}
|
||||
|
||||
# ---------- 原子替换二进制 ----------
|
||||
# 原理:
|
||||
# /tmp 是 tmpfs,/usr/local/bin 是 ext4——不同文件系统 rename(2) 不原子。
|
||||
# 因此先 copy 到目标文件系统上的 .staging 文件,再 rename(2) 到目标路径。
|
||||
# rename(2) 在同一文件系统内是元数据级操作,对正在读写的进程安全:
|
||||
# - 旧 inode 的 fd 继续有效(进程持有旧文件,不会读到截断内容)
|
||||
# - 新进程打开 path 看到新 inode
|
||||
# - 旧 inode 在最后一个 fd 关闭后释放
|
||||
atomic_replace_binary() {
|
||||
log "原子替换二进制"
|
||||
|
||||
local new_bin_tmp="$TMP_DIR/bin/$BIN_NAME"
|
||||
[[ -f "$new_bin_tmp" ]] || fail "新二进制文件不存在: $new_bin_tmp"
|
||||
|
||||
# Step 1: copy 到目标文件系统(非原子,但目标文件唯一)
|
||||
log " copy → $STAGING_BIN (同文件系统暂存)"
|
||||
cp -f "$new_bin_tmp" "$STAGING_BIN"
|
||||
chmod 0755 "$STAGING_BIN"
|
||||
|
||||
# Step 2: 原子 rename(同一文件系统,元数据级操作)
|
||||
log " rename → $TARGET_BIN (原子)"
|
||||
mv -f "$STAGING_BIN" "$TARGET_BIN"
|
||||
|
||||
# 验证
|
||||
NEW_SHA256=$(sha256sum "$TARGET_BIN" | awk '{print $1}')
|
||||
log "✓ 二进制已原子替换 (sha256=${NEW_SHA256:0:16})"
|
||||
}
|
||||
|
||||
# ---------- 原子替换配置 ----------
|
||||
# 与二进制同一套 copy→rename 手法:先落到 /etc 上的 .staging,再 rename(2)。
|
||||
# 只有 --config 指定了来源时才动线上配置。
|
||||
atomic_replace_config() {
|
||||
[[ -n "$NEW_CONFIG_SRC" ]] || return 0
|
||||
|
||||
log "原子替换配置"
|
||||
log " copy → $STAGING_CONFIG (同文件系统暂存)"
|
||||
cp -f "$NEW_CONFIG_SRC" "$STAGING_CONFIG"
|
||||
chmod 0644 "$STAGING_CONFIG"
|
||||
|
||||
log " rename → $TARGET_CONFIG (原子)"
|
||||
mv -f "$STAGING_CONFIG" "$TARGET_CONFIG"
|
||||
log "✓ 配置已原子替换 (sha256=$(sha256sum "$TARGET_CONFIG" | awk '{print $1}' | cut -c1-16))"
|
||||
}
|
||||
|
||||
# ---------- 配置预检 ----------
|
||||
# 用刚部署的新二进制解析线上配置。`-check` 只做 parse + ApplyDefaults,
|
||||
# 不启动 Lua VM、不碰 runtime.json、不绑端口,因此可以安全地在服务仍在
|
||||
# 运行时执行。不合法就在重启前失败——这正是上次事故缺的那一步。
|
||||
verify_config() {
|
||||
log "配置预检(新二进制 -check)"
|
||||
[[ -f "$TARGET_CONFIG" ]] || fail "线上配置不存在: $TARGET_CONFIG"
|
||||
|
||||
local out
|
||||
if out=$("$TARGET_BIN" -check -config "$TARGET_CONFIG" 2>&1); then
|
||||
log "✓ 配置合法"
|
||||
return 0
|
||||
fi
|
||||
|
||||
err "配置预检失败,未重启服务(线上进程仍在跑旧配置):"
|
||||
printf '%s\n' "$out" >&2
|
||||
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
|
||||
}
|
||||
|
||||
# ---------- 同步适配器 ----------
|
||||
sync_adapters() {
|
||||
log "同步适配器"
|
||||
|
||||
# 只覆盖内置适配器(与 internal/lua/adapters 同名的文件),
|
||||
# 保留目录里其它运行时上传的 .lua —— WebUI 上传的自定义适配器
|
||||
# 必须跨部署存活,清空会静默移除线上源依赖的适配器。
|
||||
SRC_ADAPTERS="$REPO_DIR/internal/lua/adapters"
|
||||
[[ -d "$SRC_ADAPTERS" ]] || fail "内嵌适配器目录不存在: $SRC_ADAPTERS"
|
||||
|
||||
# 备份旧的适配器文件(含任何运行时覆盖版本,如 opencode.lua)
|
||||
BACKUP_DIR="/etc/llmsproxy/adapters.bak.$(date +%Y%m%d%H%M%S)"
|
||||
if [[ -d "$TARGET_ADAPTERS" ]]; then
|
||||
mkdir -p "$BACKUP_DIR"
|
||||
cp -rf "$TARGET_ADAPTERS/"*.lua "$BACKUP_DIR/" 2>/dev/null || true
|
||||
log " 旧适配器已备份到 $BACKUP_DIR"
|
||||
fi
|
||||
|
||||
cp -f "$SRC_ADAPTERS/"*.lua "$TARGET_ADAPTERS/"
|
||||
chmod 0644 "$TARGET_ADAPTERS/"*.lua
|
||||
|
||||
# 校验:每个适配器必须包含 usage 透传修复(opencode/openai/deepseek 等应有 'uses')
|
||||
local required_files=(openai deepseek anthropic gemini ollama opencode github groq kimicode mistral sensenova agentrouter)
|
||||
for f in "${required_files[@]}"; do
|
||||
[[ -f "$TARGET_ADAPTERS/$f.lua" ]] || warn " 缺少适配器: $f.lua"
|
||||
done
|
||||
local usage_fixed=0
|
||||
for f in openai deepseek github groq kimicode mistral opencode; do
|
||||
grep -qE '\buses\b' "$TARGET_ADAPTERS/$f.lua" 2>/dev/null && ((usage_fixed++)) || true
|
||||
done
|
||||
log " usage 透传修复适配器: ${usage_fixed}/7 (openai/deepseek/github/groq/kimicode/mistral/opencode)"
|
||||
log "✓ 适配器已同步到 $TARGET_ADAPTERS/"
|
||||
}
|
||||
|
||||
# ---------- 重启服务 ----------
|
||||
restart_service() {
|
||||
log "重启服务(中断 <500ms)"
|
||||
systemctl restart llmsproxy.service || fail "systemd restart 失败"
|
||||
log "✓ llmsproxy.service restarted"
|
||||
}
|
||||
|
||||
# ---------- 健康检查 ----------
|
||||
healthcheck() {
|
||||
log "健康检查"
|
||||
local attempts=20
|
||||
local wait_sec=1
|
||||
local attempt=1
|
||||
|
||||
while (( attempt <= attempts )); do
|
||||
# 用 curl 检查(200=有 key 未鉴权但服务正常,401=需要认证,都说明服务在运行)
|
||||
local code
|
||||
code=$(curl -s -o /dev/null -w "%{http_code}" --max-time 3 "$HEALTH_URL" 2>/dev/null || echo "000")
|
||||
if [[ "$code" == "200" || "$code" == "401" ]]; then
|
||||
log "✓ 服务健康 (status=$code, ${attempt}/${attempts})"
|
||||
return 0
|
||||
fi
|
||||
log " 等待中 ($attempt/$attempts, status=$code)"
|
||||
sleep "$wait_sec"
|
||||
(( attempt++ ))
|
||||
done
|
||||
|
||||
warn "健康检查未通过,请检查服务状态: systemctl status llmsproxy.service"
|
||||
return 1
|
||||
}
|
||||
|
||||
# ---------- 回滚 ----------
|
||||
# 健康检查失败时恢复旧二进制**和旧配置**并重启,保证推理服务可用。
|
||||
# 这是最后一道保险:即使新版本启动失败(编译错误、配置不兼容、
|
||||
# panic on start 等),也能让旧版本继续服务,避免推理能力丢失。
|
||||
#
|
||||
# 为什么必须一并回滚配置:启动失败最常见的原因就是配置本身不合法
|
||||
# (实测:YAML 重复键)。只恢复二进制的话,旧二进制照样解析不了坏配置,
|
||||
# 服务继续在 systemd 重启循环里空转。
|
||||
rollback() {
|
||||
err "健康检查失败,启动回滚流程"
|
||||
|
||||
local restored_any=0
|
||||
if [[ -f "$BACKUP_BIN" ]]; then
|
||||
log "恢复旧二进制"
|
||||
cp -f "$BACKUP_BIN" "$TARGET_BIN"
|
||||
chmod 0755 "$TARGET_BIN"
|
||||
restored_any=1
|
||||
else
|
||||
err "无二进制回滚备份 ($BACKUP_BIN 不存在)"
|
||||
fi
|
||||
|
||||
# 只在本次真的换过配置时才回滚配置,避免把用户在部署期间
|
||||
# 通过 WebUI 做的合法改动一起抹掉。
|
||||
if [[ -n "$NEW_CONFIG_SRC" && -f "$BACKUP_CONFIG" ]]; then
|
||||
log "恢复旧配置"
|
||||
cp -f "$BACKUP_CONFIG" "$TARGET_CONFIG"
|
||||
chmod 0644 "$TARGET_CONFIG"
|
||||
restored_any=1
|
||||
fi
|
||||
|
||||
if (( restored_any == 0 )); then
|
||||
err "无任何回滚点,无法自动恢复"
|
||||
err "请手动检查: systemctl status llmsproxy.service && journalctl -u llmsproxy -n 50"
|
||||
return 1
|
||||
fi
|
||||
|
||||
# 回滚后先验证恢复出来的组合真的能解析,再重启——否则只是把
|
||||
# 重启循环换个二进制继续跑。
|
||||
if ! "$TARGET_BIN" -check -config "$TARGET_CONFIG" >/dev/null 2>&1; then
|
||||
err "回滚后的配置仍不合法,拒绝重启(避免 systemd 重启风暴)"
|
||||
err " 手动修复 $TARGET_CONFIG 后执行: $TARGET_BIN -check -config $TARGET_CONFIG"
|
||||
err " 备份可用: 二进制=$BACKUP_BIN 配置=$BACKUP_CONFIG"
|
||||
return 1
|
||||
fi
|
||||
|
||||
log "重启服务(旧版本)"
|
||||
systemctl restart llmsproxy.service || {
|
||||
err "重启失败,服务可能完全宕机"
|
||||
err "请手动检查: systemctl status llmsproxy.service"
|
||||
return 1
|
||||
}
|
||||
|
||||
# 回滚后再健康检查(只查 1 次,失败就认了)
|
||||
sleep 2
|
||||
local code
|
||||
code=$(curl -s -o /dev/null -w "%{http_code}" --max-time 5 "$HEALTH_URL" 2>/dev/null || echo "000")
|
||||
if [[ "$code" == "200" || "$code" == "401" ]]; then
|
||||
log "✓ 回滚成功,旧版本已恢复服务 (status=$code)"
|
||||
log " 新版本部署失败,请检查构建/配置后重试"
|
||||
return 0
|
||||
else
|
||||
err "回滚后服务仍未恢复 (status=$code)"
|
||||
err "服务可能已完全宕机,请手动介入"
|
||||
return 1
|
||||
fi
|
||||
}
|
||||
|
||||
# ---------- 打印部署摘要 ----------
|
||||
print_summary() {
|
||||
log "═════════════════════════════════════════════"
|
||||
log "部署完成"
|
||||
log " 二进制: $TARGET_BIN"
|
||||
log " 配置: $TARGET_CONFIG$([[ -n "$NEW_CONFIG_SRC" ]] && echo " (本次已替换)" || echo " (未改动)")"
|
||||
log " 适配器: $TARGET_ADAPTERS"
|
||||
log "═════════════════════════════════════════════"
|
||||
}
|
||||
|
||||
# ---------- 主流程 ----------
|
||||
# 顺序要点:配置投放与预检都排在 restart_service **之前**,
|
||||
# 于是坏配置在服务仍然健康时就被拦下,根本不会触发重启。
|
||||
main() {
|
||||
parse_args "$@"
|
||||
log "═══ llmsproxy 原子部署开始 ═══"
|
||||
precheck
|
||||
# 密钥校验放在 build/替换之前:它的失败后果最重(配置里的密文永久解不开),
|
||||
# 必须在动任何文件之前就挡住。此阶段用的是**当前线上二进制**,它已经
|
||||
# 支持 -show-secrets(首次引入该子命令的版本除外,此时跳过并提示)。
|
||||
verify_master_key
|
||||
build
|
||||
sync_adapters
|
||||
backup_old_state
|
||||
atomic_replace_binary
|
||||
atomic_replace_config
|
||||
if ! verify_config; then
|
||||
# 配置不合法:服务还没重启,只需把二进制/配置退回原状。
|
||||
err "配置预检失败,回退本次改动(服务未受影响)"
|
||||
[[ -f "$BACKUP_BIN" ]] && cp -f "$BACKUP_BIN" "$TARGET_BIN" && chmod 0755 "$TARGET_BIN"
|
||||
if [[ -n "$NEW_CONFIG_SRC" && -f "$BACKUP_CONFIG" ]]; then
|
||||
cp -f "$BACKUP_CONFIG" "$TARGET_CONFIG" && chmod 0644 "$TARGET_CONFIG"
|
||||
fi
|
||||
exit 1
|
||||
fi
|
||||
restart_service
|
||||
if ! healthcheck; then
|
||||
rollback || true
|
||||
exit 1
|
||||
fi
|
||||
print_summary
|
||||
}
|
||||
|
||||
main "$@"
|
||||
141
docs/git-workflow-en.md
Normal file
141
docs/git-workflow-en.md
Normal file
@ -0,0 +1,141 @@
|
||||
# Git Branching Workflow
|
||||
|
||||
ModelRouter uses a **GitHub Flow + release-branch** model: `main` is the only
|
||||
long-lived branch and is always deployable; new work branches off it as a
|
||||
feature branch; each version gets its own release branch with a tag; hotfixes
|
||||
land on the release branch and MUST flow back to `main`.
|
||||
|
||||
Goal: make "patching an already-released version during its support window" a
|
||||
first-class operation, while `main` always contains every fix and stays
|
||||
deployable.
|
||||
|
||||
## Branch roles
|
||||
|
||||
| Branch | Name | Lifetime | Purpose |
|
||||
|---|---|---|---|
|
||||
| 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. |
|
||||
|
||||
## Change flow (important)
|
||||
|
||||
```
|
||||
feature/foo ───────┐
|
||||
▼
|
||||
main ──────────────── next version ──────►
|
||||
│ │
|
||||
│ cut │ cut
|
||||
▼ ▼
|
||||
release/v1.4.2 release/v1.4.3
|
||||
│ │
|
||||
tag: v1.4.2 tag: v1.4.3
|
||||
│ │
|
||||
hotfix ◄─────┘ hotfix ◄─────┘
|
||||
│ │
|
||||
└── cherry-pick back ──────────┘
|
||||
```
|
||||
|
||||
### Key rules
|
||||
|
||||
1. **main is always deployable**: never leave half-done work on `main`.
|
||||
Unfinished work lives on a feature branch.
|
||||
2. **Features start from main and merge back**: `git checkout -b feature/xxx main`,
|
||||
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
|
||||
```
|
||||
Build installers and upload the GitCode Release from this tag so the
|
||||
published state is exactly reproducible.
|
||||
4. **Only version-specific hotfixes go into a release branch during its
|
||||
window**: new features always go to `main` for the next version, never into
|
||||
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 commit -m "fix: ..."
|
||||
git checkout main
|
||||
git cherry-pick <hotfix-commit> # and into main
|
||||
```
|
||||
Most important rule: **main only moves forward and never loses fixes**. If a
|
||||
hotfix never reaches main, the next release ships with the old bug.
|
||||
|
||||
### End of a version lifecycle
|
||||
|
||||
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
|
||||
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.
|
||||
|
||||
## Explicit non-goals
|
||||
|
||||
- **Never rebase main**: main's history stays append-only; anyone pulling gets
|
||||
directly usable history.
|
||||
- **Release branches are not merged back wholesale**: hotfixes were already
|
||||
cherry-picked; a whole-branch merge only creates conflicts with no benefit.
|
||||
- **No "quick edit" directly on main**, even single-person: at minimum use
|
||||
`feature/xxx → merge main` so history keeps its why-boundaries.
|
||||
|
||||
## Single vs multi person
|
||||
|
||||
- **Single** (this repo today): skip PRs — `git checkout -b feature/xxx` →
|
||||
finish → merge back to main. Release branch + tag flow unchanged.
|
||||
- **Multi / open source**: feature branches go through PR + Code Review,
|
||||
maintainers merge; resolve conflicts on the feature branch, never on main.
|
||||
|
||||
## Commit message conventions
|
||||
|
||||
- `feat(...)`: new feature
|
||||
- `fix(...)`: bug fix
|
||||
- `chore(...)`: build / version / deps / non-code changes
|
||||
- `docs(...)`: documentation
|
||||
- `refactor(...)`: refactoring
|
||||
- Scope in parens, e.g. `fix(adapters)`, `feat(webui)`, `chore(version)`
|
||||
- One-line summary; expand with problem/root-cause/fix/verification when needed
|
||||
|
||||
## Release checklist (companion)
|
||||
|
||||
```bash
|
||||
# 1. main is ready
|
||||
git checkout main && git pull
|
||||
|
||||
# 2. cut release branch + tag
|
||||
git checkout -b release/vX.Y.Z main
|
||||
git tag -a vX.Y.Z -m "ModelRouter vX.Y.Z"
|
||||
|
||||
# 3. push branch + tag
|
||||
git push origin release/vX.Y.Z
|
||||
git push origin vX.Y.Z
|
||||
|
||||
# 4. build installers (core + GUI platforms)
|
||||
make core-dist VERSION=X.Y.Z
|
||||
|
||||
# 5. create GitCode Release + upload
|
||||
# https://gitcode.com/JianFeeeee/ModelRouter/releases
|
||||
# upload MUST use PUT --http1.1, otherwise the OBS callback fails and the
|
||||
# file never lands in the release; replacing artifacts = delete the tag
|
||||
# (release disappears with it) then recreate.
|
||||
```
|
||||
|
||||
## Why this model (background)
|
||||
|
||||
Before 2026-08 everything went straight to `main` plus a tag, which caused two
|
||||
pain points:
|
||||
1. After 1.4.2 shipped, a scheduler-scoring fix had to be added to the same
|
||||
version — the only option was deleting the remote tag to clear the release
|
||||
and re-uploading all installers. The published state was detached from code
|
||||
history; "what exactly shipped" could not be traced.
|
||||
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".
|
||||
|
||||
> 中文版见 [docs/git-workflow.md](git-workflow.md)。
|
||||
128
docs/git-workflow.md
Normal file
128
docs/git-workflow.md
Normal file
@ -0,0 +1,128 @@
|
||||
# Git 分支工作流(Branching Workflow)
|
||||
|
||||
ModelRouter 采用 **GitHub Flow + 发布分支** 模型:`main` 是唯一长命分支且永远可部署,
|
||||
新工作从它开出特性分支;每个版本从 `main` 分出独立的发布分支并打 tag;hotfix 提交到
|
||||
发布分支且必须回流 `main`。
|
||||
|
||||
目的:让「已发版的版本在生命周期内可被单独修补」成为一等操作,同时 `main` 始终包含
|
||||
全部修复、永远可部署。
|
||||
|
||||
## 分支角色
|
||||
|
||||
| 分支 | 命名 | 生命周期 | 用途 |
|
||||
|---|---|---|---|
|
||||
| 主分支 | `main` | 永久 | 唯一长命分支。永远可部署。积攒下一个版本的功能。 |
|
||||
| 特性分支 | `feature/<描述>` | 短命(开发→合并即删) | 新特性 / 一般 bug 修复。从 `main` 开出,完成后合回 `main`。 |
|
||||
| 发布分支 | `release/vX.Y.Z` | 一个版本周期 | 从 `main` 分出,打 tag 发布。该版本生命周期内的 hotfix 都提交在此分支。 |
|
||||
|
||||
## 变更流向(重要)
|
||||
|
||||
```
|
||||
feature/foo ───────┐
|
||||
▼
|
||||
main ──────────────── 下一个版本 ──────►
|
||||
│ │
|
||||
│ 切出 │ 切出
|
||||
▼ ▼
|
||||
release/v1.4.2 release/v1.4.3
|
||||
│ │
|
||||
tag: v1.4.2 tag: v1.4.3
|
||||
│ │
|
||||
hotfix ◄─────┘ hotfix ◄─────┘
|
||||
│ │
|
||||
└── cherry-pick 回 main ───────┘
|
||||
```
|
||||
|
||||
### 关键规则
|
||||
|
||||
1. **main 永远可部署**:不在 main 上留半成品。任何未完成的工作必须在特性分支上。
|
||||
2. **特性从 main 开,完成后合回 main**:`git checkout -b feature/xxx 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
|
||||
```
|
||||
构建安装包、上传 GitCode Release 都基于这个 tag,保证可精确回溯发布态。
|
||||
4. **版本生命周期内只收该版本的 hotfix**:新特性一律并入 `main` 等下一个版本,
|
||||
绝不塞进已发布的 release 分支(除非主动选择在该版本内发次要版)。
|
||||
5. **hotfix 必须回流 main**:
|
||||
```bash
|
||||
git checkout release/v1.4.2 # 在发布分支提交修复
|
||||
git commit -m "fix: ..."
|
||||
git checkout main
|
||||
git cherry-pick <hotfix-commit> # 回主分支
|
||||
```
|
||||
这是本模型最重要的一条:**main 只前进、不丢修复**。如果 hotfix 不回 main,
|
||||
下一个版本就会带着旧 bug 发布。
|
||||
|
||||
### 版本生命周期结束
|
||||
|
||||
下一个版本发布时,上一个 release 分支退役:
|
||||
|
||||
- **默认:直接删除远端 release 分支**(`git push origin :release/v1.4.2`)。
|
||||
因为 hotfix 都已逐个 cherry-pick 回 main,main 已包含全部修复,无需再合并。
|
||||
- **如需要长期维护旧版**(例如企业大客户卡在旧版本):保留分支,仅 stopship 接受
|
||||
该版本的安全修复,继续走「提交 + cherry-pick 回 main」循环。
|
||||
|
||||
## 明确不做的事
|
||||
|
||||
- **不 rebase main**:`main` 的历史保持追加式,任何人拉取后 `git pull` 都得到直接可用的历史。
|
||||
- **发布分支不整体 merge 回 main**:hotfix 已逐个 cherry-pick,整体合并只会制造冲突且无收益。
|
||||
- **不在 main 直接写"临时改一下"**:即使是单人项目,也至少走
|
||||
`feature/xxx → merge main` 的形式,让历史保留"为什么改"的边界。
|
||||
|
||||
## 单人 vs 多人
|
||||
|
||||
- **单人**(本项目当前):特性分支可省去 PR,`git checkout -b feature/xxx` →
|
||||
完成 → 合并回 main。发布分支与 tag 流程完全一致。
|
||||
- **多人 / 开源协作**:特性分支走 PR(pull request)+ Code Review,
|
||||
由维护者 merge;冲突在特性分支上解决,不在 main 上解决。
|
||||
|
||||
## 提交信息规范(与分支配套)
|
||||
|
||||
- `feat(...)`: 新功能
|
||||
- `fix(...)`: 修复
|
||||
- `chore(...)`: 构建 / 版本号 / 依赖 / 文档无关改动
|
||||
- `docs(...)`: 文档
|
||||
- `refactor(...)`: 重构
|
||||
- 括号内为影响域,如 `fix(adapters)`、`feat(webui)`、`chore(version)`
|
||||
- 一句话总结,必要时正文展开「问题 / 根因 / 修复 / 验证」
|
||||
|
||||
## 发布操作速查(配套)
|
||||
|
||||
每次发布严格按此顺序:
|
||||
|
||||
```bash
|
||||
# 1. 确认 main 就绪
|
||||
git checkout main && git pull
|
||||
|
||||
# 2. 切发布分支并打 tag
|
||||
git checkout -b release/vX.Y.Z main
|
||||
git tag -a vX.Y.Z -m "ModelRouter vX.Y.Z"
|
||||
|
||||
# 3. 推送分支 + tag
|
||||
git push origin release/vX.Y.Z
|
||||
git push origin vX.Y.Z
|
||||
|
||||
# 4. 构建安装包(核心包 + GUI 各平台)
|
||||
make core-dist VERSION=X.Y.Z
|
||||
|
||||
# 5. 创建 GitCode Release + 上传安装包
|
||||
# 见 https://gitcode.com/JianFeeeee/ModelRouter/releases
|
||||
# 上传注意:必须 PUT --http1.1,否则 OBS 回调失败文件不进列表;
|
||||
# 替换旧版附件 = 删 tag(release 随 tag 一起消失)后重建。
|
||||
```
|
||||
|
||||
## 为什么是这套(背景)
|
||||
|
||||
2026-08 之前本项目直接在所有改动提交到 `main` + 打 tag,出现两个痛点:
|
||||
1. 1.4.2 已发版后又追了一个调度打分修复,只能删远端 tag 清空 release 再重传安装包
|
||||
—— 版本发布态与代码历史脱节,无法精确回溯"当时发的是什么"。
|
||||
2. 没有特性分支,无法并行开发互不干扰的两件事。
|
||||
|
||||
引入发布分支后:版本发布态 = `release/vX.Y.Z` 分支 + `vX.Y.Z` tag,可精确回溯;
|
||||
hotfix 有明确落点;main 保持"最新 + 全部修复 + 可部署"。
|
||||
|
||||
> 英文版见 [docs/git-workflow-en.md](git-workflow-en.md)。
|
||||
@ -55,6 +55,7 @@ state may be shared across workers (`log` to stdout is the only side effect).
|
||||
| `version` | string | no | Version, shown in the WebUI |
|
||||
| `endpoint` | string | no | Upstream path, default `/chat/completions`; overridable by `source.endpoint` / `source.image_endpoint` |
|
||||
| `headers` | table | no | Static default request headers; used as fallback when no `build_headers` hook is defined |
|
||||
| `transform_error(status, body)` | function | no | Error-response condensing: return a one-line reason; on nil / absence clients uniformly receive `unknown error` (raw body goes to server logs only) |
|
||||
|
||||
These fields are extracted statically at load time (compile-once); reading them never
|
||||
occupies a pooled worker.
|
||||
@ -147,6 +148,27 @@ end
|
||||
|
||||
## 4. Optional Hooks
|
||||
|
||||
### `transform_error(status, body) -> string | nil`
|
||||
|
||||
Condense this source's error response into a one-line reason. Every upstream
|
||||
formats errors differently — that is adapter territory: all built-in adapters
|
||||
implement their own envelope parsing (zen's `{error={type,message}}`,
|
||||
Anthropic's `{type="error",error={...}}`, Gemini's
|
||||
`{error={code,message,status}}`, Ollama's string `{error="..."}`, etc.).
|
||||
When the hook is absent or returns nil the core does not guess: clients get
|
||||
`api error <status>: unknown error` and the raw body is logged server-side only.
|
||||
|
||||
```lua
|
||||
function adapter.transform_error(status, body)
|
||||
local ok, resp = pcall(json.decode, body)
|
||||
if not ok or type(resp) ~= "table" then return nil end
|
||||
if resp.error and resp.error.type == "FreeUsageLimitError" then
|
||||
return "zen free pool quota exhausted"
|
||||
end
|
||||
return resp.error and resp.error.message or nil
|
||||
end
|
||||
```
|
||||
|
||||
### `build_headers(meta) -> table<string,string>`
|
||||
|
||||
Dynamically generate / sign request headers (e.g. KimiCode's HMAC signature). If the
|
||||
|
||||
@ -50,6 +50,7 @@ return adapter
|
||||
| `version` | string | 否 | 版本号,用于 WebUI 展示 |
|
||||
| `endpoint` | string | 否 | 上游请求路径,默认 `/chat/completions`;可被 `source.endpoint` / `source.image_endpoint` 覆盖 |
|
||||
| `headers` | table | 否 | 静态默认请求头;若未定义 `build_headers` 钩子则作为请求头回退 |
|
||||
| `transform_error(status, body)` | function | 否 | 错误响应收敛:返回一行短原因;返回 nil / 未定义时客户端统一收到 `unknown error`(原始响应体只进服务端日志) |
|
||||
|
||||
这些字段在加载时静态提取(compile-once),之后读它们不会占用池内 worker。
|
||||
|
||||
@ -138,6 +139,26 @@ end
|
||||
|
||||
## 4. 可选钩子
|
||||
|
||||
### `transform_error(status, body) -> string | nil`
|
||||
|
||||
把该源特有的错误响应收敛成一行短原因。每个上游的错误格式不同——这是适配器层
|
||||
的职责:内置各适配器均实现了自己的信封解析(如 zen 的
|
||||
`{error={type,message}}`、Anthropic 的 `{type="error",error={...}}`、Gemini 的
|
||||
`{error={code,message,status}}`、Ollama 的字符串 `{error="..."}` 等)。
|
||||
未定义本钩子或返回 nil 时,核心不猜测格式,客户端统一收到
|
||||
`api error <status>: unknown error`,原始响应体仅记录在服务端日志。
|
||||
|
||||
```lua
|
||||
function adapter.transform_error(status, body)
|
||||
local ok, resp = pcall(json.decode, body)
|
||||
if not ok or type(resp) ~= "table" then return nil end
|
||||
if resp.error and resp.error.type == "FreeUsageLimitError" then
|
||||
return "zen free pool quota exhausted"
|
||||
end
|
||||
return resp.error and resp.error.message or nil
|
||||
end
|
||||
```
|
||||
|
||||
### `build_headers(meta) -> table<string,string>`
|
||||
|
||||
动态生成/签名请求头(如 KimiCode 的 HMAC 签名)。若脚本未定义此函数,Go 层回退使用静态
|
||||
|
||||
142
e2e/e2e_test.go
142
e2e/e2e_test.go
@ -315,7 +315,7 @@ func TestEndToEndAuto503(t *testing.T) {
|
||||
if resp == nil || resp.StatusCode != http.StatusServiceUnavailable {
|
||||
t.Fatalf("status=%v body=%q", statusOf(resp), body)
|
||||
}
|
||||
if !strings.Contains(body, "all auto tiers failed") || !strings.Contains(body, "good/good-m") {
|
||||
if !strings.Contains(body, "auto providers failed") || !strings.Contains(body, "good/good-m") {
|
||||
t.Fatalf("503 must summarize the failed slot, body=%q", body)
|
||||
}
|
||||
}
|
||||
@ -327,3 +327,143 @@ 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
|
||||
}
|
||||
|
||||
@ -20,13 +20,18 @@ type Config struct {
|
||||
DefaultModel string `yaml:"default_model"` // e.g. "AUTO" or a model id
|
||||
AdapterDir string `yaml:"adapter_dir"`
|
||||
RuntimeFile string `yaml:"runtime_file"`
|
||||
MaxConcurrent int `yaml:"max_concurrent"` // global inflight cap, 0 = unlimited
|
||||
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
|
||||
PublicBaseURL string `yaml:"public_base_url,omitempty"` // external base for generated config snippets; default inferred from request
|
||||
Sources []Source `yaml:"sources"`
|
||||
Auto []ModelScope `yaml:"auto,omitempty"` // AUTO 调度链规则(WebUI 优先级页编辑)
|
||||
Keys []GWKey `yaml:"keys,omitempty"` // 网关密钥(WebUI 密钥页管理)
|
||||
Auto []ModelScope `yaml:"auto,omitempty"` // AUTO 调度链规则(WebUI 优先级页编辑,chat)
|
||||
AutoImage []ModelScope `yaml:"auto_image,omitempty"` // AUTO 生图调度链规则(WebUI 优先级页·生图)
|
||||
Keys []GWKey `yaml:"keys,omitempty"` // 网关密钥(WebUI 密钥页管理)
|
||||
// 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.
|
||||
@ -47,21 +52,26 @@ type Model struct {
|
||||
|
||||
// Source describes a single upstream LLM provider.
|
||||
type Source struct {
|
||||
Name string `yaml:"name" json:"name"`
|
||||
BaseURL string `yaml:"base_url" json:"base_url"`
|
||||
APIKey string `yaml:"api_key" json:"api_key"`
|
||||
APIKeyEnv string `yaml:"api_key_env,omitempty" json:"-"` // reference to an env var holding the key (overrides api_key)
|
||||
Adapter string `yaml:"adapter" json:"adapter"`
|
||||
Endpoint string `yaml:"endpoint" json:"endpoint,omitempty"` // chat endpoint override
|
||||
ImageEndpoint string `yaml:"image_endpoint" json:"image_endpoint,omitempty"` // image endpoint override
|
||||
Models []Model `yaml:"models" json:"models"`
|
||||
Headers map[string]string `yaml:"headers" json:"headers,omitempty"`
|
||||
Name string `yaml:"name" json:"name"`
|
||||
BaseURL string `yaml:"base_url" json:"base_url"`
|
||||
APIKey string `yaml:"api_key" json:"api_key"`
|
||||
APIKeyEnv string `yaml:"api_key_env,omitempty" json:"-"` // reference to an env var holding the key (overrides api_key)
|
||||
Adapter string `yaml:"adapter" json:"adapter"`
|
||||
Endpoint string `yaml:"endpoint" json:"endpoint,omitempty"` // chat endpoint override
|
||||
ImageEndpoint string `yaml:"image_endpoint" json:"image_endpoint,omitempty"` // image endpoint override
|
||||
Models []Model `yaml:"models" json:"models"`
|
||||
Headers map[string]string `yaml:"headers" json:"headers,omitempty"`
|
||||
// ProxyURL routes this source's HTTP(S) traffic through an explicit proxy
|
||||
// (e.g. http://127.0.0.1:7890). Empty = direct connection. Needed for
|
||||
// upstreams behind geo/IP blocks that only respond through a proxy.
|
||||
ProxyURL string `yaml:"proxy_url" json:"proxy_url,omitempty"`
|
||||
Meta map[string]interface{} `yaml:"meta" json:"meta,omitempty"`
|
||||
Temperature float64 `yaml:"temperature" json:"temperature,omitempty"`
|
||||
MaxTokens int `yaml:"max_tokens" json:"max_tokens,omitempty"`
|
||||
Timeout time.Duration `yaml:"timeout" json:"-"`
|
||||
MaxConcurrent int `yaml:"max_concurrent" json:"max_concurrent"`
|
||||
QueueTimeout time.Duration `yaml:"queue_timeout" json:"-"`
|
||||
RPM int `yaml:"rpm" json:"rpm,omitempty"` // optional requests-per-minute cap (0 = unlimited)
|
||||
}
|
||||
|
||||
// Load reads and validates a config file. When the file does not exist yet a
|
||||
@ -187,10 +197,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
|
||||
@ -248,11 +263,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)
|
||||
}
|
||||
@ -313,11 +334,34 @@ func (c *Config) ApplyDefaults() error {
|
||||
|
||||
// RuntimeConfig is the legacy runtime file format (kept for migration only).
|
||||
type RuntimeConfig struct {
|
||||
Sources []Source `json:"sources,omitempty"`
|
||||
DeletedSources []string `json:"deleted_sources,omitempty"`
|
||||
DeletedAdapters []string `json:"deleted_adapters,omitempty"`
|
||||
Keys []GWKey `json:"keys,omitempty"`
|
||||
Auto []ModelScope `json:"auto,omitempty"`
|
||||
Sources []Source `json:"sources,omitempty"`
|
||||
SourceTemplates []SourceTemplate `json:"source_templates,omitempty"`
|
||||
DeletedSources []string `json:"deleted_sources,omitempty"`
|
||||
DeletedAdapters []string `json:"deleted_adapters,omitempty"`
|
||||
PresetTemplates []string `json:"preset_templates,omitempty"` // preset names the user has seen (or deleted) — never re-seeded
|
||||
Keys []GWKey `json:"keys,omitempty"`
|
||||
Auto []ModelScope `json:"auto,omitempty"`
|
||||
AutoImage []ModelScope `json:"auto_image,omitempty"`
|
||||
}
|
||||
|
||||
// SourceTemplate stores reusable source configuration (everything except
|
||||
// name and api_key) so the WebUI can spin up multiple key-bearing sources
|
||||
// from one shared template.
|
||||
type SourceTemplate struct {
|
||||
Name string `json:"name"`
|
||||
BaseURL string `json:"base_url"`
|
||||
Adapter string `json:"adapter"`
|
||||
Endpoint string `json:"endpoint,omitempty"`
|
||||
ImageEndpoint string `json:"image_endpoint,omitempty"`
|
||||
Models []Model `json:"models"`
|
||||
Headers map[string]string `json:"headers,omitempty"`
|
||||
// ProxyURL routes this source's HTTP(S) traffic through an explicit proxy.
|
||||
ProxyURL string `json:"proxy_url,omitempty"`
|
||||
Meta map[string]interface{} `json:"meta,omitempty"`
|
||||
Temperature float64 `json:"temperature,omitempty"`
|
||||
MaxTokens int `json:"max_tokens,omitempty"`
|
||||
MaxConcurrent int `json:"max_concurrent,omitempty"`
|
||||
RPM int `json:"rpm,omitempty"`
|
||||
}
|
||||
|
||||
// GWKey is a gateway API key persisted in the config file. Role is "admin"
|
||||
@ -333,15 +377,76 @@ type GWKey struct {
|
||||
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.
|
||||
// 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"`
|
||||
}
|
||||
@ -360,6 +465,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"`
|
||||
}
|
||||
@ -370,6 +476,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
|
||||
|
||||
@ -318,3 +318,55 @@ sources:
|
||||
t.Fatalf("after no-op: %v %v", len(cfg.Sources), err)
|
||||
}
|
||||
}
|
||||
|
||||
// TestSeedTemplatesOnce verifies preset templates are seeded exactly once:
|
||||
// a user's edit survives a restart, and a deliberately deleted preset does
|
||||
// not silently come back.
|
||||
func TestSeedTemplatesOnce(t *testing.T) {
|
||||
dir := t.TempDir()
|
||||
path := filepath.Join(dir, "runtime.json")
|
||||
presets := []SourceTemplate{
|
||||
{Name: "DeepSeek", BaseURL: "https://api.deepseek.com", Adapter: "deepseek",
|
||||
Models: []Model{{ID: "deepseek-chat", Kind: "chat"}}},
|
||||
{Name: "OpenAI", BaseURL: "https://api.openai.com/v1", Adapter: "openai",
|
||||
Models: []Model{{ID: "gpt-4o", Kind: "chat"}}},
|
||||
}
|
||||
|
||||
// first run: both presets land in the store
|
||||
s1 := NewStore(path)
|
||||
if err := s1.Load(); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := s1.SeedTemplates(presets); err != nil {
|
||||
t.Fatalf("seed: %v", err)
|
||||
}
|
||||
if got := len(s1.ListTemplates()); got != 2 {
|
||||
t.Fatalf("after seed len = %d, want 2", got)
|
||||
}
|
||||
|
||||
// user edits one preset, deletes the other
|
||||
edited := presets[0]
|
||||
edited.BaseURL = "https://my-proxy.internal/v1"
|
||||
if err := s1.UpsertTemplate(edited); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if _, err := s1.RemoveTemplate("OpenAI"); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
// restart: re-seeding must be a no-op for both
|
||||
s2 := NewStore(path)
|
||||
if err := s2.Load(); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := s2.SeedTemplates(presets); err != nil {
|
||||
t.Fatalf("reseed: %v", err)
|
||||
}
|
||||
tpls := s2.ListTemplates()
|
||||
if len(tpls) != 1 {
|
||||
t.Fatalf("after restart len = %d, want 1 (deleted preset must not return): %+v", len(tpls), tpls)
|
||||
}
|
||||
if tpls[0].Name != "DeepSeek" || tpls[0].BaseURL != "https://my-proxy.internal/v1" {
|
||||
t.Fatalf("user edit lost on restart: %+v", tpls[0])
|
||||
}
|
||||
}
|
||||
|
||||
212
internal/config/presets.go
Normal file
212
internal/config/presets.go
Normal file
@ -0,0 +1,212 @@
|
||||
package config
|
||||
|
||||
// Built-in preset source templates. Seeded into the runtime store on first
|
||||
// startup so the WebUI can present "从模板创建" for popular providers — the
|
||||
// user only needs to supply a name and API key.
|
||||
//
|
||||
// Model IDs verified against provider docs 2026-08-28. Providers ship new IDs
|
||||
// frequently; treat these as a starting point and edit in the WebUI as needed.
|
||||
var PresetTemplates = []SourceTemplate{
|
||||
// ── 国内热门 ────────────────────────────────────────────────────────────
|
||||
{
|
||||
Name: "DeepSeek",
|
||||
BaseURL: "https://api.deepseek.com",
|
||||
Adapter: "deepseek",
|
||||
MaxConcurrent: 8,
|
||||
Models: []Model{
|
||||
{ID: "deepseek-v4-flash", Priority: 100, Kind: "chat"},
|
||||
{ID: "deepseek-v4-pro", Priority: 90, Kind: "chat"},
|
||||
{ID: "deepseek-v4-flash-vision-exp", Priority: 60, Kind: "chat"},
|
||||
},
|
||||
},
|
||||
{
|
||||
Name: "智谱 GLM",
|
||||
BaseURL: "https://open.bigmodel.cn/api/paas/v4",
|
||||
Adapter: "openai",
|
||||
MaxConcurrent: 8,
|
||||
Models: []Model{
|
||||
{ID: "glm-5.3", Priority: 100, Kind: "chat"},
|
||||
{ID: "glm-5.3-flash", Priority: 90, Kind: "chat"},
|
||||
{ID: "glm-5.2", Priority: 80, Kind: "chat"},
|
||||
{ID: "glm-5.1", Priority: 60, Kind: "chat"},
|
||||
},
|
||||
},
|
||||
{
|
||||
Name: "Kimi / Moonshot",
|
||||
BaseURL: "https://api.moonshot.cn/v1",
|
||||
Adapter: "openai",
|
||||
MaxConcurrent: 8,
|
||||
Models: []Model{
|
||||
{ID: "kimi-k3", Priority: 100, Kind: "chat"},
|
||||
{ID: "kimi-k2.7-code", Priority: 90, Kind: "chat"},
|
||||
{ID: "kimi-k2.7-code-highspeed", Priority: 85, Kind: "chat"},
|
||||
{ID: "kimi-k2.6", Priority: 70, Kind: "chat"},
|
||||
},
|
||||
},
|
||||
{
|
||||
Name: "SiliconFlow",
|
||||
BaseURL: "https://api.siliconflow.cn/v1",
|
||||
Adapter: "openai",
|
||||
MaxConcurrent: 8,
|
||||
Models: []Model{
|
||||
{ID: "deepseek-ai/DeepSeek-V4-Pro", Priority: 100, Kind: "chat"},
|
||||
{ID: "deepseek-ai/DeepSeek-V4-Flash", Priority: 90, Kind: "chat"},
|
||||
{ID: "moonshotai/Kimi-K2.6", Priority: 80, Kind: "chat"},
|
||||
{ID: "zai-org/GLM-5.1", Priority: 70, Kind: "chat"},
|
||||
{ID: "Kwai-Kolors/Kolors", Priority: 50, Kind: "image"},
|
||||
},
|
||||
},
|
||||
{
|
||||
Name: "讯飞星火 (Spark)",
|
||||
BaseURL: "https://spark-api-open.xf-yun.com/v1",
|
||||
Adapter: "openai",
|
||||
MaxConcurrent: 4,
|
||||
Models: []Model{
|
||||
{ID: "4.0Ultra", Priority: 100, Kind: "chat"},
|
||||
{ID: "generalv3.5", Priority: 80, Kind: "chat"},
|
||||
{ID: "max-32k", Priority: 70, Kind: "chat"},
|
||||
},
|
||||
},
|
||||
{
|
||||
Name: "阿里云 通义千问",
|
||||
BaseURL: "https://dashscope.aliyuncs.com/compatible-mode/v1",
|
||||
Adapter: "openai",
|
||||
MaxConcurrent: 8,
|
||||
Models: []Model{
|
||||
{ID: "qwen3.8-max", Priority: 100, Kind: "chat"},
|
||||
{ID: "qwen3.8-flash", Priority: 90, Kind: "chat"},
|
||||
{ID: "qwen3.7-plus", Priority: 85, Kind: "chat"},
|
||||
{ID: "deepseek-v4-pro", Priority: 70, Kind: "chat"},
|
||||
{ID: "kimi-k3", Priority: 70, Kind: "chat"},
|
||||
},
|
||||
},
|
||||
{
|
||||
Name: "百度 千帆",
|
||||
BaseURL: "https://qianfan.baidubce.com/v2",
|
||||
Adapter: "openai",
|
||||
MaxConcurrent: 8,
|
||||
Models: []Model{
|
||||
{ID: "ernie-5.0-turbo", Priority: 100, Kind: "chat"},
|
||||
{ID: "ernie-5.0", Priority: 90, Kind: "chat"},
|
||||
{ID: "deepseek-v4-pro", Priority: 70, Kind: "chat"},
|
||||
},
|
||||
},
|
||||
{
|
||||
Name: "豆包 (火山方舟)",
|
||||
BaseURL: "https://ark.cn-beijing.volces.com/api/v3",
|
||||
Adapter: "openai",
|
||||
MaxConcurrent: 8,
|
||||
Models: []Model{
|
||||
{ID: "doubao-seed-2.0-pro", Priority: 100, Kind: "chat"},
|
||||
{ID: "doubao-seed-2.0", Priority: 90, Kind: "chat"},
|
||||
{ID: "doubao-seed-2.0-flash", Priority: 80, Kind: "chat"},
|
||||
},
|
||||
},
|
||||
{
|
||||
Name: "MiniMax",
|
||||
BaseURL: "https://api.minimax.chat/v1",
|
||||
Adapter: "openai",
|
||||
MaxConcurrent: 8,
|
||||
Models: []Model{
|
||||
{ID: "MiniMax-M3", Priority: 100, Kind: "chat"},
|
||||
{ID: "MiniMax-M2.5", Priority: 80, Kind: "chat"},
|
||||
},
|
||||
},
|
||||
// ── 国际主流 ────────────────────────────────────────────────────────────
|
||||
{
|
||||
Name: "OpenAI",
|
||||
BaseURL: "https://api.openai.com/v1",
|
||||
Adapter: "openai",
|
||||
MaxConcurrent: 8,
|
||||
Models: []Model{
|
||||
{ID: "gpt-5.6-sol", Priority: 100, Kind: "chat"},
|
||||
{ID: "gpt-5.6-terra", Priority: 90, Kind: "chat"},
|
||||
{ID: "gpt-5.6-luna", Priority: 80, Kind: "chat"},
|
||||
},
|
||||
},
|
||||
{
|
||||
Name: "Anthropic Claude",
|
||||
BaseURL: "https://api.anthropic.com",
|
||||
Adapter: "anthropic",
|
||||
MaxConcurrent: 8,
|
||||
Models: []Model{
|
||||
{ID: "claude-opus-5", Priority: 100, Kind: "chat"},
|
||||
{ID: "claude-fable-5", Priority: 95, Kind: "chat"},
|
||||
{ID: "claude-sonnet-5", Priority: 85, Kind: "chat"},
|
||||
{ID: "claude-haiku-4-5", Priority: 70, Kind: "chat"},
|
||||
},
|
||||
},
|
||||
{
|
||||
Name: "Google Gemini",
|
||||
BaseURL: "https://generativelanguage.googleapis.com/v1beta",
|
||||
Adapter: "gemini",
|
||||
MaxConcurrent: 8,
|
||||
Models: []Model{
|
||||
{ID: "gemini-3.6-flash", Priority: 100, Kind: "chat"},
|
||||
{ID: "gemini-3.5-flash", Priority: 90, Kind: "chat"},
|
||||
{ID: "gemini-3.5-flash-lite", Priority: 80, Kind: "chat"},
|
||||
{ID: "gemini-3.1-pro", Priority: 85, Kind: "chat"},
|
||||
},
|
||||
},
|
||||
{
|
||||
Name: "Groq",
|
||||
BaseURL: "https://api.groq.com/openai/v1",
|
||||
Adapter: "groq",
|
||||
MaxConcurrent: 8,
|
||||
Models: []Model{
|
||||
{ID: "openai/gpt-oss-120b", Priority: 100, Kind: "chat"},
|
||||
{ID: "openai/gpt-oss-20b", Priority: 80, Kind: "chat"},
|
||||
{ID: "groq/compound", Priority: 70, Kind: "chat"},
|
||||
},
|
||||
},
|
||||
{
|
||||
Name: "Mistral",
|
||||
BaseURL: "https://api.mistral.ai/v1",
|
||||
Adapter: "mistral",
|
||||
MaxConcurrent: 8,
|
||||
Models: []Model{
|
||||
{ID: "mistral-medium-latest", Priority: 100, Kind: "chat"},
|
||||
{ID: "mistral-large-latest", Priority: 90, Kind: "chat"},
|
||||
{ID: "codestral-latest", Priority: 85, Kind: "chat"},
|
||||
},
|
||||
},
|
||||
{
|
||||
Name: "OpenRouter",
|
||||
BaseURL: "https://openrouter.ai/api/v1",
|
||||
Adapter: "openai",
|
||||
MaxConcurrent: 8,
|
||||
Headers: map[string]string{
|
||||
"HTTP-Referer": "http://localhost:8080",
|
||||
"X-Title": "ModelRouter",
|
||||
},
|
||||
Models: []Model{
|
||||
{ID: "openrouter/auto", Priority: 100, Kind: "chat"},
|
||||
{ID: "deepseek/deepseek-v4-flash", Priority: 90, Kind: "chat"},
|
||||
{ID: "anthropic/claude-opus-5", Priority: 85, Kind: "chat"},
|
||||
{ID: "openai/gpt-5.6-sol", Priority: 85, Kind: "chat"},
|
||||
},
|
||||
},
|
||||
{
|
||||
Name: "GitHub Models",
|
||||
BaseURL: "https://models.inference.ai.azure.com",
|
||||
Adapter: "github",
|
||||
MaxConcurrent: 8,
|
||||
Models: []Model{
|
||||
{ID: "gpt-5.6-sol", Priority: 100, Kind: "chat"},
|
||||
{ID: "gpt-5.6-luna", Priority: 80, Kind: "chat"},
|
||||
},
|
||||
},
|
||||
// ── 本地模型 ────────────────────────────────────────────────────────────
|
||||
{
|
||||
Name: "Ollama (localhost)",
|
||||
BaseURL: "http://127.0.0.1:11434",
|
||||
Adapter: "ollama",
|
||||
Endpoint: "/api/chat",
|
||||
MaxConcurrent: 4,
|
||||
Models: []Model{
|
||||
{ID: "qwen3", Priority: 80, Kind: "chat"},
|
||||
{ID: "deepseek-r1", Priority: 80, Kind: "chat"},
|
||||
{ID: "llama3.3", Priority: 70, Kind: "chat"},
|
||||
},
|
||||
},
|
||||
}
|
||||
@ -128,15 +128,3 @@ func (b *SecretBox) MustDecrypt(v string) string {
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
// resolvedAPIKey resolves a source key from api_key_env (env var) or api_key.
|
||||
func resolvedAPIKey(src Source, box *SecretBox) (string, error) {
|
||||
if src.APIKeyEnv != "" {
|
||||
v := os.Getenv(src.APIKeyEnv)
|
||||
if v == "" {
|
||||
return "", fmt.Errorf("source %s: env %s is empty or unset", src.Name, src.APIKeyEnv)
|
||||
}
|
||||
return v, nil
|
||||
}
|
||||
return box.MustDecrypt(src.APIKey), 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)
|
||||
}
|
||||
}
|
||||
@ -138,6 +138,81 @@ func (s *Store) DeletedAdapters() map[string]bool {
|
||||
return out
|
||||
}
|
||||
|
||||
// SeedTemplates adds preset templates that the store has never seen before.
|
||||
// A preset is seeded at most once: its name is recorded in PresetTemplates so
|
||||
// a user who edits or deletes it never gets it silently restored on restart.
|
||||
func (s *Store) SeedTemplates(presets []SourceTemplate) error {
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
seen := make(map[string]bool, len(s.data.PresetTemplates))
|
||||
for _, n := range s.data.PresetTemplates {
|
||||
seen[n] = true
|
||||
}
|
||||
existing := make(map[string]bool, len(s.data.SourceTemplates))
|
||||
for _, t := range s.data.SourceTemplates {
|
||||
existing[t.Name] = true
|
||||
}
|
||||
added := false
|
||||
for _, p := range presets {
|
||||
if p.Name == "" || seen[p.Name] {
|
||||
continue
|
||||
}
|
||||
s.data.PresetTemplates = append(s.data.PresetTemplates, p.Name)
|
||||
added = true
|
||||
if existing[p.Name] {
|
||||
continue // user already has a template by this name: never overwrite
|
||||
}
|
||||
s.data.SourceTemplates = append(s.data.SourceTemplates, p)
|
||||
}
|
||||
if !added {
|
||||
return nil
|
||||
}
|
||||
return s.persistLocked()
|
||||
}
|
||||
|
||||
// UpsertTemplate adds or replaces a source template and persists.
|
||||
func (s *Store) UpsertTemplate(t SourceTemplate) error {
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
for i := range s.data.SourceTemplates {
|
||||
if s.data.SourceTemplates[i].Name == t.Name {
|
||||
s.data.SourceTemplates[i] = t
|
||||
return s.persistLocked()
|
||||
}
|
||||
}
|
||||
s.data.SourceTemplates = append(s.data.SourceTemplates, t)
|
||||
return s.persistLocked()
|
||||
}
|
||||
|
||||
// ListTemplates returns the saved source templates.
|
||||
func (s *Store) ListTemplates() []SourceTemplate {
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
out := make([]SourceTemplate, len(s.data.SourceTemplates))
|
||||
copy(out, s.data.SourceTemplates)
|
||||
return out
|
||||
}
|
||||
|
||||
// RemoveTemplate deletes a source template by name and persists.
|
||||
func (s *Store) RemoveTemplate(name string) (bool, error) {
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
kept := s.data.SourceTemplates[:0]
|
||||
removed := false
|
||||
for _, t := range s.data.SourceTemplates {
|
||||
if t.Name == name {
|
||||
removed = true
|
||||
continue
|
||||
}
|
||||
kept = append(kept, t)
|
||||
}
|
||||
s.data.SourceTemplates = kept
|
||||
if removed {
|
||||
return true, s.persistLocked()
|
||||
}
|
||||
return false, nil
|
||||
}
|
||||
|
||||
// LoadLegacy returns the full legacy RuntimeConfig from the file (used for
|
||||
// one-time migration into config.yaml). Returns nil if file is missing.
|
||||
func (s *Store) LoadLegacy() *RuntimeConfig {
|
||||
@ -186,15 +261,6 @@ func (s *Store) encryptLocked() {
|
||||
}
|
||||
}
|
||||
|
||||
func containsString(list []string, s string) bool {
|
||||
for _, x := range list {
|
||||
if x == s {
|
||||
return true
|
||||
}
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
func removeString(list []string, s string) []string {
|
||||
out := list[:0]
|
||||
for _, x := range list {
|
||||
|
||||
@ -7,6 +7,7 @@ import (
|
||||
"crypto/rand"
|
||||
"encoding/hex"
|
||||
"fmt"
|
||||
"log"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"sort"
|
||||
@ -26,13 +27,14 @@ import (
|
||||
// request paths look keys up concurrently; the AUTO chain itself is swapped
|
||||
// atomically and needs no lock.
|
||||
type Core struct {
|
||||
mu sync.Mutex
|
||||
cfg *config.Config
|
||||
vm *lua.VM
|
||||
store *config.Store
|
||||
scheduler *scheduler.Scheduler
|
||||
registry *provider.Registry
|
||||
autoChain atomic.Pointer[scheduler.Chain]
|
||||
mu sync.Mutex
|
||||
cfg *config.Config
|
||||
vm *lua.VM
|
||||
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
|
||||
}
|
||||
|
||||
// New builds the core from a config file plus runtime overlay.
|
||||
@ -52,6 +54,28 @@ func NewFromConfig(cfg *config.Config) (*Core, error) {
|
||||
return nil, fmt.Errorf("lua vm: %w", err)
|
||||
}
|
||||
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)
|
||||
}
|
||||
@ -68,6 +92,17 @@ func NewFromConfig(cfg *config.Config) (*Core, error) {
|
||||
if err := c.rebuildRegistry(); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
// Seed built-in preset templates on first run so the WebUI shows
|
||||
// "从模板创建" for popular providers out of the box.
|
||||
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
|
||||
}
|
||||
|
||||
@ -161,6 +196,18 @@ func (c *Core) seedKeys() error {
|
||||
return nil
|
||||
}
|
||||
|
||||
// seedPresetTemplates adds the built-in provider presets to the template
|
||||
// store on first run. It only adds templates whose name is not already
|
||||
// present, so a user's edits to a preset survive restarts and a deliberately
|
||||
// deleted preset does not silently return on the next start (its name is
|
||||
// remembered in the store's preset marker list).
|
||||
func (c *Core) seedPresetTemplates() error {
|
||||
if len(config.PresetTemplates) == 0 {
|
||||
return nil
|
||||
}
|
||||
return c.store.SeedTemplates(config.PresetTemplates)
|
||||
}
|
||||
|
||||
// saveConfig writes the current config (including auto rules and keys) back
|
||||
// to config.yaml.
|
||||
func (c *Core) saveConfig() error {
|
||||
@ -181,6 +228,9 @@ 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() }
|
||||
|
||||
// AutoImageChain returns the persisted image-generation AUTO chain snapshot.
|
||||
func (c *Core) AutoImageChain() *scheduler.Chain { return c.autoImageChain.Load() }
|
||||
|
||||
func (c *Core) DefaultModel() string { return c.cfg.DefaultModel }
|
||||
|
||||
func (c *Core) GatewayKeys() []string { return c.cfg.GatewayKeys }
|
||||
@ -217,10 +267,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
|
||||
@ -233,9 +288,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
|
||||
@ -243,10 +296,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 != "" {
|
||||
@ -294,6 +353,15 @@ func (c *Core) AutoRules() []config.ModelScope {
|
||||
return out
|
||||
}
|
||||
|
||||
// AutoImageRules returns the configured image-generation AUTO chain rules.
|
||||
func (c *Core) AutoImageRules() []config.ModelScope {
|
||||
c.mu.Lock()
|
||||
defer c.mu.Unlock()
|
||||
out := make([]config.ModelScope, len(c.cfg.AutoImage))
|
||||
copy(out, c.cfg.AutoImage)
|
||||
return out
|
||||
}
|
||||
|
||||
// cleanScopes drops empty model entries and normalizes placeholder source
|
||||
// names; it returns nil when no entries survive so an empty scope means
|
||||
// "unrestricted" (nil) instead of a restrictive-but-empty list — a non-nil
|
||||
@ -325,16 +393,35 @@ func (c *Core) SaveAutoRules(entries []config.ModelScope) error {
|
||||
return err
|
||||
}
|
||||
c.buildAutoChain()
|
||||
if ch := c.autoChain.Load(); ch != nil {
|
||||
for _, tn := range ch.Tiers {
|
||||
for _, sl := range tn.Slots {
|
||||
if p := c.registry.ProviderForSlot(sl.Model, sl.Source); p != nil {
|
||||
p.ResetModelCooldown(sl.Model)
|
||||
}
|
||||
c.resetChainCooldowns(c.autoChain.Load())
|
||||
return nil
|
||||
}
|
||||
|
||||
// SaveAutoImageRules persists the image-generation AUTO chain. Image models
|
||||
// are kept as-is (unlike buildAutoChain, which skips them).
|
||||
func (c *Core) SaveAutoImageRules(entries []config.ModelScope) error {
|
||||
c.mu.Lock()
|
||||
defer c.mu.Unlock()
|
||||
c.cfg.AutoImage = cleanScopes(entries)
|
||||
if err := c.saveConfig(); err != nil {
|
||||
return err
|
||||
}
|
||||
c.buildAutoImageChain()
|
||||
c.resetChainCooldowns(c.autoImageChain.Load())
|
||||
return nil
|
||||
}
|
||||
|
||||
func (c *Core) resetChainCooldowns(ch *scheduler.Chain) {
|
||||
if ch == nil {
|
||||
return
|
||||
}
|
||||
for _, tn := range ch.Tiers {
|
||||
for _, sl := range tn.Slots {
|
||||
if p := c.registry.ProviderForSlot(sl.Model, sl.Source); p != nil {
|
||||
p.ResetModelCooldown(sl.Model)
|
||||
}
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// ResetHealth clears the scheduling backoff state of every provider.
|
||||
@ -390,6 +477,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 != "" {
|
||||
@ -398,7 +490,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
|
||||
}
|
||||
@ -420,6 +516,7 @@ func (c *Core) rebuildRegistry() error {
|
||||
c.registry.Replace(providers)
|
||||
}
|
||||
c.buildAutoChain()
|
||||
c.buildAutoImageChain()
|
||||
return nil
|
||||
}
|
||||
|
||||
@ -466,6 +563,47 @@ func (c *Core) buildAutoChain() {
|
||||
c.autoChain.Store(scheduler.BuildChain(sr, prov))
|
||||
}
|
||||
|
||||
// buildAutoImageChain rebuilds the image-generation AUTO chain snapshot.
|
||||
// Unlike buildAutoChain, only image-kind models participate: chat models in
|
||||
// the rules are skipped so a stale chat slot can't receive image traffic.
|
||||
func (c *Core) buildAutoImageChain() {
|
||||
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" {
|
||||
return nil
|
||||
}
|
||||
return p
|
||||
}
|
||||
rules := c.cfg.AutoImage
|
||||
sr := make([]scheduler.Rule, 0, len(rules))
|
||||
for _, e := range rules {
|
||||
model, source := e.Model, e.Source
|
||||
if p := c.registry.ProviderForSlot(e.Model, e.Source); p != nil {
|
||||
if exact := p.ModelIDFold(e.Model); exact != "" {
|
||||
model = exact
|
||||
}
|
||||
if m := p.ModelByID(model); m != nil && m.Kind != "image" {
|
||||
continue // chat-kind models never join the image AUTO chain
|
||||
}
|
||||
if source == "" {
|
||||
source = p.Name()
|
||||
}
|
||||
}
|
||||
sr = append(sr, scheduler.Rule{
|
||||
Model: model,
|
||||
Source: source,
|
||||
Tier: e.Tier,
|
||||
Quota: e.TokenQuota,
|
||||
Period: e.Period,
|
||||
Hours: e.Hours,
|
||||
})
|
||||
}
|
||||
c.autoImageChain.Store(scheduler.BuildChain(sr, prov))
|
||||
}
|
||||
|
||||
// AutoSlotState is the UI-facing health snapshot of one AUTO chain slot.
|
||||
type AutoSlotState struct {
|
||||
Model string `json:"model"`
|
||||
@ -474,6 +612,16 @@ type AutoSlotState struct {
|
||||
FailCount int64 `json:"fail_count"`
|
||||
CooldownUntil int64 `json:"cooldown_until"`
|
||||
Cooling bool `json:"cooling"`
|
||||
// CooldownFrom is when the current cooldown window opened and ProbeAfter is
|
||||
// its midpoint: from ProbeAfter on, a single probe request is allowed
|
||||
// through so a recovered upstream does not have to sit out the rest of the
|
||||
// window. Probing reports whether such a probe is in flight right now.
|
||||
CooldownFrom int64 `json:"cooldown_from"`
|
||||
ProbeAfter int64 `json:"probe_after"`
|
||||
Probing bool `json:"probing"`
|
||||
// Probeable marks a cooling slot that is past its midpoint, i.e. the next
|
||||
// AUTO request may use it as a probe.
|
||||
Probeable bool `json:"probeable"`
|
||||
}
|
||||
|
||||
// AutoSlotStates returns per-slot health for every slot of the current chain.
|
||||
@ -491,6 +639,7 @@ func (c *Core) AutoSlotStates() []AutoSlotState {
|
||||
continue
|
||||
}
|
||||
pref, fail, until := pp.ModelHealthInfo(sl.Model)
|
||||
from, probeAfter, probing := pp.ModelProbeInfo(sl.Model)
|
||||
out = append(out, AutoSlotState{
|
||||
Model: sl.Model,
|
||||
Source: sl.Source,
|
||||
@ -498,6 +647,10 @@ func (c *Core) AutoSlotStates() []AutoSlotState {
|
||||
FailCount: fail,
|
||||
CooldownUntil: until,
|
||||
Cooling: until > now,
|
||||
CooldownFrom: from,
|
||||
ProbeAfter: probeAfter,
|
||||
Probing: probing,
|
||||
Probeable: until > now && probeAfter > 0 && now >= probeAfter,
|
||||
})
|
||||
}
|
||||
}
|
||||
@ -549,7 +702,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.
|
||||
@ -601,6 +754,33 @@ func (c *Core) Sources() []config.Source {
|
||||
return c.mergedSources()
|
||||
}
|
||||
|
||||
// Templates returns the saved source templates.
|
||||
func (c *Core) Templates() []config.SourceTemplate {
|
||||
c.mu.Lock()
|
||||
defer c.mu.Unlock()
|
||||
return c.store.ListTemplates()
|
||||
}
|
||||
|
||||
// SaveTemplate upserts a source template.
|
||||
func (c *Core) SaveTemplate(t config.SourceTemplate) error {
|
||||
c.mu.Lock()
|
||||
defer c.mu.Unlock()
|
||||
if t.Name == "" {
|
||||
return fmt.Errorf("template requires a name")
|
||||
}
|
||||
return c.store.UpsertTemplate(t)
|
||||
}
|
||||
|
||||
// RemoveTemplate deletes a source template by name.
|
||||
func (c *Core) RemoveTemplate(name string) error {
|
||||
c.mu.Lock()
|
||||
defer c.mu.Unlock()
|
||||
if _, err := c.store.RemoveTemplate(name); err != nil {
|
||||
return err
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func normalizeSource(s *config.Source) error {
|
||||
if s.Name == "" || s.BaseURL == "" {
|
||||
return fmt.Errorf("source requires name and base_url")
|
||||
@ -622,4 +802,21 @@ func (c *Core) Close() {
|
||||
if c.vm != nil {
|
||||
c.vm.Stop()
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// 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
|
||||
}
|
||||
|
||||
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)
|
||||
}
|
||||
}
|
||||
}
|
||||
@ -4,6 +4,7 @@ import (
|
||||
"encoding/csv"
|
||||
"encoding/json"
|
||||
"io"
|
||||
"log"
|
||||
"net/http"
|
||||
"strconv"
|
||||
"strings"
|
||||
@ -31,7 +32,11 @@ func (g *Gateway) handleAdaptersAPI(w http.ResponseWriter, r *http.Request) {
|
||||
writeJSON(w, http.StatusOK, map[string]interface{}{"adapters": g.core.ListAdapters()})
|
||||
return
|
||||
}
|
||||
writeError(w, http.StatusNotFound, "not_found", "adapter code not exposed; edit in UI")
|
||||
// A GET on a specific adapter is almost always a client that meant to
|
||||
// DELETE it but let fetch default to GET; say so instead of only
|
||||
// reporting that the code is not exposed.
|
||||
writeError(w, http.StatusNotFound, "not_found",
|
||||
"adapter code not exposed; edit in UI (to remove it use DELETE /api/adapters/"+path+")")
|
||||
case http.MethodPost:
|
||||
var p adapterPayload
|
||||
if err := json.NewDecoder(r.Body).Decode(&p); err != nil {
|
||||
@ -72,6 +77,63 @@ type sourcePayload struct {
|
||||
Temperature float64 `json:"temperature"`
|
||||
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"`
|
||||
}
|
||||
|
||||
// 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) {
|
||||
@ -85,17 +147,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,
|
||||
@ -105,6 +184,30 @@ func (g *Gateway) handleSourcesAPI(w http.ResponseWriter, r *http.Request) {
|
||||
Temperature: p.Temperature,
|
||||
MaxTokens: p.MaxTokens,
|
||||
MaxConcurrent: p.MaxConcurrent,
|
||||
RPM: p.RPM,
|
||||
}
|
||||
// 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())
|
||||
@ -126,6 +229,62 @@ func (g *Gateway) handleSourcesAPI(w http.ResponseWriter, r *http.Request) {
|
||||
}
|
||||
}
|
||||
|
||||
// csvHeaders stamps the shared download headers for CSV exports.
|
||||
func csvHeaders(w http.ResponseWriter, filename string) {
|
||||
w.Header().Set("Content-Type", "text/csv; charset=utf-8")
|
||||
w.Header().Set("Content-Disposition", "attachment; filename="+filename)
|
||||
}
|
||||
|
||||
// handleSourceTemplatesAPI manages reusable source templates (a Source minus
|
||||
// name and api_key) stored in the runtime file so the WebUI can spin up
|
||||
// multiple key-bearing sources from one shared template.
|
||||
func (g *Gateway) handleSourceTemplatesAPI(w http.ResponseWriter, r *http.Request) {
|
||||
if reqRole(r.Context()) != "admin" {
|
||||
writeError(w, http.StatusForbidden, "forbidden", "admin role required")
|
||||
return
|
||||
}
|
||||
path := strings.TrimPrefix(r.URL.Path, "/api/source_templates")
|
||||
path = strings.Trim(path, "/")
|
||||
switch r.Method {
|
||||
case http.MethodGet:
|
||||
writeJSON(w, http.StatusOK, map[string]interface{}{"templates": g.core.Templates()})
|
||||
case http.MethodPost:
|
||||
var t config.SourceTemplate
|
||||
if err := json.NewDecoder(r.Body).Decode(&t); err != nil {
|
||||
writeError(w, http.StatusBadRequest, "invalid_request", "invalid json: "+err.Error())
|
||||
return
|
||||
}
|
||||
if err := g.core.SaveTemplate(t); err != nil {
|
||||
writeError(w, http.StatusBadRequest, "template_error", err.Error())
|
||||
return
|
||||
}
|
||||
writeJSON(w, http.StatusOK, map[string]interface{}{"ok": true, "templates": g.core.Templates()})
|
||||
case http.MethodDelete:
|
||||
if path == "" {
|
||||
writeError(w, http.StatusBadRequest, "invalid_request", "template name required")
|
||||
return
|
||||
}
|
||||
if err := g.core.RemoveTemplate(path); err != nil {
|
||||
writeError(w, http.StatusBadRequest, "template_error", err.Error())
|
||||
return
|
||||
}
|
||||
writeJSON(w, http.StatusOK, map[string]interface{}{"ok": true})
|
||||
default:
|
||||
writeError(w, http.StatusMethodNotAllowed, "method_not_allowed", "")
|
||||
}
|
||||
}
|
||||
|
||||
// exportKey resolves which masked key id an export covers: admins may pass
|
||||
// any key filter, user keys are always scoped to themselves (records and
|
||||
// aggregates are keyed by the masked keyID form).
|
||||
func exportKey(r *http.Request) string {
|
||||
k := r.URL.Query().Get("key")
|
||||
if reqRole(r.Context()) != "admin" {
|
||||
k = keyID(reqKey(r.Context()))
|
||||
}
|
||||
return k
|
||||
}
|
||||
|
||||
// handleStatsAPI returns per-key / per-model / per-source usage aggregates and
|
||||
// the recent request audit trail.
|
||||
func (g *Gateway) handleStatsAPI(w http.ResponseWriter, r *http.Request) {
|
||||
@ -133,33 +292,35 @@ func (g *Gateway) handleStatsAPI(w http.ResponseWriter, r *http.Request) {
|
||||
writeError(w, http.StatusMethodNotAllowed, "method_not_allowed", "use GET")
|
||||
return
|
||||
}
|
||||
limit := 500
|
||||
// Default to ONE SCREEN of records. The dashboard pages the rest through
|
||||
// /api/stats/records as the user scrolls, so a poll must not serialize the
|
||||
// whole ring buffer. An explicit limit is still honoured (capped by the ring).
|
||||
limit := firstScreenRecords
|
||||
if v := r.URL.Query().Get("limit"); v != "" {
|
||||
if n, err := strconv.Atoi(v); err == nil && n > 0 && n <= 20000 {
|
||||
limit = n
|
||||
}
|
||||
}
|
||||
key := r.URL.Query().Get("key")
|
||||
if reqRole(r.Context()) != "admin" {
|
||||
// user keys may only see their own usage - records and aggregates are
|
||||
// keyed by the masked id (keyID), so filter on that masked form.
|
||||
key = keyID(reqKey(r.Context()))
|
||||
}
|
||||
key := exportKey(r)
|
||||
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)
|
||||
if to == 0 {
|
||||
to = time.Now().UnixMilli()
|
||||
}
|
||||
w.Header().Set("Content-Type", "text/csv; charset=utf-8")
|
||||
w.Header().Set("Content-Disposition", "attachment; filename=llmsproxy-requests.csv")
|
||||
csvHeaders(w, "llmsproxy-requests.csv")
|
||||
cw := csv.NewWriter(w)
|
||||
names := map[string]string{}
|
||||
for _, k := range g.core.ListKeys() {
|
||||
names[keyID(k.Key)] = k.Name
|
||||
}
|
||||
_ = cw.Write([]string{"time", "key", "key_name", "type", "model", "source", "status", "ok", "prompt_tokens", "completion_tokens", "latency_ms", "error"})
|
||||
for _, rec := range g.stats.AuditRecords(from, to, key) {
|
||||
_ = cw.Write([]string{"time", "key", "key_name", "type", "model", "source", "status", "ok", "prompt_tokens", "completion_tokens", "latency_ms", "first_byte_ms", "cache_hit_tokens", "cache_miss_tokens", "error"})
|
||||
// Stream the window straight to the client instead of materializing it:
|
||||
// an export covering months of audit data must not be bounded by RAM, and
|
||||
// nothing is retained once the response is written.
|
||||
flushEvery := 1000
|
||||
n := 0
|
||||
err := g.stats.StreamAuditRecords(from, to, key, func(rec Req) error {
|
||||
_ = cw.Write([]string{
|
||||
time.UnixMilli(rec.Time).Format(time.RFC3339),
|
||||
rec.Key,
|
||||
@ -172,38 +333,48 @@ func (g *Gateway) handleStatsAPI(w http.ResponseWriter, r *http.Request) {
|
||||
strconv.FormatInt(rec.Prompt, 10),
|
||||
strconv.FormatInt(rec.Compl, 10),
|
||||
strconv.FormatInt(rec.LatMs, 10),
|
||||
strconv.FormatInt(rec.FirstByteMs, 10),
|
||||
strconv.FormatInt(rec.CacheHit, 10),
|
||||
strconv.FormatInt(rec.CacheMiss, 10),
|
||||
rec.Err,
|
||||
})
|
||||
}
|
||||
n++
|
||||
if n%flushEvery == 0 {
|
||||
cw.Flush()
|
||||
if err := cw.Error(); err != nil {
|
||||
return err // client went away: stop walking the audit files
|
||||
}
|
||||
if f, ok := w.(http.Flusher); ok {
|
||||
f.Flush()
|
||||
}
|
||||
}
|
||||
return nil
|
||||
})
|
||||
cw.Flush()
|
||||
if err != nil {
|
||||
log.Printf("[gateway] csv export aborted after %d rows: %v", n, err)
|
||||
}
|
||||
return
|
||||
}
|
||||
if r.URL.Query().Get("export") == "keys-csv" {
|
||||
w.Header().Set("Content-Type", "text/csv; charset=utf-8")
|
||||
w.Header().Set("Content-Disposition", "attachment; filename=llmsproxy-keys.csv")
|
||||
csvHeaders(w, "llmsproxy-keys.csv")
|
||||
cw := csv.NewWriter(w)
|
||||
_ = cw.Write([]string{"key", "key_name", "role", "models", "total_requests", "success_requests", "failed_requests", "prompt_tokens", "completion_tokens", "total_tokens", "avg_latency_ms", "max_latency_ms", "created_at"})
|
||||
// Use the same key filtering as the JSON API
|
||||
exportKey := r.URL.Query().Get("key")
|
||||
if reqRole(r.Context()) != "admin" {
|
||||
exportKey = keyID(reqKey(r.Context()))
|
||||
}
|
||||
keyFilter := exportKey(r)
|
||||
// by_key rows are keyed by the masked id (keyID); build a masked-id ->
|
||||
// record map so names/roles/models resolve for the export.
|
||||
byMasked := map[string]config.GWKey{}
|
||||
for _, k := range g.core.ListKeys() {
|
||||
byMasked[keyID(k.Key)] = k
|
||||
}
|
||||
snap := g.stats.Snapshot(0, exportKey)
|
||||
snap := g.stats.Snapshot(0, keyFilter)
|
||||
byKeyRaw, ok := snap["by_key"].([]StatsRow)
|
||||
if !ok {
|
||||
byKeyRaw = []StatsRow{}
|
||||
}
|
||||
for _, row := range byKeyRaw {
|
||||
keyInfo, found := byMasked[row.Name]
|
||||
name := ""
|
||||
role := ""
|
||||
models := ""
|
||||
name, role, models := "", "", ""
|
||||
if found {
|
||||
name = keyInfo.Name
|
||||
role = keyInfo.Role
|
||||
@ -213,17 +384,9 @@ func (g *Gateway) handleStatsAPI(w http.ResponseWriter, r *http.Request) {
|
||||
}
|
||||
models = strings.Join(modelNames, ",")
|
||||
}
|
||||
reqs := row.Stat.Reqs
|
||||
success := row.Stat.OK
|
||||
errCount := row.Stat.Err
|
||||
prompt := row.Stat.Prompt
|
||||
compl := row.Stat.Compl
|
||||
tokens := row.Stat.Prompt + row.Stat.Compl
|
||||
latSum := row.Stat.LatSum
|
||||
latMax := row.Stat.LatMax
|
||||
avgLat := int64(0)
|
||||
if reqs > 0 {
|
||||
avgLat = latSum / reqs
|
||||
if row.Stat.Reqs > 0 {
|
||||
avgLat = row.Stat.LatSum / row.Stat.Reqs
|
||||
}
|
||||
createdAt := ""
|
||||
if found && keyInfo.CreatedAt > 0 {
|
||||
@ -231,14 +394,14 @@ func (g *Gateway) handleStatsAPI(w http.ResponseWriter, r *http.Request) {
|
||||
}
|
||||
_ = cw.Write([]string{
|
||||
row.Name, name, role, models,
|
||||
strconv.FormatInt(reqs, 10),
|
||||
strconv.FormatInt(success, 10),
|
||||
strconv.FormatInt(errCount, 10),
|
||||
strconv.FormatInt(prompt, 10),
|
||||
strconv.FormatInt(compl, 10),
|
||||
strconv.FormatInt(tokens, 10),
|
||||
strconv.FormatInt(row.Stat.Reqs, 10),
|
||||
strconv.FormatInt(row.Stat.OK, 10),
|
||||
strconv.FormatInt(row.Stat.Err, 10),
|
||||
strconv.FormatInt(row.Stat.Prompt, 10),
|
||||
strconv.FormatInt(row.Stat.Compl, 10),
|
||||
strconv.FormatInt(row.Stat.Prompt+row.Stat.Compl, 10),
|
||||
strconv.FormatInt(avgLat, 10),
|
||||
strconv.FormatInt(latMax, 10),
|
||||
strconv.FormatInt(row.Stat.LatMax, 10),
|
||||
createdAt,
|
||||
})
|
||||
}
|
||||
@ -253,3 +416,38 @@ func (g *Gateway) handleStatsAPI(w http.ResponseWriter, r *http.Request) {
|
||||
snap["key_names"] = keyNames
|
||||
writeJSON(w, http.StatusOK, snap)
|
||||
}
|
||||
|
||||
// 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
|
||||
// per-client state (the cursor is the whole state, and it lives in the URL) and
|
||||
// closes every file handle before responding.
|
||||
//
|
||||
// GET /api/stats/records?before=<cursor>&limit=100[&key=<masked id>]
|
||||
// -> { records: [...newest first...], next_cursor, has_more, rotated }
|
||||
func (g *Gateway) handleStatsRecordsAPI(w http.ResponseWriter, r *http.Request) {
|
||||
if r.Method != http.MethodGet {
|
||||
writeError(w, http.StatusMethodNotAllowed, "method_not_allowed", "use GET")
|
||||
return
|
||||
}
|
||||
limit := 100
|
||||
if v := r.URL.Query().Get("limit"); v != "" {
|
||||
if n, err := strconv.Atoi(v); err == nil && n > 0 && n <= 1000 {
|
||||
limit = n
|
||||
}
|
||||
}
|
||||
// exportKey pins non-admin callers to their own key, so a user key can never
|
||||
// page another key's records.
|
||||
page := g.stats.AuditPage(r.URL.Query().Get("before"), limit, exportKey(r))
|
||||
keyNames := map[string]string{}
|
||||
for _, k := range g.core.ListKeys() {
|
||||
keyNames[keyID(k.Key)] = k.Name
|
||||
}
|
||||
writeJSON(w, http.StatusOK, map[string]interface{}{
|
||||
"records": page.Records,
|
||||
"next_cursor": page.Next,
|
||||
"has_more": page.HasMore,
|
||||
"rotated": page.Rotated,
|
||||
"key_names": keyNames,
|
||||
})
|
||||
}
|
||||
|
||||
421
internal/gateway/apiv1.go
Normal file
421
internal/gateway/apiv1.go
Normal file
@ -0,0 +1,421 @@
|
||||
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
|
||||
}
|
||||
writeJSON(w, http.StatusOK, map[string]interface{}{"source": s})
|
||||
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: "DELETE", Path: "/api/keys/{name}", Auth: "admin", Summary: "delete a gateway key",
|
||||
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: "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",
|
||||
"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 map[string]interface{}{
|
||||
"name": s.Name,
|
||||
"base_url": s.BaseURL,
|
||||
"adapter": s.Adapter,
|
||||
"endpoint": s.Endpoint,
|
||||
"image_endpoint": s.ImageEndpoint,
|
||||
"api_key": maskKey(s.APIKey),
|
||||
"api_key_set": s.APIKey != "",
|
||||
"models": s.Models,
|
||||
"headers": maskHeaders(s.Headers),
|
||||
"proxy_url": s.ProxyURL,
|
||||
"meta": s.Meta,
|
||||
"temperature": s.Temperature,
|
||||
"max_tokens": s.MaxTokens,
|
||||
"max_concurrent": s.MaxConcurrent,
|
||||
"rpm": s.RPM,
|
||||
}
|
||||
}
|
||||
|
||||
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")
|
||||
}
|
||||
}
|
||||
@ -5,7 +5,10 @@ import (
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"log"
|
||||
"net/http"
|
||||
"regexp"
|
||||
"strconv"
|
||||
"strings"
|
||||
"sync/atomic"
|
||||
"time"
|
||||
@ -27,6 +30,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.
|
||||
@ -37,6 +43,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 {
|
||||
@ -53,14 +63,18 @@ type RespMessage struct {
|
||||
}
|
||||
|
||||
type ChatChunk struct {
|
||||
ID string `json:"id"`
|
||||
Object string `json:"object"`
|
||||
Created int64 `json:"created"`
|
||||
Model string `json:"model"`
|
||||
Choices []ChunkChoice `json:"choices"`
|
||||
ID string `json:"id"`
|
||||
Object string `json:"object"`
|
||||
Created int64 `json:"created"`
|
||||
Model string `json:"model"`
|
||||
Choices []ChunkChoice `json:"choices"`
|
||||
// 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 {
|
||||
@ -114,13 +128,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
|
||||
}
|
||||
}
|
||||
@ -151,20 +167,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{}
|
||||
@ -178,48 +192,100 @@ 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)
|
||||
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" /
|
||||
@ -228,18 +294,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"), ""
|
||||
@ -286,6 +375,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
|
||||
@ -300,8 +391,8 @@ func (g *Gateway) handleChat(w http.ResponseWriter, r *http.Request) {
|
||||
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()
|
||||
@ -316,6 +407,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)),
|
||||
@ -353,8 +445,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()
|
||||
@ -371,6 +463,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)),
|
||||
@ -387,13 +480,6 @@ func (g *Gateway) handleChat(w http.ResponseWriter, r *http.Request) {
|
||||
g.singleChat(w, ctx, cands, inner, effective, rec)
|
||||
}
|
||||
|
||||
func normalizeModel(m string) string {
|
||||
if isAuto(m) {
|
||||
return ""
|
||||
}
|
||||
return m
|
||||
}
|
||||
|
||||
func firstModel(p *provider.Provider) string {
|
||||
ms := p.Models()
|
||||
if len(ms) > 0 {
|
||||
@ -523,21 +609,145 @@ func upstreamErrStatus(err error) int {
|
||||
return http.StatusBadGateway
|
||||
}
|
||||
|
||||
func (g *Gateway) singleChat(w http.ResponseWriter, ctx context.Context, cands []*provider.Provider, req *types.ChatRequest, effective string, rec *Req) {
|
||||
rec.LatMs = 0
|
||||
t0 := time.Now()
|
||||
resp, usedSrc, usedModel, err := g.core.Scheduler().Chat(ctx, toScheduler(cands), req)
|
||||
rec.LatMs = time.Since(t0).Milliseconds()
|
||||
if err != nil {
|
||||
rec.OK = false
|
||||
rec.Status = upstreamErrStatus(err)
|
||||
rec.Err = err.Error()
|
||||
g.writeRec(rec)
|
||||
writeError(w, rec.Status, "upstream_error", err.Error())
|
||||
return
|
||||
// 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
|
||||
}
|
||||
}
|
||||
rec.OK = true
|
||||
rec.Status = http.StatusOK
|
||||
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 {
|
||||
// 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, 160))
|
||||
}
|
||||
msg := strings.Join(parts, "; ")
|
||||
if len(msg) > 300 {
|
||||
msg = msg[:300] + "..."
|
||||
}
|
||||
return fmt.Sprintf("all %d auto providers failed: %s", len(parts), msg)
|
||||
}
|
||||
return types.OneLine(err.Error(), 160)
|
||||
}
|
||||
|
||||
// failChat maps a scheduling failure onto the audit record and answers the
|
||||
// client. A failed AUTO chain additionally pins its first failed tier onto
|
||||
// the record; direct-path errors never match *scheduler.ChainErr, so the
|
||||
// extraction is safely shared by all four entry points. Callers own the
|
||||
// writeRec call (inline for non-stream paths, deferred for stream paths).
|
||||
func (g *Gateway) failChat(w http.ResponseWriter, rec *Req, err error) {
|
||||
rec.OK = false
|
||||
rec.Status = upstreamErrStatus(err)
|
||||
rec.Err = err.Error()
|
||||
var ce *scheduler.ChainErr
|
||||
if errors.As(err, &ce) && len(ce.Tiers) > 0 {
|
||||
rec.Source = ce.Tiers[0].Source
|
||||
rec.Model = ce.Tiers[0].Model
|
||||
}
|
||||
writeError(w, rec.Status, "upstream_error", clientUpstreamErr(err))
|
||||
}
|
||||
|
||||
// recordChatUsage fills token accounting for a finished non-streaming
|
||||
// request: exact upstream numbers win, byte estimates fill the gaps.
|
||||
func recordChatUsage(rec *Req, req *types.ChatRequest, resp *types.UnifiedResponse) {
|
||||
rec.Prompt = int64(resp.TokenUsage.Prompt)
|
||||
if rec.Prompt == 0 {
|
||||
rec.Prompt = estimatePromptTokens(req)
|
||||
@ -546,9 +756,29 @@ func (g *Gateway) singleChat(w http.ResponseWriter, ctx context.Context, cands [
|
||||
if rec.Compl == 0 {
|
||||
rec.Compl = estimateTextTokens(resp.Content, resp.ReasoningContent, resp.ToolCalls)
|
||||
}
|
||||
rec.Source = usedSrc
|
||||
rec.Model = usedModel
|
||||
g.writeRec(rec)
|
||||
// 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.
|
||||
if d := resp.TokenUsage.PromptTokensDetails; d != nil && d.CachedTokens > 0 {
|
||||
rec.CacheHit = int64(d.CachedTokens)
|
||||
}
|
||||
if resp.TokenUsage.PromptTokensDetails != nil {
|
||||
// Upstream reported cache details (even a 0 hit) — tag the row so
|
||||
// the UI can show 0% rather than “—”.
|
||||
rec.CacheReported = true
|
||||
} else if resp.TokenUsage.PromptCacheHit > 0 {
|
||||
rec.CacheHit = int64(resp.TokenUsage.PromptCacheHit)
|
||||
rec.CacheReported = true
|
||||
}
|
||||
if resp.TokenUsage.PromptCacheMiss > 0 {
|
||||
rec.CacheMiss = int64(resp.TokenUsage.PromptCacheMiss)
|
||||
}
|
||||
}
|
||||
|
||||
// writeChatCompletion renders a unified response as an OpenAI
|
||||
// chat.completion object. modelName is the id clients see as the serving
|
||||
// model: the requested id for direct routes, the exact slot model for AUTO.
|
||||
func writeChatCompletion(w http.ResponseWriter, resp *types.UnifiedResponse, modelName string) {
|
||||
msg := RespMessage{Role: "assistant", Content: resp.Content}
|
||||
if resp.ReasoningContent != "" {
|
||||
msg.ReasoningContent = resp.ReasoningContent
|
||||
@ -560,8 +790,9 @@ func (g *Gateway) singleChat(w http.ResponseWriter, ctx context.Context, cands [
|
||||
ID: newID(),
|
||||
Object: "chat.completion",
|
||||
Created: time.Now().Unix(),
|
||||
Model: effective,
|
||||
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
|
||||
@ -569,6 +800,30 @@ func (g *Gateway) singleChat(w http.ResponseWriter, ctx context.Context, cands [
|
||||
writeJSON(w, http.StatusOK, out)
|
||||
}
|
||||
|
||||
// singleChat runs a direct (model-pinned) non-streaming request across the
|
||||
// candidate list, falling back on failure.
|
||||
func (g *Gateway) singleChat(w http.ResponseWriter, ctx context.Context, cands []*provider.Provider, req *types.ChatRequest, effective string, rec *Req) {
|
||||
rec.LatMs = 0
|
||||
t0 := time.Now()
|
||||
resp, usedSrc, usedModel, err := g.core.Scheduler().Chat(ctx, toScheduler(cands), req)
|
||||
rec.LatMs = time.Since(t0).Milliseconds()
|
||||
if err != nil {
|
||||
g.failChat(w, rec, err)
|
||||
g.writeRec(rec)
|
||||
return
|
||||
}
|
||||
rec.OK = true
|
||||
rec.Status = http.StatusOK
|
||||
recordChatUsage(rec, req, resp)
|
||||
rec.Source = usedSrc
|
||||
rec.Model = usedModel
|
||||
// Non-streaming: the whole response arrives at once, so TTFB equals
|
||||
// the total latency.
|
||||
rec.FirstByteMs = rec.LatMs
|
||||
g.writeRec(rec)
|
||||
writeChatCompletion(w, resp, effective)
|
||||
}
|
||||
|
||||
// writeRec records a finished request (audit + aggregates).
|
||||
func (g *Gateway) writeRec(rec *Req) {
|
||||
if rec == nil {
|
||||
@ -603,30 +858,28 @@ func mergeUsage(prev, cur *types.TokenUsage) *types.TokenUsage {
|
||||
out.Completion = cur.Completion
|
||||
}
|
||||
out.Total = out.Prompt + out.Completion
|
||||
if cur.PromptTokensDetails != nil {
|
||||
// Keep the details object even when CachedTokens is 0: a reported
|
||||
// zero-hit is meaningful ("cache missed") and must stay
|
||||
// distinguishable from "upstream never reported cache info".
|
||||
// Dropping it here made streaming rows show “—” instead of 0%.
|
||||
out.PromptTokensDetails = cur.PromptTokensDetails
|
||||
}
|
||||
if cur.PromptCacheHit > 0 {
|
||||
out.PromptCacheHit = cur.PromptCacheHit
|
||||
}
|
||||
if cur.PromptCacheMiss > 0 {
|
||||
out.PromptCacheMiss = cur.PromptCacheMiss
|
||||
}
|
||||
return &out
|
||||
}
|
||||
|
||||
func (g *Gateway) streamChat(w http.ResponseWriter, ctx context.Context, cands []*provider.Provider, req *types.ChatRequest, effective string, rec *Req) {
|
||||
rec.LatMs = 0
|
||||
t0 := time.Now()
|
||||
rec.OK = true
|
||||
rec.Status = http.StatusOK
|
||||
defer func() {
|
||||
rec.LatMs = time.Since(t0).Milliseconds()
|
||||
g.writeRec(rec)
|
||||
}()
|
||||
chunks, _, usedModel, err := g.core.Scheduler().ChatStream(ctx, toScheduler(cands), req)
|
||||
if err != nil {
|
||||
rec.OK = false
|
||||
rec.Status = upstreamErrStatus(err)
|
||||
rec.Err = err.Error()
|
||||
writeError(w, rec.Status, "upstream_error", err.Error())
|
||||
return
|
||||
}
|
||||
if usedModel != "" {
|
||||
rec.Model = usedModel
|
||||
}
|
||||
rec.Prompt = estimatePromptTokens(req)
|
||||
// pumpStream writes the full SSE sequence for a started stream: role
|
||||
// preamble, one chunk per unified delta, the terminating finish_reason, the
|
||||
// OpenAI-standard final usage chunk (empty choices) and [DONE]. modelName
|
||||
// follows writeChatCompletion's rule (requested id for direct routes, exact
|
||||
// slot model for AUTO).
|
||||
func (g *Gateway) pumpStream(w http.ResponseWriter, rec *Req, chunks <-chan types.UnifiedChunk, modelName string, t0 time.Time) {
|
||||
w.Header().Set("Content-Type", "text/event-stream")
|
||||
w.Header().Set("Cache-Control", "no-cache")
|
||||
w.Header().Set("Connection", "keep-alive")
|
||||
@ -650,18 +903,26 @@ func (g *Gateway) streamChat(w http.ResponseWriter, ctx context.Context, cands [
|
||||
}
|
||||
|
||||
if !send(ChatChunk{
|
||||
ID: id, Object: "chat.completion.chunk", Created: created, Model: effective,
|
||||
ID: id, Object: "chat.completion.chunk", Created: created, Model: modelName,
|
||||
Choices: []ChunkChoice{{Index: 0, Delta: RespMessage{Role: "assistant"}}},
|
||||
}) {
|
||||
return
|
||||
}
|
||||
// First SSE byte sent to the client: record time-to-first-byte for the
|
||||
// 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: effective,
|
||||
ID: id, Object: "chat.completion.chunk", Created: created, Model: modelName,
|
||||
}
|
||||
delta := RespMessage{Role: "assistant", Content: ck.Content}
|
||||
if ck.ReasoningContent != "" {
|
||||
@ -672,8 +933,12 @@ func (g *Gateway) streamChat(w http.ResponseWriter, ctx context.Context, cands [
|
||||
}
|
||||
choice := ChunkChoice{Index: 0, Delta: delta}
|
||||
if ck.Done {
|
||||
stop := "stop"
|
||||
choice.FinishReason = &stop
|
||||
fin := ck.FinishReason
|
||||
if fin == "" {
|
||||
fin = "stop"
|
||||
}
|
||||
lastFinish = fin
|
||||
choice.FinishReason = &fin
|
||||
}
|
||||
chunk.Choices = []ChunkChoice{choice}
|
||||
rec.Compl += int64(len(ck.Content)+len(ck.ReasoningContent)+len(ck.ToolCalls)) / 3
|
||||
@ -681,17 +946,47 @@ func (g *Gateway) streamChat(w http.ResponseWriter, ctx context.Context, cands [
|
||||
return
|
||||
}
|
||||
}
|
||||
stop := "stop"
|
||||
finalFinish := lastFinish
|
||||
if finalFinish == "" {
|
||||
finalFinish = "stop"
|
||||
}
|
||||
send(ChatChunk{
|
||||
ID: id, Object: "chat.completion.chunk", Created: created, Model: effective,
|
||||
Choices: []ChunkChoice{{Index: 0, Delta: RespMessage{}, FinishReason: &stop}},
|
||||
ID: id, Object: "chat.completion.chunk", Created: created, Model: modelName,
|
||||
Choices: []ChunkChoice{{Index: 0, Delta: RespMessage{}, FinishReason: &finalFinish}},
|
||||
})
|
||||
// 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)
|
||||
}
|
||||
} else if rec.Prompt+rec.Compl > 0 {
|
||||
u := types.TokenUsage{
|
||||
Prompt: int(rec.Prompt),
|
||||
@ -702,9 +997,17 @@ func (g *Gateway) streamChat(w http.ResponseWriter, ctx context.Context, cands [
|
||||
}
|
||||
if tut != nil {
|
||||
send(ChatChunk{
|
||||
ID: id, Object: "chat.completion.chunk", Created: created, Model: effective,
|
||||
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")
|
||||
@ -713,6 +1016,33 @@ func (g *Gateway) streamChat(w http.ResponseWriter, ctx context.Context, cands [
|
||||
}
|
||||
}
|
||||
|
||||
// streamChat runs a direct (model-pinned) streaming request across the
|
||||
// candidate list, falling back early on connect errors.
|
||||
func (g *Gateway) streamChat(w http.ResponseWriter, ctx context.Context, cands []*provider.Provider, req *types.ChatRequest, effective string, rec *Req) {
|
||||
rec.LatMs = 0
|
||||
t0 := time.Now()
|
||||
rec.OK = true
|
||||
rec.Status = http.StatusOK
|
||||
defer func() {
|
||||
rec.LatMs = time.Since(t0).Milliseconds()
|
||||
g.writeRec(rec)
|
||||
}()
|
||||
chunks, usedSrc, usedModel, err := g.core.Scheduler().ChatStream(ctx, toScheduler(cands), req)
|
||||
if err != nil {
|
||||
g.failChat(w, rec, err)
|
||||
return
|
||||
}
|
||||
if usedModel != "" {
|
||||
rec.Model = usedModel
|
||||
}
|
||||
// Audit accuracy: pin the source that actually served the stream (after a
|
||||
// failover it differs from the first candidate). Direct streams previously
|
||||
// discarded it.
|
||||
rec.Source = usedSrc
|
||||
rec.Prompt = estimatePromptTokens(req)
|
||||
g.pumpStream(w, rec, chunks, effective, t0)
|
||||
}
|
||||
|
||||
// singleChatAuto runs a non-streaming AUTO request down the chain (see
|
||||
// scheduler.ChainChat): tiers ascending (tier 1 = highest priority first),
|
||||
// per-tier round-robin ordered by preference, cooldown as the only hard skip,
|
||||
@ -725,48 +1055,18 @@ func (g *Gateway) singleChatAuto(w http.ResponseWriter, ctx context.Context, cha
|
||||
resp, usedSrc, usedModel, err := g.core.Scheduler().ChainChat(ctx, chain, req, quotaExhausted)
|
||||
rec.LatMs = time.Since(t0).Milliseconds()
|
||||
if err != nil {
|
||||
rec.OK = false
|
||||
rec.Status = upstreamErrStatus(err)
|
||||
rec.Err = err.Error()
|
||||
if ce, ok := err.(*scheduler.ChainErr); ok && len(ce.Tiers) > 0 {
|
||||
rec.Source = ce.Tiers[0].Source
|
||||
rec.Model = ce.Tiers[0].Model
|
||||
}
|
||||
g.failChat(w, rec, err)
|
||||
g.writeRec(rec)
|
||||
writeError(w, rec.Status, "upstream_error", err.Error())
|
||||
return
|
||||
}
|
||||
rec.OK = true
|
||||
rec.Status = http.StatusOK
|
||||
rec.Prompt = int64(resp.TokenUsage.Prompt)
|
||||
if rec.Prompt == 0 {
|
||||
rec.Prompt = estimatePromptTokens(req)
|
||||
}
|
||||
rec.Compl = int64(resp.TokenUsage.Completion)
|
||||
if rec.Compl == 0 {
|
||||
rec.Compl = estimateTextTokens(resp.Content, resp.ReasoningContent, resp.ToolCalls)
|
||||
}
|
||||
recordChatUsage(rec, req, resp)
|
||||
rec.Source = usedSrc
|
||||
rec.Model = usedModel
|
||||
rec.FirstByteMs = rec.LatMs
|
||||
g.writeRec(rec)
|
||||
msg := RespMessage{Role: "assistant", Content: resp.Content}
|
||||
if resp.ReasoningContent != "" {
|
||||
msg.ReasoningContent = resp.ReasoningContent
|
||||
}
|
||||
if len(resp.ToolCalls) > 0 {
|
||||
msg.ToolCalls = toolCallsWire(resp.ToolCalls)
|
||||
}
|
||||
out := ChatCompletion{
|
||||
ID: newID(),
|
||||
Object: "chat.completion",
|
||||
Created: time.Now().Unix(),
|
||||
Model: usedModel,
|
||||
Choices: []ChatChoice{{Index: 0, Message: msg, FinishReason: resp.FinishReason}},
|
||||
}
|
||||
if resp.TokenUsage.Total > 0 || resp.TokenUsage.Prompt > 0 || resp.TokenUsage.Completion > 0 {
|
||||
out.Usage = &resp.TokenUsage
|
||||
}
|
||||
writeJSON(w, http.StatusOK, out)
|
||||
writeChatCompletion(w, resp, usedModel)
|
||||
}
|
||||
|
||||
// streamChatAuto streams an AUTO request down the chain. A slot is abandoned
|
||||
@ -784,14 +1084,7 @@ func (g *Gateway) streamChatAuto(w http.ResponseWriter, ctx context.Context, cha
|
||||
}()
|
||||
chunks, usedSrc, usedModel, err := g.core.Scheduler().ChainChatStream(ctx, chain, req, quotaExhausted)
|
||||
if err != nil {
|
||||
rec.OK = false
|
||||
rec.Status = upstreamErrStatus(err)
|
||||
rec.Err = err.Error()
|
||||
if ce, ok := err.(*scheduler.ChainErr); ok && len(ce.Tiers) > 0 {
|
||||
rec.Source = ce.Tiers[0].Source
|
||||
rec.Model = ce.Tiers[0].Model
|
||||
}
|
||||
writeError(w, rec.Status, "upstream_error", err.Error())
|
||||
g.failChat(w, rec, err)
|
||||
return
|
||||
}
|
||||
if usedModel != "" {
|
||||
@ -799,88 +1092,7 @@ func (g *Gateway) streamChatAuto(w http.ResponseWriter, ctx context.Context, cha
|
||||
}
|
||||
rec.Source = usedSrc
|
||||
rec.Prompt = estimatePromptTokens(req)
|
||||
w.Header().Set("Content-Type", "text/event-stream")
|
||||
w.Header().Set("Cache-Control", "no-cache")
|
||||
w.Header().Set("Connection", "keep-alive")
|
||||
w.WriteHeader(http.StatusOK)
|
||||
flusher, _ := w.(http.Flusher)
|
||||
id := newID()
|
||||
created := time.Now().Unix()
|
||||
send := func(obj interface{}) bool {
|
||||
b, err := json.Marshal(obj)
|
||||
if err != nil {
|
||||
return false
|
||||
}
|
||||
if _, err := fmt.Fprintf(w, "data: %s\n\n", b); err != nil {
|
||||
return false
|
||||
}
|
||||
if flusher != nil {
|
||||
flusher.Flush()
|
||||
}
|
||||
return true
|
||||
}
|
||||
if !send(ChatChunk{
|
||||
ID: id, Object: "chat.completion.chunk", Created: created, Model: rec.Model,
|
||||
Choices: []ChunkChoice{{Index: 0, Delta: RespMessage{Role: "assistant"}}},
|
||||
}) {
|
||||
return
|
||||
}
|
||||
var lastUsage *types.TokenUsage
|
||||
for ck := range chunks {
|
||||
if ck.Usage != nil {
|
||||
lastUsage = mergeUsage(lastUsage, ck.Usage)
|
||||
}
|
||||
chunk := ChatChunk{
|
||||
ID: id, Object: "chat.completion.chunk", Created: created, Model: rec.Model,
|
||||
}
|
||||
delta := RespMessage{Role: "assistant", Content: ck.Content}
|
||||
if ck.ReasoningContent != "" {
|
||||
delta.ReasoningContent = ck.ReasoningContent
|
||||
}
|
||||
if len(ck.ToolCalls) > 0 {
|
||||
delta.ToolCalls = ck.ToolCalls
|
||||
}
|
||||
choice := ChunkChoice{Index: 0, Delta: delta}
|
||||
if ck.Done {
|
||||
stop := "stop"
|
||||
choice.FinishReason = &stop
|
||||
}
|
||||
chunk.Choices = []ChunkChoice{choice}
|
||||
rec.Compl += int64(len(ck.Content)+len(ck.ReasoningContent)+len(ck.ToolCalls)) / 3
|
||||
if !send(chunk) {
|
||||
return
|
||||
}
|
||||
}
|
||||
stop := "stop"
|
||||
send(ChatChunk{
|
||||
ID: id, Object: "chat.completion.chunk", Created: created, Model: rec.Model,
|
||||
Choices: []ChunkChoice{{Index: 0, Delta: RespMessage{}, FinishReason: &stop}},
|
||||
})
|
||||
// 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.
|
||||
var tut *types.TokenUsage
|
||||
if lastUsage != nil {
|
||||
tut = lastUsage
|
||||
} else if rec.Prompt+rec.Compl > 0 {
|
||||
u := types.TokenUsage{
|
||||
Prompt: int(rec.Prompt),
|
||||
Completion: int(rec.Compl),
|
||||
Total: int(rec.Prompt + rec.Compl),
|
||||
}
|
||||
tut = &u
|
||||
}
|
||||
if tut != nil {
|
||||
send(ChatChunk{
|
||||
ID: id, Object: "chat.completion.chunk", Created: created, Model: rec.Model,
|
||||
Choices: []ChunkChoice{},
|
||||
Usage: tut,
|
||||
})
|
||||
}
|
||||
fmt.Fprintf(w, "data: [DONE]\n\n")
|
||||
if flusher != nil {
|
||||
flusher.Flush()
|
||||
}
|
||||
g.pumpStream(w, rec, chunks, usedModel, t0)
|
||||
}
|
||||
|
||||
func (g *Gateway) handleImage(w http.ResponseWriter, r *http.Request) {
|
||||
@ -901,11 +1113,48 @@ func (g *Gateway) handleImage(w http.ResponseWriter, r *http.Request) {
|
||||
if model == "" {
|
||||
model = g.core.DefaultModel()
|
||||
}
|
||||
if !isAuto(model) {
|
||||
if allow := g.allowedModels(r.Context()); allow != nil && !g.hasScopeModel(allow, model) {
|
||||
writeError(w, http.StatusForbidden, "model_not_allowed", fmt.Sprintf("model %q is not allowed for this key", model))
|
||||
if isAuto(model) {
|
||||
if chain := g.core.AutoImageChain(); chain != nil && len(chain.Tiers) > 0 {
|
||||
if q := g.checkQuota(r.Context(), "AUTO"); q != nil {
|
||||
g.writeReject(w, q)
|
||||
return
|
||||
}
|
||||
done := g.stats.Begin()
|
||||
defer done()
|
||||
rec := &Req{Key: keyID(reqKey(r.Context())), Type: "image", Model: model, OK: false}
|
||||
t0 := time.Now()
|
||||
resp, usedSrc, usedModel, err := g.core.Scheduler().ChainImage(r.Context(), chain, &req)
|
||||
rec.LatMs = time.Since(t0).Milliseconds()
|
||||
if err != nil {
|
||||
rec.Status = upstreamErrStatus(err)
|
||||
rec.Err = err.Error()
|
||||
g.writeRec(rec)
|
||||
writeError(w, rec.Status, "upstream_error", clientUpstreamErr(err))
|
||||
return
|
||||
}
|
||||
rec.Source = usedSrc
|
||||
if usedModel != "" {
|
||||
rec.Model = usedModel // actual image model served, not "AUTO"
|
||||
}
|
||||
rec.OK = true
|
||||
rec.Status = http.StatusOK
|
||||
// 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(),
|
||||
Data: resp.ImageData,
|
||||
})
|
||||
return
|
||||
}
|
||||
// no image chain configured: fall through to legacy discovery (all
|
||||
// sources exposing an image model, tried in registry order)
|
||||
}
|
||||
if allow := g.allowedModels(r.Context()); allow != nil && !g.hasScopeModel(allow, model) {
|
||||
writeError(w, http.StatusForbidden, "model_not_allowed", fmt.Sprintf("model %q is not allowed for this key", model))
|
||||
return
|
||||
}
|
||||
cands, _ := g.resolveByModel(model)
|
||||
cands = imageOnly(cands)
|
||||
@ -916,8 +1165,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()
|
||||
@ -930,15 +1179,19 @@ func (g *Gateway) handleImage(w http.ResponseWriter, r *http.Request) {
|
||||
rec.Status = upstreamErrStatus(err)
|
||||
rec.Err = err.Error()
|
||||
g.writeRec(rec)
|
||||
writeError(w, rec.Status, "upstream_error", err.Error())
|
||||
writeError(w, rec.Status, "upstream_error", clientUpstreamErr(err))
|
||||
return
|
||||
}
|
||||
if usedSrc != "" {
|
||||
rec.Source = usedSrc
|
||||
}
|
||||
if resp.Model != "" {
|
||||
rec.Model = resp.Model // record the actual model served, not the raw request id
|
||||
}
|
||||
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(),
|
||||
|
||||
62
internal/gateway/chat_cache_test.go
Normal file
62
internal/gateway/chat_cache_test.go
Normal file
@ -0,0 +1,62 @@
|
||||
package gateway
|
||||
|
||||
import (
|
||||
"testing"
|
||||
|
||||
"llmsproxy/internal/types"
|
||||
)
|
||||
|
||||
// TestMergeUsageKeepsZeroCacheDetails locks the streaming counterpart of the
|
||||
// "distinguish missed from not reported" contract: a usage chunk that reports
|
||||
// prompt_tokens_details with 0 cached tokens must survive the merge, otherwise
|
||||
// streaming rows lose cache_reported and the UI shows "—" instead of 0%.
|
||||
func TestMergeUsageKeepsZeroCacheDetails(t *testing.T) {
|
||||
t.Run("zero-hit details survive merge", func(t *testing.T) {
|
||||
prev := &types.TokenUsage{Prompt: 10, Completion: 1, Total: 11}
|
||||
cur := &types.TokenUsage{
|
||||
Prompt: 10, Completion: 5, Total: 15,
|
||||
PromptTokensDetails: &types.PromptTokensDetails{CachedTokens: 0},
|
||||
}
|
||||
got := mergeUsage(prev, cur)
|
||||
if got.PromptTokensDetails == nil {
|
||||
t.Fatal("zero-hit prompt_tokens_details was dropped by mergeUsage")
|
||||
}
|
||||
if got.PromptTokensDetails.CachedTokens != 0 {
|
||||
t.Fatalf("cached_tokens = %d, want 0", got.PromptTokensDetails.CachedTokens)
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("non-zero hit still wins", func(t *testing.T) {
|
||||
prev := &types.TokenUsage{Prompt: 10, Completion: 1, Total: 11}
|
||||
cur := &types.TokenUsage{
|
||||
Prompt: 10, Completion: 5, Total: 15,
|
||||
PromptTokensDetails: &types.PromptTokensDetails{CachedTokens: 64},
|
||||
}
|
||||
got := mergeUsage(prev, cur)
|
||||
if got.PromptTokensDetails == nil || got.PromptTokensDetails.CachedTokens != 64 {
|
||||
t.Fatalf("cached_tokens lost: %+v", got.PromptTokensDetails)
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("absent details do not overwrite an earlier report", func(t *testing.T) {
|
||||
prev := &types.TokenUsage{
|
||||
Prompt: 10, Completion: 1, Total: 11,
|
||||
PromptTokensDetails: &types.PromptTokensDetails{CachedTokens: 32},
|
||||
}
|
||||
cur := &types.TokenUsage{Prompt: 10, Completion: 5, Total: 15}
|
||||
got := mergeUsage(prev, cur)
|
||||
if got.PromptTokensDetails == nil || got.PromptTokensDetails.CachedTokens != 32 {
|
||||
t.Fatalf("earlier cache report clobbered by a later chunk without details: %+v",
|
||||
got.PromptTokensDetails)
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("no cache data anywhere stays nil", func(t *testing.T) {
|
||||
prev := &types.TokenUsage{Prompt: 10, Completion: 1, Total: 11}
|
||||
cur := &types.TokenUsage{Prompt: 10, Completion: 5, Total: 15}
|
||||
got := mergeUsage(prev, cur)
|
||||
if got.PromptTokensDetails != nil {
|
||||
t.Fatalf("fabricated cache details: %+v", got.PromptTokensDetails)
|
||||
}
|
||||
})
|
||||
}
|
||||
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)
|
||||
}
|
||||
}
|
||||
@ -6,10 +6,13 @@ import (
|
||||
"io"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"net/url"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strconv"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"llmsproxy/internal/config"
|
||||
"llmsproxy/internal/core"
|
||||
@ -55,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)
|
||||
}
|
||||
@ -139,7 +142,7 @@ func TestChatAutoChain503Summary(t *testing.T) {
|
||||
if rr.Code != http.StatusServiceUnavailable {
|
||||
t.Fatalf("status=%d body=%s", rr.Code, rr.Body.String())
|
||||
}
|
||||
if !strings.Contains(rr.Body.String(), "all auto tiers failed") ||
|
||||
if !strings.Contains(rr.Body.String(), "auto providers failed") ||
|
||||
!strings.Contains(rr.Body.String(), "a/a-m") ||
|
||||
!strings.Contains(rr.Body.String(), "b/b-m") {
|
||||
t.Fatalf("503 must summarize every tier, body=%s", rr.Body.String())
|
||||
@ -667,7 +670,351 @@ 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")
|
||||
}
|
||||
}
|
||||
|
||||
// TestStatusRecorderFlusher pins the SSE-critical contract: the access-log
|
||||
// wrapper must implement http.Flusher, otherwise the streaming chat handlers'
|
||||
// w.(http.Flusher) assertion yields nil and chunks are never flushed until
|
||||
// the response ends (regression for the buffered-SSE bug).
|
||||
func TestStatusRecorderFlusher(t *testing.T) {
|
||||
var inner *flushRecorder
|
||||
rr := httptest.NewRecorder()
|
||||
inner = &flushRecorder{ResponseWriter: rr}
|
||||
sr := &statusRecorder{ResponseWriter: inner}
|
||||
|
||||
fl, ok := any(sr).(http.Flusher)
|
||||
if !ok {
|
||||
t.Fatal("statusRecorder does not implement http.Flusher — SSE streaming is broken")
|
||||
}
|
||||
if inner.flushed {
|
||||
t.Fatal("Flush called before Flush()")
|
||||
}
|
||||
fl.Flush()
|
||||
if !inner.flushed {
|
||||
t.Fatal("statusRecorder.Flush did not forward to the underlying writer")
|
||||
}
|
||||
}
|
||||
|
||||
// flushRecorder records whether Flush was forwarded.
|
||||
type flushRecorder struct {
|
||||
http.ResponseWriter
|
||||
flushed bool
|
||||
}
|
||||
|
||||
func (f *flushRecorder) Flush() { f.flushed = true }
|
||||
|
||||
// TestDirectStreamFailoverAuditSource: a direct streaming request whose first
|
||||
// candidate hard-fails must be served by the fallback, and the audit record's
|
||||
// Source must name the source that actually served the stream (not just the
|
||||
// first candidate).
|
||||
func TestDirectStreamFailoverAuditSource(t *testing.T) {
|
||||
aUp := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
w.WriteHeader(500)
|
||||
fmt.Fprint(w, `{"error":"boom"}`)
|
||||
}))
|
||||
defer aUp.Close()
|
||||
bUp := mockUpstream()
|
||||
defer bUp.Close()
|
||||
g := newTestGateway(t,
|
||||
// Both sources expose the same model id so the direct scheduler has a
|
||||
// real fallback candidate after a hard-fails; b is the healthy one.
|
||||
config.Source{Name: "a", BaseURL: aUp.URL, Adapter: "openai", Models: []config.Model{{ID: "m"}}},
|
||||
config.Source{Name: "b", BaseURL: bUp.URL, Adapter: "openai", Models: []config.Model{{ID: "m"}}},
|
||||
)
|
||||
body := doReq(t, g, "POST", "/v1/chat/completions",
|
||||
`{"model":"m","stream":true,"messages":[{"role":"user","content":"hi"}]}`)
|
||||
if body.Code != 200 {
|
||||
t.Fatalf("status=%d body=%s", body.Code, body.Body.String())
|
||||
}
|
||||
sb := body.Body.String()
|
||||
if !strings.Contains(sb, `"finish_reason":"stop"`) || !strings.Contains(sb, "[DONE]") {
|
||||
t.Fatalf("stream incomplete: %s", sb)
|
||||
}
|
||||
recs := g.stats.AuditRecords(0, 0, "")
|
||||
if len(recs) == 0 {
|
||||
t.Fatal("no audit records")
|
||||
}
|
||||
last := recs[len(recs)-1]
|
||||
if last.Source != "b" || !last.OK {
|
||||
t.Fatalf("audit rec source=%q ok=%v, want source=b ok=true (actual serving source)", last.Source, last.OK)
|
||||
}
|
||||
}
|
||||
|
||||
// TestStatusExposesAdapterPools checks that the elastic Lua pool sizing is
|
||||
// visible to admins (and only to them) via /api/status, so the grow/shrink
|
||||
// algorithm can be observed instead of inferred.
|
||||
func TestStatusExposesAdapterPools(t *testing.T) {
|
||||
g := newTestGateway(t)
|
||||
rr := doReq(t, g, http.MethodGet, "/api/status", "")
|
||||
if rr.Code != 200 {
|
||||
t.Fatalf("status = %d, body=%s", rr.Code, rr.Body.String())
|
||||
}
|
||||
var body map[string]interface{}
|
||||
if err := json.Unmarshal(rr.Body.Bytes(), &body); err != nil {
|
||||
t.Fatalf("decode: %v", err)
|
||||
}
|
||||
pools, ok := body["adapter_pools"].([]interface{})
|
||||
if !ok {
|
||||
t.Fatalf("adapter_pools missing or wrong type: %T", body["adapter_pools"])
|
||||
}
|
||||
if len(pools) == 0 {
|
||||
t.Fatal("adapter_pools must list the loaded adapters")
|
||||
}
|
||||
first, _ := pools[0].(map[string]interface{})
|
||||
for _, field := range []string{"name", "created", "idle", "in_use", "max", "grow_step", "shrink_step"} {
|
||||
if _, ok := first[field]; !ok {
|
||||
t.Errorf("pool entry missing %q: %v", field, first)
|
||||
}
|
||||
}
|
||||
// idle gateway: nothing booted eagerly
|
||||
if c, _ := first["created"].(float64); c != 0 {
|
||||
t.Errorf("adapter pool booted eagerly, created=%v", first["created"])
|
||||
}
|
||||
}
|
||||
|
||||
// TestStatsRecordsAPIPaging exercises the on-demand records endpoint the
|
||||
// dashboard scrolls with: a bounded first page plus a cursor that walks back
|
||||
// through the audit file.
|
||||
func TestStatsRecordsAPIPaging(t *testing.T) {
|
||||
g := newTestGateway(t)
|
||||
for i := 0; i < 250; i++ {
|
||||
g.stats.Record(Req{
|
||||
Time: time.Now().UnixMilli() - int64(250-i)*1000,
|
||||
Key: keyID("sk-test"), Type: "chat", Model: "m", Source: "s",
|
||||
Prompt: 1, Compl: 1, OK: true, Status: 200,
|
||||
})
|
||||
}
|
||||
|
||||
rr := doReq(t, g, http.MethodGet, "/api/stats/records?limit=50", "")
|
||||
if rr.Code != 200 {
|
||||
t.Fatalf("status = %d body=%s", rr.Code, rr.Body.String())
|
||||
}
|
||||
var page struct {
|
||||
Records []Req `json:"records"`
|
||||
NextCursor string `json:"next_cursor"`
|
||||
HasMore bool `json:"has_more"`
|
||||
}
|
||||
if err := json.Unmarshal(rr.Body.Bytes(), &page); err != nil {
|
||||
t.Fatalf("decode: %v", err)
|
||||
}
|
||||
if len(page.Records) != 50 {
|
||||
t.Fatalf("first page = %d records, want 50", len(page.Records))
|
||||
}
|
||||
if !page.HasMore || page.NextCursor == "" {
|
||||
t.Fatal("a long history must advertise more pages plus a cursor")
|
||||
}
|
||||
// newest first
|
||||
for i := 1; i < len(page.Records); i++ {
|
||||
if page.Records[i-1].Time < page.Records[i].Time {
|
||||
t.Fatalf("page not newest-first at %d", i)
|
||||
}
|
||||
}
|
||||
|
||||
// follow the cursor: the next page must continue strictly older
|
||||
oldest := page.Records[len(page.Records)-1].Time
|
||||
rr2 := doReq(t, g, http.MethodGet,
|
||||
"/api/stats/records?limit=50&before="+url.QueryEscape(page.NextCursor), "")
|
||||
if rr2.Code != 200 {
|
||||
t.Fatalf("page 2 status = %d", rr2.Code)
|
||||
}
|
||||
var page2 struct {
|
||||
Records []Req `json:"records"`
|
||||
}
|
||||
if err := json.Unmarshal(rr2.Body.Bytes(), &page2); err != nil {
|
||||
t.Fatalf("decode page 2: %v", err)
|
||||
}
|
||||
if len(page2.Records) == 0 {
|
||||
t.Fatal("cursor page must return records")
|
||||
}
|
||||
if page2.Records[0].Time > oldest {
|
||||
t.Fatalf("page 2 starts at %d, must continue below %d", page2.Records[0].Time, oldest)
|
||||
}
|
||||
}
|
||||
|
||||
// TestStatsSnapshotShipsOneScreen: a dashboard poll must not serialize the whole
|
||||
// ring, otherwise the "load the first screen only" contract is broken on the
|
||||
// wire even if the UI pages.
|
||||
func TestStatsSnapshotShipsOneScreen(t *testing.T) {
|
||||
g := newTestGateway(t)
|
||||
for i := 0; i < 400; i++ {
|
||||
g.stats.Record(Req{
|
||||
Time: time.Now().UnixMilli(), Key: keyID("sk-test"), Type: "chat",
|
||||
Model: "m", Source: "s", Prompt: 1, Compl: 1, OK: true, Status: 200,
|
||||
})
|
||||
}
|
||||
rr := doReq(t, g, http.MethodGet, "/api/stats", "")
|
||||
var snap struct {
|
||||
Records []Req `json:"records"`
|
||||
Total Stat `json:"total"`
|
||||
}
|
||||
if err := json.Unmarshal(rr.Body.Bytes(), &snap); err != nil {
|
||||
t.Fatalf("decode: %v", err)
|
||||
}
|
||||
if len(snap.Records) > firstScreenRecords {
|
||||
t.Fatalf("snapshot shipped %d records, want <= %d", len(snap.Records), firstScreenRecords)
|
||||
}
|
||||
// aggregates still cover every request
|
||||
if snap.Total.Reqs != 400 {
|
||||
t.Fatalf("total reqs = %d, want 400 (aggregates must not be paged away)", snap.Total.Reqs)
|
||||
}
|
||||
}
|
||||
|
||||
// TestStatsCsvExportStreams checks the export still emits every row in the
|
||||
// window now that it is written straight from the audit walk.
|
||||
func TestStatsCsvExportStreams(t *testing.T) {
|
||||
g := newTestGateway(t)
|
||||
base := time.Now().UnixMilli()
|
||||
for i := 0; i < 120; i++ {
|
||||
g.stats.Record(Req{
|
||||
Time: base + int64(i)*1000, Key: keyID("sk-test"), Type: "chat",
|
||||
Model: "m", Source: "s", Prompt: 2, Compl: 3, OK: true, Status: 200,
|
||||
})
|
||||
}
|
||||
rr := doReq(t, g, http.MethodGet,
|
||||
"/api/stats?export=csv&from="+strconv.FormatInt(base, 10)+
|
||||
"&to="+strconv.FormatInt(base+120*1000, 10), "")
|
||||
if rr.Code != 200 {
|
||||
t.Fatalf("export status = %d", rr.Code)
|
||||
}
|
||||
lines := strings.Split(strings.TrimSpace(rr.Body.String()), "\n")
|
||||
if len(lines) != 121 { // header + 120 rows
|
||||
t.Fatalf("csv had %d lines, want 121 (header + 120 records)", len(lines))
|
||||
}
|
||||
if !strings.HasPrefix(lines[0], "time,key,key_name") {
|
||||
t.Fatalf("unexpected csv header: %q", lines[0])
|
||||
}
|
||||
}
|
||||
|
||||
// TestStatsRecordsAPIScopedToOwnKey: a non-admin key must only page its own
|
||||
// records even when it asks for someone else's.
|
||||
func TestStatsRecordsAPIScopedToOwnKey(t *testing.T) {
|
||||
g := newTestGateway(t)
|
||||
mine := keyID("sk-test")
|
||||
for i := 0; i < 5; i++ {
|
||||
g.stats.Record(Req{Time: time.Now().UnixMilli(), Key: mine, Type: "chat", Model: "m", Source: "s", OK: true, Status: 200})
|
||||
}
|
||||
for i := 0; i < 5; i++ {
|
||||
g.stats.Record(Req{Time: time.Now().UnixMilli(), Key: "othermask", Type: "chat", Model: "m", Source: "s", OK: true, Status: 200})
|
||||
}
|
||||
// sk-test is an admin in the test gateway, so it legitimately sees all keys;
|
||||
// assert the explicit filter path instead.
|
||||
rr := doReq(t, g, http.MethodGet, "/api/stats/records?limit=100&key=othermask", "")
|
||||
var page struct {
|
||||
Records []Req `json:"records"`
|
||||
}
|
||||
if err := json.Unmarshal(rr.Body.Bytes(), &page); err != nil {
|
||||
t.Fatalf("decode: %v", err)
|
||||
}
|
||||
for _, r := range page.Records {
|
||||
if r.Key != "othermask" {
|
||||
t.Fatalf("key filter leaked %q", r.Key)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// TestKeyDeleteRoundTrip is the end-to-end version of the bug users hit:
|
||||
// deleting a gateway key from the admin UI. It also asserts the GET-on-a-key
|
||||
// error names the right verb, since that message is what a mis-methoded client
|
||||
// actually sees.
|
||||
func TestKeyDeleteRoundTrip(t *testing.T) {
|
||||
g := newTestGateway(t)
|
||||
|
||||
rr := doReq(t, g, http.MethodPost, "/api/keys", `{"name":"doomed","role":"user"}`)
|
||||
if rr.Code != 200 {
|
||||
t.Fatalf("create: %d %s", rr.Code, rr.Body.String())
|
||||
}
|
||||
var created struct {
|
||||
Key struct {
|
||||
Key string `json:"key"`
|
||||
Name string `json:"name"`
|
||||
} `json:"key"`
|
||||
}
|
||||
if err := json.Unmarshal(rr.Body.Bytes(), &created); err != nil {
|
||||
t.Fatalf("decode create: %v", err)
|
||||
}
|
||||
if created.Key.Key == "" {
|
||||
t.Fatalf("no key returned: %s", rr.Body.String())
|
||||
}
|
||||
|
||||
// DELETE removes it
|
||||
rr = doReq(t, g, http.MethodDelete, "/api/keys/"+url.PathEscape(created.Key.Key), "")
|
||||
if rr.Code != 200 {
|
||||
t.Fatalf("delete: %d %s", rr.Code, rr.Body.String())
|
||||
}
|
||||
for _, k := range g.core.ListKeys() {
|
||||
if k.Key == created.Key.Key {
|
||||
t.Fatal("key still present after DELETE")
|
||||
}
|
||||
}
|
||||
|
||||
// deleting again is a clean 404, not a 500
|
||||
rr = doReq(t, g, http.MethodDelete, "/api/keys/"+url.PathEscape(created.Key.Key), "")
|
||||
if rr.Code != http.StatusNotFound {
|
||||
t.Fatalf("second delete: %d, want 404", rr.Code)
|
||||
}
|
||||
}
|
||||
|
||||
// TestKeyGetOnSpecificKeyNamesTheVerb: a GET on /api/keys/{key} is what a
|
||||
// client that forgot to set method:"DELETE" ends up sending. The error must say
|
||||
// which verb to use instead of just "use GET /api/keys", which reads as if the
|
||||
// caller did nothing wrong.
|
||||
func TestKeyGetOnSpecificKeyNamesTheVerb(t *testing.T) {
|
||||
g := newTestGateway(t)
|
||||
rr := doReq(t, g, http.MethodGet, "/api/keys/sk-whatever", "")
|
||||
if rr.Code != http.StatusNotFound {
|
||||
t.Fatalf("status = %d, want 404", rr.Code)
|
||||
}
|
||||
body := rr.Body.String()
|
||||
for _, want := range []string{"DELETE /api/keys/", "PUT /api/keys/"} {
|
||||
if !strings.Contains(body, want) {
|
||||
t.Errorf("error message should mention %q, got: %s", want, body)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// TestAdapterGetOnSpecificAdapterNamesDelete is the same guard for the adapters
|
||||
// endpoint, whose delAdapter helper had the identical missing-method bug.
|
||||
func TestAdapterGetOnSpecificAdapterNamesDelete(t *testing.T) {
|
||||
g := newTestGateway(t)
|
||||
rr := doReq(t, g, http.MethodGet, "/api/adapters/openai", "")
|
||||
if rr.Code != http.StatusNotFound {
|
||||
t.Fatalf("status = %d, want 404", rr.Code)
|
||||
}
|
||||
if !strings.Contains(rr.Body.String(), "DELETE /api/adapters/openai") {
|
||||
t.Errorf("error should name the DELETE verb, got: %s", rr.Body.String())
|
||||
}
|
||||
}
|
||||
|
||||
// TestCannotDeleteOwnKey keeps the admin from locking themselves out.
|
||||
func TestCannotDeleteOwnKey(t *testing.T) {
|
||||
g := newTestGateway(t)
|
||||
rr := doReq(t, g, http.MethodDelete, "/api/keys/sk-test", "")
|
||||
if rr.Code != http.StatusBadRequest {
|
||||
t.Fatalf("status = %d, want 400", rr.Code)
|
||||
}
|
||||
if !strings.Contains(rr.Body.String(), "logged in with") {
|
||||
t.Errorf("unexpected message: %s", rr.Body.String())
|
||||
}
|
||||
}
|
||||
|
||||
// TestGatewayRingStaysSmall pins the resident record ring at the small default.
|
||||
// The gateway used to construct Stats with a 10000-record ring, which after the
|
||||
// full-history aggregate scan meant 10000 resident Req structs — putting back the
|
||||
// ~25 MB the on-demand paging work removed. Aggregates cover all of history;
|
||||
// the ring only needs the first screen plus the 5-minute status windows.
|
||||
func TestGatewayRingStaysSmall(t *testing.T) {
|
||||
g := newTestGateway(t)
|
||||
g.stats.mu.Lock()
|
||||
max := g.stats.maxRecs
|
||||
g.stats.mu.Unlock()
|
||||
if max != defaultRingSize {
|
||||
t.Fatalf("gateway record ring = %d, want %d (a big ring reintroduces the resident-memory cost)",
|
||||
max, defaultRingSize)
|
||||
}
|
||||
}
|
||||
|
||||
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)
|
||||
}
|
||||
}
|
||||
}
|
||||
@ -29,28 +29,26 @@ func (g *Gateway) handleKeysAPI(w http.ResponseWriter, r *http.Request) {
|
||||
switch r.Method {
|
||||
case http.MethodGet:
|
||||
if path != "" {
|
||||
writeError(w, http.StatusNotFound, "not_found", "use GET /api/keys")
|
||||
// A GET on a specific key is almost always a client that meant to
|
||||
// DELETE or PUT it but let fetch default to GET. Name the verbs
|
||||
// instead of only pointing back at the collection endpoint.
|
||||
writeError(w, http.StatusNotFound, "not_found",
|
||||
"no such endpoint; use GET /api/keys to list, DELETE /api/keys/{key} to remove, PUT /api/keys/{key} to update")
|
||||
return
|
||||
}
|
||||
writeJSON(w, http.StatusOK, map[string]interface{}{"keys": g.core.ListKeys()})
|
||||
case http.MethodPost:
|
||||
var body struct {
|
||||
Name string `json:"name"`
|
||||
Role string `json:"role"`
|
||||
Models []config.ModelScope `json:"models"`
|
||||
Note string `json:"note"`
|
||||
Name string `json:"name"`
|
||||
Role string `json:"role"`
|
||||
Models []config.ModelScope `json:"models"`
|
||||
Note string `json:"note"`
|
||||
}
|
||||
if err := json.NewDecoder(r.Body).Decode(&body); err != nil {
|
||||
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())
|
||||
@ -72,6 +70,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())
|
||||
@ -143,8 +144,9 @@ func (g *Gateway) allowedModels(ctx context.Context) []config.ModelScope {
|
||||
func (g *Gateway) handleAutoAPI(w http.ResponseWriter, r *http.Request) {
|
||||
if r.Method == http.MethodGet {
|
||||
writeJSON(w, http.StatusOK, map[string]interface{}{
|
||||
"rules": g.core.AutoRules(),
|
||||
"states": g.core.AutoSlotStates(),
|
||||
"rules": g.core.AutoRules(),
|
||||
"image_rules": g.core.AutoImageRules(),
|
||||
"states": g.core.AutoSlotStates(),
|
||||
})
|
||||
return
|
||||
}
|
||||
@ -155,18 +157,31 @@ func (g *Gateway) handleAutoAPI(w http.ResponseWriter, r *http.Request) {
|
||||
switch r.Method {
|
||||
case http.MethodPut, http.MethodPost:
|
||||
var body struct {
|
||||
Rules []config.ModelScope `json:"rules"`
|
||||
Rules []config.ModelScope `json:"rules"`
|
||||
ImageRules []config.ModelScope `json:"image_rules"`
|
||||
}
|
||||
if err := json.NewDecoder(r.Body).Decode(&body); err != nil {
|
||||
writeError(w, http.StatusBadRequest, "invalid_request", "invalid json: "+err.Error())
|
||||
return
|
||||
}
|
||||
if err := g.core.SaveAutoRules(body.Rules); err != nil {
|
||||
writeError(w, http.StatusBadRequest, "auto_error", err.Error())
|
||||
return
|
||||
if body.Rules != nil {
|
||||
if err := g.core.SaveAutoRules(body.Rules); err != nil {
|
||||
writeError(w, http.StatusBadRequest, "auto_error", err.Error())
|
||||
return
|
||||
}
|
||||
}
|
||||
writeJSON(w, http.StatusOK, map[string]interface{}{"ok": true, "rules": g.core.AutoRules()})
|
||||
if body.ImageRules != nil {
|
||||
if err := g.core.SaveAutoImageRules(body.ImageRules); err != nil {
|
||||
writeError(w, http.StatusBadRequest, "auto_error", err.Error())
|
||||
return
|
||||
}
|
||||
}
|
||||
writeJSON(w, http.StatusOK, map[string]interface{}{
|
||||
"ok": true,
|
||||
"rules": g.core.AutoRules(),
|
||||
"image_rules": g.core.AutoImageRules(),
|
||||
})
|
||||
default:
|
||||
writeError(w, http.StatusMethodNotAllowed, "method_not_allowed", "")
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
116
internal/gateway/login_guard_test.go
Normal file
116
internal/gateway/login_guard_test.go
Normal file
@ -0,0 +1,116 @@
|
||||
package gateway
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"strings"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestLoginBruteForceLockout(t *testing.T) {
|
||||
g := newTestGateway(t)
|
||||
body := `{"key":"wrong-key-XXXX"}`
|
||||
|
||||
// 5 failures allowed (each answered 401), 6th+ locked out
|
||||
var gotLocked bool
|
||||
for i := 1; i <= 8; i++ {
|
||||
req := httptest.NewRequest("POST", "/api/login", strings.NewReader(body))
|
||||
rr := httptest.NewRecorder()
|
||||
g.handleLoginAPI(rr, req)
|
||||
if rr.Code != http.StatusUnauthorized {
|
||||
t.Fatalf("attempt %d: status=%d want 401", i, rr.Code)
|
||||
}
|
||||
if i >= 6 {
|
||||
// after lockout the response must be identical (no info leak)
|
||||
if !strings.Contains(rr.Body.String(), "invalid gateway api key") {
|
||||
t.Fatalf("lockout reply changed body: %s", rr.Body.String())
|
||||
}
|
||||
gotLocked = true
|
||||
}
|
||||
}
|
||||
if !gotLocked {
|
||||
t.Fatal("expected lockout by attempt 6")
|
||||
}
|
||||
|
||||
// correct key from the SAME ip must also be rejected while locked
|
||||
okBody := `{"key":"sk-test"}`
|
||||
req := httptest.NewRequest("POST", "/api/login", strings.NewReader(okBody))
|
||||
rr := httptest.NewRecorder()
|
||||
g.handleLoginAPI(rr, req)
|
||||
if rr.Code != http.StatusUnauthorized {
|
||||
t.Fatalf("locked-out valid key: status=%d want 401", rr.Code)
|
||||
}
|
||||
|
||||
// a different IP is unaffected
|
||||
req = httptest.NewRequest("POST", "/api/login", strings.NewReader(okBody))
|
||||
req.RemoteAddr = "10.9.9.9:5555"
|
||||
rr = httptest.NewRecorder()
|
||||
g.handleLoginAPI(rr, req)
|
||||
if rr.Code != http.StatusOK {
|
||||
t.Fatalf("other IP valid key: status=%d want 200 body=%s", rr.Code, rr.Body.String())
|
||||
}
|
||||
var resp map[string]interface{}
|
||||
_ = json.Unmarshal(rr.Body.Bytes(), &resp)
|
||||
if resp["ok"] != true {
|
||||
t.Fatalf("login ok flag missing: %s", rr.Body.String())
|
||||
}
|
||||
}
|
||||
|
||||
func TestLoginSuccessResetsFailures(t *testing.T) {
|
||||
g := newTestGateway(t)
|
||||
mkBad := func() *http.Request {
|
||||
return httptest.NewRequest("POST", "/api/login", strings.NewReader(`{"key":"nope"}`))
|
||||
}
|
||||
for i := 0; i < loginMaxFails-1; i++ { // one below the lockout threshold
|
||||
rr := httptest.NewRecorder()
|
||||
g.handleLoginAPI(rr, mkBad())
|
||||
if rr.Code != http.StatusUnauthorized {
|
||||
t.Fatal("setup failure")
|
||||
}
|
||||
}
|
||||
// success clears the counter
|
||||
rr := httptest.NewRecorder()
|
||||
g.handleLoginAPI(rr, httptest.NewRequest("POST", "/api/login", strings.NewReader(`{"key":"sk-test"}`)))
|
||||
if rr.Code != http.StatusOK {
|
||||
t.Fatalf("valid login failed: %d", rr.Code)
|
||||
}
|
||||
// now up to loginMaxFails-1 more failures are still accepted before lock
|
||||
for i := 0; i < loginMaxFails-1; i++ {
|
||||
last := httptest.NewRecorder()
|
||||
g.handleLoginAPI(last, mkBad())
|
||||
}
|
||||
last := httptest.NewRecorder()
|
||||
g.handleLoginAPI(last, mkBad())
|
||||
// count is at loginMaxFails-1 failures since reset; this next one reaches
|
||||
// the threshold but the attempt itself is still evaluated normally (401),
|
||||
// and the FOLLOWING one must be locked.
|
||||
if last.Code != http.StatusUnauthorized {
|
||||
t.Fatalf("threshold attempt should still answer 401: %d", last.Code)
|
||||
}
|
||||
next := httptest.NewRecorder()
|
||||
g.handleLoginAPI(next, mkBad())
|
||||
if next.Code != http.StatusUnauthorized {
|
||||
t.Fatal("want 401")
|
||||
}
|
||||
// verify lock actually engaged via the guard state
|
||||
g.loginMu.Lock()
|
||||
f := g.loginFails["192.0.2.1"]
|
||||
engaged := f != nil && f.count >= loginMaxFails && f.until > 0
|
||||
g.loginMu.Unlock()
|
||||
if !engaged {
|
||||
t.Fatal("expected engaged lockout state")
|
||||
}
|
||||
}
|
||||
|
||||
func TestClientIPExtraction(t *testing.T) {
|
||||
r := httptest.NewRequest("POST", "/api/login", nil)
|
||||
r.RemoteAddr = "203.0.113.7:12345"
|
||||
if got := clientIP(r); got != "203.0.113.7" {
|
||||
t.Errorf("clientIP=%q want 203.0.113.7", got)
|
||||
}
|
||||
r.Header.Set("X-Real-IP", "198.51.100.9")
|
||||
if got := clientIP(r); got != "198.51.100.9" {
|
||||
t.Errorf("clientIP with X-Real-IP=%q want proxy value", got)
|
||||
}
|
||||
}
|
||||
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)
|
||||
}
|
||||
}
|
||||
@ -11,6 +11,7 @@ import (
|
||||
"fmt"
|
||||
"io/fs"
|
||||
"log"
|
||||
"net"
|
||||
"net/http"
|
||||
"net/url"
|
||||
"strings"
|
||||
@ -31,14 +32,101 @@ type Gateway struct {
|
||||
stats *Stats
|
||||
probeMu sync.Mutex
|
||||
lastProbe time.Time
|
||||
// loginGuard throttles /api/login brute-force attempts: per-IP failure
|
||||
// counters with exponential backoff, plus a global cap so a distributed
|
||||
// spray cannot outrun the per-IP window. Guarded by loginMu.
|
||||
loginMu sync.Mutex
|
||||
loginFails map[string]*loginFail
|
||||
loginWin int64 // unix sec of current global window start
|
||||
loginGlob int // failures in the global window
|
||||
}
|
||||
|
||||
func New(c *core.Core, gatewayKeys []string) (*Gateway, error) {
|
||||
// loginFail tracks consecutive failed logins from one source IP.
|
||||
type loginFail struct {
|
||||
count int
|
||||
until int64 // locked-out until this unix second (0 = not locked)
|
||||
lastAt int64
|
||||
}
|
||||
|
||||
const (
|
||||
loginMaxFails = 5 // failures before lockout kicks in
|
||||
loginLockBase = 30 // first lockout: 30s
|
||||
loginLockCap = 30 * 60 // repeated lockouts cap at 30min
|
||||
loginDecay = 15 * 60 // counters decay after 15min of quiet
|
||||
loginGlobalMax = 100 // global failures per minute across all IPs
|
||||
loginGlobalWindow = 60 // global window length (sec)
|
||||
)
|
||||
|
||||
// loginAllow checks (and records) a login attempt for ip. It returns false
|
||||
// when the attempt must be rejected: the IP is in exponential lockout or the
|
||||
// global failure budget for the current window is exhausted.
|
||||
func (g *Gateway) loginAllow(ip string, now int64) bool {
|
||||
g.loginMu.Lock()
|
||||
defer g.loginMu.Unlock()
|
||||
if g.loginFails == nil {
|
||||
g.loginFails = map[string]*loginFail{}
|
||||
}
|
||||
// global window rollover
|
||||
if now-g.loginWin >= loginGlobalWindow {
|
||||
g.loginWin = now
|
||||
g.loginGlob = 0
|
||||
}
|
||||
f := g.loginFails[ip]
|
||||
if f != nil && f.count > 0 && now-f.lastAt > loginDecay {
|
||||
delete(g.loginFails, ip) // quiet long enough: forgive and forget
|
||||
f = nil
|
||||
}
|
||||
if f != nil && f.until > now {
|
||||
return false // still locked out
|
||||
}
|
||||
if g.loginGlob >= loginGlobalMax {
|
||||
return false // global spray budget exhausted for this window
|
||||
}
|
||||
g.loginGlob++
|
||||
return true
|
||||
}
|
||||
|
||||
// loginRecord notes the outcome of one attempt: failures escalate toward an
|
||||
// exponential lockout; a success clears the IP's counter entirely.
|
||||
func (g *Gateway) loginRecord(ip string, ok bool, now int64) {
|
||||
g.loginMu.Lock()
|
||||
defer g.loginMu.Unlock()
|
||||
if ok {
|
||||
if g.loginFails != nil {
|
||||
delete(g.loginFails, ip)
|
||||
}
|
||||
return
|
||||
}
|
||||
f := g.loginFails[ip]
|
||||
if f == nil {
|
||||
f = &loginFail{}
|
||||
g.loginFails[ip] = f
|
||||
}
|
||||
f.count++
|
||||
f.lastAt = now
|
||||
if f.count >= loginMaxFails {
|
||||
// lock duration doubles with each extra burst of failures, capped
|
||||
epoch := int64(f.count/loginMaxFails) - 1
|
||||
dur := int64(loginLockBase) << min(epoch, 6)
|
||||
if dur > loginLockCap {
|
||||
dur = loginLockCap
|
||||
}
|
||||
f.until = now + dur
|
||||
}
|
||||
}
|
||||
|
||||
func New(c *core.Core) (*Gateway, error) {
|
||||
sub, err := fs.Sub(uiFS, "ui")
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
st := NewStats(10000)
|
||||
// The record ring is deliberately small: aggregates are built from the full
|
||||
// audit history at startup (see Stats.LoadAudit), so the ring only has to
|
||||
// cover the status page's 5-minute source windows and the dashboard's first
|
||||
// screen. Everything older is paged from disk by /api/stats/records.
|
||||
// A large ring here would put the ~25 MB of resident Req structs straight
|
||||
// back, which is exactly what the on-demand paging removed.
|
||||
st := NewStats(defaultRingSize)
|
||||
if cfg := c.Config(); cfg != nil && cfg.RuntimeFile != "" {
|
||||
st.LoadAudit(cfg.RuntimeFile + ".audit.jsonl")
|
||||
}
|
||||
@ -113,6 +201,17 @@ func (s *statusRecorder) Write(b []byte) (int, error) {
|
||||
return s.ResponseWriter.Write(b)
|
||||
}
|
||||
|
||||
// Flush forwards to the underlying writer so SSE handlers can stream
|
||||
// incrementally through the access-log wrapper. Without this method the
|
||||
// w.(http.Flusher) assertion inside the streaming chat handlers fails (the
|
||||
// embedded ResponseWriter interface does not carry Flush into the method set),
|
||||
// and every chunk stays buffered until the response ends.
|
||||
func (s *statusRecorder) Flush() {
|
||||
if f, ok := s.ResponseWriter.(http.Flusher); ok {
|
||||
f.Flush()
|
||||
}
|
||||
}
|
||||
|
||||
func (g *Gateway) routes(w http.ResponseWriter, r *http.Request) {
|
||||
switch {
|
||||
case r.URL.Path == "/v1/chat/completions":
|
||||
@ -123,14 +222,20 @@ 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)
|
||||
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/chat":
|
||||
g.handleChat(w, r)
|
||||
case r.URL.Path == "/api/status":
|
||||
g.handleStatusAPI(w, r)
|
||||
case r.URL.Path == "/api/status/reset":
|
||||
g.handleResetHealth(w, r)
|
||||
case r.URL.Path == "/api/stats/records":
|
||||
g.handleStatsRecordsAPI(w, r)
|
||||
case r.URL.Path == "/api/stats" || strings.HasPrefix(r.URL.Path, "/api/stats/"):
|
||||
g.handleStatsAPI(w, r)
|
||||
case r.URL.Path == "/api/keys" || strings.HasPrefix(r.URL.Path, "/api/keys/"):
|
||||
@ -325,23 +430,34 @@ document.getElementById('key').addEventListener('keydown',function(e){if(e.key==
|
||||
apply();
|
||||
</script></body></html>`
|
||||
|
||||
// handleLoginAPI validates the gateway key and issues a session cookie.
|
||||
// handleLoginAPI validates the gateway key and issues a session cookie,
|
||||
// guarded by per-IP + global brute-force throttling (see loginGuard).
|
||||
func (g *Gateway) handleLoginAPI(w http.ResponseWriter, r *http.Request) {
|
||||
if r.Method != http.MethodPost {
|
||||
writeError(w, http.StatusMethodNotAllowed, "method_not_allowed", "use POST")
|
||||
return
|
||||
}
|
||||
now := time.Now().Unix()
|
||||
if !g.loginAllow(clientIP(r), now) {
|
||||
// Don't reveal whether the key was right: answer the same 401 the
|
||||
// failure path uses, but skip the expensive key lookup.
|
||||
writeError(w, http.StatusUnauthorized, "invalid_api_key", "invalid gateway api key")
|
||||
return
|
||||
}
|
||||
var body struct {
|
||||
Key string `json:"key"`
|
||||
}
|
||||
if err := json.NewDecoder(r.Body).Decode(&body); err != nil {
|
||||
g.loginRecord(clientIP(r), false, now)
|
||||
writeError(w, http.StatusBadRequest, "invalid_request", "invalid json")
|
||||
return
|
||||
}
|
||||
if _, ok := g.core.FindKey(body.Key); !ok {
|
||||
g.loginRecord(clientIP(r), false, now)
|
||||
writeError(w, http.StatusUnauthorized, "invalid_api_key", "invalid gateway api key")
|
||||
return
|
||||
}
|
||||
g.loginRecord(clientIP(r), true, now)
|
||||
http.SetCookie(w, &http.Cookie{
|
||||
Name: "gw_key",
|
||||
Value: body.Key,
|
||||
@ -353,6 +469,20 @@ func (g *Gateway) handleLoginAPI(w http.ResponseWriter, r *http.Request) {
|
||||
writeJSON(w, http.StatusOK, map[string]interface{}{"ok": true})
|
||||
}
|
||||
|
||||
// clientIP returns the caller's IP for rate limiting. Behind a trusted
|
||||
// reverse proxy (nginx) the X-Real-IP header is preferred; otherwise the
|
||||
// remote address is used with the port stripped.
|
||||
func clientIP(r *http.Request) string {
|
||||
if x := r.Header.Get("X-Real-IP"); x != "" {
|
||||
return x
|
||||
}
|
||||
host, _, err := net.SplitHostPort(r.RemoteAddr)
|
||||
if err != nil {
|
||||
return r.RemoteAddr
|
||||
}
|
||||
return host
|
||||
}
|
||||
|
||||
func (g *Gateway) handleModels(w http.ResponseWriter, r *http.Request) {
|
||||
if r.Method != http.MethodGet {
|
||||
writeError(w, http.StatusMethodNotAllowed, "method_not_allowed", "use GET")
|
||||
@ -458,14 +588,23 @@ func (g *Gateway) handleStatusAPI(w http.ResponseWriter, r *http.Request) {
|
||||
if reqRole(r.Context()) == "admin" {
|
||||
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
|
||||
}
|
||||
}
|
||||
resp["sources"] = sts
|
||||
resp["adapters"] = g.core.ListAdapters()
|
||||
// Elastic Lua pool sizing: created/idle/in_use vs the ceiling, plus the
|
||||
// current grow/shrink steps, so the algorithm is inspectable in the UI
|
||||
// instead of being a black box.
|
||||
resp["adapter_pools"] = g.core.VM().PoolStats()
|
||||
}
|
||||
writeJSON(w, http.StatusOK, resp)
|
||||
}
|
||||
|
||||
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")
|
||||
}
|
||||
}
|
||||
@ -2,12 +2,16 @@ package gateway
|
||||
|
||||
import (
|
||||
"bufio"
|
||||
"bytes"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"io"
|
||||
"log"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"sort"
|
||||
"strconv"
|
||||
"strings"
|
||||
"sync"
|
||||
"time"
|
||||
)
|
||||
@ -26,23 +30,41 @@ type Req struct {
|
||||
Compl int64 `json:"completion_tokens"`
|
||||
// LatMs total handling time ms
|
||||
LatMs int64 `json:"latency_ms"`
|
||||
OK bool `json:"ok"`
|
||||
// FirstByteMs time-to-first-byte for streaming (ms from request start
|
||||
// to the first SSE chunk sent to the client); for non-streaming it
|
||||
// equals LatMs. 0 when unmeasured (legacy records).
|
||||
FirstByteMs int64 `json:"first_byte_ms,omitempty"`
|
||||
// CacheHit / CacheMiss carry the upstream prompt-cache accounting
|
||||
// (DeepSeek-style hit/miss tokens) when the upstream reports it.
|
||||
// Both 0 = upstream gave no cache data.
|
||||
CacheHit int64 `json:"cache_hit_tokens,omitempty"`
|
||||
CacheMiss int64 `json:"cache_miss_tokens,omitempty"`
|
||||
// CacheReported marks that the upstream usage reported cache
|
||||
// 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"`
|
||||
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"`
|
||||
}
|
||||
|
||||
// Stat aggregates counters for one dimension row.
|
||||
type Stat struct {
|
||||
Reqs int64 `json:"reqs"`
|
||||
OK int64 `json:"ok"`
|
||||
Err int64 `json:"err"`
|
||||
Tokens int64 `json:"tokens"`
|
||||
Prompt int64 `json:"prompt_tokens"`
|
||||
Compl int64 `json:"completion_tokens"`
|
||||
LatSum int64 `json:"latency_sum_ms"`
|
||||
LatMax int64 `json:"latency_max_ms"`
|
||||
Reqs int64 `json:"reqs"`
|
||||
OK int64 `json:"ok"`
|
||||
Err int64 `json:"err"`
|
||||
Tokens int64 `json:"tokens"`
|
||||
Prompt int64 `json:"prompt_tokens"`
|
||||
Compl int64 `json:"completion_tokens"`
|
||||
LatSum int64 `json:"latency_sum_ms"`
|
||||
LatMax int64 `json:"latency_max_ms"`
|
||||
FirstByteSum int64 `json:"first_byte_sum_ms,omitempty"`
|
||||
}
|
||||
|
||||
type agrRow struct {
|
||||
@ -53,18 +75,34 @@ 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.
|
||||
type Stats struct {
|
||||
mu sync.Mutex
|
||||
active int64
|
||||
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 // per http status code aggregates (incl. 402/400)
|
||||
recs []Req
|
||||
maxRecs int
|
||||
auditPath string
|
||||
modelHour map[string]map[int64]int64 // model -> unix-hour bucket -> tokens
|
||||
mu sync.Mutex
|
||||
active int64
|
||||
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 // per http status code aggregates (incl. 402/400)
|
||||
recs []Req
|
||||
maxRecs int
|
||||
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
|
||||
@ -73,24 +111,42 @@ const hourSec = 3600
|
||||
// file is renamed to <path>.<unix>.old and a fresh one is started); pruning
|
||||
// keeps at most auditKeepOld rotated files. Both are vars so tests can shrink
|
||||
// the threshold.
|
||||
//
|
||||
// The single-file threshold is deliberately small relative to the total budget
|
||||
// (16 MB x 16 = 256 MB): records are paged by reading a file backwards, and a
|
||||
// smaller file keeps the first page from seeking through a huge one.
|
||||
var (
|
||||
auditRotateBytes int64 = 64 << 20
|
||||
auditKeepOld = 10
|
||||
auditRotateBytes int64 = 16 << 20
|
||||
auditKeepOld = 16
|
||||
)
|
||||
|
||||
// defaultRingSize is how many recent request records stay resident. It has to
|
||||
// cover two consumers: the status page's 5-minute source windows
|
||||
// (SourceRecent/SourceAverages) and the first screen of the records table.
|
||||
// Everything beyond that is paged from the audit file.
|
||||
const defaultRingSize = 500
|
||||
|
||||
// firstScreenRecords is how many records a dashboard load ships by default — one
|
||||
// screen's worth. Scrolling pulls the rest through /api/stats/records.
|
||||
const firstScreenRecords = 100
|
||||
|
||||
func NewStats(maxRecords int) *Stats {
|
||||
if maxRecords <= 0 {
|
||||
maxRecords = 3000
|
||||
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,
|
||||
}
|
||||
}
|
||||
|
||||
@ -130,39 +186,141 @@ func incStatus(a *Stat, name string, r Req) {
|
||||
if r.LatMs > a.LatMax {
|
||||
a.LatMax = r.LatMs
|
||||
}
|
||||
if r.FirstByteMs > 0 {
|
||||
a.FirstByteSum += r.FirstByteMs
|
||||
}
|
||||
}
|
||||
|
||||
// LoadAudit primes the in-memory state from the audit file.
|
||||
//
|
||||
// The aggregates (totals, per-key/model/source rows, status counts and the
|
||||
// hourly quota buckets) are built from the FULL history: every audit file is
|
||||
// streamed oldest-first so the dashboard shows real all-time numbers rather than
|
||||
// whatever happened to fit in a replay window. This is affordable because the
|
||||
// scan keeps nothing per record — aggregate maps are keyed by key/model/source,
|
||||
// so their size is bounded by cardinality, not by request count. Measured on the
|
||||
// production host: 29 MB / 221k lines / 37k requests in ~260 ms.
|
||||
//
|
||||
// The raw-record ring is what stays bounded: only the newest maxRecs records are
|
||||
// retained, and everything older is paged from disk on demand by AuditPage /
|
||||
// StreamAuditRecords. The old behaviour — appending EVERY record into a slice
|
||||
// and truncating at the end — is what cost ~25 MB of resident memory.
|
||||
func (s *Stats) LoadAudit(path string) {
|
||||
f, err := os.Open(path)
|
||||
if err == nil {
|
||||
defer f.Close()
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
var recs []Req
|
||||
// tolerate long error summaries / oversized junk lines
|
||||
sc := bufio.NewScanner(f)
|
||||
sc.Buffer(make([]byte, 64*1024), 16*1024*1024)
|
||||
for sc.Scan() {
|
||||
var r Req
|
||||
if json.Unmarshal(sc.Bytes(), &r) != nil || r.Type == "" {
|
||||
continue // access/event rows (obj) and malformed lines are no requests
|
||||
}
|
||||
// replay EVERY request row into the aggregates so totals and
|
||||
// quota windows survive restarts; only the ring buffer view stays
|
||||
// capped at maxRecs.
|
||||
s.aggregateLocked(r)
|
||||
recs = append(recs, r)
|
||||
}
|
||||
if n := len(recs); n > s.maxRecs {
|
||||
recs = append([]Req(nil), recs[n-s.maxRecs:]...)
|
||||
}
|
||||
s.recs = recs
|
||||
s.auditPath = path
|
||||
return
|
||||
}
|
||||
s.mu.Lock()
|
||||
s.auditPath = path
|
||||
s.mu.Unlock()
|
||||
|
||||
// Oldest-first: the hourly-bucket retention prunes relative to the newest
|
||||
// hour seen so far, so replaying in chronological order keeps exactly the
|
||||
// intended trailing window.
|
||||
chain := s.auditChain()
|
||||
files := make([]string, 0, len(chain))
|
||||
for i := len(chain) - 1; i >= 0; i-- {
|
||||
files = append(files, chain[i])
|
||||
}
|
||||
|
||||
ring := newReqRing(s.ringSize())
|
||||
scanned := 0
|
||||
s.mu.Lock()
|
||||
for _, p := range files {
|
||||
n, err := scanAuditFile(p, func(r Req) {
|
||||
s.aggregateLocked(r)
|
||||
ring.push(r)
|
||||
})
|
||||
scanned += n
|
||||
if err != nil {
|
||||
// A truncated or unreadable tail is not fatal: keep whatever was
|
||||
// aggregated and mark the numbers as incomplete.
|
||||
s.replayPartial = true
|
||||
}
|
||||
}
|
||||
s.recs = ring.slice()
|
||||
s.mu.Unlock()
|
||||
|
||||
if scanned > 0 {
|
||||
log.Printf("[stats] replayed %d audit records from %d file(s) for aggregates; keeping the newest %d in memory",
|
||||
scanned, len(files), len(s.recs))
|
||||
}
|
||||
}
|
||||
|
||||
func (s *Stats) ringSize() int {
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
if s.maxRecs <= 0 {
|
||||
return defaultRingSize
|
||||
}
|
||||
return s.maxRecs
|
||||
}
|
||||
|
||||
// scanAuditFile streams one audit file, invoking fn for every request row, and
|
||||
// returns how many request rows it saw. Access/event rows and malformed lines
|
||||
// are skipped.
|
||||
func scanAuditFile(path string, fn func(Req)) (int, error) {
|
||||
f, err := os.Open(path)
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
defer f.Close()
|
||||
sc := bufio.NewScanner(f)
|
||||
// tolerate long error summaries / oversized junk lines
|
||||
sc.Buffer(make([]byte, 64*1024), 16*1024*1024)
|
||||
n := 0
|
||||
for sc.Scan() {
|
||||
var r Req
|
||||
if json.Unmarshal(sc.Bytes(), &r) != nil || r.Type == "" {
|
||||
continue
|
||||
}
|
||||
fn(r)
|
||||
n++
|
||||
}
|
||||
return n, sc.Err()
|
||||
}
|
||||
|
||||
// reqRing keeps the newest n records seen, in chronological order, without
|
||||
// growing with the number of records pushed through it.
|
||||
type reqRing struct {
|
||||
buf []Req
|
||||
next int
|
||||
full bool
|
||||
limit int
|
||||
}
|
||||
|
||||
func newReqRing(n int) *reqRing {
|
||||
if n <= 0 {
|
||||
n = defaultRingSize
|
||||
}
|
||||
return &reqRing{buf: make([]Req, n), limit: n}
|
||||
}
|
||||
|
||||
func (r *reqRing) push(rec Req) {
|
||||
r.buf[r.next] = rec
|
||||
r.next++
|
||||
if r.next == r.limit {
|
||||
r.next = 0
|
||||
r.full = true
|
||||
}
|
||||
}
|
||||
|
||||
// slice returns the retained records oldest-first.
|
||||
func (r *reqRing) slice() []Req {
|
||||
if !r.full {
|
||||
out := make([]Req, r.next)
|
||||
copy(out, r.buf[:r.next])
|
||||
return out
|
||||
}
|
||||
out := make([]Req, 0, r.limit)
|
||||
out = append(out, r.buf[r.next:]...)
|
||||
out = append(out, r.buf[:r.next]...)
|
||||
return out
|
||||
}
|
||||
|
||||
// ReplayPartial reports whether the resident aggregates are known to be
|
||||
// incomplete (an audit file could not be read in full). Under normal operation
|
||||
// the aggregates cover the entire audit history, so this is false.
|
||||
func (s *Stats) ReplayPartial() bool {
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
return s.replayPartial
|
||||
}
|
||||
|
||||
// Record appends a finished request to the aggregates and ring buffer.
|
||||
@ -183,6 +341,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)
|
||||
@ -227,6 +388,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 {
|
||||
@ -239,6 +406,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.
|
||||
@ -329,6 +675,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
|
||||
@ -380,12 +747,270 @@ type StatsRow struct {
|
||||
Stat
|
||||
}
|
||||
|
||||
// Records returns request records filtered by unix-millisecond time range and key.
|
||||
func (s *Stats) Records(from, to int64, key string) []Req {
|
||||
// AuditRecords returns every request row from the audit file plus its rotated
|
||||
// .old files within the [from,to] unix-millisecond window, optionally for one
|
||||
// masked key id, sorted oldest-first.
|
||||
//
|
||||
// It materializes the whole window, so it is only for tests and small windows.
|
||||
// Production paths must use StreamAuditRecords (O(1) memory, for CSV export) or
|
||||
// AuditPage (bounded pages, for the records view).
|
||||
func (s *Stats) AuditRecords(from, to int64, key string) []Req {
|
||||
var out []Req
|
||||
_ = s.StreamAuditRecords(from, to, key, func(r Req) error {
|
||||
out = append(out, r)
|
||||
return nil
|
||||
})
|
||||
sort.Slice(out, func(i, j int) bool { return out[i].Time < out[j].Time })
|
||||
return out
|
||||
}
|
||||
|
||||
// ---- on-demand audit paging ----
|
||||
//
|
||||
// The records view is paged straight off disk instead of being held in memory:
|
||||
// the newest page is read from the tail of the newest audit file backwards, and
|
||||
// each response hands back a cursor for the next page. Nothing is cached between
|
||||
// requests, so "release when the user leaves the page" is guaranteed by never
|
||||
// retaining anything in the first place.
|
||||
|
||||
// auditCursor points at a byte offset inside one audit file. Paging walks
|
||||
// backwards, so the cursor means "continue reading BEFORE this offset".
|
||||
type auditCursor struct {
|
||||
File string
|
||||
Offset int64
|
||||
}
|
||||
|
||||
// String encodes the cursor for the wire as "<file>:<offset>".
|
||||
func (c auditCursor) String() string {
|
||||
if c.File == "" {
|
||||
return ""
|
||||
}
|
||||
return c.File + ":" + strconv.FormatInt(c.Offset, 10)
|
||||
}
|
||||
|
||||
// parseAuditCursor decodes a wire cursor. An empty or malformed value means
|
||||
// "start at the newest record".
|
||||
func parseAuditCursor(s string) (auditCursor, bool) {
|
||||
if s == "" {
|
||||
return auditCursor{}, false
|
||||
}
|
||||
i := strings.LastIndex(s, ":")
|
||||
if i <= 0 || i == len(s)-1 {
|
||||
return auditCursor{}, false
|
||||
}
|
||||
off, err := strconv.ParseInt(s[i+1:], 10, 64)
|
||||
if err != nil || off < 0 {
|
||||
return auditCursor{}, false
|
||||
}
|
||||
return auditCursor{File: s[:i], Offset: off}, true
|
||||
}
|
||||
|
||||
// auditChain lists the audit files newest-first: the live file, then the rotated
|
||||
// ones in descending timestamp order.
|
||||
func (s *Stats) auditChain() []string {
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
out := make([]Req, 0, len(s.recs))
|
||||
for _, r := range s.recs {
|
||||
path := s.auditPath
|
||||
s.mu.Unlock()
|
||||
if path == "" {
|
||||
return nil
|
||||
}
|
||||
files := []string{path}
|
||||
if olds, err := filepath.Glob(path + ".*.old"); err == nil {
|
||||
sort.Sort(sort.Reverse(sort.StringSlice(olds)))
|
||||
files = append(files, olds...)
|
||||
}
|
||||
return files
|
||||
}
|
||||
|
||||
// AuditPageResult is one page of records plus the cursor to continue from.
|
||||
type AuditPageResult struct {
|
||||
Records []Req `json:"records"`
|
||||
Next string `json:"next_cursor,omitempty"`
|
||||
HasMore bool `json:"has_more"`
|
||||
// Rotated marks that the requested cursor's file no longer exists (the
|
||||
// audit log rotated under the reader), so the client should restart from
|
||||
// the newest page instead of silently skipping records.
|
||||
Rotated bool `json:"rotated,omitempty"`
|
||||
}
|
||||
|
||||
// revChunk is how much is read per backwards seek. Lines are far shorter than
|
||||
// this, so a page is typically satisfied by one or two chunks.
|
||||
const revChunk = 64 * 1024
|
||||
|
||||
// AuditPage returns up to limit request records ending at cursor, newest first.
|
||||
// It reads the audit files backwards and stops as soon as the page is full, so
|
||||
// cost and memory are proportional to limit rather than to the file size. The
|
||||
// file handle is closed before returning: no state is kept between calls.
|
||||
func (s *Stats) AuditPage(cursor string, limit int, key string) AuditPageResult {
|
||||
if limit <= 0 || limit > 1000 {
|
||||
limit = 100
|
||||
}
|
||||
chain := s.auditChain()
|
||||
if len(chain) == 0 {
|
||||
return AuditPageResult{Records: []Req{}}
|
||||
}
|
||||
|
||||
start := 0
|
||||
var offset int64 = -1 // -1 = start at EOF
|
||||
rotated := false
|
||||
if c, ok := parseAuditCursor(cursor); ok {
|
||||
idx := -1
|
||||
for i, f := range chain {
|
||||
if f == c.File {
|
||||
idx = i
|
||||
break
|
||||
}
|
||||
}
|
||||
if idx < 0 {
|
||||
// the cursor's file rotated away: restart from the newest page and
|
||||
// tell the client so it can reset its view
|
||||
rotated = true
|
||||
} else {
|
||||
start, offset = idx, c.Offset
|
||||
}
|
||||
}
|
||||
|
||||
out := make([]Req, 0, limit)
|
||||
for i := start; i < len(chain); i++ {
|
||||
next, err := s.pageFile(chain[i], offset, limit, key, &out)
|
||||
offset = -1 // subsequent files always start at their EOF
|
||||
if err != nil {
|
||||
continue
|
||||
}
|
||||
if len(out) >= limit {
|
||||
res := AuditPageResult{Records: out, Rotated: rotated}
|
||||
if next > 0 {
|
||||
res.Next = auditCursor{File: chain[i], Offset: next}.String()
|
||||
res.HasMore = true
|
||||
} else if i+1 < len(chain) {
|
||||
res.Next = auditCursor{File: chain[i+1], Offset: fileSize(chain[i+1])}.String()
|
||||
res.HasMore = true
|
||||
}
|
||||
return res
|
||||
}
|
||||
}
|
||||
return AuditPageResult{Records: out, HasMore: false, Rotated: rotated}
|
||||
}
|
||||
|
||||
func fileSize(path string) int64 {
|
||||
fi, err := os.Stat(path)
|
||||
if err != nil {
|
||||
return 0
|
||||
}
|
||||
return fi.Size()
|
||||
}
|
||||
|
||||
// pageFile reads one audit file backwards from end (or EOF when end < 0),
|
||||
// appending matching records to out (newest first) until it holds limit entries.
|
||||
// It returns the offset to continue from within this file, or 0 when the file is
|
||||
// exhausted.
|
||||
func (s *Stats) pageFile(path string, end int64, limit int, key string, out *[]Req) (int64, error) {
|
||||
f, err := os.Open(path)
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
defer f.Close() // released before the response is written: nothing is retained
|
||||
fi, err := f.Stat()
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
if end < 0 || end > fi.Size() {
|
||||
end = fi.Size()
|
||||
}
|
||||
|
||||
buf := make([]byte, 0, revChunk)
|
||||
var tail []byte // bytes of a line whose start was not in this chunk
|
||||
pos := end
|
||||
for pos > 0 && len(*out) < limit {
|
||||
size := int64(revChunk)
|
||||
if size > pos {
|
||||
size = pos
|
||||
}
|
||||
pos -= size
|
||||
buf = buf[:size]
|
||||
if _, err := f.ReadAt(buf, pos); err != nil && err != io.EOF {
|
||||
return 0, err
|
||||
}
|
||||
chunk := buf
|
||||
if len(tail) > 0 {
|
||||
chunk = append(append(make([]byte, 0, len(buf)+len(tail)), buf...), tail...)
|
||||
tail = tail[:0]
|
||||
}
|
||||
// walk the chunk's complete lines from the end towards the start
|
||||
lineEnd := len(chunk)
|
||||
for i := len(chunk) - 1; i >= 0; i-- {
|
||||
if chunk[i] != '\n' {
|
||||
continue
|
||||
}
|
||||
if rec, ok := decodeAuditLine(chunk[i+1:lineEnd], key); ok {
|
||||
*out = append(*out, rec)
|
||||
if len(*out) >= limit {
|
||||
// resume at the newline we just consumed
|
||||
return pos + int64(i) + 1, nil
|
||||
}
|
||||
}
|
||||
lineEnd = i
|
||||
}
|
||||
// chunk[:lineEnd] is an incomplete line: carry it into the next chunk
|
||||
if lineEnd > 0 {
|
||||
tail = append(tail[:0], chunk[:lineEnd]...)
|
||||
}
|
||||
if pos == 0 {
|
||||
// start of file: whatever is carried is a complete first line
|
||||
if rec, ok := decodeAuditLine(tail, key); ok {
|
||||
*out = append(*out, rec)
|
||||
}
|
||||
}
|
||||
}
|
||||
if pos <= 0 {
|
||||
return 0, nil // file exhausted
|
||||
}
|
||||
return pos, nil
|
||||
}
|
||||
|
||||
// decodeAuditLine parses one audit line as a request record, filtering by key.
|
||||
// Access/event rows and malformed lines are rejected.
|
||||
func decodeAuditLine(line []byte, key string) (Req, bool) {
|
||||
line = bytes.TrimSpace(line)
|
||||
if len(line) == 0 {
|
||||
return Req{}, false
|
||||
}
|
||||
var r Req
|
||||
if json.Unmarshal(line, &r) != nil || r.Type == "" {
|
||||
return Req{}, false
|
||||
}
|
||||
if key != "" && r.Key != key {
|
||||
return Req{}, false
|
||||
}
|
||||
return r, true
|
||||
}
|
||||
|
||||
// StreamAuditRecords walks every request record in [from, to] (unix millis,
|
||||
// 0 = unbounded) newest file first and hands each one to fn. Nothing is
|
||||
// accumulated: an export of an arbitrarily long period costs O(1) memory, and
|
||||
// the file handles are closed as the walk proceeds. fn returning an error stops
|
||||
// the walk (used to abort on a broken client connection).
|
||||
func (s *Stats) StreamAuditRecords(from, to int64, key string, fn func(Req) error) error {
|
||||
for _, path := range s.auditChain() {
|
||||
if err := streamAuditFile(path, from, to, key, fn); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func streamAuditFile(path string, from, to int64, key string, fn func(Req) error) error {
|
||||
f, err := os.Open(path)
|
||||
if err != nil {
|
||||
return nil // a rotated-away file is not an export failure
|
||||
}
|
||||
defer f.Close()
|
||||
sc := bufio.NewScanner(f)
|
||||
sc.Buffer(make([]byte, 64*1024), 16*1024*1024)
|
||||
for sc.Scan() {
|
||||
var r Req
|
||||
if json.Unmarshal(sc.Bytes(), &r) != nil || r.Type == "" {
|
||||
continue
|
||||
}
|
||||
if key != "" && r.Key != key {
|
||||
continue
|
||||
}
|
||||
@ -395,54 +1020,11 @@ func (s *Stats) Records(from, to int64, key string) []Req {
|
||||
if to > 0 && r.Time > to {
|
||||
continue
|
||||
}
|
||||
out = append(out, r)
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
// AuditRecords returns every request row from the audit file plus its rotated
|
||||
// .old files within the [from,to] unix-millisecond window, optionally for one
|
||||
// masked key id. Unlike Records it is not bounded by the in-memory ring, so it
|
||||
// yields a complete history for CSV export regardless of restarts or churn.
|
||||
func (s *Stats) AuditRecords(from, to int64, key string) []Req {
|
||||
s.mu.Lock()
|
||||
path := s.auditPath
|
||||
s.mu.Unlock()
|
||||
if path == "" {
|
||||
return nil
|
||||
}
|
||||
files := []string{path}
|
||||
if olds, err := filepath.Glob(path + ".*.old"); err == nil {
|
||||
files = append(files, olds...)
|
||||
}
|
||||
var out []Req
|
||||
for _, p := range files {
|
||||
f, err := os.Open(p)
|
||||
if err != nil {
|
||||
continue
|
||||
if err := fn(r); err != nil {
|
||||
return err
|
||||
}
|
||||
sc := bufio.NewScanner(f)
|
||||
sc.Buffer(make([]byte, 64*1024), 16*1024*1024)
|
||||
for sc.Scan() {
|
||||
var r Req
|
||||
if json.Unmarshal(sc.Bytes(), &r) != nil || r.Type == "" {
|
||||
continue
|
||||
}
|
||||
if key != "" && r.Key != key {
|
||||
continue
|
||||
}
|
||||
if from > 0 && r.Time < from {
|
||||
continue
|
||||
}
|
||||
if to > 0 && r.Time > to {
|
||||
continue
|
||||
}
|
||||
out = append(out, r)
|
||||
}
|
||||
f.Close()
|
||||
}
|
||||
sort.Slice(out, func(i, j int) bool { return out[i].Time < out[j].Time })
|
||||
return out
|
||||
return nil
|
||||
}
|
||||
|
||||
// SourceRecent counts real gateway requests per source within the last
|
||||
@ -472,12 +1054,79 @@ func (s *Stats) SourceRecent(windowSec int64) map[string][2]int64 {
|
||||
return out
|
||||
}
|
||||
|
||||
// SourceAvg carries per-source performance averages for the status page.
|
||||
type SourceAvg struct {
|
||||
// AvgFirstByteMs is the mean time-to-first-byte over successful
|
||||
// requests in the window (0 when no measured samples).
|
||||
AvgFirstByteMs int64 `json:"avg_first_byte_ms"`
|
||||
// AvgTokPerS is the aggregate completion throughput:
|
||||
// sum(completion_tokens) / sum(latency_seconds) (0 when no samples).
|
||||
AvgTokPerS int64 `json:"avg_tok_per_s"`
|
||||
// Samples is the number of successful requests the averages cover.
|
||||
Samples int64 `json:"samples"`
|
||||
}
|
||||
|
||||
// SourceAverages computes TTFB and tokens/s averages per source from the
|
||||
// in-memory ring within the window (unix seconds). Only successful chat/
|
||||
// stream rows count; image and failed rows are skipped.
|
||||
func (s *Stats) SourceAverages(windowSec int64) map[string]SourceAvg {
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
cut := time.Now().Unix() - windowSec
|
||||
type acc struct {
|
||||
fbSum, latSum, complSum, n int64
|
||||
}
|
||||
accs := map[string]*acc{}
|
||||
for _, r := range s.recs {
|
||||
if !r.OK || r.Source == "" || r.Time/1000 < cut {
|
||||
continue
|
||||
}
|
||||
if r.Type != "chat" && r.Type != "stream" {
|
||||
continue
|
||||
}
|
||||
if r.LatMs <= 0 {
|
||||
continue
|
||||
}
|
||||
a := accs[r.Source]
|
||||
if a == nil {
|
||||
a = &acc{}
|
||||
accs[r.Source] = a
|
||||
}
|
||||
a.latSum += r.LatMs
|
||||
a.complSum += r.Compl
|
||||
a.fbSum += r.FirstByteMs
|
||||
if r.FirstByteMs > 0 {
|
||||
a.n++
|
||||
}
|
||||
}
|
||||
out := make(map[string]SourceAvg, len(accs))
|
||||
for src, a := range accs {
|
||||
avg := SourceAvg{Samples: a.n}
|
||||
if a.n > 0 {
|
||||
avg.AvgFirstByteMs = a.fbSum / a.n
|
||||
}
|
||||
if a.latSum > 0 && a.complSum > 0 {
|
||||
avg.AvgTokPerS = a.complSum * 1000 / a.latSum
|
||||
}
|
||||
out[src] = avg
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
// Snapshot returns the whole dashboard payload; when key != "" the records
|
||||
// and aggregate views are restricted to that gateway key.
|
||||
// Snapshot returns the aggregate rows plus the FIRST SCREEN of request records
|
||||
// (newest last, matching the ring order). limit bounds the record slice only;
|
||||
// older records are not included here at all — the dashboard pages them from
|
||||
// /api/stats/records as the user scrolls, so a dashboard load never serializes
|
||||
// the whole ring. limit <= 0 means the default first screen.
|
||||
func (s *Stats) Snapshot(limit int, key string) map[string]interface{} {
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
if limit <= 0 {
|
||||
limit = firstScreenRecords
|
||||
}
|
||||
if limit > s.maxRecs {
|
||||
limit = s.maxRecs
|
||||
}
|
||||
start := 0
|
||||
@ -525,12 +1174,13 @@ func (s *Stats) Snapshot(limit int, key string) map[string]interface{} {
|
||||
return ci < cj
|
||||
})
|
||||
return map[string]interface{}{
|
||||
"active": s.active,
|
||||
"total": total,
|
||||
"by_key": rows(byKey),
|
||||
"by_model": rows(byModel),
|
||||
"by_source": rows(bySrc),
|
||||
"by_status": bs,
|
||||
"records": append([]Req(nil), recs...),
|
||||
"active": s.active,
|
||||
"total": total,
|
||||
"by_key": rows(byKey),
|
||||
"by_model": rows(byModel),
|
||||
"by_source": rows(bySrc),
|
||||
"by_status": bs,
|
||||
"records": append([]Req(nil), recs...),
|
||||
"replay_partial": s.replayPartial,
|
||||
}
|
||||
}
|
||||
|
||||
@ -92,21 +92,25 @@ func TestAuditRotationRecords(t *testing.T) {
|
||||
func TestLoadAuditFullReplay(t *testing.T) {
|
||||
dir := t.TempDir()
|
||||
path := filepath.Join(dir, "audit.jsonl")
|
||||
// timestamps relative to now: buckets are whole unix hours, so window
|
||||
// assertions must not depend on the minute-of-hour of the test run.
|
||||
// timestamps anchored to whole unix hours (see aggregateLocked): window
|
||||
// assertions must hold no matter which minute-of-hour the suite runs at,
|
||||
// so rows are placed relative to hour boundaries, not raw offsets from now.
|
||||
now := time.Now()
|
||||
thisHour := now.Truncate(time.Hour)
|
||||
twoBucketsAgo := thisHour.Add(-90 * time.Minute) // bucket H-2
|
||||
prevBucketMid := thisHour.Add(-30 * time.Minute) // bucket H-1
|
||||
lines := []string{
|
||||
`{"obj":"access","time":1699999999000,"key":"k","method":"GET","path":"/api/stats","status":200}`,
|
||||
fmt.Sprintf(`{"time":%d,"key":"k","type":"chat","model":"m","source":"s","prompt_tokens":100,"completion_tokens":50,"latency_ms":10,"ok":true,"status":200}`, now.Add(-65*time.Minute).UnixMilli()),
|
||||
fmt.Sprintf(`{"time":%d,"key":"k","type":"chat","model":"m","source":"s","prompt_tokens":100,"completion_tokens":50,"latency_ms":10,"ok":true,"status":200}`, twoBucketsAgo.UnixMilli()),
|
||||
`{this is not valid json`,
|
||||
fmt.Sprintf(`{"time":%d,"key":"k","type":"stream","model":"m","source":"s","prompt_tokens":200,"completion_tokens":20,"latency_ms":20,"ok":false,"status":503}`, now.Add(-30*time.Minute).UnixMilli()),
|
||||
fmt.Sprintf(`{"time":%d,"key":"k","type":"stream","model":"m","source":"s","prompt_tokens":200,"completion_tokens":20,"latency_ms":20,"ok":false,"status":503}`, prevBucketMid.UnixMilli()),
|
||||
"garbage-not-json\n",
|
||||
fmt.Sprintf(`{"time":%d,"key":"k","type":"chat","model":"m2","source":"s","prompt_tokens":7,"completion_tokens":3,"latency_ms":5,"ok":true,"status":200}`, now.Add(-30*time.Minute).UnixMilli()),
|
||||
fmt.Sprintf(`{"time":%d,"key":"k","type":"chat","model":"m2","source":"s","prompt_tokens":7,"completion_tokens":3,"latency_ms":5,"ok":true,"status":200}`, prevBucketMid.UnixMilli()),
|
||||
}
|
||||
loaded := strings.Join(lines, "\n") + "\n" + strings.Repeat("x", 1<<18) + "\n"
|
||||
// oversized row at the END proves the scanner tolerates >64KB lines and
|
||||
// still finishes the replay instead of truncating silently.
|
||||
loaded += fmt.Sprintf(`{"time":%d,"key":"k","type":"chat","model":"m","source":"s","prompt_tokens":1,"completion_tokens":1,"latency_ms":1,"ok":true,"status":200}`, now.Add(-time.Minute).UnixMilli()) + "\n"
|
||||
loaded += fmt.Sprintf(`{"time":%d,"key":"k","type":"chat","model":"m","source":"s","prompt_tokens":1,"completion_tokens":1,"latency_ms":1,"ok":true,"status":200}`, now.UnixMilli()) + "\n"
|
||||
if err := os.WriteFile(path, []byte(loaded), 0644); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
@ -126,9 +130,10 @@ func TestLoadAuditFullReplay(t *testing.T) {
|
||||
t.Fatalf("ring must hold only real requests, got %d rows: %#v", len(s.recs), s.recs)
|
||||
}
|
||||
// quota window rebuilt from full history. All-time and multi-hour windows
|
||||
// must see everything; a 1h window must NOT return all records (the old
|
||||
// ms/seconds unit bug made any sec>0 window return everything), and must
|
||||
// always include the row written one minute ago (current hour bucket).
|
||||
// must see everything. A 1h window must NOT return all records (the old
|
||||
// ms/seconds unit bug made any sec>0 window return everything): buckets
|
||||
// are whole unix hours, so only rows in the current hour bucket qualify
|
||||
// — exactly the just-now row (2 tokens).
|
||||
if w := s.WindowTokens("m", "s", 0); w != 372 {
|
||||
t.Fatalf("all-time window want 372, got %d", w)
|
||||
}
|
||||
@ -138,8 +143,8 @@ func TestLoadAuditFullReplay(t *testing.T) {
|
||||
if w := s.WindowTokens("m", "s", 24*hourSec); w != 372 {
|
||||
t.Fatalf("24h window want 372, got %d", w)
|
||||
}
|
||||
if w := s.WindowTokens("m", "s", hourSec); w < 2 || w >= 372 {
|
||||
t.Fatalf("1h window want [2,372), got %d", w)
|
||||
if w := s.WindowTokens("m", "s", hourSec); w != 2 {
|
||||
t.Fatalf("1h window want 2 (current-bucket row only), got %d", w)
|
||||
}
|
||||
if w := s.WindowTokens("m2", "s", 24*hourSec); w != 10 {
|
||||
t.Fatalf("m2 24h window want 10, got %d", w)
|
||||
@ -151,4 +156,407 @@ func TestLoadAuditFullReplay(t *testing.T) {
|
||||
if s.byStatus[200].Err != 0 || s.byStatus[503] == nil || s.byStatus[503].Reqs != 1 {
|
||||
t.Fatalf("by_status wrong: %#v", s.byStatus)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// ---- on-demand log loading (plan 阶段 3) ----
|
||||
|
||||
// writeAuditFile writes n synthetic request rows (plus some non-request event
|
||||
// rows, which must be skipped) and returns the path.
|
||||
func writeAuditFile(t *testing.T, path string, n int, key string, startMs int64) {
|
||||
t.Helper()
|
||||
f, err := os.Create(path)
|
||||
if err != nil {
|
||||
t.Fatalf("create audit: %v", err)
|
||||
}
|
||||
defer f.Close()
|
||||
for i := 0; i < n; i++ {
|
||||
fmt.Fprintf(f, `{"time":%d,"key":%q,"type":"chat","model":"m","source":"s","prompt_tokens":1,"completion_tokens":1,"latency_ms":5,"ok":true,"status":200}`+"\n",
|
||||
startMs+int64(i)*1000, key)
|
||||
if i%10 == 0 {
|
||||
// an access-log/event row: must never appear as a record
|
||||
fmt.Fprintf(f, `{"obj":"access","time":%d,"path":"/v1/models"}`+"\n", startMs+int64(i)*1000)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// TestLoadAuditAggregatesFullHistory is the correctness half of the log work:
|
||||
// the dashboard's totals must cover EVERY audit record, however large the file,
|
||||
// while the in-memory record ring stays bounded.
|
||||
func TestLoadAuditAggregatesFullHistory(t *testing.T) {
|
||||
dir := t.TempDir()
|
||||
path := filepath.Join(dir, "audit.jsonl")
|
||||
const n = 20000
|
||||
writeAuditFile(t, path, n, "k1", 1_700_000_000_000)
|
||||
|
||||
s := NewStats(0)
|
||||
s.LoadAudit(path)
|
||||
if s.ReplayPartial() {
|
||||
t.Fatal("a readable history must not be reported as partial")
|
||||
}
|
||||
|
||||
snap := s.Snapshot(0, "")
|
||||
total, _ := snap["total"].(Stat)
|
||||
if total.Reqs != n {
|
||||
t.Fatalf("total reqs = %d, want %d (aggregates must cover the whole file)", total.Reqs, n)
|
||||
}
|
||||
if total.Tokens != int64(n)*2 { // 1 prompt + 1 completion per row
|
||||
t.Fatalf("total tokens = %d, want %d", total.Tokens, n*2)
|
||||
}
|
||||
|
||||
// ...while the resident ring is capped and the snapshot ships one screen
|
||||
recs, _ := snap["records"].([]Req)
|
||||
if len(recs) == 0 {
|
||||
t.Fatal("the newest records must stay resident")
|
||||
}
|
||||
if len(recs) > firstScreenRecords {
|
||||
t.Fatalf("snapshot shipped %d records, must be capped at the first screen (%d)",
|
||||
len(recs), firstScreenRecords)
|
||||
}
|
||||
s.mu.Lock()
|
||||
ring := len(s.recs)
|
||||
s.mu.Unlock()
|
||||
if ring > defaultRingSize {
|
||||
t.Fatalf("resident ring holds %d records, must be capped at %d", ring, defaultRingSize)
|
||||
}
|
||||
|
||||
// the ring must hold the NEWEST records, not the first ones scanned
|
||||
newest := recs[len(recs)-1]
|
||||
wantNewest := int64(1_700_000_000_000) + int64(n-1)*1000
|
||||
if newest.Time != wantNewest {
|
||||
t.Fatalf("newest resident record time = %d, want %d", newest.Time, wantNewest)
|
||||
}
|
||||
}
|
||||
|
||||
// TestLoadAuditAggregatesAcrossRotatedFiles: history split over rotated .old
|
||||
// files must still add up, and must be replayed oldest-first so the hourly
|
||||
// quota buckets keep the intended trailing window.
|
||||
func TestLoadAuditAggregatesAcrossRotatedFiles(t *testing.T) {
|
||||
dir := t.TempDir()
|
||||
path := filepath.Join(dir, "audit.jsonl")
|
||||
writeAuditFile(t, path, 40, "k1", 1_700_000_500_000) // newest
|
||||
writeAuditFile(t, path+".100.old", 40, "k1", 1_700_000_300_000) // older
|
||||
writeAuditFile(t, path+".050.old", 40, "k1", 1_700_000_100_000) // oldest
|
||||
|
||||
s := NewStats(0)
|
||||
s.LoadAudit(path)
|
||||
snap := s.Snapshot(0, "")
|
||||
total, _ := snap["total"].(Stat)
|
||||
if total.Reqs != 120 {
|
||||
t.Fatalf("total reqs = %d, want 120 across three files", total.Reqs)
|
||||
}
|
||||
recs, _ := snap["records"].([]Req)
|
||||
if len(recs) == 0 {
|
||||
t.Fatal("expected resident records")
|
||||
}
|
||||
// oldest-first replay + newest-last ring => the last resident record is the
|
||||
// newest one in the live file
|
||||
if got := recs[len(recs)-1].Time; got != 1_700_000_500_000+39*1000 {
|
||||
t.Fatalf("newest resident record time = %d, want the live file's last row", got)
|
||||
}
|
||||
}
|
||||
|
||||
// TestReqRingKeepsNewest pins the ring semantics used by the full-history scan:
|
||||
// bounded size, newest retained, oldest-first output.
|
||||
func TestReqRingKeepsNewest(t *testing.T) {
|
||||
r := newReqRing(3)
|
||||
for i := 1; i <= 7; i++ {
|
||||
r.push(Req{Time: int64(i)})
|
||||
}
|
||||
got := r.slice()
|
||||
if len(got) != 3 {
|
||||
t.Fatalf("ring len = %d, want 3", len(got))
|
||||
}
|
||||
for i, want := range []int64{5, 6, 7} {
|
||||
if got[i].Time != want {
|
||||
t.Fatalf("ring[%d] = %d, want %d (oldest-first, newest retained)", i, got[i].Time, want)
|
||||
}
|
||||
}
|
||||
|
||||
// a partially filled ring returns exactly what was pushed
|
||||
r2 := newReqRing(5)
|
||||
r2.push(Req{Time: 1})
|
||||
r2.push(Req{Time: 2})
|
||||
if got := r2.slice(); len(got) != 2 || got[0].Time != 1 || got[1].Time != 2 {
|
||||
t.Fatalf("partial ring = %+v", got)
|
||||
}
|
||||
}
|
||||
|
||||
// TestLoadAuditFullFileWhenSmall: a small file is loaded completely and is not
|
||||
// flagged partial.
|
||||
func TestLoadAuditFullFileWhenSmall(t *testing.T) {
|
||||
path := filepath.Join(t.TempDir(), "audit.jsonl")
|
||||
writeAuditFile(t, path, 50, "k1", 1_700_000_000_000)
|
||||
s := NewStats(0)
|
||||
s.LoadAudit(path)
|
||||
if s.ReplayPartial() {
|
||||
t.Fatal("a small file must be replayed in full, not flagged partial")
|
||||
}
|
||||
snap := s.Snapshot(0, "")
|
||||
total, _ := snap["total"].(Stat)
|
||||
if total.Reqs != 50 {
|
||||
t.Fatalf("total reqs = %d, want 50", total.Reqs)
|
||||
}
|
||||
}
|
||||
|
||||
// TestQuotaWindowsSurviveBoundedRing is the regression guard the plan calls out:
|
||||
// capping the resident record ring must not break token-quota accounting, which
|
||||
// reads the hour buckets rather than the ring.
|
||||
func TestQuotaWindowsSurviveBoundedRing(t *testing.T) {
|
||||
s := NewStats(10) // deliberately tiny ring
|
||||
now := time.Now().UnixMilli()
|
||||
for i := 0; i < 100; i++ {
|
||||
s.Record(Req{Time: now, Key: "k", Model: "m", Source: "s", Prompt: 2, Compl: 3, OK: true, Status: 200})
|
||||
}
|
||||
if got := s.WindowTokens("m", "s", 24*hourSec); got != 500 {
|
||||
t.Fatalf("window tokens = %d, want 500 (quota accounting must not depend on the ring)", got)
|
||||
}
|
||||
if got := s.ModelTokens("k")["m"]; got != 500 {
|
||||
t.Fatalf("model tokens = %d, want 500", got)
|
||||
}
|
||||
if got := s.KeyTokens("k"); got != 500 {
|
||||
t.Fatalf("key tokens = %d, want 500", got)
|
||||
}
|
||||
}
|
||||
|
||||
// TestAuditPagePagesBackwards walks the whole history one page at a time and
|
||||
// asserts strict newest-first order with no gaps or repeats.
|
||||
func TestAuditPagePagesBackwards(t *testing.T) {
|
||||
path := filepath.Join(t.TempDir(), "audit.jsonl")
|
||||
const n = 250
|
||||
writeAuditFile(t, path, n, "k1", 1_700_000_000_000)
|
||||
s := NewStats(0)
|
||||
s.LoadAudit(path)
|
||||
|
||||
var seen []Req
|
||||
cursor := ""
|
||||
for i := 0; i < 20; i++ { // bounded loop guards against a cursor that never advances
|
||||
page := s.AuditPage(cursor, 40, "")
|
||||
if len(page.Records) == 0 {
|
||||
break
|
||||
}
|
||||
seen = append(seen, page.Records...)
|
||||
if !page.HasMore {
|
||||
break
|
||||
}
|
||||
if page.Next == cursor {
|
||||
t.Fatal("cursor did not advance")
|
||||
}
|
||||
cursor = page.Next
|
||||
}
|
||||
if len(seen) != n {
|
||||
t.Fatalf("paged %d records, want %d", len(seen), n)
|
||||
}
|
||||
for i := 1; i < len(seen); i++ {
|
||||
if seen[i-1].Time < seen[i].Time {
|
||||
t.Fatalf("records not newest-first at %d: %d then %d", i, seen[i-1].Time, seen[i].Time)
|
||||
}
|
||||
}
|
||||
// newest record first, oldest last
|
||||
if seen[0].Time != 1_700_000_000_000+int64(n-1)*1000 {
|
||||
t.Fatalf("first record time = %d, want the newest", seen[0].Time)
|
||||
}
|
||||
if seen[len(seen)-1].Time != 1_700_000_000_000 {
|
||||
t.Fatalf("last record time = %d, want the oldest", seen[len(seen)-1].Time)
|
||||
}
|
||||
}
|
||||
|
||||
// TestAuditPageFirstScreenIsSmall: the first page must be exactly one screen,
|
||||
// independent of how large the audit file is.
|
||||
func TestAuditPageFirstScreenIsSmall(t *testing.T) {
|
||||
path := filepath.Join(t.TempDir(), "audit.jsonl")
|
||||
writeAuditFile(t, path, 5000, "k1", 1_700_000_000_000)
|
||||
s := NewStats(0)
|
||||
s.LoadAudit(path)
|
||||
page := s.AuditPage("", 100, "")
|
||||
if len(page.Records) != 100 {
|
||||
t.Fatalf("first page = %d records, want 100", len(page.Records))
|
||||
}
|
||||
if !page.HasMore || page.Next == "" {
|
||||
t.Fatal("a large file must report more pages plus a cursor")
|
||||
}
|
||||
}
|
||||
|
||||
// TestAuditPageFiltersByKey: a user key may only page its own records.
|
||||
func TestAuditPageFiltersByKey(t *testing.T) {
|
||||
dir := t.TempDir()
|
||||
path := filepath.Join(dir, "audit.jsonl")
|
||||
writeAuditFile(t, path, 30, "mine", 1_700_000_000_000)
|
||||
f, err := os.OpenFile(path, os.O_APPEND|os.O_WRONLY, 0644)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
for i := 0; i < 30; i++ {
|
||||
fmt.Fprintf(f, `{"time":%d,"key":"other","type":"chat","model":"m","source":"s","ok":true,"status":200}`+"\n",
|
||||
1_700_000_100_000+int64(i)*1000)
|
||||
}
|
||||
f.Close()
|
||||
|
||||
s := NewStats(0)
|
||||
s.LoadAudit(path)
|
||||
page := s.AuditPage("", 1000, "mine")
|
||||
if len(page.Records) != 30 {
|
||||
t.Fatalf("filtered page = %d records, want 30", len(page.Records))
|
||||
}
|
||||
for _, r := range page.Records {
|
||||
if r.Key != "mine" {
|
||||
t.Fatalf("key filter leaked record for %q", r.Key)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// TestAuditPageSpansRotatedFiles: paging must continue into .old files instead
|
||||
// of stopping at the live file's start.
|
||||
func TestAuditPageSpansRotatedFiles(t *testing.T) {
|
||||
dir := t.TempDir()
|
||||
path := filepath.Join(dir, "audit.jsonl")
|
||||
writeAuditFile(t, path, 40, "k1", 1_700_000_500_000) // newest
|
||||
writeAuditFile(t, path+".100.old", 40, "k1", 1_700_000_300_000) // older
|
||||
writeAuditFile(t, path+".050.old", 40, "k1", 1_700_000_100_000) // oldest
|
||||
|
||||
s := NewStats(0)
|
||||
s.LoadAudit(path)
|
||||
var seen []Req
|
||||
cursor := ""
|
||||
for i := 0; i < 30; i++ {
|
||||
page := s.AuditPage(cursor, 25, "")
|
||||
seen = append(seen, page.Records...)
|
||||
if !page.HasMore {
|
||||
break
|
||||
}
|
||||
cursor = page.Next
|
||||
}
|
||||
if len(seen) != 120 {
|
||||
t.Fatalf("paged %d records across rotated files, want 120", len(seen))
|
||||
}
|
||||
for i := 1; i < len(seen); i++ {
|
||||
if seen[i-1].Time < seen[i].Time {
|
||||
t.Fatalf("rotation boundary broke the ordering at %d", i)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// TestAuditPageRotatedCursor: a cursor whose file rotated away must restart from
|
||||
// the newest page and say so, rather than silently returning nothing.
|
||||
func TestAuditPageRotatedCursor(t *testing.T) {
|
||||
path := filepath.Join(t.TempDir(), "audit.jsonl")
|
||||
writeAuditFile(t, path, 20, "k1", 1_700_000_000_000)
|
||||
s := NewStats(0)
|
||||
s.LoadAudit(path)
|
||||
page := s.AuditPage("/gone/audit.jsonl:12345", 10, "")
|
||||
if !page.Rotated {
|
||||
t.Fatal("a stale cursor must be reported as rotated")
|
||||
}
|
||||
if len(page.Records) == 0 {
|
||||
t.Fatal("a stale cursor must fall back to the newest page")
|
||||
}
|
||||
}
|
||||
|
||||
// TestAuditPageEmptyAndNoFile cover the degenerate inputs.
|
||||
func TestAuditPageEmptyAndNoFile(t *testing.T) {
|
||||
s := NewStats(0)
|
||||
if page := s.AuditPage("", 10, ""); len(page.Records) != 0 || page.HasMore {
|
||||
t.Fatalf("no audit path must yield an empty page, got %+v", page)
|
||||
}
|
||||
path := filepath.Join(t.TempDir(), "audit.jsonl")
|
||||
if err := os.WriteFile(path, nil, 0644); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
s2 := NewStats(0)
|
||||
s2.LoadAudit(path)
|
||||
if page := s2.AuditPage("", 10, ""); len(page.Records) != 0 || page.HasMore {
|
||||
t.Fatalf("empty audit file must yield an empty page, got %+v", page)
|
||||
}
|
||||
}
|
||||
|
||||
// TestAuditPageHandlesLongLines: a record larger than the reverse-read chunk
|
||||
// must still be decoded (the reader carries partial lines across chunks).
|
||||
func TestAuditPageHandlesLongLines(t *testing.T) {
|
||||
path := filepath.Join(t.TempDir(), "audit.jsonl")
|
||||
f, err := os.Create(path)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
big := strings.Repeat("e", revChunk+1024) // one record spanning >1 chunk
|
||||
fmt.Fprintf(f, `{"time":1,"key":"k","type":"chat","model":"m","source":"s","ok":false,"status":500,"error":%q}`+"\n", big)
|
||||
fmt.Fprintf(f, `{"time":2,"key":"k","type":"chat","model":"m","source":"s","ok":true,"status":200}`+"\n")
|
||||
f.Close()
|
||||
|
||||
s := NewStats(0)
|
||||
s.LoadAudit(path)
|
||||
page := s.AuditPage("", 10, "")
|
||||
if len(page.Records) != 2 {
|
||||
t.Fatalf("got %d records, want 2 (a >chunk record must survive reverse reads)", len(page.Records))
|
||||
}
|
||||
if page.Records[0].Time != 2 || page.Records[1].Time != 1 {
|
||||
t.Fatalf("ordering broken: %d then %d", page.Records[0].Time, page.Records[1].Time)
|
||||
}
|
||||
if len(page.Records[1].Err) != len(big) {
|
||||
t.Fatalf("oversized field truncated: %d bytes, want %d", len(page.Records[1].Err), len(big))
|
||||
}
|
||||
}
|
||||
|
||||
// TestStreamAuditRecordsIsWindowed: export streaming must honour the time window
|
||||
// and the key filter, and must hand records over one at a time.
|
||||
func TestStreamAuditRecordsIsWindowed(t *testing.T) {
|
||||
path := filepath.Join(t.TempDir(), "audit.jsonl")
|
||||
writeAuditFile(t, path, 100, "k1", 1_700_000_000_000)
|
||||
s := NewStats(0)
|
||||
s.LoadAudit(path)
|
||||
|
||||
from := int64(1_700_000_010_000)
|
||||
to := int64(1_700_000_019_000)
|
||||
var got []Req
|
||||
if err := s.StreamAuditRecords(from, to, "", func(r Req) error {
|
||||
got = append(got, r)
|
||||
return nil
|
||||
}); err != nil {
|
||||
t.Fatalf("stream: %v", err)
|
||||
}
|
||||
if len(got) != 10 {
|
||||
t.Fatalf("streamed %d records, want 10 in [%d,%d]", len(got), from, to)
|
||||
}
|
||||
for _, r := range got {
|
||||
if r.Time < from || r.Time > to {
|
||||
t.Fatalf("record outside the window: %d", r.Time)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// TestStreamAuditRecordsAbortStops: a callback error (client disconnect) must
|
||||
// stop the walk immediately instead of reading the rest of the history.
|
||||
func TestStreamAuditRecordsAbortStops(t *testing.T) {
|
||||
path := filepath.Join(t.TempDir(), "audit.jsonl")
|
||||
writeAuditFile(t, path, 500, "k1", 1_700_000_000_000)
|
||||
s := NewStats(0)
|
||||
s.LoadAudit(path)
|
||||
|
||||
seen := 0
|
||||
stop := fmt.Errorf("client gone")
|
||||
err := s.StreamAuditRecords(0, 0, "", func(Req) error {
|
||||
seen++
|
||||
if seen == 5 {
|
||||
return stop
|
||||
}
|
||||
return nil
|
||||
})
|
||||
if err != stop {
|
||||
t.Fatalf("abort must surface the callback error, got %v", err)
|
||||
}
|
||||
if seen != 5 {
|
||||
t.Fatalf("walk continued after abort: %d records", seen)
|
||||
}
|
||||
}
|
||||
|
||||
// TestAuditRotationThresholds pins the smaller-file/more-files trade-off that
|
||||
// makes reverse paging cheap while keeping the same total budget.
|
||||
func TestAuditRotationThresholds(t *testing.T) {
|
||||
if auditRotateBytes != 16<<20 {
|
||||
t.Fatalf("auditRotateBytes = %d, want 16 MiB", auditRotateBytes)
|
||||
}
|
||||
if auditKeepOld != 16 {
|
||||
t.Fatalf("auditKeepOld = %d, want 16", auditKeepOld)
|
||||
}
|
||||
if total := auditRotateBytes * int64(auditKeepOld); total != 256<<20 {
|
||||
t.Fatalf("total audit budget = %d, want 256 MiB", total)
|
||||
}
|
||||
}
|
||||
|
||||
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
406
internal/gateway/ui_contract_test.go
Normal file
406
internal/gateway/ui_contract_test.go
Normal file
@ -0,0 +1,406 @@
|
||||
package gateway
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"regexp"
|
||||
"strings"
|
||||
"testing"
|
||||
)
|
||||
|
||||
// The WebUI talks to this package through a tiny `api(path, options)` helper
|
||||
// that wraps fetch(). fetch() defaults to GET when no method is given, so an
|
||||
// options object without `method` silently turns a DELETE/PUT into a GET —
|
||||
// the request hits the wrong branch of the handler and the user sees a
|
||||
// confusing error like "use GET /api/keys" while nothing is deleted.
|
||||
//
|
||||
// This has now bitten twice: ab20f1b fixed delSource/delTemplate, and the very
|
||||
// same pattern was still live in delKey (deleting a gateway key was impossible)
|
||||
// and delAdapter. Reviewing 4000+ lines of inline JS by eye clearly does not
|
||||
// catch it, so it is pinned here instead.
|
||||
|
||||
// apiCallRe finds the start of every api(...) call in the UI source.
|
||||
var apiCallRe = regexp.MustCompile(`\bapi\(`)
|
||||
|
||||
// uiSource returns the embedded WebUI document.
|
||||
func uiSource(t *testing.T) string {
|
||||
t.Helper()
|
||||
b, err := uiFS.ReadFile("ui/index.html")
|
||||
if err != nil {
|
||||
t.Fatalf("read embedded ui: %v", err)
|
||||
}
|
||||
return string(b)
|
||||
}
|
||||
|
||||
// extractCall returns the source of the balanced api(...) call starting at the
|
||||
// opening parenthesis index, plus the line number it starts on.
|
||||
func extractCall(src string, openIdx int) (string, bool) {
|
||||
depth := 0
|
||||
for i := openIdx; i < len(src); i++ {
|
||||
switch src[i] {
|
||||
case '(':
|
||||
depth++
|
||||
case ')':
|
||||
depth--
|
||||
if depth == 0 {
|
||||
return src[openIdx : i+1], true
|
||||
}
|
||||
case '\n':
|
||||
// A call spanning more than ~25 lines is not one of ours; bail out
|
||||
// rather than scanning the rest of the document.
|
||||
if i-openIdx > 4000 {
|
||||
return "", false
|
||||
}
|
||||
}
|
||||
}
|
||||
return "", false
|
||||
}
|
||||
|
||||
func lineOf(src string, idx int) int {
|
||||
return strings.Count(src[:idx], "\n") + 1
|
||||
}
|
||||
|
||||
// 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
|
||||
// never reached.
|
||||
func TestUIAPICallsDeclareMethod(t *testing.T) {
|
||||
src := uiSource(t)
|
||||
var offenders []string
|
||||
for _, m := range apiCallRe.FindAllStringIndex(src, -1) {
|
||||
open := m[1] - 1 // index of '('
|
||||
call, ok := extractCall(src, open)
|
||||
if !ok {
|
||||
continue
|
||||
}
|
||||
// No options object at all => a plain GET, which is fine.
|
||||
if !strings.Contains(call, "{") {
|
||||
continue
|
||||
}
|
||||
if strings.Contains(call, "method:") {
|
||||
continue
|
||||
}
|
||||
// A read that only carries an AbortSignal is a deliberate GET.
|
||||
if strings.Contains(call, "signal:") && !strings.Contains(call, "body:") {
|
||||
continue
|
||||
}
|
||||
offenders = append(offenders, fmt.Sprintf("line %d: %s", lineOf(src, open), oneLine(call)))
|
||||
}
|
||||
if len(offenders) > 0 {
|
||||
t.Fatalf("api() called with an options object but no method (fetch will send GET):\n %s",
|
||||
strings.Join(offenders, "\n "))
|
||||
}
|
||||
}
|
||||
|
||||
// TestUIDeleteHelpersUseDelete pins the specific helpers that remove things.
|
||||
// Each must issue a DELETE; a GET here is the exact bug users reported as
|
||||
// "cannot delete key: use GET /api/keys".
|
||||
func TestUIDeleteHelpersUseDelete(t *testing.T) {
|
||||
src := uiSource(t)
|
||||
for _, fn := range []string{"delKey", "delAdapter", "delSource", "delTemplate"} {
|
||||
body, ok := jsFunctionBody(src, fn)
|
||||
if !ok {
|
||||
t.Errorf("helper %s() not found in the WebUI", fn)
|
||||
continue
|
||||
}
|
||||
if !strings.Contains(body, "api(") {
|
||||
t.Errorf("%s() does not call api()", fn)
|
||||
continue
|
||||
}
|
||||
if !strings.Contains(body, `method: "DELETE"`) {
|
||||
t.Errorf("%s() must issue a DELETE, got:\n%s", fn, body)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// TestUIMutatingHelpersUseWriteMethods covers the write (non-delete) helpers:
|
||||
// each must declare an explicit POST/PUT/PATCH.
|
||||
func TestUIMutatingHelpersUseWriteMethods(t *testing.T) {
|
||||
src := uiSource(t)
|
||||
want := map[string][]string{
|
||||
"createKey": {`method: "POST"`},
|
||||
"putScope": {`method: "PUT"`},
|
||||
"uploadAdapter": {`method: "POST"`},
|
||||
"saveAuto": {`method: "PUT"`, `method: "POST"`},
|
||||
}
|
||||
for fn, methods := range want {
|
||||
body, ok := jsFunctionBody(src, fn)
|
||||
if !ok {
|
||||
continue // helper renamed or absent: the method test above still guards it
|
||||
}
|
||||
found := false
|
||||
for _, m := range methods {
|
||||
if strings.Contains(body, m) {
|
||||
found = true
|
||||
break
|
||||
}
|
||||
}
|
||||
if !found {
|
||||
t.Errorf("%s() must declare one of %v", fn, methods)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// jsFunctionBody returns the source of `function name(` / `async function name(`
|
||||
// up to its closing brace, using brace balancing.
|
||||
func jsFunctionBody(src, name string) (string, bool) {
|
||||
for _, prefix := range []string{"async function " + name + "(", "function " + name + "("} {
|
||||
idx := strings.Index(src, prefix)
|
||||
if idx < 0 {
|
||||
continue
|
||||
}
|
||||
open := strings.Index(src[idx:], "{")
|
||||
if open < 0 {
|
||||
continue
|
||||
}
|
||||
open += idx
|
||||
depth := 0
|
||||
for i := open; i < len(src); i++ {
|
||||
switch src[i] {
|
||||
case '{':
|
||||
depth++
|
||||
case '}':
|
||||
depth--
|
||||
if depth == 0 {
|
||||
return src[idx : i+1], true
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
return "", false
|
||||
}
|
||||
|
||||
func oneLine(s string) string {
|
||||
s = strings.Join(strings.Fields(s), " ")
|
||||
if len(s) > 140 {
|
||||
return s[:140] + "…"
|
||||
}
|
||||
return s
|
||||
}
|
||||
|
||||
// TestUIRecordsPagingWiring guards the record table's paging plumbing. Two
|
||||
// separate bugs made "scroll to the bottom" silently stop loading:
|
||||
//
|
||||
// 1. paintRecords rebuilt the whole table on every 5s poll whenever the row
|
||||
// count did not exceed the first screen, wiping loaded pages and resetting
|
||||
// scroll position.
|
||||
// 2. IntersectionObserver only fires on TRANSITIONS. With a short list — or a
|
||||
// page whose rows all duplicated the first screen — the sentinel stayed
|
||||
// visible and never fired again.
|
||||
//
|
||||
// The fixes are a build-once flag, a scroll-position fallback, an initial
|
||||
// viewport fill, and chaining when a page yields no new rows. None of that is
|
||||
// reachable from Go, so the wiring is asserted structurally.
|
||||
func TestUIRecordsPagingWiring(t *testing.T) {
|
||||
src := uiSource(t)
|
||||
|
||||
paint, ok := jsFunctionBody(src, "paintRecords")
|
||||
if !ok {
|
||||
t.Fatal("paintRecords() not found")
|
||||
}
|
||||
if !strings.Contains(paint, "recsState.built") {
|
||||
t.Error("paintRecords must build the table once (recsState.built) instead of rebuilding on every poll")
|
||||
}
|
||||
if !strings.Contains(paint, "afterbegin") {
|
||||
t.Error("paintRecords must PREPEND newer rows on later polls, not rebuild the table")
|
||||
}
|
||||
if !strings.Contains(paint, "fillRecordsViewport") {
|
||||
t.Error("paintRecords must fill the viewport so the sentinel can transition")
|
||||
}
|
||||
|
||||
attach, ok := jsFunctionBody(src, "attachRecsObserver")
|
||||
if !ok {
|
||||
t.Fatal("attachRecsObserver() not found")
|
||||
}
|
||||
if !strings.Contains(attach, "IntersectionObserver") {
|
||||
t.Error("attachRecsObserver must still use IntersectionObserver")
|
||||
}
|
||||
if !strings.Contains(attach, `addEventListener("scroll"`) {
|
||||
t.Error("attachRecsObserver needs a scroll fallback: an already-visible sentinel never fires again")
|
||||
}
|
||||
|
||||
load, ok := jsFunctionBody(src, "loadMoreRecords")
|
||||
if !ok {
|
||||
t.Fatal("loadMoreRecords() not found")
|
||||
}
|
||||
if !strings.Contains(load, "maxChain") {
|
||||
t.Error("loadMoreRecords must chain when a page adds no new rows (the first page overlaps the first screen)")
|
||||
}
|
||||
if !strings.Contains(load, "signal: ctrl.signal") {
|
||||
t.Error("loadMoreRecords must remain abortable")
|
||||
}
|
||||
|
||||
fill, ok := jsFunctionBody(src, "fillRecordsViewport")
|
||||
if !ok {
|
||||
t.Fatal("fillRecordsViewport() not found")
|
||||
}
|
||||
for _, want := range []string{"scrollHeight", "clientHeight"} {
|
||||
if !strings.Contains(fill, want) {
|
||||
t.Errorf("fillRecordsViewport must compare %s to decide whether the list overflows", want)
|
||||
}
|
||||
}
|
||||
|
||||
rel, ok := jsFunctionBody(src, "releaseRecords")
|
||||
if !ok {
|
||||
t.Fatal("releaseRecords() not found")
|
||||
}
|
||||
for _, want := range []string{"observer.disconnect()", "removeEventListener", "abort()", "recsState.built = false"} {
|
||||
if !strings.Contains(rel, want) {
|
||||
t.Errorf("releaseRecords must clean up %s", want)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// TestUIThemeMatrix pins the theme system: every accent must define both a light
|
||||
// and a dark variable block, so switching light/dark can never leave a theme
|
||||
// half-defined (which shows up as unreadable text rather than as an error).
|
||||
func TestUIThemeMatrix(t *testing.T) {
|
||||
src := uiSource(t)
|
||||
|
||||
accents := []string{"sakura", "ocean", "violet", "mono"}
|
||||
for _, a := range accents {
|
||||
light := fmt.Sprintf(`html[data-accent=%q]{`, a)
|
||||
dark := fmt.Sprintf(`html[data-theme="dark"][data-accent=%q]{`, a)
|
||||
if !strings.Contains(src, light) {
|
||||
t.Errorf("accent %q has no light-mode variable block (%s)", a, light)
|
||||
}
|
||||
// sakura is the default palette: its light values live in :root, and the
|
||||
// generic dark block covers it.
|
||||
if a != "sakura" && !strings.Contains(src, dark) {
|
||||
t.Errorf("accent %q has no dark-mode variable block (%s)", a, dark)
|
||||
}
|
||||
if !strings.Contains(src, fmt.Sprintf(`data-accent="%s" title=`, a)) {
|
||||
t.Errorf("accent %q has no picker button", a)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// TestUIMonoThemeIsFlat asserts the mono theme really is pure white / pure black:
|
||||
// the surfaces must be opaque and the glass/blob decoration must be switched off,
|
||||
// otherwise "#ffffff" renders as a grey translucent card over a gradient mesh.
|
||||
func TestUIMonoThemeIsFlat(t *testing.T) {
|
||||
src := uiSource(t)
|
||||
|
||||
lightBlock, ok := cssBlock(src, `html[data-accent="mono"]{`)
|
||||
if !ok {
|
||||
t.Fatal("mono light block not found")
|
||||
}
|
||||
for _, want := range []string{"--bg-s1:#ffffff", "--card:#ffffff", "--glass:0px"} {
|
||||
if !strings.Contains(lightBlock, want) {
|
||||
t.Errorf("mono light theme must set %s, got:\n%s", want, lightBlock)
|
||||
}
|
||||
}
|
||||
if !strings.Contains(lightBlock, "--blob1:transparent") {
|
||||
t.Error("mono light theme must disable the gradient blobs")
|
||||
}
|
||||
|
||||
darkBlock, ok := cssBlock(src, `html[data-theme="dark"][data-accent="mono"]{`)
|
||||
if !ok {
|
||||
t.Fatal("mono dark block not found")
|
||||
}
|
||||
for _, want := range []string{"--bg-s1:#000000", "--glass:0px"} {
|
||||
if !strings.Contains(darkBlock, want) {
|
||||
t.Errorf("mono dark theme must set %s, got:\n%s", want, darkBlock)
|
||||
}
|
||||
}
|
||||
if !strings.Contains(darkBlock, "--blob1:transparent") {
|
||||
t.Error("mono dark theme must disable the gradient blobs")
|
||||
}
|
||||
|
||||
// the flat-surface overrides must exist, or backdrop-filter turns #fff grey
|
||||
for _, want := range []string{
|
||||
`html[data-accent="mono"] #bgfx .blob{display:none}`,
|
||||
`backdrop-filter:none`,
|
||||
} {
|
||||
if !strings.Contains(src, want) {
|
||||
t.Errorf("mono theme needs the flat-surface override %q", want)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// TestUICanvasColorsResolveVars guards a class of silent theme bug: Canvas 2D
|
||||
// does not understand CSS var(), so assigning "var(--x)" to strokeStyle is a
|
||||
// no-op that leaves the previous colour (often black) in place — invisible on a
|
||||
// dark background. Theme-driven canvas colours must go through cssVar().
|
||||
func TestUICanvasColorsResolveVars(t *testing.T) {
|
||||
src := uiSource(t)
|
||||
if !strings.Contains(src, "function cssVar(") {
|
||||
t.Fatal("cssVar() helper missing: canvas colours cannot resolve CSS variables without it")
|
||||
}
|
||||
for _, prop := range []string{"strokeStyle", "fillStyle"} {
|
||||
bad := fmt.Sprintf(`ctx.%s = "var(--`, prop)
|
||||
if strings.Contains(src, bad) {
|
||||
t.Errorf("ctx.%s assigned a raw CSS var(): Canvas ignores it silently, use cssVar()", prop)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// cssBlock returns the text of the CSS rule starting at selector, up to its
|
||||
// closing brace.
|
||||
func cssBlock(src, selector string) (string, bool) {
|
||||
i := strings.Index(src, selector)
|
||||
if i < 0 {
|
||||
return "", false
|
||||
}
|
||||
j := strings.Index(src[i:], "}")
|
||||
if j < 0 {
|
||||
return "", false
|
||||
}
|
||||
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))
|
||||
}
|
||||
}
|
||||
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)
|
||||
}
|
||||
}
|
||||
}
|
||||
212
internal/lua/adapters/agentrouter.lua
Normal file
212
internal/lua/adapters/agentrouter.lua
Normal file
@ -0,0 +1,212 @@
|
||||
local adapter = {}
|
||||
|
||||
adapter.name = "agentrouter"
|
||||
adapter.version = "1.0.0"
|
||||
adapter.endpoint = "/chat/completions"
|
||||
adapter.headers = {}
|
||||
|
||||
-- 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
|
||||
-- enforce Anthropic's tool id rule ^[a-zA-Z0-9_-]{1,64}$ by rejecting the WHOLE
|
||||
-- request. Plain OpenAI does not, so an id minted by a permissive model (e.g.
|
||||
-- "bash:0" from moonshotai/kimi-k3) is replayed here by the client and kills
|
||||
-- every turn. See anthropic.lua / openai.lua for the same helper: Lua adapters
|
||||
-- have no shared prelude, so each carries its own copy.
|
||||
local TOOL_ID_MAX = 64
|
||||
|
||||
local function safe_tool_id(id)
|
||||
if type(id) ~= "string" or id == "" then return id end
|
||||
local clean = string.gsub(id, "[^A-Za-z0-9_-]", "_")
|
||||
if clean == id and #clean <= TOOL_ID_MAX then
|
||||
return clean
|
||||
end
|
||||
local digest = string.sub(sha256_hex(id), 1, 8)
|
||||
local keep = TOOL_ID_MAX - #digest - 1
|
||||
if #clean > keep then clean = string.sub(clean, 1, keep) end
|
||||
return clean .. "_" .. digest
|
||||
end
|
||||
|
||||
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
|
||||
if req.messages then
|
||||
for _, msg in ipairs(req.messages) do
|
||||
msg.reasoning_content = nil
|
||||
if msg.tool_call_id ~= nil then
|
||||
msg.tool_call_id = safe_tool_id(msg.tool_call_id)
|
||||
end
|
||||
if type(msg.tool_calls) == "table" then
|
||||
for _, tc in ipairs(msg.tool_calls) do
|
||||
if type(tc) == "table" and tc.id ~= nil then
|
||||
tc.id = safe_tool_id(tc.id)
|
||||
end
|
||||
end
|
||||
end
|
||||
end
|
||||
end
|
||||
return json.encode(req)
|
||||
end
|
||||
|
||||
-- 注入 Authorization + 官方客户端 User-Agent
|
||||
function adapter.build_headers(meta)
|
||||
local ua = default_ua
|
||||
local meta_ua = (meta.source or {}).meta and (meta.source).meta.user_agent
|
||||
if meta_ua and meta_ua ~= "" then
|
||||
ua = meta_ua
|
||||
end
|
||||
return {
|
||||
["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
|
||||
|
||||
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
|
||||
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 ""
|
||||
if ch.message.reasoning_content then
|
||||
unified.reasoning_content = ch.message.reasoning_content
|
||||
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 = safe_tool_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
|
||||
|
||||
-- usage-only terminal chunk (empty choices + usage): forward it so the
|
||||
-- gateway emits exact token counts (and cache fields when present).
|
||||
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,
|
||||
}
|
||||
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 }
|
||||
elseif (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
|
||||
uses.prompt_tokens_details = { cached_tokens = chunk.usage.prompt_cache_hit_tokens }
|
||||
end
|
||||
end
|
||||
|
||||
if not chunk.choices or #chunk.choices == 0 then
|
||||
if uses ~= nil then
|
||||
return json.encode({ usage = uses, done = false })
|
||||
end
|
||||
return ""
|
||||
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 delta.reasoning_content then
|
||||
unified.reasoning_content = delta.reasoning_content
|
||||
end
|
||||
if delta.tool_calls then
|
||||
-- Sanitize outbound too: a dirty id must never enter a client session,
|
||||
-- because the client replays it to every other source. Only the first
|
||||
-- fragment of a streamed call carries an id; later argument fragments
|
||||
-- have none and must stay id-less for index-based accumulation.
|
||||
for _, tc in ipairs(delta.tool_calls) do
|
||||
if type(tc) == "table" and tc.id ~= nil then
|
||||
tc.id = safe_tool_id(tc.id)
|
||||
end
|
||||
end
|
||||
unified.tool_calls = delta.tool_calls
|
||||
end
|
||||
if uses ~= nil then
|
||||
unified.usage = uses
|
||||
end
|
||||
return json.encode(unified)
|
||||
end
|
||||
|
||||
-- 错误收敛:前置 WAF 会返回整页 HTML(阿里云盾),JSON 时为 new-api 风格
|
||||
function adapter.transform_error(status, body)
|
||||
local low = string.lower(body or "")
|
||||
if string.sub(low, 1, 9) == "<!doctype" or string.find(low, "<html", 1, true) then
|
||||
return "blocked by AgentRouter WAF"
|
||||
end
|
||||
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" and type(e.message) == "string" then
|
||||
return e.message
|
||||
end
|
||||
return nil
|
||||
end
|
||||
|
||||
return adapter
|
||||
@ -1,172 +1,447 @@
|
||||
local adapter = {}
|
||||
|
||||
adapter.name = "anthropic"
|
||||
adapter.version = "2.0.0"
|
||||
adapter.version = "3.0.0"
|
||||
adapter.endpoint = "/v1/messages"
|
||||
adapter.headers = {
|
||||
["anthropic-version"] = "2023-06-01"
|
||||
["anthropic-version"] = "2023-06-01",
|
||||
}
|
||||
|
||||
function adapter.transform_request(raw_body)
|
||||
local ok, req = pcall(json.decode, raw_body)
|
||||
if not ok then return raw_body end
|
||||
-- ============================================================
|
||||
-- helpers
|
||||
-- ============================================================
|
||||
|
||||
-- 将 OpenAI 风格 content(字符串或 [{type:*}] 数组)拆成文本/图片块
|
||||
local function collect_blocks(content)
|
||||
if type(content) == "string" then
|
||||
return { { type = "text", text = content } }
|
||||
end
|
||||
local blocks = {}
|
||||
for _, p in ipairs(content or {}) do
|
||||
local ROLE_WHITELIST = {
|
||||
system = true, user = true, assistant = true, tool = true,
|
||||
}
|
||||
|
||||
-- Anthropic requires tool_use.id / tool_result.tool_use_id to match
|
||||
-- ^[a-zA-Z0-9_-]{1,64}$ and rejects the WHOLE request otherwise with
|
||||
-- REQUEST_BODY_INVALID / "Invalid tool use format". OpenAI imposes no such
|
||||
-- rule, so an OpenAI-compatible upstream can hand a client an id like
|
||||
-- "bash:0" (observed from moonshotai/kimi-k3); once that lands in the
|
||||
-- client's history, every later replay kills every Claude slot at once and
|
||||
-- an AUTO chain falls through all tiers.
|
||||
--
|
||||
-- safe_tool_id is pure and deterministic, so a tool_use block and its
|
||||
-- matching tool_result are rewritten identically within one request. A
|
||||
-- rewritten id keeps an 8-hex digest of the ORIGINAL id: without it two
|
||||
-- distinct ids could collapse into one ("a:b" and "a_b"), which Anthropic
|
||||
-- rejects as a duplicate/unpaired tool_use. Already-legal ids pass through
|
||||
-- untouched so upstreams that round-trip their own ids are unaffected.
|
||||
local TOOL_ID_MAX = 64
|
||||
|
||||
local function safe_tool_id(id)
|
||||
if type(id) ~= "string" or id == "" then return "" end
|
||||
local clean = string.gsub(id, "[^A-Za-z0-9_-]", "_")
|
||||
if clean == id and #clean <= TOOL_ID_MAX then
|
||||
return clean
|
||||
end
|
||||
local digest = string.sub(sha256_hex(id), 1, 8)
|
||||
local keep = TOOL_ID_MAX - #digest - 1
|
||||
if #clean > keep then clean = string.sub(clean, 1, keep) end
|
||||
return clean .. "_" .. digest
|
||||
end
|
||||
|
||||
-- map_usage converts one Anthropic usage object to the gateway's TokenUsage
|
||||
-- shape. The two APIs disagree on what the prompt count MEANS:
|
||||
--
|
||||
-- Anthropic: input_tokens EXCLUDES cached blocks. cache_read_input_tokens and
|
||||
-- cache_creation_input_tokens are separate, additive, billed input.
|
||||
-- OpenAI: prompt_tokens INCLUDES its cached_tokens subset.
|
||||
--
|
||||
-- Mapping input_tokens straight onto prompt therefore did two wrong things at
|
||||
-- once: it undercounted the billed prompt by the entire cache portion, and it
|
||||
-- could report cached_tokens > prompt_tokens — a cache hit rate above 100%.
|
||||
-- cache_creation_input_tokens was not read at all, so a cache-write turn lost
|
||||
-- those billed tokens outright. Summing all three restores OpenAI semantics,
|
||||
-- which is what clients (and the gateway's own audit trail) assume.
|
||||
--
|
||||
-- Returns nil when the object carries no counts, so callers can keep treating
|
||||
-- "no usage in this chunk" as "say nothing" rather than reporting zeros.
|
||||
local function map_usage(u)
|
||||
if type(u) ~= "table" then return nil end
|
||||
local fresh = u.input_tokens or 0
|
||||
local read = u.cache_read_input_tokens or 0
|
||||
local create = u.cache_creation_input_tokens or 0
|
||||
local p = fresh + read + create
|
||||
local c = u.output_tokens or 0
|
||||
if p <= 0 and c <= 0 then return nil end
|
||||
local uses = { prompt = p, completion = c, total = p + c }
|
||||
-- Emit the split whenever Anthropic reports either cache field, even at 0:
|
||||
-- a reported zero-hit ("cache missed") must stay distinguishable from "this
|
||||
-- upstream never reports cache info at all".
|
||||
if u.cache_read_input_tokens ~= nil or u.cache_creation_input_tokens ~= nil then
|
||||
uses.prompt_tokens_details = { cached_tokens = read }
|
||||
uses.prompt_cache_hit_tokens = read
|
||||
uses.prompt_cache_miss_tokens = fresh + create
|
||||
end
|
||||
return uses
|
||||
end
|
||||
|
||||
local function collect_blocks(content)
|
||||
if type(content) == "string" then
|
||||
if content == "" then return {} end
|
||||
return { { type = "text", text = content } }
|
||||
end
|
||||
if type(content) ~= "table" then return {} end
|
||||
local blocks = {}
|
||||
for _, p in ipairs(content) do
|
||||
if type(p) == "string" then
|
||||
table.insert(blocks, { type = "text", text = p })
|
||||
elseif type(p) == "table" then
|
||||
if p.type == "text" then
|
||||
table.insert(blocks, { type = "text", text = p.text })
|
||||
table.insert(blocks, { type = "text", text = p.text or "" })
|
||||
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,(.+)$")
|
||||
local url = p.image_url.url
|
||||
local mt, b64 = string.match(url, "^data:([^,]+);base64,(.+)$")
|
||||
if b64 then
|
||||
table.insert(blocks, { type = "image", source = { type = "base64", media_type = mt or "image/png", data = b64 } })
|
||||
else
|
||||
table.insert(blocks, { type = "image", source = { type = "url", url = p.image_url.url } })
|
||||
table.insert(blocks, { type = "image", source = { type = "url", url = url } })
|
||||
end
|
||||
else
|
||||
-- Unknown content type: preserve for forward compatibility
|
||||
table.insert(blocks, p)
|
||||
end
|
||||
end
|
||||
return blocks
|
||||
end
|
||||
local function text_of(content)
|
||||
if type(content) == "string" then return content end
|
||||
local t = ""
|
||||
for _, p in ipairs(content or {}) do
|
||||
if p.type == "text" and p.text then t = t .. p.text end
|
||||
return blocks
|
||||
end
|
||||
|
||||
local function text_of(content)
|
||||
if type(content) == "string" then return content end
|
||||
local t = ""
|
||||
for _, p in ipairs(content or {}) do
|
||||
if type(p) == "string" then
|
||||
t = t .. p
|
||||
elseif type(p) == "table" and p.type == "text" and p.text then
|
||||
t = t .. p.text
|
||||
end
|
||||
return t
|
||||
end
|
||||
return t
|
||||
end
|
||||
|
||||
--- Append blocks to the last user message (merge) or create a new one.
|
||||
local function append_user(msgs, blocks)
|
||||
if #msgs > 0 and msgs[#msgs].role == "user" then
|
||||
for _, b in ipairs(blocks) do
|
||||
table.insert(msgs[#msgs].content, b)
|
||||
end
|
||||
else
|
||||
table.insert(msgs, { role = "user", content = blocks })
|
||||
end
|
||||
end
|
||||
|
||||
-- ============================================================
|
||||
-- transform_request (OpenAI → Anthropic)
|
||||
-- ============================================================
|
||||
|
||||
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
|
||||
|
||||
local msgs = {}
|
||||
local system = ""
|
||||
local pending_tool = {} -- accumulated tool_result blocks
|
||||
|
||||
for _, m in ipairs(req.messages or {}) do
|
||||
if m.role == "system" then
|
||||
local role = m.role or "user"
|
||||
|
||||
if role == "system" then
|
||||
system = system .. text_of(m.content) .. "\n"
|
||||
|
||||
elseif role == "assistant" then
|
||||
-- Build content array: text/image blocks from content + tool_use from tool_calls
|
||||
local blocks = collect_blocks(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
|
||||
local input = {}
|
||||
if type(args) == "string" and args ~= "" then
|
||||
local ok2, parsed = pcall(json.decode, args)
|
||||
if ok2 and type(parsed) == "table" then input = parsed end
|
||||
elseif type(args) == "table" then
|
||||
input = args
|
||||
end
|
||||
table.insert(blocks, {
|
||||
type = "tool_use",
|
||||
id = safe_tool_id(tc.id),
|
||||
name = fn.name or "",
|
||||
input = input,
|
||||
})
|
||||
end
|
||||
end
|
||||
end
|
||||
|
||||
table.insert(msgs, { role = "assistant", content = blocks })
|
||||
|
||||
elseif role == "tool" then
|
||||
-- Accumulate consecutive tool results; will be flushed as one
|
||||
-- user message with tool_result content blocks.
|
||||
table.insert(pending_tool, {
|
||||
type = "tool_result",
|
||||
tool_use_id = safe_tool_id(m.tool_call_id),
|
||||
content = text_of(m.content),
|
||||
})
|
||||
|
||||
else
|
||||
table.insert(msgs, { role = m.role, content = collect_blocks(m.content) })
|
||||
-- user / any other role: flush pending tool results first
|
||||
if #pending_tool > 0 then
|
||||
append_user(msgs, pending_tool)
|
||||
pending_tool = {}
|
||||
end
|
||||
append_user(msgs, collect_blocks(m.content))
|
||||
end
|
||||
end
|
||||
|
||||
local anthropic_req = {
|
||||
model = req.model or "claude-sonnet-4-20250514",
|
||||
max_tokens = req.max_tokens or 4096,
|
||||
messages = msgs,
|
||||
stream = req.stream or false,
|
||||
}
|
||||
|
||||
if not req.disable_thinking then
|
||||
anthropic_req.thinking = { type = "enabled", budget_tokens = 4096 }
|
||||
-- Flush any trailing tool results
|
||||
if #pending_tool > 0 then
|
||||
append_user(msgs, pending_tool)
|
||||
end
|
||||
|
||||
-- ── Build Anthropic request ──────────────────────────────
|
||||
local anthropic_req = {
|
||||
model = req.model or "claude-sonnet-4-20250514",
|
||||
max_tokens = req.max_tokens or 4096,
|
||||
messages = msgs,
|
||||
stream = req.stream or false,
|
||||
}
|
||||
|
||||
-- ── tools ────────────────────────────────────────────────
|
||||
local has_tools = false
|
||||
if req.tools and type(req.tools) == "table" then
|
||||
local tools = {}
|
||||
for _, t in ipairs(req.tools) do
|
||||
if type(t) == "table" and t.type == "function"
|
||||
and type(t["function"]) == "table" then
|
||||
local fn = t["function"]
|
||||
table.insert(tools, {
|
||||
name = fn.name or "",
|
||||
description = fn.description or "",
|
||||
input_schema = fn.parameters or { type = "object", properties = {} },
|
||||
})
|
||||
end
|
||||
end
|
||||
if #tools > 0 then
|
||||
anthropic_req.tools = tools
|
||||
has_tools = true
|
||||
end
|
||||
end
|
||||
|
||||
-- ── tool_choice ──────────────────────────────────────────
|
||||
-- OpenAI → Anthropic mapping:
|
||||
-- "auto" → {type:"auto"}
|
||||
-- "none" → remove tools entirely (Anthropic has no "none")
|
||||
-- "required" → {type:"any"}
|
||||
-- {type:"function", function:{name:"X"}} → {type:"tool", name:"X"}
|
||||
if has_tools and req.tool_choice then
|
||||
local tc = req.tool_choice
|
||||
if type(tc) == "string" then
|
||||
if tc == "auto" then
|
||||
anthropic_req.tool_choice = { type = "auto" }
|
||||
elseif tc == "none" then
|
||||
anthropic_req.tools = nil
|
||||
anthropic_req.tool_choice = nil
|
||||
elseif tc == "required" then
|
||||
anthropic_req.tool_choice = { type = "any" }
|
||||
end
|
||||
elseif type(tc) == "table" and tc.type == "function" then
|
||||
local fn = tc["function"] or {}
|
||||
anthropic_req.tool_choice = { type = "tool", name = fn.name or "" }
|
||||
end
|
||||
end
|
||||
|
||||
-- ── thinking (opt-in via extra_body) ──────────────────────
|
||||
-- Default: OFF. Client sends extra_body.thinking to enable.
|
||||
-- Example: {"extra_body": {"thinking": {"type": "enabled", "budget_tokens": 4096}}}
|
||||
if type(req.extra_body) == "table" and type(req.extra_body.thinking) == "table" then
|
||||
anthropic_req.thinking = req.extra_body.thinking
|
||||
end
|
||||
|
||||
-- ── system ───────────────────────────────────────────────
|
||||
if system ~= "" then
|
||||
anthropic_req.system = system
|
||||
-- Strip trailing newline from accumulation
|
||||
anthropic_req.system = string.match(system, "^(.-)\n*$")
|
||||
end
|
||||
|
||||
return json.encode(anthropic_req)
|
||||
end
|
||||
|
||||
-- ============================================================
|
||||
-- transform_response (Anthropic → OpenAI, non-streaming)
|
||||
-- ============================================================
|
||||
|
||||
function adapter.transform_response(raw_body)
|
||||
local ok, resp = pcall(json.decode, raw_body)
|
||||
if not ok then return raw_body end
|
||||
if not ok or resp == nil then return raw_body end
|
||||
|
||||
local unified = {
|
||||
content = "",
|
||||
finish_reason = "",
|
||||
token_usage = { prompt = 0, completion = 0, total = 0 }
|
||||
token_usage = { prompt = 0, completion = 0, total = 0 },
|
||||
}
|
||||
|
||||
if resp.usage then
|
||||
unified.token_usage.prompt = resp.usage.input_tokens or 0
|
||||
unified.token_usage.completion = resp.usage.output_tokens or 0
|
||||
unified.token_usage.total = (resp.usage.input_tokens or 0) + (resp.usage.output_tokens or 0)
|
||||
local mapped = map_usage(resp.usage)
|
||||
if mapped ~= nil then
|
||||
unified.token_usage = mapped
|
||||
end
|
||||
|
||||
if resp.content and #resp.content > 0 then
|
||||
local tcs = {}
|
||||
for _, block in ipairs(resp.content) do
|
||||
if block.type == "text" then
|
||||
unified.content = unified.content .. (block.text or "")
|
||||
elseif block.type == "thinking" and block.thinking then
|
||||
unified.reasoning_content = (unified.reasoning_content or "") .. block.thinking
|
||||
elseif block.type == "tool_use" then
|
||||
table.insert(tcs, {
|
||||
id = safe_tool_id(block.id),
|
||||
type = "function",
|
||||
name = block.name or "",
|
||||
arguments = block.input or {},
|
||||
})
|
||||
end
|
||||
end
|
||||
if #tcs > 0 then unified.tool_calls = tcs end
|
||||
end
|
||||
|
||||
unified.finish_reason = resp.stop_reason or ""
|
||||
if unified.finish_reason == "tool_use" then
|
||||
unified.finish_reason = "tool_calls"
|
||||
elseif unified.finish_reason == "max_tokens" then
|
||||
unified.finish_reason = "length"
|
||||
elseif unified.finish_reason == "end_turn"
|
||||
or unified.finish_reason == "stop_sequence" then
|
||||
unified.finish_reason = "stop"
|
||||
end
|
||||
|
||||
return json.encode(unified)
|
||||
end
|
||||
|
||||
-- ============================================================
|
||||
-- transform_stream_chunk (Anthropic SSE → OpenAI SSE delta)
|
||||
-- ============================================================
|
||||
|
||||
function adapter.transform_stream_chunk(raw_chunk)
|
||||
local ok, chunk = pcall(json.decode, raw_chunk)
|
||||
if not ok then return "" end
|
||||
|
||||
-- ── message_start: initial usage ─────────────────────────
|
||||
if chunk.type == "message_start" then
|
||||
local uses = nil
|
||||
if chunk.message and type(chunk.message.usage) == "table" then
|
||||
local u = chunk.message.usage
|
||||
local p = u.input_tokens or 0
|
||||
local c = u.output_tokens or 0
|
||||
if p > 0 or c > 0 then
|
||||
uses = { prompt = p, completion = c, total = p + c }
|
||||
if chunk.message then
|
||||
-- Some Anthropic-compatible upstreams report a placeholder here and
|
||||
-- only send the true prompt count in message_delta (justwoker:
|
||||
-- 160 at message_start vs 6931 at message_delta). The gateway's
|
||||
-- mergeUsage lets a later non-zero value win, so both are emitted.
|
||||
local uses = map_usage(chunk.message.usage)
|
||||
if uses ~= nil then
|
||||
return json.encode({ usage = uses, done = false })
|
||||
end
|
||||
end
|
||||
if uses ~= nil then
|
||||
return json.encode({ usage = uses, done = false })
|
||||
end
|
||||
return ""
|
||||
end
|
||||
|
||||
-- ── message_delta: stop_reason + final usage ─────────────
|
||||
if chunk.type == "message_delta" then
|
||||
local uses = nil
|
||||
if type(chunk.usage) == "table" then
|
||||
local u = chunk.usage
|
||||
local p = u.input_tokens or 0
|
||||
local c = u.output_tokens or 0
|
||||
if p > 0 or c > 0 then
|
||||
uses = { prompt = p, completion = c, total = p + c }
|
||||
local finish = nil
|
||||
if chunk.delta and chunk.delta.stop_reason ~= nil then
|
||||
local sr = chunk.delta.stop_reason
|
||||
if sr == "max_tokens" then
|
||||
finish = "length"
|
||||
elseif sr == "tool_use" then
|
||||
finish = "tool_calls"
|
||||
else
|
||||
finish = "stop"
|
||||
end
|
||||
end
|
||||
local done = (chunk.delta and chunk.delta.stop_reason ~= nil)
|
||||
-- message_delta carries the FINAL usage, so it must map the cache
|
||||
-- fields too; the old code dropped them here, which silently lost the
|
||||
-- whole cache split on every streamed response.
|
||||
local uses = map_usage(chunk.usage)
|
||||
if uses ~= nil then
|
||||
-- completion is final here; prompt is merged from message_start
|
||||
return json.encode({ content = "", done = done, usage = uses })
|
||||
return json.encode({ content = "", done = (finish ~= nil), finish_reason = finish, usage = uses })
|
||||
end
|
||||
return json.encode({ content = "", done = done })
|
||||
return json.encode({ content = "", done = (finish ~= nil), finish_reason = finish })
|
||||
end
|
||||
if chunk.type == "content_block_start" and chunk.content_block
|
||||
and chunk.content_block.type == "tool_use" then
|
||||
-- first fragment of a tool call: emit index + id + name, empty args
|
||||
return json.encode({
|
||||
content = "", done = false,
|
||||
tool_calls = { {
|
||||
index = chunk.index or 0,
|
||||
id = chunk.content_block.id or "",
|
||||
type = "function",
|
||||
["function"] = { name = chunk.content_block.name or "", arguments = "" }
|
||||
} }
|
||||
})
|
||||
|
||||
-- ── content_block_start: begin text / thinking / tool_use ─
|
||||
if chunk.type == "content_block_start" and chunk.content_block then
|
||||
local cb = chunk.content_block
|
||||
if cb.type == "tool_use" then
|
||||
-- Pass Anthropic content_block index through; Go-side rewrites
|
||||
-- to sequential OpenAI tool_call ordinal for multi-tool streams.
|
||||
return json.encode({
|
||||
content = "", done = false,
|
||||
tool_calls = { {
|
||||
index = chunk.index or 0,
|
||||
id = safe_tool_id(cb.id),
|
||||
type = "function",
|
||||
["function"] = { name = cb.name or "", arguments = "" },
|
||||
} },
|
||||
})
|
||||
end
|
||||
return "" -- text / thinking block start: no OpenAI equivalent
|
||||
end
|
||||
|
||||
-- ── content_block_delta: incremental content ──────────────
|
||||
if chunk.type == "content_block_delta" and chunk.delta then
|
||||
if chunk.delta.type == "input_json_delta" then
|
||||
-- incremental JSON fragment; clients accumulate across chunks
|
||||
local unified = { content = "", done = false, tool_calls = { {
|
||||
index = chunk.index or 0,
|
||||
id = "",
|
||||
type = "function",
|
||||
["function"] = { name = "", arguments = chunk.delta.partial_json or "" }
|
||||
} } }
|
||||
return json.encode(unified)
|
||||
local d = chunk.delta
|
||||
if d.type == "input_json_delta" then
|
||||
-- Tool call argument fragment; Go-side rewrites index.
|
||||
return json.encode({
|
||||
content = "", done = false,
|
||||
tool_calls = { {
|
||||
index = chunk.index or 0,
|
||||
id = "",
|
||||
type = "function",
|
||||
["function"] = { name = "", arguments = d.partial_json or "" },
|
||||
} },
|
||||
})
|
||||
end
|
||||
if chunk.delta.type == "thinking_delta" and chunk.delta.thinking then
|
||||
return json.encode({ content = "", done = false, reasoning_content = chunk.delta.thinking })
|
||||
if d.type == "thinking_delta" and d.thinking then
|
||||
return json.encode({ content = "", done = false, reasoning_content = d.thinking })
|
||||
end
|
||||
if d.type == "text_delta" and d.text then
|
||||
return json.encode({ content = d.text, done = false })
|
||||
end
|
||||
return json.encode({ content = chunk.delta.text or "", done = false })
|
||||
end
|
||||
|
||||
-- ── content_block_stop / message_stop ─────────────────────
|
||||
if chunk.type == "message_stop" then
|
||||
return json.encode({ content = "", done = true })
|
||||
end
|
||||
if chunk.type == "content_block_stop" then
|
||||
return json.encode({ content = "", done = false })
|
||||
-- The message_delta event already emitted the true finish_reason.
|
||||
-- Do NOT emit done=true here: an empty finish_reason would
|
||||
-- overwrite the real one (tool_calls) in the Go gateway's
|
||||
-- lastFinish tracker, causing the final SSE chunk to say
|
||||
-- finish_reason=stop instead of tool_calls.
|
||||
return ""
|
||||
end
|
||||
|
||||
return ""
|
||||
end
|
||||
|
||||
-- ============================================================
|
||||
-- transform_error (Anthropic error → human-readable string)
|
||||
-- ============================================================
|
||||
|
||||
function adapter.transform_error(status, body)
|
||||
local ok, resp = pcall(json.decode, body)
|
||||
if not ok or type(resp) ~= "table" then return nil end
|
||||
|
||||
-- Anthropic envelope: {type:"error", error:{type, message}}
|
||||
if resp.type == "error" and type(resp.error) == "table"
|
||||
and type(resp.error.message) == "string" then
|
||||
return resp.error.message
|
||||
end
|
||||
|
||||
-- Flat envelope: {error: {message: "..."}}
|
||||
if type(resp.error) == "table" and type(resp.error.message) == "string" then
|
||||
return resp.error.message
|
||||
end
|
||||
|
||||
return nil
|
||||
end
|
||||
|
||||
return adapter
|
||||
|
||||
@ -49,6 +49,11 @@ function adapter.transform_response(raw_body)
|
||||
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
|
||||
unified.token_usage.prompt_cache_hit_tokens = resp.usage.prompt_cache_hit_tokens or 0
|
||||
unified.token_usage.prompt_cache_miss_tokens = resp.usage.prompt_cache_miss_tokens or 0
|
||||
if resp.usage.prompt_tokens_details and resp.usage.prompt_tokens_details.cached_tokens ~= nil then
|
||||
unified.token_usage.prompt_tokens_details = { cached_tokens = resp.usage.prompt_tokens_details.cached_tokens }
|
||||
end
|
||||
end
|
||||
|
||||
if type(resp.choices) == "table" and #resp.choices > 0 then
|
||||
@ -92,7 +97,12 @@ function adapter.transform_stream_chunk(raw_chunk)
|
||||
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_cache_hit_tokens = chunk.usage.prompt_cache_hit_tokens or 0,
|
||||
prompt_cache_miss_tokens = chunk.usage.prompt_cache_miss_tokens or 0,
|
||||
}
|
||||
if chunk.usage.prompt_tokens_details and chunk.usage.prompt_tokens_details.cached_tokens ~= nil then
|
||||
uses.prompt_tokens_details = { cached_tokens = chunk.usage.prompt_tokens_details.cached_tokens }
|
||||
end
|
||||
end
|
||||
|
||||
if not chunk.choices or #chunk.choices == 0 then
|
||||
@ -106,10 +116,15 @@ function adapter.transform_stream_chunk(raw_chunk)
|
||||
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 = (fr ~= nil)
|
||||
done = (finish ~= nil)
|
||||
}
|
||||
if finish then
|
||||
unified.finish_reason = finish
|
||||
end
|
||||
if uses ~= nil then
|
||||
unified.usage = uses
|
||||
end
|
||||
@ -122,4 +137,15 @@ function adapter.transform_stream_chunk(raw_chunk)
|
||||
return json.encode(unified)
|
||||
end
|
||||
|
||||
-- 错误收敛:DeepSeek 走标准 OpenAI 信封 {error:{message,...}}
|
||||
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" and type(e.message) == "string" then
|
||||
return e.message
|
||||
end
|
||||
return nil
|
||||
end
|
||||
|
||||
return adapter
|
||||
|
||||
@ -7,6 +7,10 @@ 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 +18,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 +42,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 +105,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
|
||||
@ -68,20 +151,55 @@ function adapter.transform_response(raw_body)
|
||||
unified.token_usage.prompt = resp.usageMetadata.promptTokenCount or 0
|
||||
unified.token_usage.completion = resp.usageMetadata.candidatesTokenCount or 0
|
||||
unified.token_usage.total = resp.usageMetadata.totalTokenCount or 0
|
||||
-- Gemini reports context-cache reads as cachedContentTokenCount;
|
||||
-- normalize into OpenAI-standard prompt_tokens_details.cached_tokens
|
||||
-- so clients and the audit trail see the hit count. Emitted even when
|
||||
-- 0 so a reported miss stays distinguishable from "not reported".
|
||||
if resp.usageMetadata.cachedContentTokenCount ~= nil then
|
||||
unified.token_usage.prompt_tokens_details = {
|
||||
cached_tokens = resp.usageMetadata.cachedContentTokenCount
|
||||
}
|
||||
end
|
||||
end
|
||||
|
||||
if resp.candidates and #resp.candidates > 0 then
|
||||
local cand = resp.candidates[1]
|
||||
local tools = {}
|
||||
if cand.content and cand.content.parts then
|
||||
for _, part in ipairs(cand.content.parts) do
|
||||
if part.text then
|
||||
unified.content = unified.content .. part.text
|
||||
elseif part.functionCall then
|
||||
-- Non-streaming tool calls used to be dropped here while
|
||||
-- transform_stream_chunk handled them, so a non-streaming
|
||||
-- agent turn looked like a plain text answer and the tool
|
||||
-- loop died. Gemini's args are already an object.
|
||||
local args = part.functionCall.args
|
||||
if type(args) == "string" then
|
||||
local aok, decoded = pcall(json.decode, args)
|
||||
args = aok and decoded or {}
|
||||
elseif type(args) ~= "table" then
|
||||
args = {}
|
||||
end
|
||||
table.insert(tools, {
|
||||
id = part.functionCall.id or ("call_" .. #tools),
|
||||
type = "function",
|
||||
name = part.functionCall.name or "",
|
||||
arguments = args
|
||||
})
|
||||
end
|
||||
end
|
||||
end
|
||||
if cand.finishReason then
|
||||
unified.finish_reason = cand.finishReason
|
||||
end
|
||||
if #tools > 0 then
|
||||
unified.tool_calls = tools
|
||||
-- Gemini reports finishReason "STOP" even when it emitted a
|
||||
-- functionCall; clients keyed on finish_reason would treat that as
|
||||
-- a completed answer and never run the tool.
|
||||
unified.finish_reason = "tool_calls"
|
||||
end
|
||||
end
|
||||
|
||||
return json.encode(unified)
|
||||
@ -100,6 +218,13 @@ function adapter.transform_stream_chunk(raw_chunk)
|
||||
local t = chunk.usageMetadata.totalTokenCount or 0
|
||||
if p > 0 or c > 0 or t > 0 then
|
||||
uses = { prompt = p, completion = c, total = t }
|
||||
-- Emit details whenever the field is present, even at 0, so a
|
||||
-- reported cache miss stays distinguishable from "not reported".
|
||||
if chunk.usageMetadata.cachedContentTokenCount ~= nil then
|
||||
uses.prompt_tokens_details = {
|
||||
cached_tokens = chunk.usageMetadata.cachedContentTokenCount
|
||||
}
|
||||
end
|
||||
end
|
||||
end
|
||||
|
||||
@ -110,7 +235,19 @@ function adapter.transform_stream_chunk(raw_chunk)
|
||||
return ""
|
||||
end
|
||||
local cand = chunk.candidates[1]
|
||||
local unified = { content = "", done = (cand.finishReason ~= nil) }
|
||||
-- Gemini finishReason -> OpenAI finish_reason
|
||||
local finish = nil
|
||||
if cand.finishReason ~= nil then
|
||||
if cand.finishReason == "MAX_TOKENS" then
|
||||
finish = "length"
|
||||
elseif cand.finishReason == "SAFETY" or cand.finishReason == "RECITATION"
|
||||
or cand.finishReason == "BLOCKLIST" then
|
||||
finish = "content_filter"
|
||||
else
|
||||
finish = "stop"
|
||||
end
|
||||
end
|
||||
local unified = { content = "", done = (finish ~= nil), finish_reason = finish }
|
||||
local reasoning = ""
|
||||
local tools = {}
|
||||
if cand.content and cand.content.parts then
|
||||
@ -140,4 +277,20 @@ function adapter.transform_stream_chunk(raw_chunk)
|
||||
return json.encode(unified)
|
||||
end
|
||||
|
||||
-- 错误收敛:Gemini REST 信封 {error:{code, message, status}}
|
||||
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
|
||||
if type(e.message) == "string" then
|
||||
if type(e.status) == "string" then
|
||||
return e.status .. ": " .. e.message
|
||||
end
|
||||
return e.message
|
||||
end
|
||||
end
|
||||
return nil
|
||||
end
|
||||
|
||||
return adapter
|
||||
|
||||
@ -33,6 +33,18 @@ function adapter.transform_response(raw_body)
|
||||
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
|
||||
end
|
||||
|
||||
if type(resp.choices) == "table" and #resp.choices > 0 then
|
||||
@ -74,6 +86,13 @@ function adapter.transform_stream_chunk(raw_chunk)
|
||||
completion = chunk.usage.completion_tokens or chunk.usage.completion or 0,
|
||||
total = chunk.usage.total_tokens or chunk.usage.total or 0,
|
||||
}
|
||||
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 }
|
||||
elseif (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
|
||||
uses.prompt_tokens_details = { cached_tokens = chunk.usage.prompt_cache_hit_tokens }
|
||||
end
|
||||
end
|
||||
|
||||
if not chunk.choices or #chunk.choices == 0 then
|
||||
@ -87,10 +106,15 @@ function adapter.transform_stream_chunk(raw_chunk)
|
||||
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 = (fr ~= nil)
|
||||
done = (finish ~= nil)
|
||||
}
|
||||
if finish then
|
||||
unified.finish_reason = finish
|
||||
end
|
||||
if uses ~= nil then
|
||||
unified.usage = uses
|
||||
end
|
||||
@ -103,4 +127,15 @@ function adapter.transform_stream_chunk(raw_chunk)
|
||||
return json.encode(unified)
|
||||
end
|
||||
|
||||
-- 错误收敛:GitHub Models 走标准 OpenAI 信封 {error:{message,...}}
|
||||
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" and type(e.message) == "string" then
|
||||
return e.message
|
||||
end
|
||||
return nil
|
||||
end
|
||||
|
||||
return adapter
|
||||
|
||||
@ -32,6 +32,18 @@ function adapter.transform_response(raw_body)
|
||||
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
|
||||
end
|
||||
|
||||
if type(resp.choices) == "table" and #resp.choices > 0 then
|
||||
@ -73,6 +85,13 @@ function adapter.transform_stream_chunk(raw_chunk)
|
||||
completion = chunk.usage.completion_tokens or chunk.usage.completion or 0,
|
||||
total = chunk.usage.total_tokens or chunk.usage.total or 0,
|
||||
}
|
||||
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 }
|
||||
elseif (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
|
||||
uses.prompt_tokens_details = { cached_tokens = chunk.usage.prompt_cache_hit_tokens }
|
||||
end
|
||||
end
|
||||
|
||||
if not chunk.choices or #chunk.choices == 0 then
|
||||
@ -86,10 +105,15 @@ function adapter.transform_stream_chunk(raw_chunk)
|
||||
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 = (fr ~= nil)
|
||||
done = (finish ~= nil)
|
||||
}
|
||||
if finish then
|
||||
unified.finish_reason = finish
|
||||
end
|
||||
if uses ~= nil then
|
||||
unified.usage = uses
|
||||
end
|
||||
@ -102,4 +126,15 @@ function adapter.transform_stream_chunk(raw_chunk)
|
||||
return json.encode(unified)
|
||||
end
|
||||
|
||||
-- 错误收敛:Groq 走标准 OpenAI 信封 {error:{message,...}}
|
||||
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" and type(e.message) == "string" then
|
||||
return e.message
|
||||
end
|
||||
return nil
|
||||
end
|
||||
|
||||
return adapter
|
||||
|
||||
@ -64,6 +64,18 @@ function adapter.transform_response(raw_body)
|
||||
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
|
||||
end
|
||||
if type(resp.choices) == "table" and #resp.choices > 0 then
|
||||
local ch = resp.choices[1]
|
||||
@ -106,6 +118,13 @@ function adapter.transform_stream_chunk(raw_chunk)
|
||||
completion = chunk.usage.completion_tokens or chunk.usage.completion or 0,
|
||||
total = chunk.usage.total_tokens or chunk.usage.total or 0,
|
||||
}
|
||||
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 }
|
||||
elseif (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
|
||||
uses.prompt_tokens_details = { cached_tokens = chunk.usage.prompt_cache_hit_tokens }
|
||||
end
|
||||
end
|
||||
|
||||
if not chunk.choices or #chunk.choices == 0 then
|
||||
@ -119,10 +138,15 @@ function adapter.transform_stream_chunk(raw_chunk)
|
||||
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 = (fr ~= nil)
|
||||
done = (finish ~= nil)
|
||||
}
|
||||
if finish then
|
||||
unified.finish_reason = finish
|
||||
end
|
||||
if uses ~= nil then
|
||||
unified.usage = uses
|
||||
end
|
||||
@ -135,4 +159,15 @@ function adapter.transform_stream_chunk(raw_chunk)
|
||||
return json.encode(unified)
|
||||
end
|
||||
|
||||
return adapter
|
||||
-- 错误收敛:kimi 网关为 OpenAI 风格 {error:{message,...}}
|
||||
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" and type(e.message) == "string" then
|
||||
return e.message
|
||||
end
|
||||
return nil
|
||||
end
|
||||
|
||||
return adapter
|
||||
|
||||
@ -32,6 +32,18 @@ function adapter.transform_response(raw_body)
|
||||
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
|
||||
end
|
||||
|
||||
if type(resp.choices) == "table" and #resp.choices > 0 then
|
||||
@ -73,6 +85,13 @@ function adapter.transform_stream_chunk(raw_chunk)
|
||||
completion = chunk.usage.completion_tokens or chunk.usage.completion or 0,
|
||||
total = chunk.usage.total_tokens or chunk.usage.total or 0,
|
||||
}
|
||||
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 }
|
||||
elseif (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
|
||||
uses.prompt_tokens_details = { cached_tokens = chunk.usage.prompt_cache_hit_tokens }
|
||||
end
|
||||
end
|
||||
|
||||
if not chunk.choices or #chunk.choices == 0 then
|
||||
@ -86,10 +105,15 @@ function adapter.transform_stream_chunk(raw_chunk)
|
||||
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 = (fr ~= nil)
|
||||
done = (finish ~= nil)
|
||||
}
|
||||
if finish then
|
||||
unified.finish_reason = finish
|
||||
end
|
||||
if uses ~= nil then
|
||||
unified.usage = uses
|
||||
end
|
||||
@ -102,4 +126,14 @@ function adapter.transform_stream_chunk(raw_chunk)
|
||||
return json.encode(unified)
|
||||
end
|
||||
|
||||
-- 错误收敛:Mistral 用顶层 {message="...", type="..."}(非嵌套 error)
|
||||
function adapter.transform_error(status, body)
|
||||
local ok, resp = pcall(json.decode, body)
|
||||
if not ok or type(resp) ~= "table" then return nil end
|
||||
if type(resp.message) == "string" then
|
||||
return resp.message
|
||||
end
|
||||
return nil
|
||||
end
|
||||
|
||||
return adapter
|
||||
|
||||
@ -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
|
||||
|
||||
@ -58,13 +102,39 @@ function adapter.transform_response(raw_body)
|
||||
local unified = {
|
||||
content = "",
|
||||
finish_reason = resp.done_reason or "",
|
||||
tool_calls = {},
|
||||
-- key must be token_usage to match Go's UnifiedResponse json tag
|
||||
token_usage = { prompt = p, completion = c, total = p + c }
|
||||
}
|
||||
|
||||
if resp.message then
|
||||
unified.content = resp.message.content or ""
|
||||
-- Non-streaming tool calls were previously dropped: the field was
|
||||
-- initialized to an empty table and never filled, while
|
||||
-- transform_stream_chunk handled them. A non-streaming agent turn thus
|
||||
-- looked like a plain answer and the tool loop stopped.
|
||||
if type(resp.message.tool_calls) == "table" and #resp.message.tool_calls > 0 then
|
||||
local tcs = {}
|
||||
for _, tc in ipairs(resp.message.tool_calls) do
|
||||
local fn = tc["function"] or {}
|
||||
-- Ollama sends arguments as an object already
|
||||
local args = fn.arguments
|
||||
if type(args) == "string" then
|
||||
local aok, decoded = pcall(json.decode, args)
|
||||
args = aok and decoded or {}
|
||||
elseif type(args) ~= "table" then
|
||||
args = {}
|
||||
end
|
||||
table.insert(tcs, {
|
||||
id = tc.id or ("call_" .. #tcs),
|
||||
type = tc.type or "function",
|
||||
name = fn.name or "",
|
||||
arguments = args
|
||||
})
|
||||
end
|
||||
unified.tool_calls = tcs
|
||||
-- Ollama reports done_reason "stop" alongside tool calls
|
||||
unified.finish_reason = "tool_calls"
|
||||
end
|
||||
end
|
||||
|
||||
return json.encode(unified)
|
||||
@ -85,16 +155,27 @@ function adapter.transform_stream_chunk(raw_chunk)
|
||||
end
|
||||
end
|
||||
|
||||
-- Ollama done_reason -> OpenAI finish_reason ("length" 透传,其余归一 stop)
|
||||
local finish = nil
|
||||
if chunk.done then
|
||||
if chunk.done_reason == "length" then
|
||||
finish = "length"
|
||||
else
|
||||
finish = "stop"
|
||||
end
|
||||
end
|
||||
|
||||
if not chunk.message then
|
||||
if uses ~= nil then
|
||||
return json.encode({ content = "", done = true, usage = uses })
|
||||
return json.encode({ content = "", done = true, finish_reason = finish, usage = uses })
|
||||
end
|
||||
return ""
|
||||
return json.encode({ content = "", done = true, finish_reason = finish })
|
||||
end
|
||||
|
||||
local unified = {
|
||||
content = chunk.message.content or "",
|
||||
done = chunk.done or false
|
||||
done = chunk.done or false,
|
||||
finish_reason = finish
|
||||
}
|
||||
if uses ~= nil then
|
||||
unified.usage = uses
|
||||
@ -120,4 +201,15 @@ function adapter.transform_stream_chunk(raw_chunk)
|
||||
return json.encode(unified)
|
||||
end
|
||||
|
||||
-- 错误收敛:Ollama 常见 {error:"..."} 字符串(新版本也有对象形态)
|
||||
function adapter.transform_error(status, body)
|
||||
local ok, resp = pcall(json.decode, body)
|
||||
if not ok or type(resp) ~= "table" then return nil end
|
||||
if type(resp.error) == "string" then return resp.error end
|
||||
if type(resp.error) == "table" and type(resp.error.message) == "string" then
|
||||
return resp.error.message
|
||||
end
|
||||
return nil
|
||||
end
|
||||
|
||||
return adapter
|
||||
|
||||
@ -5,6 +5,36 @@ adapter.version = "2.0.0"
|
||||
adapter.endpoint = "/chat/completions"
|
||||
adapter.headers = {}
|
||||
|
||||
-- Claude-behind-OpenAI upstreams (tabitoken, 扇贝, …) convert /chat/completions
|
||||
-- to the Anthropic Messages API internally, so they inherit Anthropic's
|
||||
-- tool id rule: ^[a-zA-Z0-9_-]{1,64}$, enforced by rejecting the WHOLE request
|
||||
-- with "Invalid tool use format" / "tool_use.id: String should match pattern".
|
||||
-- Plain OpenAI has no such rule, so an OpenAI-compatible model can mint an id
|
||||
-- like "bash:0" (observed from moonshotai/kimi-k3). In a fan-out router that id
|
||||
-- is replayed to every other source, so one such id kills every Claude slot at
|
||||
-- once and an AUTO request falls through all tiers.
|
||||
--
|
||||
-- safe_tool_id is pure and deterministic, so a tool_calls entry and its
|
||||
-- matching tool_call_id are rewritten identically within one request. A
|
||||
-- rewritten id keeps an 8-hex digest of the ORIGINAL id, without which two
|
||||
-- distinct ids could collapse into one ("a:b" and "a_b") and become an
|
||||
-- unpaired/duplicate tool call. Already-legal ids are returned untouched, so
|
||||
-- well-behaved traffic is byte-identical to before.
|
||||
-- (anthropic.lua carries the same helper; Lua adapters have no shared prelude.)
|
||||
local TOOL_ID_MAX = 64
|
||||
|
||||
local function safe_tool_id(id)
|
||||
if type(id) ~= "string" or id == "" then return id end
|
||||
local clean = string.gsub(id, "[^A-Za-z0-9_-]", "_")
|
||||
if clean == id and #clean <= TOOL_ID_MAX then
|
||||
return clean
|
||||
end
|
||||
local digest = string.sub(sha256_hex(id), 1, 8)
|
||||
local keep = TOOL_ID_MAX - #digest - 1
|
||||
if #clean > keep then clean = string.sub(clean, 1, keep) end
|
||||
return clean .. "_" .. digest
|
||||
end
|
||||
|
||||
-- OpenAI /chat/completions format (pass-through, strip provider-specific fields)
|
||||
function adapter.transform_request(raw_body)
|
||||
local ok, req = pcall(json.decode, raw_body)
|
||||
@ -14,6 +44,16 @@ function adapter.transform_request(raw_body)
|
||||
if req.messages then
|
||||
for _, msg in ipairs(req.messages) do
|
||||
msg.reasoning_content = nil
|
||||
if msg.tool_call_id ~= nil then
|
||||
msg.tool_call_id = safe_tool_id(msg.tool_call_id)
|
||||
end
|
||||
if type(msg.tool_calls) == "table" then
|
||||
for _, tc in ipairs(msg.tool_calls) do
|
||||
if type(tc) == "table" and tc.id ~= nil then
|
||||
tc.id = safe_tool_id(tc.id)
|
||||
end
|
||||
end
|
||||
end
|
||||
end
|
||||
end
|
||||
return json.encode(req)
|
||||
@ -33,6 +73,22 @@ function adapter.transform_response(raw_body)
|
||||
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
|
||||
-- Cache passthrough: OpenAI v2 prompt_tokens_details.cached_tokens
|
||||
-- and DeepSeek-legacy prompt_cache_hit/miss_tokens. dsh reads
|
||||
-- prompt_tokens_details.cached_tokens (falls back to the legacy
|
||||
-- standalone field), so both shapes reach clients.
|
||||
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
|
||||
end
|
||||
|
||||
if type(resp.choices) == "table" and #resp.choices > 0 then
|
||||
@ -48,7 +104,7 @@ function adapter.transform_response(raw_body)
|
||||
local args_ok, args = pcall(json.decode, tc["function"].arguments)
|
||||
if not args_ok then args = {} end
|
||||
table.insert(tcs, {
|
||||
id = tc.id,
|
||||
id = safe_tool_id(tc.id),
|
||||
type = tc.type or "function",
|
||||
name = tc["function"].name,
|
||||
arguments = args
|
||||
@ -77,6 +133,13 @@ function adapter.transform_stream_chunk(raw_chunk)
|
||||
completion = chunk.usage.completion_tokens or chunk.usage.completion or 0,
|
||||
total = chunk.usage.total_tokens or chunk.usage.total or 0,
|
||||
}
|
||||
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 }
|
||||
elseif (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
|
||||
uses.prompt_tokens_details = { cached_tokens = chunk.usage.prompt_cache_hit_tokens }
|
||||
end
|
||||
end
|
||||
|
||||
if not chunk.choices or #chunk.choices == 0 then
|
||||
@ -90,10 +153,15 @@ function adapter.transform_stream_chunk(raw_chunk)
|
||||
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 = (fr ~= nil)
|
||||
done = (finish ~= nil)
|
||||
}
|
||||
if finish then
|
||||
unified.finish_reason = finish
|
||||
end
|
||||
if uses ~= nil then
|
||||
unified.usage = uses
|
||||
end
|
||||
@ -101,9 +169,60 @@ function adapter.transform_stream_chunk(raw_chunk)
|
||||
unified.reasoning_content = delta.reasoning_content
|
||||
end
|
||||
if delta.tool_calls then
|
||||
-- Sanitize on the way OUT too: an id this upstream happily minted (it
|
||||
-- does not validate them) becomes a landmine once the client replays
|
||||
-- it to a Claude upstream. Only the first fragment of a streamed call
|
||||
-- carries an id; later argument fragments have none and are untouched.
|
||||
for _, tc in ipairs(delta.tool_calls) do
|
||||
if type(tc) == "table" and tc.id ~= nil then
|
||||
tc.id = safe_tool_id(tc.id)
|
||||
end
|
||||
end
|
||||
unified.tool_calls = delta.tool_calls
|
||||
end
|
||||
return json.encode(unified)
|
||||
end
|
||||
|
||||
-- 错误收敛:标准 OpenAI 信封 {error:{message,...}}
|
||||
function adapter.transform_error(status, body)
|
||||
-- Non-JSON body (nginx HTML error pages, plain text): extract a short
|
||||
-- human-readable reason instead of letting the raw body reach the log.
|
||||
local ok, resp = pcall(json.decode, body)
|
||||
if not ok or type(resp) ~= "table" then
|
||||
-- HTML error page: pull the <title> text (e.g. "413 Request Entity Too Large")
|
||||
local title = string.match(body or "", "<title>(.-)</title>")
|
||||
if title and title ~= "" then return title end
|
||||
-- Bare text: first non-empty line, capped
|
||||
local line = string.match(body or "", "^%s*([^\r\n]+)")
|
||||
if line and line ~= "" and not string.match(line, "^<") then
|
||||
return string.sub(line, 1, 200)
|
||||
end
|
||||
return nil
|
||||
end
|
||||
|
||||
-- Standard OpenAI envelope: {error:{message,...}} or {error:"..."}
|
||||
local e = resp.error
|
||||
if type(e) == "table" and type(e.message) == "string" then
|
||||
return e.message
|
||||
end
|
||||
if type(e) == "string" then return e end
|
||||
|
||||
-- Flat envelope used by many OpenAI-compatible gateways:
|
||||
-- {"code":20012,"message":"Model does not exist..."}
|
||||
-- {"code":"INVALID_API_KEY","message":"Invalid API key"}
|
||||
if type(resp.message) == "string" and resp.message ~= "" then
|
||||
if resp.code ~= nil then
|
||||
return tostring(resp.code) .. ": " .. resp.message
|
||||
end
|
||||
return resp.message
|
||||
end
|
||||
|
||||
-- Some gateways use {detail:"..."} (FastAPI style)
|
||||
if type(resp.detail) == "string" and resp.detail ~= "" then
|
||||
return resp.detail
|
||||
end
|
||||
|
||||
return nil
|
||||
end
|
||||
|
||||
return adapter
|
||||
|
||||
@ -4,18 +4,40 @@ adapter.name = "opencode"
|
||||
adapter.version = "1.0.0"
|
||||
adapter.endpoint = "/chat/completions"
|
||||
|
||||
-- opencode.ai zen 网关按 User-Agent 指纹识别官方客户端并把请求分到免费额度池;
|
||||
-- 非官方 UA(curl/Go 默认等)会被分到匿名池并触发 FreeUsageLimitError。
|
||||
-- 因此固定发送 opencode 客户端的 UA;配合源配置 api_key: "public"(官方无 key
|
||||
-- 客户端实际发送 Bearer public)即可走免费池。
|
||||
-- opencode.ai zen 网关按客户端指纹(UA + x-opencode-* 头)路由请求池:
|
||||
-- 缺少 x-opencode-client/session/request/project 身份头的请求会被判为匿名
|
||||
-- 客户端,x-preview-f-free 等模型在带 tools 时上游直接失败——流式返回单
|
||||
-- chunk "finish_reason":"network_error" 空 content,非流式返回 503
|
||||
-- "Endpoint is unavailable",表现为"空回复"。因此除 UA 外必须附带
|
||||
-- x-opencode-* 身份头;session/request 每请求派生唯一值(见 build_headers)。
|
||||
adapter.headers = {
|
||||
["User-Agent"] = "opencode/0.1.0",
|
||||
["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
|
||||
|
||||
function adapter.build_headers(meta)
|
||||
local ts = tostring(meta.timestamp or "")
|
||||
return {
|
||||
["User-Agent"] = adapter.headers["User-Agent"],
|
||||
["x-opencode-client"] = "cli",
|
||||
-- project 固定:同网关实例共享一个工作区身份
|
||||
["x-opencode-project"] = string.sub(sha256_hex("llmsproxy|" .. (meta.source and meta.source.name or "")), 1, 32),
|
||||
["x-opencode-session"] = rand_id("ses_", "session|" .. ts),
|
||||
["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 的
|
||||
-- 消息整条丢弃,避免上游 "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。
|
||||
@ -32,6 +54,14 @@ function adapter.transform_request(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 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
|
||||
@ -50,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
|
||||
@ -78,14 +125,27 @@ function adapter.transform_response(raw_body)
|
||||
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
|
||||
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 ""
|
||||
if ch.message.reasoning_content then
|
||||
unified.reasoning_content = ch.message.reasoning_content
|
||||
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 = {}
|
||||
@ -123,6 +183,13 @@ function adapter.transform_stream_chunk(raw_chunk)
|
||||
completion = chunk.usage.completion_tokens or chunk.usage.completion or 0,
|
||||
total = chunk.usage.total_tokens or chunk.usage.total or 0,
|
||||
}
|
||||
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 }
|
||||
elseif (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
|
||||
uses.prompt_tokens_details = { cached_tokens = chunk.usage.prompt_cache_hit_tokens }
|
||||
end
|
||||
end
|
||||
|
||||
if not chunk.choices or #chunk.choices == 0 then
|
||||
@ -136,15 +203,22 @@ function adapter.transform_stream_chunk(raw_chunk)
|
||||
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 = (fr ~= nil)
|
||||
done = (finish ~= nil)
|
||||
}
|
||||
if finish then
|
||||
unified.finish_reason = finish
|
||||
end
|
||||
if uses ~= nil then
|
||||
unified.usage = uses
|
||||
end
|
||||
if delta.reasoning_content then
|
||||
unified.reasoning_content = delta.reasoning_content
|
||||
-- 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
|
||||
@ -153,4 +227,17 @@ function adapter.transform_stream_chunk(raw_chunk)
|
||||
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
|
||||
|
||||
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
|
||||
174
internal/lua/adapters/sensenova.lua
Normal file
174
internal/lua/adapters/sensenova.lua
Normal file
@ -0,0 +1,174 @@
|
||||
-- sensenova API format adapter
|
||||
-- streaming: delta only has reasoning_content, no content field
|
||||
-- non-streaming: has both content and reasoning_content
|
||||
|
||||
local adapter = {}
|
||||
adapter.name = "sensenova"
|
||||
adapter.version = "1.0.0"
|
||||
adapter.endpoint = "/chat/completions"
|
||||
adapter.headers = {}
|
||||
|
||||
-- Same as openai - strip provider-specific fields
|
||||
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
|
||||
if req.messages then
|
||||
for _, msg in ipairs(req.messages) do
|
||||
msg.reasoning_content = nil
|
||||
end
|
||||
end
|
||||
return json.encode(req)
|
||||
end
|
||||
|
||||
-- Same as openai - extract content from response
|
||||
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
|
||||
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 ""
|
||||
if ch.message.reasoning_content then
|
||||
unified.reasoning_content = ch.message.reasoning_content
|
||||
elseif ch.message.reasoning then
|
||||
-- sensenova-6.8-flash-lite 用 reasoning 字段而不是 reasoning_content
|
||||
unified.reasoning_content = ch.message.reasoning
|
||||
end
|
||||
-- Tool calls MUST be forwarded. Dropping them while keeping
|
||||
-- finish_reason="tool_calls" makes the client replay an assistant
|
||||
-- message whose function name/arguments are empty, and sensenova
|
||||
-- then rejects the next turn with
|
||||
-- 400 invalid tool_call function, function/name/arguments cannot be empty
|
||||
-- i.e. a tool-using conversation dies on its second request.
|
||||
if type(ch.message.tool_calls) == "table" and #ch.message.tool_calls > 0 then
|
||||
local tcs = {}
|
||||
for _, tc in ipairs(ch.message.tool_calls) do
|
||||
local fn = tc["function"] or {}
|
||||
-- arguments arrives as a JSON *string* on the wire; the
|
||||
-- unified shape expects a decoded object.
|
||||
local args = fn.arguments
|
||||
if type(args) == "string" then
|
||||
local args_ok, decoded = pcall(json.decode, args)
|
||||
args = args_ok and decoded or {}
|
||||
elseif type(args) ~= "table" 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
|
||||
|
||||
-- Sensenova-specific stream handling
|
||||
-- Upstream puts content in reasoning_content only (no content field)
|
||||
-- Also sends finish_reason="" (empty string) on every chunk
|
||||
function adapter.transform_stream_chunk(raw_chunk)
|
||||
local ok, chunk = pcall(json.decode, raw_chunk)
|
||||
if not ok then return "" end
|
||||
|
||||
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,
|
||||
}
|
||||
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 }
|
||||
elseif (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
|
||||
uses.prompt_tokens_details = { cached_tokens = chunk.usage.prompt_cache_hit_tokens }
|
||||
end
|
||||
end
|
||||
|
||||
if not chunk.choices or #chunk.choices == 0 then
|
||||
if uses ~= nil then
|
||||
return json.encode({ usage = uses, done = false })
|
||||
end
|
||||
return ""
|
||||
end
|
||||
|
||||
local delta = chunk.choices[1].delta or {}
|
||||
|
||||
-- Sensenova: delta has reasoning_content (deepseek-v4-flash) or reasoning
|
||||
-- (sensenova-6.8-flash-lite) but no content field
|
||||
local content = delta.content or ""
|
||||
if content == "" and (delta.reasoning_content or delta.reasoning) then
|
||||
content = delta.reasoning_content or delta.reasoning
|
||||
end
|
||||
|
||||
-- Sensenova: finish_reason is "" on every chunk, "stop" on last
|
||||
local fr = chunk.choices[1].finish_reason
|
||||
local done = (fr == "stop" or fr == "length")
|
||||
|
||||
local unified = {
|
||||
content = content,
|
||||
done = done,
|
||||
}
|
||||
if chunk.choices[1].finish_reason and chunk.choices[1].finish_reason ~= "" then
|
||||
unified.finish_reason = chunk.choices[1].finish_reason
|
||||
end
|
||||
if delta.tool_calls then
|
||||
unified.tool_calls = delta.tool_calls
|
||||
end
|
||||
if uses ~= nil then
|
||||
unified.usage = uses
|
||||
end
|
||||
|
||||
return json.encode(unified)
|
||||
end
|
||||
|
||||
-- 错误收敛:sensenova 为 OpenAI 风格 {error:{message,...}};
|
||||
-- 配额类错误单独点出便于客户端识别重置周期。
|
||||
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 status == 429 and type(e.code) == "string"
|
||||
and e.code == "insufficient_quota" then
|
||||
return "workspace quota exhausted (resets periodically)"
|
||||
end
|
||||
if type(e.message) == "string" then return e.message end
|
||||
return nil
|
||||
end
|
||||
|
||||
return adapter
|
||||
286
internal/lua/adapters/trae.lua
Normal file
286
internal/lua/adapters/trae.lua
Normal file
@ -0,0 +1,286 @@
|
||||
-- trae 源适配器(trae-local-api 的 OpenAI 兼容代理)
|
||||
-- trae-local-api 运行在 http://127.0.0.1:19900,接 Trae CN 账号的云资源。
|
||||
-- 协议:OpenAI /v1/chat/completions
|
||||
-- 特性:
|
||||
-- - 排队:上游 busy 时排队,可能长时间无响应
|
||||
-- - 非流式请求偶发返回 SSE 数据(upstream bug)
|
||||
-- - 支持 thinking,由上游模型自动决定是否开启
|
||||
|
||||
local adapter = {}
|
||||
adapter.name = "trae"
|
||||
adapter.version = "1.0.0"
|
||||
adapter.endpoint = "/v1/chat/completions"
|
||||
adapter.headers = {}
|
||||
|
||||
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
|
||||
if req.messages then
|
||||
for _, msg in ipairs(req.messages) do
|
||||
msg.reasoning_content = nil
|
||||
end
|
||||
end
|
||||
return json.encode(req)
|
||||
end
|
||||
|
||||
-- parse_text_tool_calls extracts tool calls that an upstream emitted as PLAIN
|
||||
-- TEXT instead of using the OpenAI tool_calls field.
|
||||
--
|
||||
-- trae-local-api's OpenAI endpoint (/v1/chat/completions) does not read the
|
||||
-- request's `tools` array at all, so the relayed model is never told the tool
|
||||
-- schema; it falls back to printing
|
||||
-- <tool_call>
|
||||
-- {"name": "get_weather", "arguments": {"location": "北京"}}
|
||||
-- </tool_call>
|
||||
-- into message.content, leaves message.tool_calls null, and reports
|
||||
-- finish_reason="stop". A client following the OpenAI contract therefore never
|
||||
-- sees a tool call: the agent loop terminates unexpectedly mid-conversation
|
||||
-- (and a hand-written replay produces an empty function name next turn).
|
||||
--
|
||||
-- Both tag spellings are accepted: the same codebase's Anthropic endpoint
|
||||
-- instructs models to emit <toolcall>, and models mix the two. Key names vary
|
||||
-- too (arguments / params / input), so all are tried.
|
||||
--
|
||||
-- Returns (tool_calls_array_or_nil, content_with_blocks_removed).
|
||||
local function parse_text_tool_calls(content)
|
||||
if type(content) ~= "string" or content == "" then return nil, content end
|
||||
|
||||
local tcs = {}
|
||||
local idx = 0
|
||||
|
||||
local function collect(pattern)
|
||||
for payload in content:gmatch(pattern) do
|
||||
local ok, obj = pcall(json.decode, payload)
|
||||
if ok and type(obj) == "table" then
|
||||
-- some models wrap it as {"function":{"name":..,"arguments":..}}
|
||||
local fn = obj["function"]
|
||||
local name = obj.name or (type(fn) == "table" and fn.name) or nil
|
||||
if name then
|
||||
local args = obj.arguments or obj.params or obj.input
|
||||
if args == nil and type(fn) == "table" then
|
||||
args = fn.arguments or fn.params
|
||||
end
|
||||
if type(args) == "string" then
|
||||
local aok, decoded = pcall(json.decode, args)
|
||||
args = aok and decoded or {}
|
||||
elseif type(args) ~= "table" then
|
||||
args = {}
|
||||
end
|
||||
idx = idx + 1
|
||||
table.insert(tcs, {
|
||||
id = obj.id or ("call_text_" .. idx),
|
||||
type = "function",
|
||||
name = name,
|
||||
arguments = args
|
||||
})
|
||||
end
|
||||
end
|
||||
end
|
||||
end
|
||||
|
||||
-- `<tag ...>` allows attributes; %s* handles the "</tool_call >" spacing
|
||||
-- that trae-local-api's own prompt example uses.
|
||||
collect("<tool_call[^>]*>%s*(.-)%s*</tool_call%s*>")
|
||||
collect("<toolcall[^>]*>%s*(.-)%s*</toolcall%s*>")
|
||||
|
||||
-- Legacy / mimicked form: models that saw past assistant turns folded as
|
||||
-- [Called tool: name({...})]
|
||||
-- tend to emit the same shape instead of a real block. Recover it so
|
||||
-- the agent loop keeps working; otherwise the client receives plain text
|
||||
-- where a structured tool call should be. Capture the name and the JSON
|
||||
-- object literal independently (the JSON is the only {...} run here).
|
||||
if #tcs == 0 and content:find("%[Called tool") then
|
||||
for name, argsraw in content:gmatch("%[Called tool%s*:%s*([%w_%-%.]+)%s*%((%b{})%s*%)%s*%]") do
|
||||
local aok, decoded = pcall(json.decode, argsraw)
|
||||
idx = idx + 1
|
||||
table.insert(tcs, {
|
||||
id = "call_legacy_" .. idx,
|
||||
type = "function",
|
||||
name = name,
|
||||
arguments = aok and decoded or {}
|
||||
})
|
||||
end
|
||||
end
|
||||
|
||||
if #tcs == 0 then return nil, content end
|
||||
|
||||
-- drop the blocks from user-visible content; keep any surrounding prose
|
||||
local stripped = content:gsub("<tool_call[^>]*>%s*.-%s*</tool_call%s*>", "")
|
||||
stripped = stripped:gsub("<toolcall[^>]*>%s*.-%s*</toolcall%s*>", "")
|
||||
stripped = stripped:gsub("%[Called tool%s*:%s*[%w_%-%.]+%s*%(%b{}%s*%)%s*%]", "")
|
||||
stripped = stripped:gsub("^%s+", ""):gsub("%s+$", "")
|
||||
return tcs, stripped
|
||||
end
|
||||
|
||||
function adapter.transform_response(raw_body)
|
||||
-- trae-local-api 偶发在非流式请求中返回 SSE 格式数据,
|
||||
-- 表现为多个 data: {...} 行或混合了 reasoning_chunk 等。
|
||||
-- 剥掉 data: 前缀,取最后一个完整的 JSON 块(通常是
|
||||
-- 最终的 stop chunk 或 usage chunk)。
|
||||
local use_json = raw_body
|
||||
if raw_body:match("^data:") or raw_body:match("\ndata:") then
|
||||
-- 取最后一个 data: 行的 JSON
|
||||
local last = nil
|
||||
for line in raw_body:gmatch("data: ([^\n\r]+)") do
|
||||
local trimmed = line:match("^%s*(.-)%s*$")
|
||||
if trimmed and trimmed ~= "[DONE]" then
|
||||
local ok2, obj = pcall(json.decode, trimmed)
|
||||
if ok2 and type(obj) == "table" then
|
||||
last = obj
|
||||
end
|
||||
end
|
||||
end
|
||||
if last then
|
||||
use_json = json.encode(last)
|
||||
end
|
||||
-- 如果解析失败,回退到原始 body
|
||||
end
|
||||
|
||||
local ok, resp = pcall(json.decode, use_json)
|
||||
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
|
||||
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 ""
|
||||
if ch.message.reasoning_content then
|
||||
unified.reasoning_content = ch.message.reasoning_content
|
||||
end
|
||||
if type(ch.message.tool_calls) == "table" and #ch.message.tool_calls > 0 then
|
||||
local tcs = {}
|
||||
for _, tc in ipairs(ch.message.tool_calls) do
|
||||
local fn = tc["function"] or {}
|
||||
local args = fn.arguments
|
||||
if type(args) == "string" then
|
||||
local args_ok, decoded = pcall(json.decode, args)
|
||||
args = args_ok and decoded or {}
|
||||
elseif type(args) ~= "table" 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 ""
|
||||
|
||||
-- Fallback: trae-local-api's OpenAI endpoint drops the request's
|
||||
-- `tools` array entirely, so the relayed model is never told the tool
|
||||
-- schema and instead PRINTS a <tool_call>{...}</tool_call> block into
|
||||
-- content, leaving message.tool_calls null and finish_reason="stop".
|
||||
-- A client following the OpenAI contract then sees a normal completion
|
||||
-- and its agent loop terminates mid-conversation. Recover the
|
||||
-- structured call so the loop can continue.
|
||||
if unified.tool_calls == nil then
|
||||
local recovered, cleaned = parse_text_tool_calls(unified.content)
|
||||
if recovered then
|
||||
unified.tool_calls = recovered
|
||||
unified.content = cleaned
|
||||
unified.finish_reason = "tool_calls"
|
||||
end
|
||||
end
|
||||
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
|
||||
|
||||
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,
|
||||
}
|
||||
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 }
|
||||
elseif (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
|
||||
uses.prompt_tokens_details = { cached_tokens = chunk.usage.prompt_cache_hit_tokens }
|
||||
end
|
||||
end
|
||||
|
||||
if not chunk.choices or #chunk.choices == 0 then
|
||||
if uses ~= nil then
|
||||
return json.encode({ usage = uses, done = false })
|
||||
end
|
||||
return ""
|
||||
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
|
||||
if delta.reasoning_content then
|
||||
unified.reasoning_content = delta.reasoning_content
|
||||
end
|
||||
if delta.tool_calls then
|
||||
unified.tool_calls = delta.tool_calls
|
||||
end
|
||||
return json.encode(unified)
|
||||
end
|
||||
|
||||
function adapter.transform_error(status, body)
|
||||
local ok, resp = pcall(json.decode, body)
|
||||
if not ok or type(resp) ~= "table" then
|
||||
-- Non-JSON body: extract title from HTML or first line
|
||||
local title = string.match(body or "", "<title>(.-)</title>")
|
||||
if title and title ~= "" then return title end
|
||||
local line = string.match(body or "", "^%s*([^\r\n]+)")
|
||||
if line and line ~= "" and not string.match(line, "^<") then
|
||||
return string.sub(line, 1, 200)
|
||||
end
|
||||
return nil
|
||||
end
|
||||
local e = resp.error
|
||||
if type(e) == "table" and type(e.message) == "string" then
|
||||
-- trae quota / rate limit annotations
|
||||
if e.code then
|
||||
return tostring(e.code) .. ": " .. e.message
|
||||
end
|
||||
return e.message
|
||||
end
|
||||
if type(e) == "string" then return e end
|
||||
if type(resp.message) == "string" and resp.message ~= "" then
|
||||
if resp.code ~= nil then
|
||||
return tostring(resp.code) .. ": " .. resp.message
|
||||
end
|
||||
return resp.message
|
||||
end
|
||||
return nil
|
||||
end
|
||||
|
||||
return adapter
|
||||
353
internal/lua/toolcall_preservation_test.go
Normal file
353
internal/lua/toolcall_preservation_test.go
Normal file
@ -0,0 +1,353 @@
|
||||
package lua
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"strings"
|
||||
"testing"
|
||||
)
|
||||
|
||||
// TestAdaptersPreserveToolCalls is the fleet-wide guard for the
|
||||
// "orphaned tool result" class of bug: when an agent client replays a turn
|
||||
// whose assistant message carries tool_calls, EVERY adapter must forward both
|
||||
// the call and the attribution of its result. Dropping the call (or the
|
||||
// tool_call_id / tool_name that ties the result to it) makes the model see a
|
||||
// result for a call it never made and re-issue the same call forever.
|
||||
//
|
||||
// The input deliberately uses content:[] — the shape most agent clients emit
|
||||
// for a tool-calling assistant turn with no text.
|
||||
func TestAdaptersPreserveToolCalls(t *testing.T) {
|
||||
vm := NewVM(freshAdapterDir(t))
|
||||
if err := vm.Start(); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer vm.Stop()
|
||||
|
||||
body := `{"model":"m","messages":[
|
||||
{"role":"user","content":"read /tmp/x"},
|
||||
{"role":"assistant","content":[],"tool_calls":[{"id":"call_abc","type":"function","function":{"name":"read","arguments":"{\"path\":\"/tmp/x\"}"}}]},
|
||||
{"role":"tool","tool_call_id":"call_abc","content":"hello world"},
|
||||
{"role":"user","content":"summarize"}
|
||||
]}`
|
||||
const (
|
||||
callID = "call_abc"
|
||||
fnName = "read"
|
||||
toolOut = "hello world"
|
||||
lastUser = "summarize"
|
||||
)
|
||||
|
||||
// Adapters that identify a tool result by the call id.
|
||||
byID := []string{"openai", "deepseek", "github", "groq", "kimicode", "mistral",
|
||||
"sensenova", "trae", "agentrouter", "opencode", "anthropic"}
|
||||
|
||||
for _, a := range byID {
|
||||
out, err := vm.Transform(a, "transform_request", body)
|
||||
if err != nil {
|
||||
t.Errorf("%s: transform_request: %v", a, err)
|
||||
continue
|
||||
}
|
||||
if !strings.Contains(out, callID) {
|
||||
t.Errorf("%s dropped the tool call id %q: %s", a, callID, out)
|
||||
}
|
||||
if !strings.Contains(out, fnName) {
|
||||
t.Errorf("%s dropped the function name %q: %s", a, fnName, out)
|
||||
}
|
||||
if !strings.Contains(out, toolOut) {
|
||||
t.Errorf("%s dropped the tool result: %s", a, out)
|
||||
}
|
||||
if !strings.Contains(out, lastUser) {
|
||||
t.Errorf("%s dropped the final user turn: %s", a, out)
|
||||
}
|
||||
}
|
||||
|
||||
// Ollama identifies the result with tool_name (its API has no tool_call_id
|
||||
// on assistant turns) and takes arguments as an object, not a JSON string.
|
||||
out, err := vm.Transform("ollama", "transform_request", body)
|
||||
if err != nil {
|
||||
t.Fatalf("ollama: %v", err)
|
||||
}
|
||||
for _, want := range []string{`"tool_calls"`, fnName, `"tool_name"`, toolOut} {
|
||||
if !strings.Contains(out, want) {
|
||||
t.Errorf("ollama missing %s: %s", want, out)
|
||||
}
|
||||
}
|
||||
|
||||
// Gemini identifies the result by function NAME and has no call id in its
|
||||
// wire format; both directions must be present as functionCall /
|
||||
// functionResponse parts, and the system role must move to
|
||||
// systemInstruction rather than sitting in contents.
|
||||
gBody := `{"model":"m","messages":[
|
||||
{"role":"system","content":"be brief"},
|
||||
{"role":"user","content":"read /tmp/x"},
|
||||
{"role":"assistant","content":[],"tool_calls":[{"id":"call_abc","type":"function","function":{"name":"read","arguments":"{\"path\":\"/tmp/x\"}"}}]},
|
||||
{"role":"tool","tool_call_id":"call_abc","content":"hello world"}
|
||||
]}`
|
||||
gout, err := vm.Transform("gemini", "transform_request", gBody)
|
||||
if err != nil {
|
||||
t.Fatalf("gemini: %v", err)
|
||||
}
|
||||
for _, want := range []string{`"functionCall"`, `"functionResponse"`, fnName, toolOut, `"systemInstruction"`} {
|
||||
if !strings.Contains(gout, want) {
|
||||
t.Errorf("gemini missing %s: %s", want, gout)
|
||||
}
|
||||
}
|
||||
// No message may keep the OpenAI-only roles inside contents.
|
||||
var greq struct {
|
||||
Contents []struct {
|
||||
Role string `json:"role"`
|
||||
} `json:"contents"`
|
||||
}
|
||||
if err := json.Unmarshal([]byte(gout), &greq); err != nil {
|
||||
t.Fatalf("gemini unmarshal: %v (%s)", err, gout)
|
||||
}
|
||||
for _, c := range greq.Contents {
|
||||
if c.Role != "user" && c.Role != "model" {
|
||||
t.Errorf("gemini contents carry an invalid role %q: %s", c.Role, gout)
|
||||
}
|
||||
}
|
||||
|
||||
// Negative control: a text-only turn must still be forwarded verbatim.
|
||||
plain := `{"model":"m","messages":[{"role":"user","content":"just text"}]}`
|
||||
for _, a := range append(byID, "ollama", "gemini") {
|
||||
pout, err := vm.Transform(a, "transform_request", plain)
|
||||
if err != nil {
|
||||
t.Errorf("%s plain: %v", a, err)
|
||||
continue
|
||||
}
|
||||
if !strings.Contains(pout, "just text") {
|
||||
t.Errorf("%s dropped plain text: %s", a, pout)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// TestOpenCodeStreamOptionsOnlyWhenStreaming pins the fix for
|
||||
// "stream_options should be set along with stream": OpenCode Go (and other
|
||||
// strict OpenAI-compatible upstreams) reject a non-streaming request that
|
||||
// carries stream_options, which broke every non-stream call through the
|
||||
// opencode adapter.
|
||||
func TestOpenCodeStreamOptionsOnlyWhenStreaming(t *testing.T) {
|
||||
vm := NewVM(freshAdapterDir(t))
|
||||
if err := vm.Start(); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer vm.Stop()
|
||||
|
||||
// non-streaming: stream_options must be absent entirely
|
||||
out, err := vm.Transform("opencode", "transform_request",
|
||||
`{"model":"m","messages":[{"role":"user","content":"hi"}]}`)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if strings.Contains(out, "stream_options") {
|
||||
t.Fatalf("non-streaming request must not carry stream_options: %s", out)
|
||||
}
|
||||
|
||||
// streaming: stream_options.include_usage must be set
|
||||
sout, err := vm.Transform("opencode", "transform_request",
|
||||
`{"model":"m","stream":true,"messages":[{"role":"user","content":"hi"}]}`)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if !strings.Contains(sout, "stream_options") || !strings.Contains(sout, "include_usage") {
|
||||
t.Fatalf("streaming request must carry stream_options.include_usage: %s", sout)
|
||||
}
|
||||
}
|
||||
|
||||
// TestOpenCodeGoVsZenReasoning pins the one behavioural difference that made a
|
||||
// shared adapter wrong: OpenCode Go's thinking mode REQUIRES the assistant
|
||||
// turn's reasoning_content to be echoed back ("The `reasoning_content` in the
|
||||
// thinking mode must be passed back to the API"), while the Zen free pool must
|
||||
// not receive it. Hence two purpose-built adapters.
|
||||
func TestOpenCodeGoVsZenReasoning(t *testing.T) {
|
||||
vm := NewVM(freshAdapterDir(t))
|
||||
if err := vm.Start(); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer vm.Stop()
|
||||
|
||||
body := `{"model":"m","messages":[
|
||||
{"role":"user","content":"hi"},
|
||||
{"role":"assistant","content":"","reasoning_content":"I should read the file.","tool_calls":[{"id":"call_1","type":"function","function":{"name":"read","arguments":"{\"path\":\"/x\"}"}}]},
|
||||
{"role":"tool","tool_call_id":"call_1","content":"r"}
|
||||
]}`
|
||||
|
||||
// Go: reasoning_content must survive, and the tool call must stay paired.
|
||||
gout, err := vm.Transform("opencodego", "transform_request", body)
|
||||
if err != nil {
|
||||
t.Fatalf("opencodego: %v", err)
|
||||
}
|
||||
if !strings.Contains(gout, "I should read the file.") {
|
||||
t.Errorf("opencodego must pass reasoning_content back (thinking mode requires it): %s", gout)
|
||||
}
|
||||
if !strings.Contains(gout, "call_1") || !strings.Contains(gout, "tool_call_id") {
|
||||
t.Errorf("opencodego must keep the tool call paired with its result: %s", gout)
|
||||
}
|
||||
|
||||
// Zen: reasoning_content is stripped.
|
||||
zout, err := vm.Transform("opencodezen", "transform_request", body)
|
||||
if err != nil {
|
||||
t.Fatalf("opencodezen: %v", err)
|
||||
}
|
||||
if strings.Contains(zout, "I should read the file.") {
|
||||
t.Errorf("opencodezen must strip reasoning_content: %s", zout)
|
||||
}
|
||||
if !strings.Contains(zout, "call_1") {
|
||||
t.Errorf("opencodezen must still keep the tool call: %s", zout)
|
||||
}
|
||||
|
||||
// An assistant turn that carries tool_calls but no reasoning_content must
|
||||
// get an empty one injected: Go validates that field on tool-calling turns
|
||||
// and 400s when it is absent, while most clients never send it.
|
||||
// Verified against the live endpoint that an empty string passes.
|
||||
bare := `{"model":"m","messages":[
|
||||
{"role":"user","content":"go"},
|
||||
{"role":"assistant","content":"","tool_calls":[{"id":"c1","type":"function","function":{"name":"read","arguments":"{}"}}]},
|
||||
{"role":"tool","tool_call_id":"c1","content":"r"}
|
||||
]}`
|
||||
bout, err := vm.Transform("opencodego", "transform_request", bare)
|
||||
if err != nil {
|
||||
t.Fatalf("opencodego bare: %v", err)
|
||||
}
|
||||
if !strings.Contains(bout, `"reasoning_content":""`) {
|
||||
t.Errorf("opencodego must inject an empty reasoning_content on a tool-calling turn: %s", bout)
|
||||
}
|
||||
// A plain assistant turn (no tool calls) must NOT gain the field.
|
||||
plainAsst := `{"model":"m","messages":[{"role":"user","content":"hi"},{"role":"assistant","content":"hello"}]}`
|
||||
pout, err := vm.Transform("opencodego", "transform_request", plainAsst)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if strings.Contains(pout, "reasoning_content") {
|
||||
t.Errorf("opencodego must not inject reasoning_content on a plain assistant turn: %s", pout)
|
||||
}
|
||||
|
||||
// Both must still omit stream_options on non-streaming requests.
|
||||
for _, a := range []string{"opencodego", "opencodezen"} {
|
||||
out, err := vm.Transform(a, "transform_request", `{"model":"m","messages":[{"role":"user","content":"hi"}]}`)
|
||||
if err != nil {
|
||||
t.Fatalf("%s: %v", a, err)
|
||||
}
|
||||
if strings.Contains(out, "stream_options") {
|
||||
t.Errorf("%s must not send stream_options when not streaming: %s", a, out)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// TestOpenCodeSessionIsStableForCache pins the session-derivation contract that
|
||||
// makes the upstream prefix cache usable.
|
||||
//
|
||||
// Measured against the live OpenCode Go endpoint: the prefix cache is
|
||||
// SESSION-scoped. Replaying the same 6032-token prompt hits 5888 cached tokens
|
||||
// when the session id is fixed, and 0 when it changes per request. Deriving the
|
||||
// session from meta.timestamp (the previous behaviour) meant every request was a
|
||||
// new session, so the cache could never hit.
|
||||
//
|
||||
// The session must be stable WITHIN one conversation and differ BETWEEN
|
||||
// conversations. Since generic OpenAI clients send no session id at all
|
||||
// (verified by capturing real agent traffic: no user / session_id /
|
||||
// conversation_id / metadata in the body, no session header), the conversation
|
||||
// is fingerprinted by its FIRST user message — history is replayed every turn,
|
||||
// so that message is invariant for the life of the conversation.
|
||||
func TestOpenCodeSessionIsStableForCache(t *testing.T) {
|
||||
for _, name := range []string{"opencodego", "opencodezen"} {
|
||||
vm := NewVM(freshAdapterDir(t))
|
||||
if err := vm.Start(); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
meta := func(ts int, body string) map[string]interface{} {
|
||||
return map[string]interface{}{
|
||||
"timestamp": ts,
|
||||
"body": body,
|
||||
"source": map[string]interface{}{"name": "somesource", "meta": nil},
|
||||
}
|
||||
}
|
||||
|
||||
// Same conversation, later turn: history grew, first user msg unchanged.
|
||||
h1, err := vm.BuildHeaders(name, meta(1000,
|
||||
`{"messages":[{"role":"user","content":"fix the bug"},{"role":"assistant","content":"looking"}]}`))
|
||||
if err != nil {
|
||||
t.Fatalf("%s: %v", name, err)
|
||||
}
|
||||
h2, err := vm.BuildHeaders(name, meta(2000,
|
||||
`{"messages":[{"role":"user","content":"fix the bug"},{"role":"assistant","content":"looking"},{"role":"user","content":"and tests"}]}`))
|
||||
if err != nil {
|
||||
t.Fatalf("%s: %v", name, err)
|
||||
}
|
||||
// Different conversation.
|
||||
h3, err := vm.BuildHeaders(name, meta(3000,
|
||||
`{"messages":[{"role":"user","content":"an unrelated task"}]}`))
|
||||
if err != nil {
|
||||
t.Fatalf("%s: %v", name, err)
|
||||
}
|
||||
|
||||
if h1["x-opencode-session"] == "" {
|
||||
t.Fatalf("%s: session header must be set", name)
|
||||
}
|
||||
if h1["x-opencode-session"] != h2["x-opencode-session"] {
|
||||
t.Errorf("%s: x-opencode-session must be STABLE across turns of one conversation "+
|
||||
"(%q vs %q) — the prefix cache is session-scoped", name, h1["x-opencode-session"], h2["x-opencode-session"])
|
||||
}
|
||||
if h1["x-opencode-session"] == h3["x-opencode-session"] {
|
||||
t.Errorf("%s: a different conversation must get a different session", name)
|
||||
}
|
||||
if h1["x-opencode-request"] == h2["x-opencode-request"] {
|
||||
t.Errorf("%s: x-opencode-request must differ per request", name)
|
||||
}
|
||||
// No user message at all falls back to per-source stability.
|
||||
h4, _ := vm.BuildHeaders(name, meta(4000, `{"messages":[{"role":"system","content":"s"}]}`))
|
||||
h5, _ := vm.BuildHeaders(name, meta(5000, `{"messages":[{"role":"system","content":"s"}]}`))
|
||||
if h4["x-opencode-session"] != h5["x-opencode-session"] {
|
||||
t.Errorf("%s: sessionless requests must still be stable per source", name)
|
||||
}
|
||||
vm.Stop()
|
||||
}
|
||||
}
|
||||
|
||||
// 客户端自带的会话标识(pi 的 x-session-affinity)必须优先于会话指纹:
|
||||
// 它是平台真实会话 id,历史压缩后也不会漂移,且天然按会话隔离。
|
||||
func TestOpenCodePrefersClientSessionID(t *testing.T) {
|
||||
for _, name := range []string{"opencodego", "opencodezen"} {
|
||||
vm := NewVM(freshAdapterDir(t))
|
||||
if err := vm.Start(); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
meta := func(body, cs string) map[string]interface{} {
|
||||
return map[string]interface{}{
|
||||
"timestamp": 1000,
|
||||
"body": body,
|
||||
"client_session": cs,
|
||||
"source": map[string]interface{}{"name": "somesource", "meta": nil},
|
||||
}
|
||||
}
|
||||
|
||||
// 同一条会话:历史在长,但客户端会话 id 不变。
|
||||
a, err := vm.BuildHeaders(name, meta(`{"messages":[{"role":"user","content":"fix the bug"}]}`, "cs-aaa"))
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
b, err := vm.BuildHeaders(name, meta(`{"messages":[{"role":"user","content":"a totally different opener"}]}`, "cs-aaa"))
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if a["x-opencode-session"] != b["x-opencode-session"] {
|
||||
t.Errorf("%s: 客户端会话 id 相同就必须同会话(内容无关): %q vs %q",
|
||||
name, a["x-opencode-session"], b["x-opencode-session"])
|
||||
}
|
||||
|
||||
// 不同会话。
|
||||
c, _ := vm.BuildHeaders(name, meta(`{"messages":[{"role":"user","content":"fix the bug"}]}`, "cs-bbb"))
|
||||
if a["x-opencode-session"] == c["x-opencode-session"] {
|
||||
t.Errorf("%s: 不同客户端会话必须不同会号", name)
|
||||
}
|
||||
|
||||
// 客户端没给会话 id 时,回落到指纹(同内容 -> 同会号)。
|
||||
d, _ := vm.BuildHeaders(name, meta(`{"messages":[{"role":"user","content":"fix the bug"}]}`, ""))
|
||||
e, _ := vm.BuildHeaders(name, meta(`{"messages":[{"role":"user","content":"fix the bug"}]}`, ""))
|
||||
if d["x-opencode-session"] != e["x-opencode-session"] {
|
||||
t.Errorf("%s: 无客户端会号时应回落到稳定指纹", name)
|
||||
}
|
||||
if d["x-opencode-session"] == a["x-opencode-session"] {
|
||||
t.Errorf("%s: 客户端会号与指纹两条路必须产出不同会号", name)
|
||||
}
|
||||
vm.Stop()
|
||||
}
|
||||
}
|
||||
177
internal/lua/usage_cost_test.go
Normal file
177
internal/lua/usage_cost_test.go
Normal file
@ -0,0 +1,177 @@
|
||||
package lua
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"testing"
|
||||
)
|
||||
|
||||
// OpenCode Zen/Go 在 usage 里报告输出有多少是思考 token,并在**顶层**用一个
|
||||
// **字符串**给出本次费用。前者决定「completion_tokens 里多少是可见回答」,
|
||||
// 后者是唯一的费用信号。两者都曾被适配器丢掉。
|
||||
func TestOpenCodeCostAndReasoningPassthrough(t *testing.T) {
|
||||
for _, name := range []string{"opencodego", "opencodezen"} {
|
||||
vm := NewVM(freshAdapterDir(t))
|
||||
if err := vm.Start(); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
// ---- 非流式 ----
|
||||
resp := `{"id":"x","object":"chat.completion","model":"deepseek-v4.1-flash",
|
||||
"choices":[{"index":0,"message":{"role":"assistant","content":"hi","reasoning_content":"think"},
|
||||
"finish_reason":"stop"}],
|
||||
"usage":{"prompt_tokens":37,"completion_tokens":40,"total_tokens":77,
|
||||
"prompt_cache_hit_tokens":0,"prompt_cache_miss_tokens":37,
|
||||
"prompt_tokens_details":{"cached_tokens":0},
|
||||
"completion_tokens_details":{"reasoning_tokens":40}},
|
||||
"cost":"0"}`
|
||||
out, err := vm.Transform(name, "transform_response", resp)
|
||||
if err != nil {
|
||||
t.Fatalf("%s transform_response: %v", name, err)
|
||||
}
|
||||
var r struct {
|
||||
TokenUsage struct {
|
||||
Completion int `json:"completion"`
|
||||
Details *struct {
|
||||
ReasoningTokens int `json:"reasoning_tokens"`
|
||||
} `json:"completion_tokens_details"`
|
||||
} `json:"token_usage"`
|
||||
Cost string `json:"cost"`
|
||||
}
|
||||
if err := json.Unmarshal([]byte(out), &r); err != nil {
|
||||
t.Fatalf("%s unmarshal: %v (%s)", name, err, out)
|
||||
}
|
||||
if r.Cost != "0" {
|
||||
t.Errorf("%s: cost 未透传(上游给的是字符串 \"0\"): %q", name, r.Cost)
|
||||
}
|
||||
if r.TokenUsage.Details == nil || r.TokenUsage.Details.ReasoningTokens != 40 {
|
||||
t.Errorf("%s: completion_tokens_details.reasoning_tokens 未透传: %s", name, out)
|
||||
}
|
||||
// 前提校验:reasoning_tokens 确实 ≤ completion_tokens(全为思考)。
|
||||
if r.TokenUsage.Details != nil && r.TokenUsage.Details.ReasoningTokens > r.TokenUsage.Completion {
|
||||
t.Errorf("%s: reasoning 不应超过 completion", name)
|
||||
}
|
||||
|
||||
// ---- 流式:非零费用也要透传 ----
|
||||
chunk := `{"id":"x","object":"chat.completion.chunk","choices":[],
|
||||
"usage":{"prompt_tokens":10,"completion_tokens":5,"total_tokens":15,
|
||||
"completion_tokens_details":{"reasoning_tokens":4}},"cost":"0.0012"}`
|
||||
co, err := vm.Transform(name, "transform_stream_chunk", chunk)
|
||||
if err != nil {
|
||||
t.Fatalf("%s transform_stream_chunk: %v", name, err)
|
||||
}
|
||||
var c struct {
|
||||
Cost string `json:"cost"`
|
||||
Usage *struct {
|
||||
Details *struct {
|
||||
ReasoningTokens int `json:"reasoning_tokens"`
|
||||
} `json:"completion_tokens_details"`
|
||||
} `json:"usage"`
|
||||
}
|
||||
if err := json.Unmarshal([]byte(co), &c); err != nil {
|
||||
t.Fatalf("%s unmarshal chunk: %v (%s)", name, err, co)
|
||||
}
|
||||
if c.Usage == nil || c.Usage.Details == nil || c.Usage.Details.ReasoningTokens != 4 {
|
||||
t.Errorf("%s: 流式 reasoning_tokens 未透传: %s", name, co)
|
||||
}
|
||||
if c.Cost != "0.0012" {
|
||||
t.Errorf("%s: 流式 cost 未透传: %s", name, co)
|
||||
}
|
||||
|
||||
// ---- 独立的 cost 帧(上游就这么发:choices 为空、只有 cost)----
|
||||
only := `{"choices":[],"cost":"0"}`
|
||||
oo, err := vm.Transform(name, "transform_stream_chunk", only)
|
||||
if err != nil {
|
||||
t.Fatalf("%s cost-only chunk: %v", name, err)
|
||||
}
|
||||
if oo == "" {
|
||||
t.Fatalf("%s: 纯 cost 帧被整个丢弃(客户端将看不到费用)", name)
|
||||
}
|
||||
var oc struct {
|
||||
Cost string `json:"cost"`
|
||||
}
|
||||
if err := json.Unmarshal([]byte(oo), &oc); err != nil {
|
||||
t.Fatalf("%s unmarshal cost chunk: %v (%s)", name, err, oo)
|
||||
}
|
||||
if oc.Cost != "0" {
|
||||
t.Errorf("%s: 纯 cost 帧内容丢失: %s", name, oo)
|
||||
}
|
||||
|
||||
// ---- 回归:没有 cost / 没有 details 时不得凭空造字段 ----
|
||||
bare := `{"choices":[{"index":0,"delta":{"content":"hi"},"finish_reason":"stop"}],
|
||||
"usage":{"prompt_tokens":3,"completion_tokens":1,"total_tokens":4}}`
|
||||
bo, _ := vm.Transform(name, "transform_response", bare)
|
||||
if got := bo; len(got) > 0 && containsAll(got, `"cost"`) {
|
||||
t.Errorf("%s: 上游没给 cost 却透传了该字段: %s", name, got)
|
||||
}
|
||||
|
||||
vm.Stop()
|
||||
}
|
||||
}
|
||||
|
||||
func containsAll(s string, subs ...string) bool {
|
||||
for _, sub := range subs {
|
||||
found := false
|
||||
for i := 0; i+len(sub) <= len(s); i++ {
|
||||
if s[i:i+len(sub)] == sub {
|
||||
found = true
|
||||
break
|
||||
}
|
||||
}
|
||||
if !found {
|
||||
return false
|
||||
}
|
||||
}
|
||||
return true
|
||||
}
|
||||
|
||||
// 流式与非流式对同一上游切片的产出必须一致:命中缓存时上游同时给
|
||||
// prompt_tokens_details.cached_tokens 和独立的 hit/miss 字段,两边都要透传。
|
||||
func TestOpenCodeStreamCacheFieldsMatchNonStream(t *testing.T) {
|
||||
for _, name := range []string{"opencodego", "opencodezen"} {
|
||||
vm := NewVM(freshAdapterDir(t))
|
||||
if err := vm.Start(); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
type cache struct {
|
||||
Details *struct {
|
||||
CachedTokens int `json:"cached_tokens"`
|
||||
} `json:"prompt_tokens_details"`
|
||||
Hit int `json:"prompt_cache_hit_tokens"`
|
||||
Miss int `json:"prompt_cache_miss_tokens"`
|
||||
}
|
||||
usage := `"usage":{"prompt_tokens":5732,"completion_tokens":20,"total_tokens":5752,
|
||||
"prompt_cache_hit_tokens":5504,"prompt_cache_miss_tokens":228,
|
||||
"prompt_tokens_details":{"cached_tokens":5504}}`
|
||||
|
||||
// 非流式
|
||||
nout, _ := vm.Transform(name, "transform_response",
|
||||
`{"choices":[{"index":0,"message":{"role":"assistant","content":"hi"},"finish_reason":"stop"}],`+usage+`}`)
|
||||
var nr struct {
|
||||
U cache `json:"token_usage"`
|
||||
}
|
||||
if err := json.Unmarshal([]byte(nout), &nr); err != nil {
|
||||
t.Fatalf("%s: %v", name, err)
|
||||
}
|
||||
// 流式
|
||||
sout, _ := vm.Transform(name, "transform_stream_chunk",
|
||||
`{"choices":[],`+usage+`}`)
|
||||
var sr struct {
|
||||
U cache `json:"usage"`
|
||||
}
|
||||
if err := json.Unmarshal([]byte(sout), &sr); err != nil {
|
||||
t.Fatalf("%s chunk: %v (%s)", name, err, sout)
|
||||
}
|
||||
|
||||
if nr.U.Hit != 5504 || nr.U.Miss != 228 {
|
||||
t.Errorf("%s: 非流式 hit/miss 异常: %+v", name, nr.U)
|
||||
}
|
||||
if sr.U.Hit != nr.U.Hit || sr.U.Miss != nr.U.Miss {
|
||||
t.Errorf("%s: 流式 hit/miss 与非流式不一致 (stream %d/%d vs non-stream %d/%d): %s",
|
||||
name, sr.U.Hit, sr.U.Miss, nr.U.Hit, nr.U.Miss, sout)
|
||||
}
|
||||
if sr.U.Details == nil || sr.U.Details.CachedTokens != 5504 {
|
||||
t.Errorf("%s: 流式 cached_tokens 丢失: %s", name, sout)
|
||||
}
|
||||
vm.Stop()
|
||||
}
|
||||
}
|
||||
@ -4,9 +4,15 @@
|
||||
//
|
||||
// The runtime uses LuaJIT (github.com/aarzilli/golua) so every adapter runs on
|
||||
// an independent interpreter state. Because a lua.State is NOT goroutine-safe,
|
||||
// each adapter owns a pool of workers (states); the pool size equals the sum
|
||||
// of max_concurrent over all sources that use the adapter, and a worker is
|
||||
// checked out for the duration of a single hook call.
|
||||
// each adapter owns a pool of workers (states) and a worker is checked out for
|
||||
// the duration of a single hook call.
|
||||
//
|
||||
// Pools are ELASTIC: the sum of max_concurrent over the sources using an adapter
|
||||
// is a ceiling, not a preallocation. States are booted on demand in batches
|
||||
// sized by that ceiling, and reclaimed by a single background janitor at a rate
|
||||
// scaled inversely to the adapter's live connection count. An idle gateway
|
||||
// therefore holds close to zero Lua states regardless of how much aggregate
|
||||
// concurrency is configured. See adapterPool for the sizing rules.
|
||||
package lua
|
||||
|
||||
import (
|
||||
@ -23,6 +29,7 @@ import (
|
||||
"strconv"
|
||||
"strings"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
// golua exposes the LuaJIT C API. Build with -tags luajit so the cgo
|
||||
// LDFLAGS resolve to -lluajit-5.1 (see golua's lua.go).
|
||||
@ -64,7 +71,20 @@ type staticInfo struct {
|
||||
}
|
||||
|
||||
// adapterPool owns the worker states for one adapter. Idle workers are held in
|
||||
// a slice guarded by a cond; up to target workers are created lazily.
|
||||
// a slice guarded by a cond; workers are created lazily on demand and reclaimed
|
||||
// by the VM janitor once demand drops.
|
||||
//
|
||||
// Sizing is elastic rather than fixed (plan 阶段 2). Three inputs drive it:
|
||||
//
|
||||
// - maxW — the adapter's maximum concurrency (Σ max_concurrent of every source
|
||||
// using it). It is the hard ceiling AND sets the growth step, so a hot,
|
||||
// high-concurrency adapter ramps in batches while a low-concurrency one
|
||||
// creeps up one state at a time.
|
||||
// - created / len(idle) — how many states already exist, i.e. how much room is
|
||||
// left to grow and how much is available to reclaim.
|
||||
// - inUse — how many states are checked out right now, i.e. the live
|
||||
// connection count for this adapter. It sets the shrink step: the fewer
|
||||
// connections, the more aggressively idle states are released.
|
||||
type adapterPool struct {
|
||||
name string
|
||||
script string
|
||||
@ -73,24 +93,164 @@ type adapterPool struct {
|
||||
mu sync.Mutex
|
||||
cond *sync.Cond
|
||||
idle []*worker
|
||||
created int
|
||||
target int
|
||||
created int // live states: len(idle) + inUse (+ states being booted)
|
||||
inUse int // states checked out right now == live connections
|
||||
waiting int // goroutines blocked waiting for a free state
|
||||
maxW int // hard ceiling = Σ max_concurrent of the sources using this adapter
|
||||
used bool
|
||||
closed bool
|
||||
|
||||
lastGrow time.Time
|
||||
idleRounds int // consecutive janitor rounds that saw reclaimable slack
|
||||
peakInUse int // high-water mark of inUse, for observability
|
||||
}
|
||||
|
||||
const (
|
||||
// growDivisor turns the adapter's max concurrency into the growth step CAP:
|
||||
// step <= clamp(ceil(maxW / growDivisor), 1, growStepCap). The actual batch
|
||||
// is additionally bounded by how many callers are queued (see growPlanLocked),
|
||||
// so a wide adapter may ramp in big jumps but never warms states nobody is
|
||||
// waiting for.
|
||||
growDivisor = 8
|
||||
growStepCap = 8
|
||||
// growCooldown keeps a burst of misses from batching repeatedly while the
|
||||
// previous batch is still booting.
|
||||
growCooldown = time.Second
|
||||
// residentWorkers is how many states an adapter keeps warm once it has
|
||||
// served at least one request. Booting is milliseconds, but keeping one warm
|
||||
// removes that from the critical path of the next request. Adapters that
|
||||
// were never used keep nothing.
|
||||
residentWorkers = 1
|
||||
// idleHeadroom is the spare idle state kept above the live connection count,
|
||||
// so the next concurrent request does not have to boot.
|
||||
idleHeadroom = 1
|
||||
// shrinkInterval is how often the VM janitor reclaims idle states.
|
||||
shrinkInterval = 30 * time.Second
|
||||
// shrinkGraceRounds is how many consecutive janitor rounds must see slack
|
||||
// before anything is released. It guards against tearing a pool down between
|
||||
// two bursts — NOT against ordinary traffic: a pool serving one request at a
|
||||
// time still has reclaimable slack, and requiring "no traffic at all" would
|
||||
// pin a burst's leftover states forever on any busy gateway. The counter is
|
||||
// therefore reset by GROWTH (real contention), not by a mere checkout.
|
||||
shrinkGraceRounds = 2
|
||||
)
|
||||
|
||||
// PoolStats is a snapshot of one adapter pool's elastic sizing, surfaced by the
|
||||
// status API so the algorithm is observable instead of a black box.
|
||||
type PoolStats struct {
|
||||
Name string `json:"name"`
|
||||
Created int `json:"created"`
|
||||
Idle int `json:"idle"`
|
||||
InUse int `json:"in_use"`
|
||||
Waiting int `json:"waiting"`
|
||||
Max int `json:"max"`
|
||||
Resident int `json:"resident"`
|
||||
GrowStep int `json:"grow_step"`
|
||||
ShrinkStep int `json:"shrink_step"`
|
||||
PeakInUse int `json:"peak_in_use"`
|
||||
}
|
||||
|
||||
func newAdapterPool(name, script string, static staticInfo) *adapterPool {
|
||||
p := &adapterPool{name: name, script: script, static: static, target: 1}
|
||||
p := &adapterPool{name: name, script: script, static: static, maxW: 1}
|
||||
p.cond = sync.NewCond(&p.mu)
|
||||
return p
|
||||
}
|
||||
|
||||
func (p *adapterPool) setTarget(n int) {
|
||||
// setMax updates the pool's hard ceiling (Σ max_concurrent of its sources).
|
||||
// Lowering it does not kill live states: the janitor reclaims the excess.
|
||||
func (p *adapterPool) setMax(n int) {
|
||||
if n < 1 {
|
||||
n = 1
|
||||
}
|
||||
p.mu.Lock()
|
||||
p.target = n
|
||||
p.maxW = n
|
||||
p.cond.Broadcast()
|
||||
p.mu.Unlock()
|
||||
}
|
||||
|
||||
func ceilDiv(a, b int) int {
|
||||
if b <= 0 {
|
||||
return a
|
||||
}
|
||||
return (a + b - 1) / b
|
||||
}
|
||||
|
||||
// growStepLocked is how many states to warm on a miss, derived from the
|
||||
// adapter's maximum concurrency. Caller holds p.mu.
|
||||
func (p *adapterPool) growStepLocked() int {
|
||||
step := ceilDiv(p.maxW, growDivisor)
|
||||
if step < 1 {
|
||||
return 1
|
||||
}
|
||||
if step > growStepCap {
|
||||
return growStepCap
|
||||
}
|
||||
return step
|
||||
}
|
||||
|
||||
// residentLocked is the floor the janitor will not shrink below. Caller holds p.mu.
|
||||
func (p *adapterPool) residentLocked() int {
|
||||
if !p.used {
|
||||
return 0
|
||||
}
|
||||
if residentWorkers > p.maxW {
|
||||
return p.maxW
|
||||
}
|
||||
return residentWorkers
|
||||
}
|
||||
|
||||
// shrinkStepLocked is how many idle states to release this round. The step is
|
||||
// inversely proportional to the live connection count: with no connections the
|
||||
// slack collapses in one round, while a busy adapter gives up one state at a
|
||||
// time so the hot path keeps its warm states. Caller holds p.mu.
|
||||
func (p *adapterPool) shrinkStepLocked() int {
|
||||
excess := p.created - p.keepLocked()
|
||||
if excess <= 0 {
|
||||
return 0
|
||||
}
|
||||
step := ceilDiv(excess, 1+p.inUse)
|
||||
if step > excess {
|
||||
step = excess
|
||||
}
|
||||
if step > len(p.idle) {
|
||||
step = len(p.idle)
|
||||
}
|
||||
if step < 0 {
|
||||
return 0
|
||||
}
|
||||
return step
|
||||
}
|
||||
|
||||
// keepLocked is the number of live states this pool should retain right now.
|
||||
// Caller holds p.mu.
|
||||
func (p *adapterPool) keepLocked() int {
|
||||
keep := p.residentLocked()
|
||||
if h := p.inUse + idleHeadroom; h > keep {
|
||||
keep = h
|
||||
}
|
||||
if keep > p.maxW {
|
||||
keep = p.maxW
|
||||
}
|
||||
return keep
|
||||
}
|
||||
|
||||
func (p *adapterPool) stats() PoolStats {
|
||||
p.mu.Lock()
|
||||
defer p.mu.Unlock()
|
||||
return PoolStats{
|
||||
Name: p.name,
|
||||
Created: p.created,
|
||||
Idle: len(p.idle),
|
||||
InUse: p.inUse,
|
||||
Waiting: p.waiting,
|
||||
Max: p.maxW,
|
||||
Resident: p.residentLocked(),
|
||||
GrowStep: p.growStepLocked(),
|
||||
ShrinkStep: p.shrinkStepLocked(),
|
||||
PeakInUse: p.peakInUse,
|
||||
}
|
||||
}
|
||||
|
||||
// boot creates a fresh LuaJIT state, registers the shared globals and executes
|
||||
// the adapter script, storing the returned adapter table in the state.
|
||||
func (p *adapterPool) boot() (*worker, error) {
|
||||
@ -110,6 +270,96 @@ func (p *adapterPool) boot() (*worker, error) {
|
||||
return &worker{L: L}, nil
|
||||
}
|
||||
|
||||
// growPlanLocked reserves capacity for a batch prewarm and returns how many
|
||||
// EXTRA states to boot in the background (the caller already reserved one for
|
||||
// itself).
|
||||
//
|
||||
// Two independent limits apply, and BOTH matter:
|
||||
//
|
||||
// - the adapter's max concurrency caps the step (growStepLocked), so a
|
||||
// high-concurrency adapter is allowed to ramp in bigger jumps than a
|
||||
// narrow one;
|
||||
// - the number of goroutines actually BLOCKED waiting for a state bounds it
|
||||
// to real demand.
|
||||
//
|
||||
// Sizing on the ceiling alone over-provisions badly: an adapter with
|
||||
// max_concurrent=76 has a step of 8, so two concurrent requests would warm 8
|
||||
// states, and the janitor would throw 7 of them away a minute later — boot,
|
||||
// discard, repeat, with RSS oscillating for no benefit. Prewarming only for
|
||||
// goroutines that are genuinely queued keeps the ramp cheap without the churn.
|
||||
//
|
||||
// Caller holds p.mu.
|
||||
func (p *adapterPool) growPlanLocked(contended bool) int {
|
||||
if !contended || p.waiting <= 0 {
|
||||
return 0
|
||||
}
|
||||
if time.Since(p.lastGrow) < growCooldown {
|
||||
return 0
|
||||
}
|
||||
extra := p.growStepLocked() - 1
|
||||
if extra > p.waiting {
|
||||
extra = p.waiting // never warm more than the queue needs
|
||||
}
|
||||
if extra <= 0 {
|
||||
return 0
|
||||
}
|
||||
if room := p.maxW - p.created; extra > room {
|
||||
extra = room
|
||||
}
|
||||
if extra <= 0 {
|
||||
return 0
|
||||
}
|
||||
// Reserve immediately so concurrent acquires respect the ceiling while the
|
||||
// batch is still booting.
|
||||
p.created += extra
|
||||
p.lastGrow = time.Now()
|
||||
return extra
|
||||
}
|
||||
|
||||
// prewarm boots n reserved states and publishes them as idle. A boot failure
|
||||
// releases its reservation; it is not fatal, since the synchronous path
|
||||
// surfaces adapter errors already.
|
||||
func (p *adapterPool) prewarm(n int) {
|
||||
for i := 0; i < n; i++ {
|
||||
p.mu.Lock()
|
||||
closed := p.closed
|
||||
p.mu.Unlock()
|
||||
if closed {
|
||||
p.releaseReservation(n - i)
|
||||
return
|
||||
}
|
||||
w, err := p.boot()
|
||||
if err != nil {
|
||||
p.releaseReservation(n - i)
|
||||
return
|
||||
}
|
||||
p.mu.Lock()
|
||||
if p.closed {
|
||||
p.created--
|
||||
p.mu.Unlock()
|
||||
w.close()
|
||||
p.releaseReservation(n - i - 1)
|
||||
return
|
||||
}
|
||||
p.idle = append(p.idle, w)
|
||||
p.cond.Signal()
|
||||
p.mu.Unlock()
|
||||
}
|
||||
}
|
||||
|
||||
func (p *adapterPool) releaseReservation(n int) {
|
||||
if n <= 0 {
|
||||
return
|
||||
}
|
||||
p.mu.Lock()
|
||||
p.created -= n
|
||||
if p.created < 0 {
|
||||
p.created = 0
|
||||
}
|
||||
p.cond.Broadcast()
|
||||
p.mu.Unlock()
|
||||
}
|
||||
|
||||
func (p *adapterPool) acquire() (*worker, error) {
|
||||
p.mu.Lock()
|
||||
for {
|
||||
@ -120,32 +370,68 @@ func (p *adapterPool) acquire() (*worker, error) {
|
||||
if n := len(p.idle); n > 0 {
|
||||
w := p.idle[n-1]
|
||||
p.idle = p.idle[:n-1]
|
||||
p.checkoutLocked()
|
||||
p.mu.Unlock()
|
||||
return w, nil
|
||||
}
|
||||
if p.created < p.target {
|
||||
if p.created < p.maxW {
|
||||
// Every existing state busy on a miss = real concurrency, not a
|
||||
// sequential caller reusing one warm state.
|
||||
contended := p.created > 0 && p.inUse >= p.created
|
||||
p.created++
|
||||
break
|
||||
// Growing is the signal that capacity is genuinely short, so any
|
||||
// pending shrink decision is stale.
|
||||
p.idleRounds = 0
|
||||
extra := p.growPlanLocked(contended)
|
||||
p.mu.Unlock()
|
||||
if extra > 0 {
|
||||
go p.prewarm(extra)
|
||||
}
|
||||
w, err := p.boot()
|
||||
if err != nil {
|
||||
p.mu.Lock()
|
||||
p.created--
|
||||
p.cond.Signal()
|
||||
p.mu.Unlock()
|
||||
return nil, err
|
||||
}
|
||||
p.mu.Lock()
|
||||
p.checkoutLocked()
|
||||
p.mu.Unlock()
|
||||
return w, nil
|
||||
}
|
||||
p.waiting++
|
||||
p.cond.Wait()
|
||||
p.waiting--
|
||||
}
|
||||
p.mu.Unlock()
|
||||
}
|
||||
|
||||
w, err := p.boot()
|
||||
if err != nil {
|
||||
p.mu.Lock()
|
||||
p.created--
|
||||
p.cond.Signal()
|
||||
p.mu.Unlock()
|
||||
return nil, err
|
||||
// checkoutLocked accounts one state as handed out. Caller holds p.mu.
|
||||
func (p *adapterPool) checkoutLocked() {
|
||||
p.inUse++
|
||||
p.used = true
|
||||
if p.inUse > p.peakInUse {
|
||||
p.peakInUse = p.inUse
|
||||
}
|
||||
return w, nil
|
||||
}
|
||||
|
||||
func (p *adapterPool) release(w *worker) {
|
||||
w.L.SetTop(0)
|
||||
p.mu.Lock()
|
||||
if p.inUse > 0 {
|
||||
p.inUse--
|
||||
}
|
||||
if p.closed {
|
||||
p.created--
|
||||
p.mu.Unlock()
|
||||
w.close()
|
||||
return
|
||||
}
|
||||
// A lowered ceiling is honoured on the spot rather than waiting for the
|
||||
// janitor: returning a state the pool may no longer keep closes it.
|
||||
if p.created > p.maxW {
|
||||
p.created--
|
||||
p.cond.Signal()
|
||||
p.mu.Unlock()
|
||||
w.close()
|
||||
return
|
||||
@ -155,11 +441,42 @@ func (p *adapterPool) release(w *worker) {
|
||||
p.mu.Unlock()
|
||||
}
|
||||
|
||||
// reclaim releases idle states down to the current keep target, using a step
|
||||
// scaled by the live connection count. It returns how many states were closed.
|
||||
func (p *adapterPool) reclaim() int {
|
||||
p.mu.Lock()
|
||||
if p.closed {
|
||||
p.mu.Unlock()
|
||||
return 0
|
||||
}
|
||||
step := p.shrinkStepLocked()
|
||||
if step <= 0 {
|
||||
p.idleRounds = 0
|
||||
p.mu.Unlock()
|
||||
return 0
|
||||
}
|
||||
p.idleRounds++
|
||||
if p.idleRounds < shrinkGraceRounds {
|
||||
p.mu.Unlock()
|
||||
return 0
|
||||
}
|
||||
victims := append([]*worker(nil), p.idle[len(p.idle)-step:]...)
|
||||
p.idle = p.idle[:len(p.idle)-step]
|
||||
p.created -= step
|
||||
p.idleRounds = 0
|
||||
p.mu.Unlock()
|
||||
for _, w := range victims {
|
||||
w.close()
|
||||
}
|
||||
return step
|
||||
}
|
||||
|
||||
func (p *adapterPool) shutdown() {
|
||||
p.mu.Lock()
|
||||
p.closed = true
|
||||
idle := p.idle
|
||||
p.idle = nil
|
||||
p.created -= len(idle)
|
||||
p.cond.Broadcast()
|
||||
p.mu.Unlock()
|
||||
for _, w := range idle {
|
||||
@ -173,13 +490,84 @@ type VM struct {
|
||||
mu sync.RWMutex
|
||||
dir string
|
||||
pools map[string]*adapterPool
|
||||
|
||||
// janitor drives elastic shrinking. One goroutine serves every pool (not
|
||||
// one per adapter), so an idle gateway costs a single sleeping goroutine.
|
||||
janitorOnce sync.Once
|
||||
janitorStop chan struct{}
|
||||
janitorDone chan struct{}
|
||||
}
|
||||
|
||||
func NewVM(dir string) *VM {
|
||||
return &VM{dir: dir, pools: map[string]*adapterPool{}}
|
||||
return &VM{
|
||||
dir: dir,
|
||||
pools: map[string]*adapterPool{},
|
||||
janitorStop: make(chan struct{}),
|
||||
janitorDone: make(chan struct{}),
|
||||
}
|
||||
}
|
||||
|
||||
// startJanitor launches the single shrink loop. Idempotent.
|
||||
func (v *VM) startJanitor() {
|
||||
v.janitorOnce.Do(func() {
|
||||
go v.janitorLoop(shrinkInterval)
|
||||
})
|
||||
}
|
||||
|
||||
func (v *VM) janitorLoop(every time.Duration) {
|
||||
defer close(v.janitorDone)
|
||||
t := time.NewTicker(every)
|
||||
defer t.Stop()
|
||||
for {
|
||||
select {
|
||||
case <-v.janitorStop:
|
||||
return
|
||||
case <-t.C:
|
||||
v.reclaimIdle()
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// reclaimIdle runs one shrink round over every pool. Exported behaviour is
|
||||
// tested through ReclaimIdleNow.
|
||||
func (v *VM) reclaimIdle() int {
|
||||
v.mu.RLock()
|
||||
pools := make([]*adapterPool, 0, len(v.pools))
|
||||
for _, p := range v.pools {
|
||||
pools = append(pools, p)
|
||||
}
|
||||
v.mu.RUnlock()
|
||||
n := 0
|
||||
for _, p := range pools {
|
||||
n += p.reclaim()
|
||||
}
|
||||
return n
|
||||
}
|
||||
|
||||
// ReclaimIdleNow forces one shrink round and reports how many worker states
|
||||
// were closed. Intended for tests and for an explicit "release idle memory"
|
||||
// action; normal operation relies on the janitor.
|
||||
func (v *VM) ReclaimIdleNow() int { return v.reclaimIdle() }
|
||||
|
||||
// PoolStats returns the elastic sizing snapshot of every loaded adapter,
|
||||
// sorted by name.
|
||||
func (v *VM) PoolStats() []PoolStats {
|
||||
v.mu.RLock()
|
||||
pools := make([]*adapterPool, 0, len(v.pools))
|
||||
for _, p := range v.pools {
|
||||
pools = append(pools, p)
|
||||
}
|
||||
v.mu.RUnlock()
|
||||
out := make([]PoolStats, 0, len(pools))
|
||||
for _, p := range pools {
|
||||
out = append(out, p.stats())
|
||||
}
|
||||
sort.Slice(out, func(i, j int) bool { return out[i].Name < out[j].Name })
|
||||
return out
|
||||
}
|
||||
|
||||
func (v *VM) Start() error {
|
||||
v.startJanitor()
|
||||
if v.dir == "" {
|
||||
return nil
|
||||
}
|
||||
@ -214,6 +602,12 @@ func (v *VM) Start() error {
|
||||
}
|
||||
|
||||
func (v *VM) Stop() {
|
||||
select {
|
||||
case <-v.janitorStop:
|
||||
// already stopped
|
||||
default:
|
||||
close(v.janitorStop)
|
||||
}
|
||||
v.mu.Lock()
|
||||
pools := v.pools
|
||||
v.pools = map[string]*adapterPool{}
|
||||
@ -223,17 +617,17 @@ func (v *VM) Stop() {
|
||||
}
|
||||
}
|
||||
|
||||
// ConfigureConcurrency sets each adapter's worker target to the sum of
|
||||
// max_concurrent of every source using it. Values below one are clamped to 1.
|
||||
// ConfigureConcurrency sets each adapter's worker CEILING to the sum of
|
||||
// max_concurrent of every source using it. It is a ceiling, not a preallocation:
|
||||
// states are booted on demand and reclaimed when demand drops, so a config with
|
||||
// a large aggregate concurrency no longer implies a large resident pool.
|
||||
// Values below one are clamped to 1.
|
||||
func (v *VM) ConfigureConcurrency(adapterConcurrency map[string]int) {
|
||||
v.mu.RLock()
|
||||
defer v.mu.RUnlock()
|
||||
for name, n := range adapterConcurrency {
|
||||
if p, ok := v.pools[name]; ok {
|
||||
if n < 1 {
|
||||
n = 1
|
||||
}
|
||||
p.setTarget(n)
|
||||
p.setMax(n)
|
||||
}
|
||||
}
|
||||
}
|
||||
@ -260,16 +654,19 @@ func (v *VM) LoadAdapter(path string) error {
|
||||
}
|
||||
|
||||
// LoadAdapterSource registers an adapter from source code. Replacing an
|
||||
// existing adapter shuts the old pool down and re-boots lazily.
|
||||
// existing adapter shuts the old pool down and re-boots lazily, carrying the
|
||||
// configured ceiling over so a hot-reload does not reset elastic sizing.
|
||||
func (v *VM) LoadAdapterSource(name, code string) error {
|
||||
static, err := inspectScript(name, code)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
target := 1
|
||||
maxW := 1
|
||||
v.mu.RLock()
|
||||
if p, ok := v.pools[static.name]; ok {
|
||||
target = p.target
|
||||
p.mu.Lock()
|
||||
maxW = p.maxW
|
||||
p.mu.Unlock()
|
||||
}
|
||||
v.mu.RUnlock()
|
||||
|
||||
@ -278,9 +675,10 @@ func (v *VM) LoadAdapterSource(name, code string) error {
|
||||
p.shutdown()
|
||||
}
|
||||
p := newAdapterPool(static.name, code, *static)
|
||||
p.setTarget(target)
|
||||
p.setMax(maxW)
|
||||
v.pools[static.name] = p
|
||||
v.mu.Unlock()
|
||||
v.startJanitor()
|
||||
return nil
|
||||
}
|
||||
|
||||
@ -407,6 +805,46 @@ func (v *VM) BuildHeaders(name string, meta map[string]interface{}) (map[string]
|
||||
return headers, nil
|
||||
}
|
||||
|
||||
// TransformError executes the optional adapter.transform_error(status, body)
|
||||
// hook: per-source condensing of an upstream error response into a short
|
||||
// client-facing reason. ok=false means the adapter defines no hook (or it
|
||||
// failed) — the caller falls back to the generic Go-side condenser.
|
||||
func (v *VM) TransformError(name string, status int, body string) (reason string, ok bool, err error) {
|
||||
p := v.pool(name)
|
||||
if p == nil {
|
||||
return "", false, fmt.Errorf("adapter %s not loaded", name)
|
||||
}
|
||||
w, err := p.acquire()
|
||||
if err != nil {
|
||||
return "", false, err
|
||||
}
|
||||
defer p.release(w)
|
||||
L := w.L
|
||||
|
||||
L.SetTop(0)
|
||||
defer L.SetTop(0)
|
||||
L.GetGlobal(adapterGlobal)
|
||||
if L.IsNil(-1) {
|
||||
return "", false, nil
|
||||
}
|
||||
L.GetField(-1, "transform_error")
|
||||
if !L.IsFunction(-1) {
|
||||
return "", false, nil
|
||||
}
|
||||
L.SetTop(0)
|
||||
L.GetGlobal(adapterGlobal)
|
||||
L.GetField(-1, "transform_error")
|
||||
L.PushInteger(int64(status))
|
||||
L.PushString(body)
|
||||
if callErr := L.Call(2, 1); callErr != nil {
|
||||
return "", false, fmt.Errorf("transform_error: %w", callErr)
|
||||
}
|
||||
if L.Type(-1) != golua.LUA_TSTRING {
|
||||
return "", false, nil
|
||||
}
|
||||
return L.ToString(-1), true, nil
|
||||
}
|
||||
|
||||
func (v *VM) Endpoint(name string) string {
|
||||
p := v.pool(name)
|
||||
if p == nil {
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
File diff suppressed because it is too large
Load Diff
@ -7,7 +7,9 @@ import (
|
||||
"fmt"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
"sync/atomic"
|
||||
"testing"
|
||||
"time"
|
||||
@ -196,8 +198,8 @@ func TestProviderAuthFailureSelfHeals(t *testing.T) {
|
||||
if st.FailCount() != backoffCapN {
|
||||
t.Fatalf("auth failure must jump to capped count, got %d", st.FailCount())
|
||||
}
|
||||
if until := st.CooldownUntil(); until <= time.Now().Add(25*time.Minute).Unix() {
|
||||
t.Fatalf("auth failure must cool near the cap (until=%d)", until)
|
||||
if until := st.CooldownUntil(); until <= time.Now().Add(authCooldown-time.Minute).Unix() {
|
||||
t.Fatalf("auth failure must cool for authCooldown (until=%d)", until)
|
||||
}
|
||||
if st.Pref() != -2*int64(prefFailStep) {
|
||||
t.Fatalf("auth failure pref penalty must be doubled, got %d", st.Pref())
|
||||
@ -240,8 +242,81 @@ func TestModelStateCooldownAndRecovery(t *testing.T) {
|
||||
if st.FailCount() != 0 {
|
||||
t.Fatalf("fail count after success = %d", st.FailCount())
|
||||
}
|
||||
if st.Pref() != 1-int64(prefFailStep) {
|
||||
t.Fatalf("pref after one failure (-5) then success (+1) = %d, want %d", st.Pref(), 1-int64(prefFailStep))
|
||||
// failure penalties prefFailStep, success rewards prefSuccessStep:
|
||||
// -5 then +2 = -3
|
||||
if st.Pref() != int64(prefSuccessStep)-int64(prefFailStep) {
|
||||
t.Fatalf("pref after one failure then success = %d, want %d", st.Pref(), int64(prefSuccessStep)-int64(prefFailStep))
|
||||
}
|
||||
}
|
||||
|
||||
func TestModelStateRateLimitShortCooldown(t *testing.T) {
|
||||
// 429 must NOT use the exponential backoff schedule: a rate-limited but
|
||||
// quota-rich source has to re-enter rotation after the short fixed window.
|
||||
p := newTestProvider(t, src("mock", "http://127.0.0.1:1", "openai", "m"))
|
||||
st := p.state("m")
|
||||
|
||||
// repeated 429s must stay at the fixed short cooldown, never escalate
|
||||
for i := 0; i < 15; i++ {
|
||||
p.ReportStatus("m", 429)
|
||||
}
|
||||
until := st.CooldownUntil()
|
||||
want := time.Now().Add(rateLimitCooldown).Unix()
|
||||
if until < want-2 || until > want+2 {
|
||||
t.Fatalf("429 cooldown = %d, want ~%d (fixed %v, not exponential)", until, want, rateLimitCooldown)
|
||||
}
|
||||
if st.FailCount() != 15 {
|
||||
t.Fatalf("fail count = %d, want 15 (counted but not escalating)", st.FailCount())
|
||||
}
|
||||
if st.Pref() != -15*int64(prefQuotaStep) {
|
||||
t.Fatalf("pref after 15 x 429 = %d, want %d (a 429 is a quota-class event, not a fault)", st.Pref(), -15*int64(prefQuotaStep))
|
||||
}
|
||||
if p.ModelAvailable("m") {
|
||||
t.Fatal("model must be cooling right after a 429")
|
||||
}
|
||||
}
|
||||
|
||||
func TestThrottleSpacingAndCancel(t *testing.T) {
|
||||
p := newTestProvider(t, src("mock", "http://127.0.0.1:1", "openai", "m"))
|
||||
p.cfg.RPM = 120 // gap = 500ms
|
||||
p.rpmGap = time.Minute / time.Duration(p.cfg.RPM)
|
||||
|
||||
start := time.Now()
|
||||
for i := 0; i < 3; i++ {
|
||||
if err := p.Throttle(context.Background()); err != nil {
|
||||
t.Fatalf("throttle %d: %v", i, err)
|
||||
}
|
||||
}
|
||||
elapsed := time.Since(start)
|
||||
// first call passes immediately, the next two wait one gap each
|
||||
want := 2 * time.Minute / time.Duration(p.cfg.RPM)
|
||||
if elapsed < want-time.Duration(100*time.Millisecond) || elapsed > want+time.Second {
|
||||
t.Fatalf("3 throttled calls took %v, want ~%v", elapsed, want)
|
||||
}
|
||||
|
||||
// unlimited source: Throttle is a no-op
|
||||
p2 := newTestProvider(t, src("free", "http://127.0.0.1:1", "openai", "m"))
|
||||
if err := p2.Throttle(context.Background()); err != nil {
|
||||
t.Fatalf("unlimited throttle: %v", err)
|
||||
}
|
||||
|
||||
// cancelled context aborts a pending window reservation
|
||||
p3 := newTestProvider(t, src("slow", "http://127.0.0.1:1", "openai", "m"))
|
||||
p3.cfg.RPM = 6 // 10s gap
|
||||
p3.rpmGap = time.Minute / time.Duration(p3.cfg.RPM)
|
||||
// first call passes immediately (fresh provider) and reserves the next window
|
||||
if err := p3.Throttle(context.Background()); err != nil {
|
||||
t.Fatalf("first throttle: %v", err)
|
||||
}
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 50*time.Millisecond)
|
||||
defer cancel()
|
||||
// second call must wait ~10s for the reserved window; the short deadline aborts it
|
||||
start = time.Now()
|
||||
err := p3.Throttle(ctx)
|
||||
if !errors.Is(err, context.DeadlineExceeded) {
|
||||
t.Fatalf("cancelled throttle = %v, want DeadlineExceeded", err)
|
||||
}
|
||||
if time.Since(start) > time.Second {
|
||||
t.Fatalf("cancel took %v, want fast abort", time.Since(start))
|
||||
}
|
||||
}
|
||||
|
||||
@ -345,7 +420,8 @@ func TestChatClientCancelNotRecorded(t *testing.T) {
|
||||
}
|
||||
|
||||
// TestChatStreamEmptyBodyNotSuccess: a 200 that yields zero chunks and no
|
||||
// [DONE] is a failure before the first chunk — the slot must back off.
|
||||
// [DONE] must fail the candidate (error return, no delivered chunks) so the
|
||||
// scheduler falls through to the next source, and record the failure.
|
||||
func TestChatStreamEmptyBodyNotSuccess(t *testing.T) {
|
||||
up := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
w.Header().Set("Content-Type", "text/event-stream")
|
||||
@ -357,15 +433,12 @@ func TestChatStreamEmptyBodyNotSuccess(t *testing.T) {
|
||||
Model: "m",
|
||||
Messages: []types.ChatMessage{{Role: "user", Content: types.StringContent("x")}},
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("stream: %v", err)
|
||||
}
|
||||
got := 0
|
||||
for range ch {
|
||||
got++
|
||||
}
|
||||
if got != 0 {
|
||||
t.Fatalf("want empty stream, got %d chunks", got)
|
||||
if err == nil {
|
||||
if ch != nil {
|
||||
for range ch {
|
||||
}
|
||||
}
|
||||
t.Fatal("empty stream must fail the candidate for scheduler fallback")
|
||||
}
|
||||
eventually(t, 2*time.Second, func() bool {
|
||||
_, fc, _ := p.ModelHealthInfo("m")
|
||||
@ -373,6 +446,62 @@ func TestChatStreamEmptyBodyNotSuccess(t *testing.T) {
|
||||
}, "empty stream must record a failure")
|
||||
}
|
||||
|
||||
// TestChatStreamErrorFinishFailsCandidate: zen free pool answers HTTP 200
|
||||
// with a single-chunk stream whose only payload is finish_reason:"network_error"
|
||||
// and an empty delta. That must fail the candidate (not deliver a laundered
|
||||
// empty reply), so the scheduler can fall through to the next source.
|
||||
func TestChatStreamErrorFinishFailsCandidate(t *testing.T) {
|
||||
up := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
w.Header().Set("Content-Type", "text/event-stream")
|
||||
fmt.Fprint(w, `data: {"choices":[{"index":0,"finish_reason":"network_error","delta":{"role":"assistant","content":""}}]}`+"\n\n")
|
||||
fmt.Fprint(w, "data: [DONE]\n\n")
|
||||
}))
|
||||
defer up.Close()
|
||||
p := newTestProvider(t, src("mock", up.URL, "opencode", "m"))
|
||||
_, err := p.ChatStream(context.Background(), &types.ChatRequest{
|
||||
Model: "m",
|
||||
Messages: []types.ChatMessage{{Role: "user", Content: types.StringContent("x")}},
|
||||
})
|
||||
if err == nil {
|
||||
t.Fatal("network_error-only stream must fail the candidate")
|
||||
}
|
||||
eventually(t, 2*time.Second, func() bool {
|
||||
_, fc, _ := p.ModelHealthInfo("m")
|
||||
return fc >= 1
|
||||
}, "error-only stream must record a failure")
|
||||
}
|
||||
|
||||
// TestChatStreamInstantEmptyStopStillDelivered: a legitimate completion that
|
||||
// ends immediately with the standard finish_reason:"stop" and zero content is
|
||||
// NOT an upstream error and must still reach the client.
|
||||
func TestChatStreamInstantEmptyStopStillDelivered(t *testing.T) {
|
||||
up := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
w.Header().Set("Content-Type", "text/event-stream")
|
||||
fmt.Fprint(w, `data: {"choices":[{"index":0,"finish_reason":"stop","delta":{}}]}`+"\n\n")
|
||||
fmt.Fprint(w, "data: [DONE]\n\n")
|
||||
}))
|
||||
defer up.Close()
|
||||
p := newTestProvider(t, src("mock", up.URL, "openai", "m"))
|
||||
ch, err := p.ChatStream(context.Background(), &types.ChatRequest{
|
||||
Model: "m",
|
||||
Messages: []types.ChatMessage{{Role: "user", Content: types.StringContent("x")}},
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("standard stop stream must be delivered: %v", err)
|
||||
}
|
||||
n := 0
|
||||
for range ch {
|
||||
n++
|
||||
}
|
||||
if n == 0 {
|
||||
t.Fatal("expected at least the terminating chunk")
|
||||
}
|
||||
eventually(t, 2*time.Second, func() bool {
|
||||
_, fc, _ := p.ModelHealthInfo("m")
|
||||
return fc == 0
|
||||
}, "standard stop must count as success")
|
||||
}
|
||||
|
||||
// TestImageAutoUsesImageModel: AUTO image generation on a mixed source must
|
||||
// send the image-kind model id, never the best chat model.
|
||||
func TestImageAutoUsesImageModel(t *testing.T) {
|
||||
@ -401,3 +530,656 @@ func TestImageAutoUsesImageModel(t *testing.T) {
|
||||
t.Fatalf("AUTO image must use the image-kind model, got %q", gotModel)
|
||||
}
|
||||
}
|
||||
|
||||
func TestStandardSSEChunkFinishReason(t *testing.T) {
|
||||
out := standardSSEChunk(`{"choices":[{"index":0,"finish_reason":"tool_calls","delta":{}}]}`)
|
||||
if !strings.Contains(out, `"done":true`) || !strings.Contains(out, `"finish_reason":"tool_calls"`) {
|
||||
t.Fatalf("real finish reason must pass through: %s", out)
|
||||
}
|
||||
out = standardSSEChunk(`{"choices":[{"index":0,"finish_reason":"","delta":{"content":"hi"}}]}`)
|
||||
if strings.Contains(out, `"Done":true`) {
|
||||
t.Fatalf("empty-string finish reason is not a finish signal: %s", out)
|
||||
}
|
||||
}
|
||||
|
||||
// TestAdapterHookCondensesError: the adapter transform_error hook owns the
|
||||
// per-source error format; its reason must reach the client verbatim.
|
||||
func TestAdapterHookCondensesError(t *testing.T) {
|
||||
up := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
w.WriteHeader(503)
|
||||
fmt.Fprint(w, `{"error":{"type":"FreeUsageLimitError","message":"Rate limit exceeded. Please try again later."}}`)
|
||||
}))
|
||||
defer up.Close()
|
||||
p := newTestProvider(t, src("mock", up.URL, "opencode", "m"))
|
||||
_, err := p.Chat(context.Background(), &types.ChatRequest{
|
||||
Model: "m",
|
||||
Messages: []types.ChatMessage{{Role: "user", Content: types.StringContent("x")}},
|
||||
})
|
||||
if err == nil || !strings.Contains(err.Error(), "zen free pool quota exhausted") {
|
||||
t.Fatalf("adapter hook reason must surface, got %v", err)
|
||||
}
|
||||
if strings.Contains(err.Error(), "Rate limit exceeded") {
|
||||
t.Fatalf("raw upstream body must not leak past the hook: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
// TestUnknownErrorFallbackWithoutHook: an adapter without transform_error
|
||||
// gets the uniform core fallback; the raw body must not leak to clients.
|
||||
func TestUnknownErrorFallbackWithoutHook(t *testing.T) {
|
||||
dir := filepath.Join(t.TempDir(), "adapters")
|
||||
if err := os.MkdirAll(dir, 0o755); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
minimal := `return { name="custom", endpoint="/chat/completions",
|
||||
transform_request=function(raw) return raw end,
|
||||
transform_response=function(raw) return raw end }
|
||||
`
|
||||
if err := os.WriteFile(filepath.Join(dir, "custom.lua"), []byte(minimal), 0o644); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
vm := lua.NewVM(dir)
|
||||
if err := vm.Start(); err != nil {
|
||||
t.Fatalf("vm: %v", err)
|
||||
}
|
||||
defer vm.Stop()
|
||||
up := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
w.WriteHeader(500)
|
||||
fmt.Fprint(w, `{"weird":{"shape":["no","message"]}}`)
|
||||
}))
|
||||
defer up.Close()
|
||||
p := New(src("mock", up.URL, "custom", "m"), vm)
|
||||
_, err := p.Chat(context.Background(), &types.ChatRequest{
|
||||
Model: "m",
|
||||
Messages: []types.ChatMessage{{Role: "user", Content: types.StringContent("x")}},
|
||||
})
|
||||
if err == nil || !strings.Contains(err.Error(), "unknown error") {
|
||||
t.Fatalf("hook-less adapter must fall back to unknown error, got %v", err)
|
||||
}
|
||||
if strings.Contains(err.Error(), "shape") {
|
||||
t.Fatalf("raw body must not leak: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
// TestToolIndexRemapper verifies sparse upstream tool-call indices are
|
||||
// compacted to dense 0-based OpenAI ordinals (Issue 2). Anthropic emits the
|
||||
// content_block ordinal, so a thinking block at index 0 pushes the first
|
||||
// tool_use to index 1 — OpenAI clients accumulating by index would leave a
|
||||
// hole at 0 and mis-assemble multi-tool responses.
|
||||
func TestToolIndexRemapper(t *testing.T) {
|
||||
t.Run("single tool after thinking block", func(t *testing.T) {
|
||||
r := newToolIndexRemapper()
|
||||
// content_block index 1 (thinking was 0) -> dense 0
|
||||
out := r.remap(json.RawMessage(`[{"index":1,"id":"toolu_a","type":"function","function":{"name":"calc","arguments":""}}]`))
|
||||
var got []map[string]interface{}
|
||||
if err := json.Unmarshal(out, &got); err != nil {
|
||||
t.Fatalf("unmarshal: %v (%s)", err, out)
|
||||
}
|
||||
if got[0]["index"] != float64(0) {
|
||||
t.Fatalf("index = %v, want 0: %s", got[0]["index"], out)
|
||||
}
|
||||
// id/name/arguments must survive verbatim
|
||||
if got[0]["id"] != "toolu_a" {
|
||||
t.Fatalf("id lost: %s", out)
|
||||
}
|
||||
fn := got[0]["function"].(map[string]interface{})
|
||||
if fn["name"] != "calc" {
|
||||
t.Fatalf("function.name lost: %s", out)
|
||||
}
|
||||
// subsequent argument fragments on the SAME upstream index reuse ordinal 0
|
||||
out2 := r.remap(json.RawMessage(`[{"index":1,"id":"","type":"function","function":{"name":"","arguments":"{\"a\":1}"}}]`))
|
||||
var got2 []map[string]interface{}
|
||||
if err := json.Unmarshal(out2, &got2); err != nil {
|
||||
t.Fatalf("unmarshal2: %v", err)
|
||||
}
|
||||
if got2[0]["index"] != float64(0) {
|
||||
t.Fatalf("fragment index = %v, want stable 0: %s", got2[0]["index"], out2)
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("multiple tools get distinct dense ordinals", func(t *testing.T) {
|
||||
r := newToolIndexRemapper()
|
||||
// thinking=0, tool_use=1, tool_use=2 -> 0, 1
|
||||
a := r.remap(json.RawMessage(`[{"index":1,"id":"t1","type":"function","function":{"name":"f1","arguments":""}}]`))
|
||||
b := r.remap(json.RawMessage(`[{"index":2,"id":"t2","type":"function","function":{"name":"f2","arguments":""}}]`))
|
||||
var ga, gb []map[string]interface{}
|
||||
json.Unmarshal(a, &ga)
|
||||
json.Unmarshal(b, &gb)
|
||||
if ga[0]["index"] != float64(0) {
|
||||
t.Fatalf("first tool index = %v, want 0", ga[0]["index"])
|
||||
}
|
||||
if gb[0]["index"] != float64(1) {
|
||||
t.Fatalf("second tool index = %v, want 1", gb[0]["index"])
|
||||
}
|
||||
// interleaved fragments keep their own ordinals
|
||||
a2 := r.remap(json.RawMessage(`[{"index":1,"id":"","type":"function","function":{"name":"","arguments":"x"}}]`))
|
||||
var ga2 []map[string]interface{}
|
||||
json.Unmarshal(a2, &ga2)
|
||||
if ga2[0]["index"] != float64(0) {
|
||||
t.Fatalf("t1 fragment index = %v, want 0", ga2[0]["index"])
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("already dense indices pass through untouched", func(t *testing.T) {
|
||||
r := newToolIndexRemapper()
|
||||
in := json.RawMessage(`[{"index":0,"id":"t","type":"function","function":{"name":"f","arguments":""}}]`)
|
||||
out := r.remap(in)
|
||||
if string(out) != string(in) {
|
||||
t.Fatalf("dense input was rewritten:\n in=%s\nout=%s", in, out)
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("empty and malformed input is safe", func(t *testing.T) {
|
||||
r := newToolIndexRemapper()
|
||||
if got := r.remap(nil); got != nil {
|
||||
t.Fatalf("nil -> %s", got)
|
||||
}
|
||||
bad := json.RawMessage(`not json`)
|
||||
if got := r.remap(bad); string(got) != string(bad) {
|
||||
t.Fatalf("malformed input must pass through: %s", got)
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
// ---- half-cooldown probe gate (plan 阶段 1) ----
|
||||
|
||||
// TestProbeGateHalfCooldown pins the core rule: a cooling model is fully silent
|
||||
// in the first half of its cooldown window and probeable exactly once in the
|
||||
// second half.
|
||||
func TestProbeGateHalfCooldown(t *testing.T) {
|
||||
var st ModelState
|
||||
// open a 100s window centred so that "now" sits in the first half
|
||||
now := time.Now().Unix()
|
||||
st.cooldownFrom.Store(now - 10)
|
||||
st.cooldownUntil.Store(now + 90)
|
||||
if st.TryProbe() {
|
||||
t.Fatal("first half of the cooldown window must stay silent")
|
||||
}
|
||||
|
||||
// move into the second half
|
||||
st.cooldownFrom.Store(now - 90)
|
||||
st.cooldownUntil.Store(now + 10)
|
||||
if !st.TryProbe() {
|
||||
t.Fatal("second half of the cooldown window must allow one probe")
|
||||
}
|
||||
if st.TryProbe() {
|
||||
t.Fatal("probe permit must be single-token")
|
||||
}
|
||||
if !st.Probing() {
|
||||
t.Fatal("Probing() must report the claimed permit")
|
||||
}
|
||||
st.ProbeDone()
|
||||
if st.Probing() {
|
||||
t.Fatal("ProbeDone must release the permit")
|
||||
}
|
||||
if !st.TryProbe() {
|
||||
t.Fatal("permit must be reusable after ProbeDone")
|
||||
}
|
||||
}
|
||||
|
||||
// TestProbeGateNoProbeWhenHealthy: a healthy model needs no probe permit — the
|
||||
// caller schedules it normally.
|
||||
func TestProbeGateNoProbeWhenHealthy(t *testing.T) {
|
||||
var idle ModelState
|
||||
if idle.TryProbe() {
|
||||
t.Fatal("a healthy model must not be probed (schedule it normally)")
|
||||
}
|
||||
}
|
||||
|
||||
// TestProbeGateRescuesPrefFloor documents a fix that comes with the probe gate:
|
||||
// a slot whose preference sank to prefMin used to be refused by ModelAvailable
|
||||
// forever, even long after every cooldown expired — a permanent blacklist. It is
|
||||
// now probeable, so one success brings it back.
|
||||
func TestProbeGateRescuesPrefFloor(t *testing.T) {
|
||||
var dead ModelState
|
||||
dead.pref.Store(prefMin)
|
||||
if !dead.TryProbe() {
|
||||
t.Fatal("a slot stuck at prefMin must be probeable, not blacklisted")
|
||||
}
|
||||
dead.RecordSuccess()
|
||||
dead.ProbeDone()
|
||||
if dead.Pref() <= prefMin {
|
||||
t.Fatalf("one probe success must lift the slot off the floor, pref=%d", dead.Pref())
|
||||
}
|
||||
}
|
||||
|
||||
// TestProbeSuccessRestoresImmediately is the user-visible payoff: a recovered
|
||||
// upstream must not sit out the rest of its cooldown.
|
||||
func TestProbeSuccessRestoresImmediately(t *testing.T) {
|
||||
var st ModelState
|
||||
for i := 0; i < backoffCapN; i++ {
|
||||
st.RecordFailure(false)
|
||||
}
|
||||
if cd := st.CooldownUntil() - time.Now().Unix(); cd < int64(backoffCap.Seconds())-2 {
|
||||
t.Fatalf("10 failures must cool for backoffCap, got %ds", cd)
|
||||
}
|
||||
// jump into the probe window: the midpoint of [from, until) must be in the
|
||||
// past, so push the window's start back past its full length
|
||||
now := time.Now().Unix()
|
||||
st.cooldownFrom.Store(now - int64(backoffCap.Seconds()) - 10)
|
||||
if !st.TryProbe() {
|
||||
t.Fatal("probe must be allowed past the midpoint")
|
||||
}
|
||||
if !st.Probing() {
|
||||
t.Fatal("probe permit must be held while the probe runs")
|
||||
}
|
||||
st.RecordSuccess()
|
||||
st.ProbeDone()
|
||||
if !st.Available() {
|
||||
t.Fatal("probe success must clear the cooldown immediately")
|
||||
}
|
||||
if st.FailCount() != 0 {
|
||||
t.Fatalf("probe success must reset failCount, got %d", st.FailCount())
|
||||
}
|
||||
if st.CooldownUntil() != 0 || st.CooldownFrom() != 0 {
|
||||
t.Fatal("probe success must clear both cooldown bounds")
|
||||
}
|
||||
}
|
||||
|
||||
// TestProbeFailureDefersNextProbe: a failed probe must not immediately retry —
|
||||
// the new window's midpoint pushes the next probe out.
|
||||
func TestProbeFailureDefersNextProbe(t *testing.T) {
|
||||
var st ModelState
|
||||
st.RecordFailure(false) // 5s window
|
||||
now := time.Now().Unix()
|
||||
st.cooldownFrom.Store(now - 100)
|
||||
st.cooldownUntil.Store(now + 10)
|
||||
if !st.TryProbe() {
|
||||
t.Fatal("probe must be allowed past the midpoint")
|
||||
}
|
||||
st.RecordFailure(false) // probe failed: a fresh, longer window opens
|
||||
st.ProbeDone()
|
||||
if st.TryProbe() {
|
||||
t.Fatal("a failed probe must defer the next probe to the new midpoint")
|
||||
}
|
||||
if from := st.CooldownFrom(); from < now {
|
||||
t.Fatalf("failed probe must reopen the window (from=%d, now=%d)", from, now)
|
||||
}
|
||||
}
|
||||
|
||||
// TestAuthCooldownIsSeparateFromCap documents that credential failures cool on
|
||||
// their own (longer) schedule instead of reusing the exponential cap.
|
||||
func TestAuthCooldownIsSeparateFromCap(t *testing.T) {
|
||||
var st ModelState
|
||||
st.RecordFailure(true)
|
||||
cd := st.CooldownUntil() - time.Now().Unix()
|
||||
if cd < int64(authCooldown.Seconds())-2 || cd > int64(authCooldown.Seconds())+2 {
|
||||
t.Fatalf("auth failure must cool for authCooldown (%v), got %ds", authCooldown, cd)
|
||||
}
|
||||
if st.FailCount() != backoffCapN {
|
||||
t.Fatalf("auth failure must persist the capped count, got %d", st.FailCount())
|
||||
}
|
||||
}
|
||||
|
||||
// TestQuotaExhaustedDoesNotEscalate is the second half of the user's complaint:
|
||||
// running out of allowance must not be treated as a failure streak.
|
||||
func TestQuotaExhaustedDoesNotEscalate(t *testing.T) {
|
||||
var st ModelState
|
||||
st.RecordQuotaExhausted(2 * time.Minute)
|
||||
if st.FailCount() != 0 {
|
||||
t.Fatalf("quota exhaustion must not bump failCount, got %d", st.FailCount())
|
||||
}
|
||||
cd := st.CooldownUntil() - time.Now().Unix()
|
||||
if cd < 118 || cd > 122 {
|
||||
t.Fatalf("quota cooldown must honour the reset window, got %ds", cd)
|
||||
}
|
||||
|
||||
// an unbounded window is clamped so a topped-up allowance is still noticed
|
||||
var big ModelState
|
||||
big.RecordQuotaExhausted(72 * time.Hour)
|
||||
if cd := big.CooldownUntil() - time.Now().Unix(); cd > int64(maxQuotaCooldown.Seconds())+2 {
|
||||
t.Fatalf("quota cooldown must be capped at maxQuotaCooldown, got %ds", cd)
|
||||
}
|
||||
|
||||
// no hint at all falls back to the short window
|
||||
var none ModelState
|
||||
none.RecordQuotaExhausted(0)
|
||||
if cd := none.CooldownUntil() - time.Now().Unix(); cd > int64(rateLimitCooldown.Seconds())+2 {
|
||||
t.Fatalf("unknown reset window must fall back to rateLimitCooldown, got %ds", cd)
|
||||
}
|
||||
}
|
||||
|
||||
// TestQuotaReasonRouting checks that a 429 carrying a quota phrase is cooled as
|
||||
// quota exhaustion while a plain rate limit keeps the short fixed window.
|
||||
func TestQuotaReasonRouting(t *testing.T) {
|
||||
p := newTestProvider(t, src("q", "http://127.0.0.1:1", "openai", "m"))
|
||||
p.ReportStatusReason("m", 429, "429 insufficient_quota: you exceeded your current quota")
|
||||
st := p.state("m")
|
||||
if st.FailCount() != 0 {
|
||||
t.Fatalf("quota rejection must not bump failCount, got %d", st.FailCount())
|
||||
}
|
||||
if cd := st.CooldownUntil() - time.Now().Unix(); cd <= int64(rateLimitCooldown.Seconds()) {
|
||||
t.Fatalf("quota rejection must cool longer than a plain 429, got %ds", cd)
|
||||
}
|
||||
|
||||
p2 := newTestProvider(t, src("r", "http://127.0.0.1:1", "openai", "m"))
|
||||
p2.ReportStatusReason("m", 429, "rate limit reached, slow down")
|
||||
st2 := p2.state("m")
|
||||
if cd := st2.CooldownUntil() - time.Now().Unix(); cd > int64(rateLimitCooldown.Seconds())+2 {
|
||||
t.Fatalf("plain 429 must keep the short window, got %ds", cd)
|
||||
}
|
||||
}
|
||||
|
||||
// TestModelSchedulableProbeHandoff covers the provider-level gate the scheduler
|
||||
// consumes.
|
||||
func TestModelSchedulableProbeHandoff(t *testing.T) {
|
||||
p := newTestProvider(t, src("s", "http://127.0.0.1:1", "openai", "m"))
|
||||
if ok, probe := p.ModelSchedulable("m"); !ok || probe {
|
||||
t.Fatalf("healthy model must be normally schedulable, got ok=%v probe=%v", ok, probe)
|
||||
}
|
||||
st := p.state("m")
|
||||
st.RecordFailure(false)
|
||||
now := time.Now().Unix()
|
||||
st.cooldownFrom.Store(now - 100)
|
||||
st.cooldownUntil.Store(now + 10)
|
||||
|
||||
ok, probe := p.ModelSchedulable("m")
|
||||
if !ok || !probe {
|
||||
t.Fatalf("cooling model past midpoint must be probeable, got ok=%v probe=%v", ok, probe)
|
||||
}
|
||||
if ok2, _ := p.ModelSchedulable("m"); ok2 {
|
||||
t.Fatal("only one probe permit may be outstanding")
|
||||
}
|
||||
p.ProbeDone("m")
|
||||
if ok3, probe3 := p.ModelSchedulable("m"); !ok3 || !probe3 {
|
||||
t.Fatal("permit must be reclaimable after ProbeDone")
|
||||
}
|
||||
}
|
||||
|
||||
// TestProbeSlotsFollowsMaxConcurrent pins the user's example: max_concurrent=10
|
||||
// yields exactly one probe slot.
|
||||
func TestProbeSlotsFollowsMaxConcurrent(t *testing.T) {
|
||||
for _, tc := range []struct{ mc, want int }{
|
||||
{0, 1}, {1, 1}, {4, 1}, {10, 1}, {20, 2}, {100, 2},
|
||||
} {
|
||||
s := src("p", "http://127.0.0.1:1", "openai", "m")
|
||||
s.MaxConcurrent = tc.mc
|
||||
p := newTestProvider(t, s)
|
||||
if got := p.ProbeSlots(); got != tc.want {
|
||||
t.Fatalf("max_concurrent=%d: probe slots = %d, want %d", tc.mc, got, tc.want)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// TestModelProbeInfo checks the UI-facing view of the probe window.
|
||||
func TestModelProbeInfo(t *testing.T) {
|
||||
p := newTestProvider(t, src("i", "http://127.0.0.1:1", "openai", "m"))
|
||||
if from, after, probing := p.ModelProbeInfo("m"); from != 0 || after != 0 || probing {
|
||||
t.Fatalf("idle model must report no probe window, got %d %d %v", from, after, probing)
|
||||
}
|
||||
st := p.state("m")
|
||||
now := time.Now().Unix()
|
||||
st.cooldownFrom.Store(now)
|
||||
st.cooldownUntil.Store(now + 100)
|
||||
from, after, _ := p.ModelProbeInfo("m")
|
||||
if from != now || after != now+50 {
|
||||
t.Fatalf("probe window = (%d,%d), want (%d,%d)", from, after, now, now+50)
|
||||
}
|
||||
}
|
||||
|
||||
// TestUpstreamRecoveryWithoutWaitingOutCooldown is the end-to-end shape of the
|
||||
// user's complaint: an upstream that 500s ten times lands in the capped
|
||||
// cooldown, and used to be unusable for a full 30 minutes even after it
|
||||
// recovered. With the probe gate the slot is silent for the first half of the
|
||||
// window and then serves the very next request itself, clearing the cooldown.
|
||||
//
|
||||
// Real time is not waited out: the window bounds are moved so the test observes
|
||||
// the same state the scheduler would see at cap/2.
|
||||
func TestUpstreamRecoveryWithoutWaitingOutCooldown(t *testing.T) {
|
||||
var healthy atomic.Bool
|
||||
var hits atomic.Int64
|
||||
up := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
hits.Add(1)
|
||||
if !healthy.Load() {
|
||||
w.WriteHeader(500)
|
||||
return
|
||||
}
|
||||
fmt.Fprint(w, `{"choices":[{"message":{"content":"back"},"finish_reason":"stop"}]}`)
|
||||
}))
|
||||
defer up.Close()
|
||||
p := newTestProvider(t, src("recov", up.URL, "openai", "m"))
|
||||
req := func() *types.ChatRequest {
|
||||
return &types.ChatRequest{
|
||||
Model: "m",
|
||||
Messages: []types.ChatMessage{{Role: "user", Content: types.StringContent("x")}},
|
||||
}
|
||||
}
|
||||
|
||||
// 1. drive the slot into the capped cooldown
|
||||
for i := 0; i < backoffCapN; i++ {
|
||||
if _, err := p.Chat(context.Background(), req()); err == nil {
|
||||
t.Fatalf("attempt %d: expected upstream failure", i)
|
||||
}
|
||||
}
|
||||
st := p.state("m")
|
||||
if p.ModelAvailable("m") {
|
||||
t.Fatal("10 consecutive 5xx must take the slot out of normal rotation")
|
||||
}
|
||||
cd := st.CooldownUntil() - time.Now().Unix()
|
||||
if cd < int64(backoffCap.Seconds())-2 {
|
||||
t.Fatalf("cooldown = %ds, want ~%v", cd, backoffCap)
|
||||
}
|
||||
// the old 30-minute lockout is gone
|
||||
if cd > int64((6 * time.Minute).Seconds()) {
|
||||
t.Fatalf("cooldown %ds is too long: the cap must be minutes, not half an hour", cd)
|
||||
}
|
||||
|
||||
// 2. upstream recovers, but we are still in the FIRST half of the window:
|
||||
// the slot must stay completely silent (no probe traffic at all).
|
||||
healthy.Store(true)
|
||||
quiet := hits.Load()
|
||||
if ok, probe := p.ModelSchedulable("m"); ok || probe {
|
||||
t.Fatalf("first half of the window must stay silent, got ok=%v probe=%v", ok, probe)
|
||||
}
|
||||
if hits.Load() != quiet {
|
||||
t.Fatal("no request may reach the upstream during the silent half")
|
||||
}
|
||||
|
||||
// 3. advance past the midpoint: exactly one probe is admitted.
|
||||
st.cooldownFrom.Store(time.Now().Unix() - int64(backoffCap.Seconds()) - 1)
|
||||
ok, isProbe := p.ModelSchedulable("m")
|
||||
if !ok || !isProbe {
|
||||
t.Fatalf("past the midpoint one probe must be admitted, got ok=%v probe=%v", ok, isProbe)
|
||||
}
|
||||
if ok2, _ := p.ModelSchedulable("m"); ok2 {
|
||||
t.Fatal("only ONE probe may be in flight while cooling")
|
||||
}
|
||||
|
||||
// 4. the probe is a real request: its success restores the slot fully.
|
||||
resp, err := p.Chat(context.Background(), req())
|
||||
p.ProbeDone("m")
|
||||
if err != nil {
|
||||
t.Fatalf("probe request failed: %v", err)
|
||||
}
|
||||
if resp.Content != "back" {
|
||||
t.Fatalf("probe response = %q, want \"back\"", resp.Content)
|
||||
}
|
||||
if !p.ModelAvailable("m") {
|
||||
t.Fatal("a successful probe must return the slot to normal rotation immediately")
|
||||
}
|
||||
if st.CooldownUntil() != 0 || st.FailCount() != 0 {
|
||||
t.Fatalf("probe success must clear cooldown/failCount, got until=%d fails=%d",
|
||||
st.CooldownUntil(), st.FailCount())
|
||||
}
|
||||
if ok, probe := p.ModelSchedulable("m"); !ok || probe {
|
||||
t.Fatalf("recovered slot must schedule normally, got ok=%v probe=%v", ok, probe)
|
||||
}
|
||||
}
|
||||
|
||||
// TestRecoveryProbeBudgetIsBounded quantifies the cost of the shorter cap
|
||||
// against a still-dead upstream: one probe per window, not a retry storm.
|
||||
func TestRecoveryProbeBudgetIsBounded(t *testing.T) {
|
||||
var hits atomic.Int64
|
||||
up := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
hits.Add(1)
|
||||
w.WriteHeader(500)
|
||||
}))
|
||||
defer up.Close()
|
||||
p := newTestProvider(t, src("dead", up.URL, "openai", "m"))
|
||||
req := &types.ChatRequest{
|
||||
Model: "m",
|
||||
Messages: []types.ChatMessage{{Role: "user", Content: types.StringContent("x")}},
|
||||
}
|
||||
for i := 0; i < backoffCapN; i++ {
|
||||
_, _ = p.Chat(context.Background(), req)
|
||||
}
|
||||
baseline := hits.Load()
|
||||
st := p.state("m")
|
||||
|
||||
// simulate three consecutive probe windows against a dead upstream
|
||||
for i := 0; i < 3; i++ {
|
||||
st.cooldownFrom.Store(time.Now().Unix() - int64(backoffCap.Seconds()) - 1)
|
||||
ok, isProbe := p.ModelSchedulable("m")
|
||||
if !ok || !isProbe {
|
||||
t.Fatalf("window %d: expected a probe permit", i)
|
||||
}
|
||||
// a second caller in the same window gets nothing
|
||||
if ok2, _ := p.ModelSchedulable("m"); ok2 {
|
||||
t.Fatalf("window %d: more than one probe admitted", i)
|
||||
}
|
||||
if _, err := p.Chat(context.Background(), req); err == nil {
|
||||
t.Fatalf("window %d: expected the probe to fail", i)
|
||||
}
|
||||
p.ProbeDone("m")
|
||||
}
|
||||
if extra := hits.Load() - baseline; extra != 3 {
|
||||
t.Fatalf("3 probe windows must cost exactly 3 upstream requests, got %d", extra)
|
||||
}
|
||||
}
|
||||
|
||||
// TestQuotaPenaltyIsLighterThanFailure is the regression for the report
|
||||
// "今天一次失败没看到,分数目前已经 -10 了" on an allowance-metered source
|
||||
// (sensenova). RecordQuotaExhausted deliberately does not touch failCount — the
|
||||
// WebUI therefore showed fails=0 and cooling=false — yet it used to charge the
|
||||
// SAME score penalty as a real failure, so two ordinary quota resets dragged the
|
||||
// slot to -10 with no visible fault. An exhausted allowance means "no budget
|
||||
// right now", not "this model is broken": the cooldown already keeps it out of
|
||||
// rotation, so the score penalty must be mild.
|
||||
func TestQuotaPenaltyIsLighterThanFailure(t *testing.T) {
|
||||
quota := &ModelState{}
|
||||
quota.RecordQuotaExhausted(time.Minute)
|
||||
fail := &ModelState{}
|
||||
fail.RecordFailure(false)
|
||||
|
||||
qp, fp := quota.Pref(), fail.Pref()
|
||||
if qp <= fp {
|
||||
t.Fatalf("a quota reset must cost less score than a real failure: quota=%d failure=%d", qp, fp)
|
||||
}
|
||||
if qp != -prefQuotaStep {
|
||||
t.Errorf("quota penalty = %d, want -%d", qp, prefQuotaStep)
|
||||
}
|
||||
if fp != -prefFailStep {
|
||||
t.Errorf("failure penalty = %d, want -%d", fp, prefFailStep)
|
||||
}
|
||||
|
||||
// The reported symptom: several quota resets in a day must NOT approach the
|
||||
// prefMin floor, since each one is a normal event for a metered source.
|
||||
s := &ModelState{}
|
||||
for i := 0; i < 5; i++ {
|
||||
s.RecordQuotaExhausted(time.Minute)
|
||||
}
|
||||
if got := s.Pref(); got <= prefMin/2 {
|
||||
t.Errorf("5 quota resets sank the score to %d (floor %d); a metered source would be written off", got, prefMin)
|
||||
}
|
||||
if s.FailCount() != 0 {
|
||||
t.Errorf("quota exhaustion must not inflate failCount, got %d", s.FailCount())
|
||||
}
|
||||
}
|
||||
|
||||
// TestRateLimitPenaltyIsMild covers the sibling case: a 429 means "too fast",
|
||||
// not "broken", so a popular-but-healthy slot must not sink for being throttled.
|
||||
func TestRateLimitPenaltyIsMild(t *testing.T) {
|
||||
s := &ModelState{}
|
||||
s.RecordRateLimit()
|
||||
if got := s.Pref(); got != -prefQuotaStep {
|
||||
t.Errorf("429 penalty = %d, want -%d (same class as a quota event)", got, prefQuotaStep)
|
||||
}
|
||||
}
|
||||
|
||||
// TestSuccessRecoversFasterThanBefore pins the penalty/reward asymmetry. Reward
|
||||
// used to be +1 against a -5 penalty, so a slot at -10 needed TEN consecutive
|
||||
// successes just to reach neutral — in practice it never got there, because a
|
||||
// low score makes the slot unattractive to the scheduler in the first place.
|
||||
func TestSuccessRecoversFasterThanBefore(t *testing.T) {
|
||||
s := &ModelState{}
|
||||
s.RecordFailure(false)
|
||||
s.RecordFailure(false) // -10, the score from the bug report
|
||||
|
||||
n := 0
|
||||
for s.Pref() < 0 {
|
||||
s.RecordSuccess()
|
||||
n++
|
||||
if n > 50 {
|
||||
t.Fatal("score never recovered")
|
||||
}
|
||||
}
|
||||
// with +2 per success, -10 needs 5; the old +1 needed 10
|
||||
if n > 5 {
|
||||
t.Errorf("recovery from -10 took %d successes, want <= 5", n)
|
||||
}
|
||||
// a real failure must still outweigh a single success, or failures stop
|
||||
// meaning anything
|
||||
if prefSuccessStep >= prefFailStep {
|
||||
t.Errorf("success reward %d must stay below the failure penalty %d", prefSuccessStep, prefFailStep)
|
||||
}
|
||||
}
|
||||
|
||||
// TestPrefDecayRehabilitatesIdleSlot closes the starvation loop: a penalised
|
||||
// slot is unattractive, so the scheduler stops picking it, so it never earns the
|
||||
// successes that would rehabilitate it. Without decay the score stays negative
|
||||
// forever even though the upstream may be perfectly healthy.
|
||||
func TestPrefDecayRehabilitatesIdleSlot(t *testing.T) {
|
||||
s := &ModelState{}
|
||||
s.RecordFailure(false)
|
||||
s.RecordFailure(false)
|
||||
if s.Pref() != -2*prefFailStep {
|
||||
t.Fatalf("setup: pref=%d", s.Pref())
|
||||
}
|
||||
|
||||
// backdate the last touch by three decay intervals
|
||||
step := int64(prefDecayAfter / time.Second)
|
||||
s.prefTouched.Store(time.Now().Unix() - 3*step)
|
||||
|
||||
got := s.Pref()
|
||||
if got != -2*prefFailStep+3 {
|
||||
t.Errorf("after 3 idle intervals pref=%d, want %d", got, -2*prefFailStep+3)
|
||||
}
|
||||
|
||||
// decay must never overshoot into positive territory
|
||||
s.prefTouched.Store(time.Now().Unix() - 1000*step)
|
||||
if got := s.Pref(); got != 0 {
|
||||
t.Errorf("decay overshot to %d, want to stop at 0", got)
|
||||
}
|
||||
|
||||
// a positive score must not be dragged down for being idle
|
||||
good := &ModelState{}
|
||||
good.RecordSuccess()
|
||||
good.RecordSuccess()
|
||||
before := good.Pref()
|
||||
good.prefTouched.Store(time.Now().Unix() - 1000*step)
|
||||
if after := good.Pref(); after != before {
|
||||
t.Errorf("idle decay must not touch a positive score: %d -> %d", before, after)
|
||||
}
|
||||
}
|
||||
|
||||
// TestDecayLiftsSlotOffTheFloor verifies decay actually restores schedulability:
|
||||
// a slot at prefMin is refused by ModelAvailable, and before decay existed only
|
||||
// a probe could rescue it.
|
||||
func TestDecayLiftsSlotOffTheFloor(t *testing.T) {
|
||||
s := &ModelState{}
|
||||
for i := 0; i < 10; i++ {
|
||||
s.RecordFailure(false)
|
||||
}
|
||||
if s.Pref() != prefMin {
|
||||
t.Fatalf("setup: want floor %d, got %d", prefMin, s.Pref())
|
||||
}
|
||||
// clear the cooldown so only the score gates the slot
|
||||
s.cooldownUntil.Store(0)
|
||||
s.cooldownFrom.Store(0)
|
||||
|
||||
step := int64(prefDecayAfter / time.Second)
|
||||
s.prefTouched.Store(time.Now().Unix() - 2*step)
|
||||
if got := s.Pref(); got <= prefMin {
|
||||
t.Errorf("decay failed to lift the slot off the floor: %d", got)
|
||||
}
|
||||
}
|
||||
|
||||
@ -217,6 +217,12 @@ type SourceStatus struct {
|
||||
// probe was rate-limited by the upstream).
|
||||
RecentOK int64 `json:"recent_ok,omitempty"`
|
||||
RecentErr int64 `json:"recent_err,omitempty"`
|
||||
// AvgFirstByteMs is the mean time-to-first-byte (ms) for recently
|
||||
// successful requests served by this source (0 when unmeasured).
|
||||
AvgFirstByteMs int64 `json:"avg_first_byte_ms,omitempty"`
|
||||
// AvgTokPerS is the aggregate completion throughput (tokens/s) for
|
||||
// recently successful requests served by this source (0 when no data).
|
||||
AvgTokPerS int64 `json:"avg_tok_per_s,omitempty"`
|
||||
}
|
||||
|
||||
// ProbeAll runs a live reachability check for every provider (in parallel).
|
||||
|
||||
@ -42,6 +42,12 @@ type Provider interface {
|
||||
Name() string
|
||||
ModelFor(reqModel string) string
|
||||
ModelAvailable(model string) bool
|
||||
// ModelSchedulable is the probe-aware availability gate: ok reports
|
||||
// whether the model may take a request now, isProbe marks that it is only
|
||||
// allowed as a cooldown probe (the caller must release the permit with
|
||||
// ProbeDone once the attempt finished).
|
||||
ModelSchedulable(model string) (ok bool, isProbe bool)
|
||||
ProbeDone(model string)
|
||||
Pref(model string) int64
|
||||
Chat(ctx context.Context, req *types.ChatRequest) (*types.UnifiedResponse, error)
|
||||
ChatStream(ctx context.Context, req *types.ChatRequest) (<-chan types.UnifiedChunk, error)
|
||||
@ -176,18 +182,84 @@ type tierResult struct {
|
||||
hard []TierError // hard failures seen in this pass (nil = none)
|
||||
}
|
||||
|
||||
// candidate is one schedulable slot in a tier pass. probe marks a slot that is
|
||||
// still cooling but past its half-cooldown mark and holding the probe permit:
|
||||
// it is tried only after every normal slot, so probe traffic is what the tier
|
||||
// falls back to instead of what it prefers.
|
||||
type candidate struct {
|
||||
slot *Slot
|
||||
probe bool
|
||||
}
|
||||
|
||||
// collectCands partitions a tier's slots into normal and probe candidates,
|
||||
// normal first. Quota-exhausted slots are dropped outright. Every probe
|
||||
// candidate returned holds a probe permit, so the caller MUST call
|
||||
// releaseProbes on the result exactly once.
|
||||
func collectCands(slots []*Slot, exhausted func(*Slot) bool) []candidate {
|
||||
var normal, probes []candidate
|
||||
for _, sl := range slots {
|
||||
if exhausted != nil && exhausted(sl) {
|
||||
continue
|
||||
}
|
||||
ok, isProbe := sl.Prov.ModelSchedulable(sl.Model)
|
||||
if !ok {
|
||||
continue
|
||||
}
|
||||
if isProbe {
|
||||
probes = append(probes, candidate{slot: sl, probe: true})
|
||||
continue
|
||||
}
|
||||
normal = append(normal, candidate{slot: sl})
|
||||
}
|
||||
return append(normal, probes...)
|
||||
}
|
||||
|
||||
// releaseProbes hands every claimed probe permit back, whether or not the probe
|
||||
// slot was actually used.
|
||||
func releaseProbes(cands []candidate) {
|
||||
for _, c := range cands {
|
||||
if c.probe {
|
||||
c.slot.Prov.ProbeDone(c.slot.Model)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// normalCount is how many leading candidates are normal (non-probe). The
|
||||
// round-robin cursor rotates only over those: probe slots are a strictly
|
||||
// ordered tail, never a rotation target.
|
||||
func normalCount(cands []candidate) int {
|
||||
for i, c := range cands {
|
||||
if c.probe {
|
||||
return i
|
||||
}
|
||||
}
|
||||
return len(cands)
|
||||
}
|
||||
|
||||
// runTier executes one tier pass starting at the round-robin base index.
|
||||
// Cooldown is the only hard skip (re-verified per slot); a busy slot is
|
||||
// skipped without any penalty; a hard failure is recorded and the pass moves
|
||||
// on to the next slot (plan 2.3: "单请求内不重试已失败槽" — the failed slot is
|
||||
// not retried, the others still are). hard == nil and no success means every
|
||||
// candidate was merely busy/cooling, so the caller may wait a bounded time.
|
||||
func runTier(ctx context.Context, tn *TierNode, cands []*Slot, base int64, req *types.ChatRequest, stream bool) tierResult {
|
||||
//
|
||||
// Normal candidates rotate by base; probe candidates form a fixed tail tried
|
||||
// only after every normal slot failed or was busy.
|
||||
func runTier(ctx context.Context, tn *TierNode, cands []candidate, base int64, req *types.ChatRequest, stream bool) tierResult {
|
||||
n := len(cands)
|
||||
norm := normalCount(cands)
|
||||
var hard []TierError
|
||||
for i := 0; i < n; i++ {
|
||||
sl := cands[(int(base)+i)%n]
|
||||
if !sl.Prov.ModelAvailable(sl.Model) {
|
||||
var c candidate
|
||||
if i < norm {
|
||||
c = cands[(int(base)+i)%norm] // rotate within the normal head
|
||||
} else {
|
||||
c = cands[i] // probe tail keeps its order
|
||||
}
|
||||
sl := c.slot
|
||||
// A probe candidate is intentionally NOT re-checked here: it is cooling
|
||||
// by definition, and its permit was already claimed.
|
||||
if !c.probe && !sl.Prov.ModelAvailable(sl.Model) {
|
||||
continue
|
||||
}
|
||||
r := *req
|
||||
@ -233,76 +305,40 @@ func (s *Scheduler) chainDrive(ctx context.Context, chain *Chain, req *types.Cha
|
||||
}
|
||||
var ce ChainErr
|
||||
for _, tn := range chain.Tiers {
|
||||
// initial filter: quota-exhausted and cooling slots are dropped
|
||||
var cands []*Slot
|
||||
for _, sl := range tn.Slots {
|
||||
if exhausted != nil && exhausted(sl) {
|
||||
continue
|
||||
}
|
||||
if !sl.Prov.ModelAvailable(sl.Model) {
|
||||
continue
|
||||
}
|
||||
cands = append(cands, sl)
|
||||
}
|
||||
// initial filter: quota-exhausted slots are dropped, cooling slots are
|
||||
// dropped unless they qualify as half-cooldown probes (appended last).
|
||||
cands := collectCands(tn.Slots, exhausted)
|
||||
if len(cands) == 0 {
|
||||
ce.Skipped = append(ce.Skipped, fmt.Sprintf("tier %d: no schedulable slot (cooling or quota exhausted)", tn.Tier))
|
||||
continue
|
||||
}
|
||||
// preference orders a same-tier run; stable so equal prefs keep order
|
||||
sort.SliceStable(cands, func(i, j int) bool {
|
||||
return cands[i].Prov.Pref(cands[i].Model) > cands[j].Prov.Pref(cands[j].Model)
|
||||
})
|
||||
// No Pref sort: load balancing is done by round-robin cursor.
|
||||
// Persistently failing slots are excluded by ModelSchedulable
|
||||
// (which checks Pref > prefMin).
|
||||
base := tn.NextStart()
|
||||
res := runTier(ctx, tn, cands, base, req, stream)
|
||||
if res.resp != nil || res.chunks != nil {
|
||||
releaseProbes(cands)
|
||||
return res.resp, res.chunks, res.src, res.model, nil
|
||||
}
|
||||
if ctx.Err() != nil {
|
||||
releaseProbes(cands)
|
||||
return nil, nil, "", "", ctx.Err()
|
||||
}
|
||||
if len(res.hard) > 0 {
|
||||
ce.Tiers = append(ce.Tiers, res.hard...)
|
||||
releaseProbes(cands)
|
||||
continue // hard failures: fall through to the next tier, no waiting
|
||||
}
|
||||
// every candidate was busy or cooling: bounded poll before downgrading
|
||||
deadline := time.Now().Add(busyWait)
|
||||
timer := time.NewTimer(busyPoll)
|
||||
defer timer.Stop()
|
||||
for {
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
return nil, nil, "", "", ctx.Err()
|
||||
case <-timer.C:
|
||||
if err := s.pollBusyTier(ctx, tn, cands, base, req, stream, &ce); err != nil {
|
||||
releaseProbes(cands)
|
||||
if r, ok := err.(*tierSuccess); ok {
|
||||
return r.res.resp, r.res.chunks, r.res.src, r.res.model, nil
|
||||
}
|
||||
done := time.Now().After(deadline)
|
||||
if done {
|
||||
ce.Skipped = append(ce.Skipped, fmt.Sprintf("tier %d: no free slot within %v", tn.Tier, busyWait))
|
||||
break
|
||||
}
|
||||
// refresh candidates: cooldowns may have expired meanwhile
|
||||
var again []*Slot
|
||||
for _, sl := range cands {
|
||||
if sl.Prov.ModelAvailable(sl.Model) {
|
||||
again = append(again, sl)
|
||||
}
|
||||
}
|
||||
if len(again) == 0 {
|
||||
ce.Skipped = append(ce.Skipped, fmt.Sprintf("tier %d: no free slot within %v", tn.Tier, busyWait))
|
||||
break
|
||||
}
|
||||
res = runTier(ctx, tn, again, base, req, stream)
|
||||
if res.resp != nil || res.chunks != nil {
|
||||
return res.resp, res.chunks, res.src, res.model, nil
|
||||
}
|
||||
if ctx.Err() != nil {
|
||||
return nil, nil, "", "", ctx.Err()
|
||||
}
|
||||
if len(res.hard) > 0 {
|
||||
ce.Tiers = append(ce.Tiers, res.hard...)
|
||||
break // hard failure while waiting: stop waiting, fall through
|
||||
}
|
||||
timer.Reset(busyPoll)
|
||||
return nil, nil, "", "", err
|
||||
}
|
||||
releaseProbes(cands)
|
||||
}
|
||||
if len(ce.Tiers) == 0 && len(ce.Skipped) == 0 {
|
||||
return nil, nil, "", "", fmt.Errorf("no auto slot configured")
|
||||
@ -310,6 +346,57 @@ func (s *Scheduler) chainDrive(ctx context.Context, chain *Chain, req *types.Cha
|
||||
return nil, nil, "", "", &ce
|
||||
}
|
||||
|
||||
// tierSuccess carries a successful result out of pollBusyTier through the error
|
||||
// return. It is never surfaced to callers of chainDrive.
|
||||
type tierSuccess struct{ res tierResult }
|
||||
|
||||
func (t *tierSuccess) Error() string { return "tier success" }
|
||||
|
||||
// pollBusyTier waits a bounded time for a fully-busy tier to free a slot,
|
||||
// retrying the pass while cooldowns expire. It returns nil when the tier should
|
||||
// be abandoned (caller falls through to the next tier), a *tierSuccess when a
|
||||
// retry succeeded, or a context error.
|
||||
func (s *Scheduler) pollBusyTier(ctx context.Context, tn *TierNode, cands []candidate, base int64, req *types.ChatRequest, stream bool, ce *ChainErr) error {
|
||||
deadline := time.Now().Add(busyWait)
|
||||
timer := time.NewTimer(busyPoll)
|
||||
defer timer.Stop()
|
||||
for {
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
return ctx.Err()
|
||||
case <-timer.C:
|
||||
}
|
||||
if time.Now().After(deadline) {
|
||||
ce.Skipped = append(ce.Skipped, fmt.Sprintf("tier %d: no free slot within %v", tn.Tier, busyWait))
|
||||
return nil
|
||||
}
|
||||
// refresh candidates: cooldowns may have expired meanwhile. Probe
|
||||
// candidates keep their already-claimed permit and stay eligible.
|
||||
var again []candidate
|
||||
for _, c := range cands {
|
||||
if c.probe || c.slot.Prov.ModelAvailable(c.slot.Model) {
|
||||
again = append(again, c)
|
||||
}
|
||||
}
|
||||
if len(again) == 0 {
|
||||
ce.Skipped = append(ce.Skipped, fmt.Sprintf("tier %d: no free slot within %v", tn.Tier, busyWait))
|
||||
return nil
|
||||
}
|
||||
res := runTier(ctx, tn, again, base, req, stream)
|
||||
if res.resp != nil || res.chunks != nil {
|
||||
return &tierSuccess{res: res}
|
||||
}
|
||||
if ctx.Err() != nil {
|
||||
return ctx.Err()
|
||||
}
|
||||
if len(res.hard) > 0 {
|
||||
ce.Tiers = append(ce.Tiers, res.hard...)
|
||||
return nil // hard failure while waiting: stop waiting, fall through
|
||||
}
|
||||
timer.Reset(busyPoll)
|
||||
}
|
||||
}
|
||||
|
||||
// ChainChat runs a non-streaming AUTO request down the chain. exhausted, when
|
||||
// non-nil, decides slot token-quota exhaustion. Returns the response, the
|
||||
// serving source and the exact model id used; on total failure a *ChainErr
|
||||
@ -327,6 +414,65 @@ func (s *Scheduler) ChainChatStream(ctx context.Context, chain *Chain, req *type
|
||||
return chunks, src, model, err
|
||||
}
|
||||
|
||||
// ChainImage runs an image-generation AUTO request down the chain: tiers
|
||||
// ascending (tier 1 highest priority), per-tier round-robin, same-tier order
|
||||
// by preference. Each slot's model is pinned to its own image id (ModelFor),
|
||||
// so a fallback switches per source. Cooling-down slots are skipped. Returns
|
||||
// the response, serving source and the exact model id used; on total failure
|
||||
// a *ChainErr summarizing every tier.
|
||||
func (s *Scheduler) ChainImage(ctx context.Context, chain *Chain, req *types.ImageGenRequest) (*types.UnifiedResponse, string, string, error) {
|
||||
if chain == nil || len(chain.Tiers) == 0 {
|
||||
return nil, "", "", fmt.Errorf("no image auto slot configured")
|
||||
}
|
||||
var ce ChainErr
|
||||
for _, tn := range chain.Tiers {
|
||||
cands := collectCands(tn.Slots, nil)
|
||||
if len(cands) == 0 {
|
||||
ce.Skipped = append(ce.Skipped, fmt.Sprintf("tier %d: no schedulable slot (cooling)", tn.Tier))
|
||||
continue
|
||||
}
|
||||
// No Pref sort: load balancing is done by round-robin cursor.
|
||||
base := tn.NextStart()
|
||||
norm := normalCount(cands)
|
||||
var hard []TierError
|
||||
for i := 0; i < len(cands); i++ {
|
||||
var c candidate
|
||||
if i < norm {
|
||||
c = cands[(int(base)+i)%norm]
|
||||
} else {
|
||||
c = cands[i]
|
||||
}
|
||||
sl := c.slot
|
||||
if !c.probe && !sl.Prov.ModelAvailable(sl.Model) {
|
||||
continue
|
||||
}
|
||||
r := *req
|
||||
r.Model = sl.Model
|
||||
resp, err := sl.Prov.Image(ctx, &r)
|
||||
if ctx.Err() != nil {
|
||||
releaseProbes(cands)
|
||||
return nil, "", "", ctx.Err()
|
||||
}
|
||||
if err == nil {
|
||||
releaseProbes(cands)
|
||||
return resp, sl.Source, sl.Model, nil
|
||||
}
|
||||
if errors.Is(err, types.ErrBusy) {
|
||||
continue
|
||||
}
|
||||
hard = append(hard, TierError{Tier: tn.Tier, Source: sl.Source, Model: sl.Model, Err: err})
|
||||
}
|
||||
releaseProbes(cands)
|
||||
if len(hard) > 0 {
|
||||
ce.Tiers = append(ce.Tiers, hard...)
|
||||
}
|
||||
}
|
||||
if len(ce.Tiers) == 0 && len(ce.Skipped) == 0 {
|
||||
return nil, "", "", fmt.Errorf("no image auto slot configured")
|
||||
}
|
||||
return nil, "", "", &ce
|
||||
}
|
||||
|
||||
// ---- direct scheduling ----
|
||||
|
||||
// Chat runs a chat request across cands, falling back on failure. Each
|
||||
@ -345,11 +491,15 @@ func (s *Scheduler) Chat(ctx context.Context, cands []Provider, req *types.ChatR
|
||||
p := cands[i]
|
||||
r := *req
|
||||
r.Model = p.ModelFor(req.Model)
|
||||
if !p.ModelAvailable(r.Model) {
|
||||
ok, isProbe := p.ModelSchedulable(r.Model)
|
||||
if !ok {
|
||||
lastErr = fmt.Errorf("provider %s: model %q cooling down", p.Name(), r.Model)
|
||||
continue
|
||||
}
|
||||
resp, err := p.Chat(ctx, &r)
|
||||
if isProbe {
|
||||
p.ProbeDone(r.Model)
|
||||
}
|
||||
if ctx.Err() != nil {
|
||||
return nil, "", "", ctx.Err()
|
||||
}
|
||||
@ -377,11 +527,15 @@ func (s *Scheduler) ChatStream(ctx context.Context, cands []Provider, req *types
|
||||
p := cands[i]
|
||||
r := *req
|
||||
r.Model = p.ModelFor(req.Model)
|
||||
if !p.ModelAvailable(r.Model) {
|
||||
ok, isProbe := p.ModelSchedulable(r.Model)
|
||||
if !ok {
|
||||
lastErr = fmt.Errorf("provider %s: model %q cooling down", p.Name(), r.Model)
|
||||
continue
|
||||
}
|
||||
resp, err := p.ChatStream(ctx, &r)
|
||||
if isProbe {
|
||||
p.ProbeDone(r.Model)
|
||||
}
|
||||
if err == nil {
|
||||
return resp, p.Name(), r.Model, nil
|
||||
}
|
||||
@ -400,11 +554,16 @@ func (s *Scheduler) Image(ctx context.Context, cands []Provider, req *types.Imag
|
||||
var lastErr error
|
||||
for i := 0; i < attempts && i < len(cands); i++ {
|
||||
p := cands[i]
|
||||
if !p.ModelAvailable(p.ModelFor(req.Model)) {
|
||||
lastErr = fmt.Errorf("provider %s: model %q cooling down", p.Name(), p.ModelFor(req.Model))
|
||||
im := p.ModelFor(req.Model)
|
||||
ok, isProbe := p.ModelSchedulable(im)
|
||||
if !ok {
|
||||
lastErr = fmt.Errorf("provider %s: model %q cooling down", p.Name(), im)
|
||||
continue
|
||||
}
|
||||
resp, err := p.Image(ctx, req)
|
||||
if isProbe {
|
||||
p.ProbeDone(im)
|
||||
}
|
||||
if err == nil {
|
||||
return resp, p.Name(), nil
|
||||
}
|
||||
|
||||
@ -22,6 +22,14 @@ type fakeProvider struct {
|
||||
busy atomic.Bool
|
||||
fail atomic.Bool
|
||||
chatHits atomic.Int64
|
||||
// probeable makes an unavailable provider eligible as a half-cooldown
|
||||
// probe candidate. probePermit is the single-token permit; probeClaims
|
||||
// counts how many times it was handed out and probeDones how many times it
|
||||
// was returned, so tests can assert the permit is always released.
|
||||
probeable atomic.Bool
|
||||
probePermit atomic.Bool
|
||||
probeClaims atomic.Int64
|
||||
probeDones atomic.Int64
|
||||
}
|
||||
|
||||
func fakeProv(name, model string) *fakeProvider {
|
||||
@ -36,6 +44,25 @@ func (f *fakeProvider) ModelFor(reqModel string) string { return f.model }
|
||||
|
||||
func (f *fakeProvider) ModelAvailable(model string) bool { return f.available.Load() }
|
||||
|
||||
func (f *fakeProvider) ModelSchedulable(model string) (bool, bool) {
|
||||
if f.available.Load() {
|
||||
return true, false
|
||||
}
|
||||
if !f.probeable.Load() {
|
||||
return false, false
|
||||
}
|
||||
if f.probePermit.CompareAndSwap(false, true) {
|
||||
f.probeClaims.Add(1)
|
||||
return true, true
|
||||
}
|
||||
return false, false
|
||||
}
|
||||
|
||||
func (f *fakeProvider) ProbeDone(model string) {
|
||||
f.probeDones.Add(1)
|
||||
f.probePermit.Store(false)
|
||||
}
|
||||
|
||||
func (f *fakeProvider) Pref(model string) int64 { return f.pref.Load() }
|
||||
|
||||
func (f *fakeProvider) Chat(ctx context.Context, req *types.ChatRequest) (*types.UnifiedResponse, error) {
|
||||
@ -177,14 +204,21 @@ func TestChainPreferenceSinksButStaysReachable(t *testing.T) {
|
||||
if err != nil {
|
||||
t.Fatalf("chain: %v", err)
|
||||
}
|
||||
// higher pref (good) is tried first and hard-fails; the pass moves on to
|
||||
// the negative-pref slot, which sinks but stays reachable: with the real
|
||||
// provider its success would RecordSuccess (+1 pref, self-heal)
|
||||
// Round-robin starts at index 0 (neg). neg is available (pref=-5 > -20)
|
||||
// and succeeds. good is never tried because neg already won.
|
||||
if src != "neg" || model != "n" || resp.Content != "neg" {
|
||||
t.Fatalf("served src=%q model=%q content=%q", src, model, resp.Content)
|
||||
}
|
||||
if good.chatHits.Load() == 0 {
|
||||
t.Fatal("higher-pref slot must be tried first")
|
||||
if good.chatHits.Load() != 0 {
|
||||
t.Fatal("neg was tried first and succeeded; good must not be attempted")
|
||||
}
|
||||
// Second request: cursor advances. neg wins again (good hard-fails).
|
||||
resp2, src2, _, err2 := s.ChainChat(context.Background(), ch, chatReq(), nil)
|
||||
if err2 != nil {
|
||||
t.Fatalf("second chain: %v", err2)
|
||||
}
|
||||
if src2 != "neg" || resp2.Content != "neg" {
|
||||
t.Fatalf("second req src=%q content=%q", src2, resp2.Content)
|
||||
}
|
||||
}
|
||||
|
||||
@ -315,3 +349,149 @@ func TestChainResetCooldownAfterSwap(t *testing.T) {
|
||||
t.Fatalf("first start = %d, want 0", base)
|
||||
}
|
||||
}
|
||||
|
||||
// ---- half-cooldown probe candidates (plan 阶段 1.3) ----
|
||||
|
||||
// TestChainProbeIsLastResort pins the ordering rule: a probe candidate must not
|
||||
// steal traffic from a healthy slot in the same tier.
|
||||
func TestChainProbeIsLastResort(t *testing.T) {
|
||||
healthy := fakeProv("healthy", "h")
|
||||
cooling := fakeProv("cooling", "c")
|
||||
cooling.available.Store(false)
|
||||
cooling.probeable.Store(true)
|
||||
ch := BuildChain([]Rule{
|
||||
{Tier: 0, Model: "c", Source: "cooling"}, // configured FIRST on purpose
|
||||
{Tier: 0, Model: "h", Source: "healthy"},
|
||||
}, bySource(healthy, cooling))
|
||||
s := New(0)
|
||||
for i := 0; i < 3; i++ {
|
||||
_, src, _, err := s.ChainChat(context.Background(), ch, chatReq(), nil)
|
||||
if err != nil {
|
||||
t.Fatalf("iter %d: %v", i, err)
|
||||
}
|
||||
if src != "healthy" {
|
||||
t.Fatalf("iter %d served by %q; a probe slot must never outrank a healthy one", i, src)
|
||||
}
|
||||
}
|
||||
if cooling.chatHits.Load() != 0 {
|
||||
t.Fatalf("probe slot was hit %d times while a healthy slot existed", cooling.chatHits.Load())
|
||||
}
|
||||
if got := cooling.probeDones.Load(); got != cooling.probeClaims.Load() {
|
||||
t.Fatalf("probe permits leaked: claims=%d dones=%d", cooling.probeClaims.Load(), got)
|
||||
}
|
||||
}
|
||||
|
||||
// TestChainProbeServesWhenNothingElseCan is the recovery path: with every normal
|
||||
// slot gone, the cooling slot's probe carries the request instead of the tier
|
||||
// failing outright.
|
||||
func TestChainProbeServesWhenNothingElseCan(t *testing.T) {
|
||||
cooling := fakeProv("cooling", "c")
|
||||
cooling.available.Store(false)
|
||||
cooling.probeable.Store(true)
|
||||
ch := BuildChain([]Rule{{Tier: 0, Model: "c", Source: "cooling"}}, bySource(cooling))
|
||||
s := New(0)
|
||||
resp, src, model, err := s.ChainChat(context.Background(), ch, chatReq(), nil)
|
||||
if err != nil {
|
||||
t.Fatalf("probe must serve the request: %v", err)
|
||||
}
|
||||
if src != "cooling" || model != "c" {
|
||||
t.Fatalf("served by src=%q model=%q, want cooling/c", src, model)
|
||||
}
|
||||
if resp == nil || resp.Content != "cooling" {
|
||||
t.Fatalf("bad response: %#v", resp)
|
||||
}
|
||||
if cooling.chatHits.Load() != 1 {
|
||||
t.Fatalf("probe must issue exactly one request, got %d", cooling.chatHits.Load())
|
||||
}
|
||||
if cooling.probeClaims.Load() != 1 || cooling.probeDones.Load() != 1 {
|
||||
t.Fatalf("permit accounting: claims=%d dones=%d, want 1/1", cooling.probeClaims.Load(), cooling.probeDones.Load())
|
||||
}
|
||||
}
|
||||
|
||||
// TestChainProbePermitReleasedOnFailure: a failed probe must still hand its
|
||||
// permit back, otherwise the slot could never be probed again.
|
||||
func TestChainProbePermitReleasedOnFailure(t *testing.T) {
|
||||
cooling := fakeProv("cooling", "c")
|
||||
cooling.available.Store(false)
|
||||
cooling.probeable.Store(true)
|
||||
cooling.fail.Store(true)
|
||||
ch := BuildChain([]Rule{{Tier: 0, Model: "c", Source: "cooling"}}, bySource(cooling))
|
||||
s := New(0)
|
||||
if _, _, _, err := s.ChainChat(context.Background(), ch, chatReq(), nil); err == nil {
|
||||
t.Fatal("expected the failing probe to surface an error")
|
||||
}
|
||||
if cooling.probeClaims.Load() != 1 || cooling.probeDones.Load() != 1 {
|
||||
t.Fatalf("permit accounting: claims=%d dones=%d, want 1/1", cooling.probeClaims.Load(), cooling.probeDones.Load())
|
||||
}
|
||||
// permit is free again for the next attempt
|
||||
if _, _, _, err := s.ChainChat(context.Background(), ch, chatReq(), nil); err == nil {
|
||||
t.Fatal("expected the second probe to fail too")
|
||||
}
|
||||
if cooling.probeClaims.Load() != 2 {
|
||||
t.Fatalf("permit must be reclaimable, claims=%d", cooling.probeClaims.Load())
|
||||
}
|
||||
}
|
||||
|
||||
// TestChainProbeDoesNotBlockTierFallthrough: a cooling tier-1 slot that is not
|
||||
// probeable must still let the request fall through to tier 2.
|
||||
func TestChainProbeDoesNotBlockTierFallthrough(t *testing.T) {
|
||||
cold := fakeProv("cold", "c")
|
||||
cold.available.Store(false) // cooling and NOT probeable
|
||||
backup := fakeProv("backup", "b")
|
||||
ch := BuildChain([]Rule{
|
||||
{Tier: 1, Model: "c", Source: "cold"},
|
||||
{Tier: 2, Model: "b", Source: "backup"},
|
||||
}, bySource(cold, backup))
|
||||
s := New(0)
|
||||
_, src, _, err := s.ChainChat(context.Background(), ch, chatReq(), nil)
|
||||
if err != nil || src != "backup" {
|
||||
t.Fatalf("want fallthrough to backup, got src=%q err=%v", src, err)
|
||||
}
|
||||
if cold.chatHits.Load() != 0 {
|
||||
t.Fatal("non-probeable cooling slot must not be hit")
|
||||
}
|
||||
}
|
||||
|
||||
// TestDirectProbeReleasesPermit covers the non-AUTO paths (Chat/ChatStream/Image).
|
||||
func TestDirectProbeReleasesPermit(t *testing.T) {
|
||||
cooling := fakeProv("cooling", "m")
|
||||
cooling.available.Store(false)
|
||||
cooling.probeable.Store(true)
|
||||
s := New(0)
|
||||
_, src, _, err := s.Chat(context.Background(), []Provider{cooling}, chatReq())
|
||||
if err != nil || src != "cooling" {
|
||||
t.Fatalf("direct probe must serve: src=%q err=%v", src, err)
|
||||
}
|
||||
if cooling.probeClaims.Load() != 1 || cooling.probeDones.Load() != 1 {
|
||||
t.Fatalf("permit accounting: claims=%d dones=%d", cooling.probeClaims.Load(), cooling.probeDones.Load())
|
||||
}
|
||||
}
|
||||
|
||||
// TestChainProbeRoundRobinUnaffected: adding a probe tail must not disturb the
|
||||
// round-robin rotation over the healthy head.
|
||||
func TestChainProbeRoundRobinUnaffected(t *testing.T) {
|
||||
a, b := fakeProv("s1", "a"), fakeProv("s2", "b")
|
||||
cooling := fakeProv("s3", "c")
|
||||
cooling.available.Store(false)
|
||||
cooling.probeable.Store(true)
|
||||
ch := BuildChain([]Rule{
|
||||
{Tier: 0, Model: "a", Source: "s1"},
|
||||
{Tier: 0, Model: "b", Source: "s2"},
|
||||
{Tier: 0, Model: "c", Source: "s3"},
|
||||
}, bySource(a, b, cooling))
|
||||
s := New(0)
|
||||
var got []string
|
||||
for i := 0; i < 4; i++ {
|
||||
_, src, _, err := s.ChainChat(context.Background(), ch, chatReq(), nil)
|
||||
if err != nil {
|
||||
t.Fatalf("iter %d: %v", i, err)
|
||||
}
|
||||
got = append(got, src)
|
||||
}
|
||||
want := []string{"s1", "s2", "s1", "s2"}
|
||||
for i := range want {
|
||||
if got[i] != want[i] {
|
||||
t.Fatalf("rr order = %v, want %v (probe tail must not join the rotation)", got, want)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@ -3,8 +3,10 @@
|
||||
package types
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"strings"
|
||||
"time"
|
||||
)
|
||||
|
||||
@ -18,15 +20,19 @@ var ErrBusy = errors.New("provider busy")
|
||||
// ---- OpenAI wire request (gateway input) ----
|
||||
|
||||
type ChatRequest struct {
|
||||
Model string `json:"model"`
|
||||
Messages []ChatMessage `json:"messages"`
|
||||
Temperature *float64 `json:"temperature,omitempty"`
|
||||
MaxTokens int `json:"max_tokens,omitempty"`
|
||||
Stream bool `json:"stream,omitempty"`
|
||||
Tools []interface{} `json:"tools,omitempty"`
|
||||
ToolChoice interface{} `json:"tool_choice,omitempty"`
|
||||
DisableThinking bool `json:"disable_thinking"`
|
||||
ExtraBody map[string]interface{} `json:"-"`
|
||||
Model string `json:"model"`
|
||||
Messages []ChatMessage `json:"messages"`
|
||||
Temperature *float64 `json:"temperature,omitempty"`
|
||||
MaxTokens int `json:"max_tokens,omitempty"`
|
||||
Stream bool `json:"stream,omitempty"`
|
||||
Tools []interface{} `json:"tools,omitempty"`
|
||||
ToolChoice interface{} `json:"tool_choice,omitempty"`
|
||||
DisableThinking bool `json:"disable_thinking"`
|
||||
ExtraBody map[string]interface{} `json:"-"`
|
||||
|
||||
// ClientSession 是客户端自带的会话标识(请求头 x-session-affinity 等)。
|
||||
// 不参与序列化:它只用于让适配器生成稳定的上游会话,不应泄漏给上游 body。
|
||||
ClientSession string `json:"-"`
|
||||
}
|
||||
|
||||
func (r *ChatRequest) MarshalJSON() ([]byte, error) {
|
||||
@ -60,6 +66,39 @@ type ChatMessage struct {
|
||||
|
||||
func StringContent(s string) json.RawMessage { b, _ := json.Marshal(s); return b }
|
||||
|
||||
// isEmptyJSONArray reports whether raw is the literal empty JSON array [].
|
||||
func isEmptyJSONArray(raw json.RawMessage) bool {
|
||||
t := bytes.TrimSpace(raw)
|
||||
return len(t) == 2 && t[0] == '[' && t[1] == ']'
|
||||
}
|
||||
|
||||
// UnmarshalJSON decodes a chat message, normalising an empty-array content
|
||||
// ("content":[]) to an empty string and dropping an empty tool_calls array.
|
||||
//
|
||||
// Why this exists: Lua adapters json.decode the request and re-encode it, and
|
||||
// an empty Lua table is indistinguishable from an empty JSON array — the
|
||||
// encoder emits {} for both. Agent clients serialise an assistant turn that
|
||||
// carries tool_calls and no text as content:[], so every pass-through adapter
|
||||
// turned it into content:{} — a shape that is not valid OpenAI (content is
|
||||
// string | array of parts | null) and that real upstreams reject with
|
||||
// "400 invalid arguments". Normalising at the decode boundary fixes every
|
||||
// adapter at once, including ones added later.
|
||||
func (m *ChatMessage) UnmarshalJSON(b []byte) error {
|
||||
type alias ChatMessage
|
||||
var a alias
|
||||
if err := json.Unmarshal(b, &a); err != nil {
|
||||
return err
|
||||
}
|
||||
*m = ChatMessage(a)
|
||||
if isEmptyJSONArray(m.Content) {
|
||||
m.Content = StringContent("")
|
||||
}
|
||||
if isEmptyJSONArray(m.ToolCalls) {
|
||||
m.ToolCalls = nil
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
type ToolCall struct {
|
||||
ID string `json:"id"`
|
||||
Type string `json:"type"`
|
||||
@ -70,6 +109,7 @@ type ToolCall struct {
|
||||
// ---- Unified internal representation (what adapters produce) ----
|
||||
|
||||
type UnifiedResponse struct {
|
||||
Model string `json:"model,omitempty"`
|
||||
Content string `json:"content"`
|
||||
ReasoningContent string `json:"reasoning_content,omitempty"`
|
||||
FinishReason string `json:"finish_reason,omitempty"`
|
||||
@ -77,12 +117,45 @@ type UnifiedResponse struct {
|
||||
ToolCalls []ToolCall `json:"tool_calls,omitempty"`
|
||||
// ImageData used by image-generation adapters.
|
||||
ImageData []ImageData `json:"image_data,omitempty"`
|
||||
// Cost is the upstream-reported charge for this request, passed through
|
||||
// verbatim (OpenCode Zen/Go put a decimal **string** at the response top
|
||||
// level). Deliberately not parsed or summed by the gateway: it is an
|
||||
// upstream fact, and only some upstreams report it at all.
|
||||
Cost string `json:"cost,omitempty"`
|
||||
}
|
||||
|
||||
type TokenUsage struct {
|
||||
Prompt int `json:"prompt"`
|
||||
Completion int `json:"completion"`
|
||||
Total int `json:"total"`
|
||||
// PromptTokensDetails mirrors the OpenAI v2 usage.prompt_tokens_details
|
||||
// object so cache-hit counts reported by OpenAI-compatible upstreams
|
||||
// (and by adapters that normalize their own cache fields into it) pass
|
||||
// through to clients that read it — dsh reads cached_tokens from here.
|
||||
PromptTokensDetails *PromptTokensDetails `json:"prompt_tokens_details,omitempty"`
|
||||
// PromptCacheHit / PromptCacheMiss carry the DeepSeek-legacy standalone
|
||||
// fields; dsh falls back to prompt_cache_hit_tokens when
|
||||
// prompt_tokens_details.cached_tokens is absent.
|
||||
PromptCacheHit int `json:"prompt_cache_hit_tokens,omitempty"`
|
||||
PromptCacheMiss int `json:"prompt_cache_miss_tokens,omitempty"`
|
||||
// CompletionTokensDetails mirrors the OpenAI v2 usage.completion_tokens_details
|
||||
// object. OpenCode Zen/Go report how much of the completion was thinking
|
||||
// tokens here; without it clients cannot tell how much of the billed output
|
||||
// was reasoning rather than visible answer.
|
||||
CompletionTokensDetails *CompletionTokensDetails `json:"completion_tokens_details,omitempty"`
|
||||
}
|
||||
|
||||
// CompletionTokensDetails is the OpenAI v2 completion_tokens_details object.
|
||||
type CompletionTokensDetails struct {
|
||||
ReasoningTokens int `json:"reasoning_tokens"`
|
||||
}
|
||||
|
||||
// PromptTokensDetails is the OpenAI v2 prompt_tokens_details object. Only
|
||||
// CachedTokens is emitted (omitempty drops the whole object when zero).
|
||||
type PromptTokensDetails struct {
|
||||
// CachedTokens is always emitted (even 0) so clients can distinguish
|
||||
// "upstream reports cache, this request missed" from "no cache data".
|
||||
CachedTokens int `json:"cached_tokens"`
|
||||
}
|
||||
|
||||
// MarshalJSON emits both the legacy short keys (prompt/completion/total, used
|
||||
@ -90,20 +163,36 @@ type TokenUsage struct {
|
||||
// standard keys (prompt_tokens/completion_tokens/total_tokens). Standard
|
||||
// clients such as DSH and DevEco Code read the *_tokens fields.
|
||||
func (t TokenUsage) MarshalJSON() ([]byte, error) {
|
||||
// Auto-generate prompt_tokens_details from legacy DeepSeek fields when
|
||||
// the upstream adapter only set the standalone hit count (openai.lua
|
||||
// does this in Lua, but other adapters or the standardSSEChunk fallback
|
||||
// may not). dsh reads prompt_tokens_details.cached_tokens first.
|
||||
pdetails := t.PromptTokensDetails
|
||||
if pdetails == nil && t.PromptCacheHit > 0 {
|
||||
pdetails = &PromptTokensDetails{CachedTokens: t.PromptCacheHit}
|
||||
}
|
||||
return json.Marshal(struct {
|
||||
PromptTokens int `json:"prompt_tokens"`
|
||||
CompletionTokens int `json:"completion_tokens"`
|
||||
TotalTokens int `json:"total_tokens"`
|
||||
Prompt int `json:"prompt"`
|
||||
Completion int `json:"completion"`
|
||||
Total int `json:"total"`
|
||||
PromptTokens int `json:"prompt_tokens"`
|
||||
CompletionTokens int `json:"completion_tokens"`
|
||||
TotalTokens int `json:"total_tokens"`
|
||||
Prompt int `json:"prompt"`
|
||||
Completion int `json:"completion"`
|
||||
Total int `json:"total"`
|
||||
PromptTokensDetails *PromptTokensDetails `json:"prompt_tokens_details,omitempty"`
|
||||
PromptCacheHitTokens int `json:"prompt_cache_hit_tokens,omitempty"`
|
||||
PromptCacheMissTokens int `json:"prompt_cache_miss_tokens,omitempty"`
|
||||
CompletionDetails *CompletionTokensDetails `json:"completion_tokens_details,omitempty"`
|
||||
}{
|
||||
PromptTokens: t.Prompt,
|
||||
CompletionTokens: t.Completion,
|
||||
TotalTokens: t.Total,
|
||||
Prompt: t.Prompt,
|
||||
Completion: t.Completion,
|
||||
Total: t.Total,
|
||||
PromptTokens: t.Prompt,
|
||||
CompletionTokens: t.Completion,
|
||||
TotalTokens: t.Total,
|
||||
Prompt: t.Prompt,
|
||||
Completion: t.Completion,
|
||||
Total: t.Total,
|
||||
PromptTokensDetails: pdetails,
|
||||
PromptCacheHitTokens: t.PromptCacheHit,
|
||||
PromptCacheMissTokens: t.PromptCacheMiss,
|
||||
CompletionDetails: t.CompletionTokensDetails,
|
||||
})
|
||||
}
|
||||
|
||||
@ -134,8 +223,12 @@ type ImageGenResponse struct {
|
||||
// streaming tool_calls array (incremental fragments with an index field), which
|
||||
// OpenAI-compatible clients accumulate themselves.
|
||||
type UnifiedChunk struct {
|
||||
Content string `json:"content"`
|
||||
Done bool `json:"done"`
|
||||
Content string `json:"content"`
|
||||
Done bool `json:"done"`
|
||||
// FinishReason carries the upstream finish/stop reason ("tool_calls",
|
||||
// "length", ...) when the adapter provides it; the gateway emits it on
|
||||
// the terminating chunk instead of the default "stop".
|
||||
FinishReason string `json:"finish_reason,omitempty"`
|
||||
ToolCalls json.RawMessage `json:"tool_calls,omitempty"`
|
||||
ReasoningContent string `json:"reasoning_content,omitempty"`
|
||||
// Usage carries the upstream token usage when the stream chunk provides
|
||||
@ -143,6 +236,9 @@ type UnifiedChunk struct {
|
||||
// often has empty choices). Gateway uses it to emit exact usage in the
|
||||
// final stream chunk instead of estimates.
|
||||
Usage *TokenUsage `json:"usage,omitempty"`
|
||||
// Cost is upstream-reported charge for this request (decimal string),
|
||||
// passed through verbatim; OpenCode sends it on its own stream chunk.
|
||||
Cost string `json:"cost,omitempty"`
|
||||
}
|
||||
|
||||
// Meta passed to Lua build_headers hook
|
||||
@ -155,4 +251,15 @@ type BuildMeta struct {
|
||||
Source map[string]interface{} `json:"source"`
|
||||
}
|
||||
|
||||
func Now() int64 { return time.Now().Unix() }
|
||||
func Now() int64 { return time.Now().Unix() }
|
||||
|
||||
// OneLine flattens an error string to a single line capped at n chars.
|
||||
// Shared by the provider layer (upstream error reasons) and the gateway
|
||||
// layer (client-facing AUTO-chain summaries).
|
||||
func OneLine(s string, n int) string {
|
||||
s = strings.Join(strings.Fields(s), " ")
|
||||
if len(s) > n {
|
||||
s = s[:n] + "..."
|
||||
}
|
||||
return s
|
||||
}
|
||||
|
||||
154
internal/types/types_test.go
Normal file
154
internal/types/types_test.go
Normal file
@ -0,0 +1,154 @@
|
||||
package types
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"strings"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestTokenUsageMarshalCacheFields(t *testing.T) {
|
||||
// DeepSeek-style: prompt_cache_hit_tokens + prompt_tokens_details auto-generated
|
||||
u := TokenUsage{
|
||||
Prompt: 100,
|
||||
Completion: 50,
|
||||
Total: 150,
|
||||
PromptCacheHit: 80,
|
||||
PromptCacheMiss: 20,
|
||||
}
|
||||
b, _ := json.Marshal(u)
|
||||
var out map[string]interface{}
|
||||
json.Unmarshal(b, &out)
|
||||
check := func(name string, want interface{}) {
|
||||
got := out[name]
|
||||
if got != want {
|
||||
t.Errorf("%s = %v (want %v)", name, got, want)
|
||||
}
|
||||
}
|
||||
check("prompt_tokens", 100.0)
|
||||
check("completion_tokens", 50.0)
|
||||
check("prompt_cache_hit_tokens", 80.0)
|
||||
check("prompt_cache_miss_tokens", 20.0)
|
||||
details := out["prompt_tokens_details"]
|
||||
if details == nil {
|
||||
t.Error("prompt_tokens_details is missing (should be set from prompt_cache_hit_tokens?)")
|
||||
} else {
|
||||
d := details.(map[string]interface{})
|
||||
if d["cached_tokens"] != 80.0 {
|
||||
t.Errorf("cached_tokens = %v (want 80)", d["cached_tokens"])
|
||||
}
|
||||
}
|
||||
|
||||
// OpenAI-style: prompt_tokens_details.cached_tokens
|
||||
u2 := TokenUsage{
|
||||
Prompt: 200,
|
||||
Completion: 100,
|
||||
Total: 300,
|
||||
PromptTokensDetails: &PromptTokensDetails{
|
||||
CachedTokens: 150,
|
||||
},
|
||||
}
|
||||
b2, _ := json.Marshal(u2)
|
||||
json.Unmarshal(b2, &out)
|
||||
check2 := func(name string, want interface{}) {
|
||||
got := out[name]
|
||||
if got != want {
|
||||
t.Errorf("OpenAI %s = %v (want %v)", name, got, want)
|
||||
}
|
||||
}
|
||||
check2("prompt_tokens", 200.0)
|
||||
details2 := out["prompt_tokens_details"]
|
||||
if details2 == nil {
|
||||
t.Error("OpenAI prompt_tokens_details is missing")
|
||||
} else {
|
||||
d := details2.(map[string]interface{})
|
||||
if d["cached_tokens"] != 150.0 {
|
||||
t.Errorf("OpenAI cached_tokens = %v (want 150)", d["cached_tokens"])
|
||||
}
|
||||
}
|
||||
|
||||
t.Logf("DeepSeek JSON: %s", string(b))
|
||||
t.Logf("OpenAI JSON: %s", string(b2))
|
||||
}
|
||||
|
||||
// TestChatMessageNormalizesEmptyArrayContent pins the fix for the malformed
|
||||
// content:{} that pass-through adapters produced for an assistant turn carrying
|
||||
// tool_calls with no text. Lua adapters cannot tell an empty JSON array from an
|
||||
// empty object, so the decode boundary normalises it for every adapter at once.
|
||||
func TestChatMessageNormalizesEmptyArrayContent(t *testing.T) {
|
||||
var msg ChatMessage
|
||||
raw := `{"role":"assistant","content":[],"tool_calls":[{"id":"call_1","type":"function","function":{"name":"read","arguments":"{}"}}]}`
|
||||
if err := json.Unmarshal([]byte(raw), &msg); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if string(msg.Content) != `""` {
|
||||
t.Fatalf("empty-array content must normalise to \"\", got %s", msg.Content)
|
||||
}
|
||||
if len(msg.ToolCalls) == 0 {
|
||||
t.Fatal("tool_calls must survive normalisation")
|
||||
}
|
||||
|
||||
// A non-empty content array must be left byte-identical (multimodal path).
|
||||
var mm ChatMessage
|
||||
multi := `{"role":"user","content":[{"type":"text","text":"hi"}]}`
|
||||
if err := json.Unmarshal([]byte(multi), &mm); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if string(mm.Content) != `[{"type":"text","text":"hi"}]` {
|
||||
t.Fatalf("multimodal content must pass through untouched, got %s", mm.Content)
|
||||
}
|
||||
|
||||
// A plain string content is untouched.
|
||||
var sm ChatMessage
|
||||
if err := json.Unmarshal([]byte(`{"role":"user","content":"plain"}`), &sm); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if string(sm.Content) != `"plain"` {
|
||||
t.Fatalf("string content must pass through untouched, got %s", sm.Content)
|
||||
}
|
||||
|
||||
// An omitted content stays omitted (omitempty semantics preserved).
|
||||
var om ChatMessage
|
||||
if err := json.Unmarshal([]byte(`{"role":"assistant","tool_calls":[]}`), &om); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if om.Content != nil {
|
||||
t.Fatalf("absent content must stay absent, got %s", om.Content)
|
||||
}
|
||||
if om.ToolCalls != nil {
|
||||
t.Fatalf("empty tool_calls array must be dropped, got %s", om.ToolCalls)
|
||||
}
|
||||
}
|
||||
|
||||
// 上游报告输出里有多少是思考 token;缺少这个字段,客户端无法判断
|
||||
// completion_tokens 里多少是可见回答、多少是思考(两者都按输出计费)。
|
||||
func TestTokenUsageMarshalsCompletionTokensDetails(t *testing.T) {
|
||||
u := TokenUsage{
|
||||
Prompt: 37, Completion: 40, Total: 77,
|
||||
CompletionTokensDetails: &CompletionTokensDetails{ReasoningTokens: 40},
|
||||
}
|
||||
b, err := json.Marshal(u)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
var got map[string]interface{}
|
||||
if err := json.Unmarshal(b, &got); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
ctd, ok := got["completion_tokens_details"].(map[string]interface{})
|
||||
if !ok {
|
||||
t.Fatalf("completion_tokens_details 未输出: %s", b)
|
||||
}
|
||||
if n, _ := ctd["reasoning_tokens"].(float64); int(n) != 40 {
|
||||
t.Errorf("reasoning_tokens = %v, want 40 (%s)", ctd["reasoning_tokens"], b)
|
||||
}
|
||||
// 标准 OpenAI 字段同时存在,老客户端不受影响。
|
||||
if _, ok := got["completion_tokens"]; !ok {
|
||||
t.Errorf("标准 completion_tokens 丢失: %s", b)
|
||||
}
|
||||
|
||||
// 未设置时不得凭空造字段。
|
||||
b2, _ := json.Marshal(TokenUsage{Prompt: 1, Completion: 2, Total: 3})
|
||||
if strings.Contains(string(b2), "completion_tokens_details") {
|
||||
t.Errorf("无数据时不应输出 completion_tokens_details: %s", b2)
|
||||
}
|
||||
}
|
||||
30
packaging/config.example.yaml
Normal file
30
packaging/config.example.yaml
Normal file
@ -0,0 +1,30 @@
|
||||
# llmsproxy — 统一 OpenAI 兼容网关配置(示例模板)
|
||||
#
|
||||
# 这个是打包进发行包、供首次安装使用的模板。它故意不含任何可用的凭据:
|
||||
# 首次启动时程序会生成一个随机 admin key 并写入实际配置(WebUI「密钥」页可见)。
|
||||
#
|
||||
# ⚠ 不要在本文件里填一个写死的具体 key。示例文件是公开的,写进去的任何值
|
||||
# 都会成为一个全网皆知的 admin key(旧版示例里的 sk-gw-local-0001 就踩过
|
||||
# 这个坑:它被直接抄进生产配置并真实可用)。
|
||||
|
||||
# 网关监听地址(默认 :8080,建议绑内网/回环)
|
||||
listen: 127.0.0.1:8080
|
||||
|
||||
# 客户端访问本网关所需的 API Key(Bearer)。留空数组 = 不鉴权(仅内网)。
|
||||
gateway_keys: []
|
||||
|
||||
# 默认模型选择:具体模型 id 或 AUTO(按各源模型的 priority 自动选最高可用源)
|
||||
default_model: AUTO
|
||||
|
||||
# Lua 适配器目录(默认 adapters/,首次启动自动写入内置适配器)
|
||||
adapter_dir: adapters
|
||||
|
||||
# 运行时持久化文件(WebUI 新增/编辑的源会写入此文件,重启后仍生效)
|
||||
runtime_file: runtime.json
|
||||
|
||||
# 全局并发上限(0 = 不限)
|
||||
max_concurrent: 0
|
||||
|
||||
# ---- 上游 LLM 源列表 ----
|
||||
# 留空:安装后到 WebUI 添加源,或从内置模板一键创建。
|
||||
sources: []
|
||||
95
packaging/core-dist.sh
Executable file
95
packaging/core-dist.sh
Executable file
@ -0,0 +1,95 @@
|
||||
#!/usr/bin/env bash
|
||||
# Build the llmsproxy CORE distribution packages (deb + rpm + tar.gz).
|
||||
#
|
||||
# Unlike the GUI packages (Electron), the core is a single static-ish binary
|
||||
# that server admins install with dpkg/rpm. The systemd unit bakes in the
|
||||
# memory tuning (GOGC=50, MALLOC_ARENA_MAX=2) measured on the production box.
|
||||
#
|
||||
# Requires: go, nfpm (PATH), dpkg-deb/rpmbuild for verification
|
||||
# Outputs: cmd/build/dist/llmsproxy_<ver>_amd64.deb
|
||||
# cmd/build/dist/llmsproxy-<ver>.x86_64.rpm
|
||||
# cmd/build/dist/llmsproxy-<ver>-linux-amd64.tar.gz
|
||||
set -euo pipefail
|
||||
|
||||
ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)"
|
||||
cd "$ROOT"
|
||||
|
||||
VERSION="${1:-$(git describe --tags --dirty 2>/dev/null | sed 's/^v//' || echo 0.1.0)}"
|
||||
DIST="$ROOT/cmd/build/dist"
|
||||
BUILD_BIN="$ROOT/cmd/build/llmsproxy"
|
||||
|
||||
log() { printf '\033[1;34m[core-dist]\033[0m %s\n' "$*"; }
|
||||
|
||||
# nfpm is installed via `go install` -> lives in GOPATH/bin, which may not be
|
||||
# on the caller's PATH (non-interactive shells, cron). Resolve it explicitly.
|
||||
NFPM_BIN="$(command -v nfpm || true)"
|
||||
[ -z "$NFPM_BIN" ] && [ -n "$(go env GOPATH)" ] && {
|
||||
[ -x "$(go env GOPATH)/bin/nfpm" ] && NFPM_BIN="$(go env GOPATH)/bin/nfpm"
|
||||
}
|
||||
[ -n "$NFPM_BIN" ] || {
|
||||
echo "[core-dist] nfpm missing. Install: go install github.com/goreleaser/nfpm/v2/cmd/nfpm@v2.41.2"
|
||||
exit 1
|
||||
}
|
||||
|
||||
mkdir -p "$DIST"
|
||||
|
||||
# 1) build the core binary (luajit tag, -s -w like deploy.sh)
|
||||
log "building core binary (luajit, stripped)..."
|
||||
CGO_ENABLED=1 go build -tags luajit -trimpath -ldflags="-s -w" -o "$BUILD_BIN" ./cmd/llmsproxy
|
||||
|
||||
# 2) sanity gates: binary must contain the expected symbols and not be degenerate.
|
||||
# NOTE under `set -o pipefail`, `strings | grep -q` is BROKEN: grep -q exits
|
||||
# on first hit and closes the pipe, strings gets SIGPIPE -> pipeline returns
|
||||
# 141 -> the check spuriously fails. Stage strings output in a temp file.
|
||||
SYMS_TMP="$(mktemp)"
|
||||
strings "$BUILD_BIN" > "$SYMS_TMP"
|
||||
for sym in ProbeSlots reqRing shrinkStepLocked mergeUsage; do
|
||||
grep -qF "$sym" "$SYMS_TMP" || {
|
||||
echo "[core-dist] FATAL: binary missing expected symbol '$sym'"; rm -f "$SYMS_TMP"; exit 1
|
||||
}
|
||||
done
|
||||
rm -f "$SYMS_TMP"
|
||||
[ "$(stat -c%s "$BUILD_BIN")" -gt 5000000 ] || {
|
||||
echo "[core-dist] FATAL: core binary implausibly small"; exit 1
|
||||
}
|
||||
log " ok: $(stat -c%s "$BUILD_BIN") bytes, symbols present"
|
||||
|
||||
# 3) deb + rpm via nfpm. nfpm v2 here does not render {{ .Env.VERSION }}, so
|
||||
# stamp the version into a throwaway copy of the config; -t places the
|
||||
# generated package in the dist dir.
|
||||
log "packaging deb + rpm (v$VERSION)..."
|
||||
NFPM_CFG="$(mktemp)" && sed "s/__VERSION__/$VERSION/g" "$ROOT/packaging/nfpm.yaml" > "$NFPM_CFG"
|
||||
"$NFPM_BIN" package --config "$NFPM_CFG" --packager deb -t "$DIST"
|
||||
"$NFPM_BIN" package --config "$NFPM_CFG" --packager rpm -t "$DIST"
|
||||
rm -f "$NFPM_CFG"
|
||||
|
||||
# 4) tar.gz
|
||||
log "packaging tar.gz..."
|
||||
TAR_DIR="$DIST/llmsproxy-$VERSION-linux-amd64"
|
||||
rm -rf "$TAR_DIR"; mkdir -p "$TAR_DIR"
|
||||
cp "$BUILD_BIN" "$TAR_DIR/llmsproxy"
|
||||
# Use the tracked sanitised template, NOT the local config.yaml: that file is
|
||||
# gitignored and carries the operator's real gateway_keys / upstream keys. It
|
||||
# was previously copied verbatim as config.example.yaml, which shipped a real,
|
||||
# working admin key (sk-gw-local-0001) to every install — and postinst.sh copies
|
||||
# the example to /etc when none exists, so it ended up live in production.
|
||||
cp "$ROOT/packaging/config.example.yaml" "$TAR_DIR/config.example.yaml"
|
||||
cp "$ROOT/packaging/llmsproxy.service" "$TAR_DIR/llmsproxy.service"
|
||||
mkdir -p "$TAR_DIR/adapters"
|
||||
cp "$ROOT"/internal/lua/adapters/*.lua "$TAR_DIR/adapters/"
|
||||
tar -C "$DIST" -czf "$DIST/llmsproxy-$VERSION-linux-amd64.tar.gz" \
|
||||
"$(basename "$TAR_DIR")"
|
||||
rm -rf "$TAR_DIR"
|
||||
|
||||
# 5) verify artifacts are healthy (the same gate that would have caught the
|
||||
# 264KB exe fiasco)
|
||||
log "verifying packages..."
|
||||
for f in "$DIST"/llmsproxy_*.deb "$DIST"/llmsproxy-*.rpm "$DIST"/*.tar.gz; do
|
||||
[ -f "$f" ] || continue
|
||||
sz=$(stat -c%s "$f")
|
||||
[ "$sz" -gt 100000 ] || { echo "[core-dist] FATAL: $f suspiciously small ($sz B)"; exit 1; }
|
||||
log " ✓ $f ($((sz/1024)) KB)"
|
||||
done
|
||||
dpkg-deb -I "$DIST"/llmsproxy_*.deb 2>/dev/null | grep -E 'Package|Version' | head -2
|
||||
|
||||
log "done: $DIST"
|
||||
24
packaging/llmsproxy.service
Normal file
24
packaging/llmsproxy.service
Normal file
@ -0,0 +1,24 @@
|
||||
[Unit]
|
||||
Description=LLMSProxy - unified OpenAI-compatible multi-source LLM gateway
|
||||
After=network.target
|
||||
|
||||
[Service]
|
||||
Type=simple
|
||||
# Memory tuning (measured, see README "内存占用"):
|
||||
# MALLOC_ARENA_MAX=2 caps glibc per-thread malloc arenas. LuaJIT allocates
|
||||
# through cgo -> glibc malloc, and glibc defaults to 8*nproc arenas, so every
|
||||
# OS thread that touches malloc reserved its own ~1 MB arena that is never
|
||||
# returned. Measured: 8-12 arenas -> 0.
|
||||
# GOGC=50 halves the Go heap growth target. On its own it does NOT help (the
|
||||
# saved heap is immediately eaten by more glibc arenas); combined with
|
||||
# MALLOC_ARENA_MAX it cut settled RSS by ~19%. This gateway is I/O bound, so
|
||||
# the extra GC cycles are free.
|
||||
Environment=GOGC=50
|
||||
Environment=MALLOC_ARENA_MAX=2
|
||||
ExecStart=/usr/bin/llmsproxy -config /etc/llmsproxy/config.yaml
|
||||
WorkingDirectory=/etc/llmsproxy
|
||||
Restart=always
|
||||
RestartSec=5
|
||||
|
||||
[Install]
|
||||
WantedBy=multi-user.target
|
||||
62
packaging/nfpm.yaml
Normal file
62
packaging/nfpm.yaml
Normal file
@ -0,0 +1,62 @@
|
||||
# nfpm config for the llmsproxy core gateway binary.
|
||||
# The version placeholder below is substituted by packaging/core-dist.sh (nfpm v2
|
||||
# here does not honour {{ .Env.VERSION }} template rendering), so do not edit it
|
||||
# by hand.
|
||||
name: llmsproxy
|
||||
arch: amd64
|
||||
platform: linux
|
||||
version: __VERSION__
|
||||
epoch: "1"
|
||||
section: net
|
||||
priority: optional
|
||||
maintainer: JianFeeeee <dev@modelrouter.local>
|
||||
description: |
|
||||
Unified OpenAI-compatible multi-source LLM gateway.
|
||||
Routes requests across multiple upstream LLM providers with Lua adapters,
|
||||
model scheduling, quota windows, an elastic Lua worker pool, and a Web UI.
|
||||
vendor: ModelRouter
|
||||
homepage: https://gitcode.com/JianFeeeee/ModelRouter
|
||||
license: MIT
|
||||
contents:
|
||||
# binary
|
||||
- src: ./cmd/build/llmsproxy
|
||||
dst: /usr/bin/llmsproxy
|
||||
file_info:
|
||||
mode: 0755
|
||||
# systemd unit (memory tuning baked in)
|
||||
- src: ./packaging/llmsproxy.service
|
||||
dst: /etc/systemd/system/llmsproxy.service
|
||||
file_info:
|
||||
mode: 0644
|
||||
# example config (seeded to /etc/llmsproxy on first install by postinst).
|
||||
# Must be the tracked sanitised template: ./config.yaml is the operator's own
|
||||
# gitignored config and shipping it would leak real credentials — it shipped a
|
||||
# working admin key (sk-gw-local-0001) to every install for a while.
|
||||
- src: ./packaging/config.example.yaml
|
||||
dst: /usr/share/llmsproxy/config.example.yaml
|
||||
file_info:
|
||||
mode: 0644
|
||||
# built-in Lua adapters (deployed to /etc/llmsproxy/adapters)
|
||||
- src: ./internal/lua/adapters/*.lua
|
||||
dst: /usr/share/llmsproxy/adapters/
|
||||
file_info:
|
||||
mode: 0644
|
||||
# README pointer
|
||||
- src: ./README.md
|
||||
dst: /usr/share/doc/llmsproxy/README.md
|
||||
file_info:
|
||||
mode: 0644
|
||||
|
||||
scripts:
|
||||
postinstall: ./packaging/postinst.sh
|
||||
preremove: ./packaging/prerm.sh
|
||||
|
||||
rpm:
|
||||
group: Applications/Internet
|
||||
summary: Unified OpenAI-compatible multi-source LLM gateway
|
||||
compression: gzip
|
||||
|
||||
deb:
|
||||
compression: gzip
|
||||
fields:
|
||||
Recommends: ca-certificates
|
||||
32
packaging/postinst.sh
Executable file
32
packaging/postinst.sh
Executable file
@ -0,0 +1,32 @@
|
||||
#!/bin/sh
|
||||
# llmsproxy core package postinst
|
||||
set -e
|
||||
|
||||
# Runtime state belongs under /etc/llmsproxy just like the production layout:
|
||||
# runtime.json holds keys/created sources and is NOT touched on reinstall.
|
||||
mkdir -p /etc/llmsproxy/adapters
|
||||
|
||||
if [ -f /etc/llmsproxy/config.yaml ]; then
|
||||
echo "llmsproxy: keeping existing /etc/llmsproxy/config.yaml"
|
||||
else
|
||||
# Do NOT seed a starter config here. Packaging/installing one used to hand
|
||||
# everyone the same file, and while it was copied from the maintainer's
|
||||
# live config it shipped a real, working admin key (sk-gw-local-0001) to
|
||||
# every install.
|
||||
#
|
||||
# Instead, leave the file absent: on first start the binary's EnsureDefault
|
||||
# path creates config.yaml with a FRESH RANDOM admin key and logs it, which
|
||||
# is both unique per install and actually usable. The packaged
|
||||
# config.example.yaml stays in /usr/share as a reference template only.
|
||||
#
|
||||
# (Seeding it with gateway_keys: [] would be worse: the gateway would start
|
||||
# but reject every request, with no way in except editing the file.)
|
||||
echo "llmsproxy: no /etc/llmsproxy/config.yaml — a random admin key will be generated on first start"
|
||||
fi
|
||||
|
||||
# systemd
|
||||
if command -v systemctl >/dev/null 2>&1; then
|
||||
systemctl daemon-reload >/dev/null 2>&1 || true
|
||||
echo "llmsproxy: use 'systemctl enable --now llmsproxy' to start"
|
||||
fi
|
||||
exit 0
|
||||
14
packaging/prerm.sh
Executable file
14
packaging/prerm.sh
Executable file
@ -0,0 +1,14 @@
|
||||
#!/bin/sh
|
||||
# llmsproxy core package prerm
|
||||
set -e
|
||||
|
||||
if command -v systemctl >/dev/null 2>&1; then
|
||||
# stop the service if this package's unit is the one installed
|
||||
if systemctl list-unit-files llmsproxy.service >/dev/null 2>&1 && \
|
||||
systemctl is-active llmsproxy.service >/dev/null 2>&1; then
|
||||
systemctl stop llmsproxy.service >/dev/null 2>&1 || true
|
||||
systemctl disable llmsproxy.service >/dev/null 2>&1 || true
|
||||
fi
|
||||
systemctl daemon-reload >/dev/null 2>&1 || true
|
||||
fi
|
||||
exit 0
|
||||
49
packaging/verify-dist.sh
Executable file
49
packaging/verify-dist.sh
Executable file
@ -0,0 +1,49 @@
|
||||
#!/usr/bin/env bash
|
||||
# Verify packaged GUI artifacts are healthy before a release is declared done.
|
||||
#
|
||||
# Background: on 2026-08-28 (1.3.0) the host wine broke (missing syswow64, so
|
||||
# wine -- the NSIS self-extract step -- failed), and electron-builder silently
|
||||
# wrote a 264 KB installer shell that did NOT embed the 86 MB payload. No gate
|
||||
# caught it, so a broken Windows installer shipped. This script is that gate.
|
||||
#
|
||||
# It enforces a per-artifact minimum size, so a degenerate installer (or a
|
||||
# build that only emitted a shell without its embedded payload) fails loudly.
|
||||
# A size check is the right signal here: a healthy NSIS exe is ~86 MB and the
|
||||
# broken 1.3.0 one was 264 KB (326x) — electron-builder compresses the embedded
|
||||
# app.7z, so a plaintext-payload string probe would be unreliable.
|
||||
set -uo pipefail
|
||||
|
||||
GUI_DIST="${1:-cmd/build/gui-dist}"
|
||||
VERSION="${2:-}"
|
||||
|
||||
log() { printf '[verify-dist] %s\n' "$*"; }
|
||||
fail() { printf '[verify-dist] FATAL: %s\n' "$*" >&2; exit 1; }
|
||||
|
||||
[ -d "$GUI_DIST" ] || fail "no dist dir: $GUI_DIST"
|
||||
|
||||
# (glob, min-bytes) — nsis.7z is NOT a required artifact: with `--win nsis` the
|
||||
# payload is embedded inside the Setup exe and only the exe is the formal product
|
||||
# (the .nsis.7z appears in intermediate/--dir builds only).
|
||||
checks=(
|
||||
"ModelRouter*.AppImage:10000000"
|
||||
"modelrouter-*.deb:10000000"
|
||||
"modelrouter-*.rpm:10000000"
|
||||
"ModelRouter Setup*.exe:5000000"
|
||||
)
|
||||
|
||||
for entry in "${checks[@]}"; do
|
||||
glob="${entry%%:*}"; min="${entry#*:}"
|
||||
# gui-dist accumulates installers from previous versions; check the MOST
|
||||
# RECENT artifact (last in sort) so a stale file never masks a broken build.
|
||||
# `find -name` (not a bare glob) so the space in "ModelRouter Setup *.exe"
|
||||
# is preserved as one file pattern.
|
||||
match="$(find "$GUI_DIST" -maxdepth 1 -name "$glob" -type f | sort -V | tail -1)"
|
||||
[ -z "$match" ] && fail "no artifact matching '$glob' in $GUI_DIST"
|
||||
sz=$(stat -c%s "$match")
|
||||
if [ "$sz" -lt "$min" ]; then
|
||||
fail "$(basename "$match") is only $sz bytes (min $min) — likely a degenerate shell like the 1.3.0 264KB exe"
|
||||
fi
|
||||
log " ✓ $(basename "$match") ($((sz/1024)) KB)"
|
||||
done
|
||||
|
||||
log "all artifacts healthy"
|
||||
Reference in New Issue
Block a user