Compare commits

...

32 Commits

Author SHA1 Message Date
504c5ac9a0 chore(version): 1.7.4 -> 1.7.5
启动密钥警告误报的修复。
2026-09-28 23:42:48 +08:00
70f1c879bd fix(startup): 密钥警告改读真实生效的 key 集合
启动时那条「gateway_keys is EMPTY — without a key every request is rejected」
读的是 legacy 的 cfg.GatewayKeys 段,而鉴权实际用 cfg.Keys(core.ListKeys)。
seedKeys 首次启动把 gateway_keys 搬进 keys[] 之后,YAML 里那个列表就不再
被鉴权使用。于是在它被清空(例如轮换掉 starter key 之后)而 keys[] 仍有
7 把可用 key(含 admin)时,进程每次启动都谎报「所有请求都会被拒绝」。

实测:生产日志出现该警告,而同一个 key 请求 /v1/models 返回 200。

- main.go 改为检查 c.ListKeys(),文案改成不绑定字段名。
- 顺带删掉 gateway.New 的 gatewayKeys 参数:函数体从未使用它,
  只读 ListKeys(),留着会继续诱导人以为鉴权来自那个列表。

判据:e2e/TestStartupWarningReflectsRealKeysNotLegacyList —— 构造
「gateway_keys 空 + keys[] 有 key」的真实形态,先断言该 key 确实能鉴权,
再断言日志里不再出现那句谎报。变异验证:回退成 GatewayKeys() 即变红。
2026-09-28 23:42:48 +08:00
04e544c823 chore(version): 1.7.3 -> 1.7.4
发行包不再内置可用 admin key 的修复,走 patch 发布。
2026-09-28 23:27:54 +08:00
7e33d11d15 fix(packaging): 发行包不再内置可用的 admin key
打包时把本地 config.yaml(gitignored,含运维真实密钥)原样复制成
config.example.yaml,而 postinst 在首次安装且 /etc 无配置时又把它
cp 成生产配置 ⇒ 每次安装都得到一个同值的、公开已知的 admin key。
实测该 key(sk-gw-local-0001)在生产上真实有效(/v1/models 返回 200,
而网关监听 0.0.0.0)。

三处修正:
- 新增 packaging/config.example.yaml(受 git 跟踪的净化模板),
  gateway_keys 留空、sources 留空,并写明不要填死值。
- core-dist.sh / nfpm.yaml 改为打包该模板,不再碰本地 config.yaml。
- postinst.sh 不再投递示例配置:留空文件会让网关启动但拒绝所有请求
  (无门可入)。改为让二进制首启时自行生成随机 admin key 并打印 ——
  每次安装都不同,且开箱可用。示例文件仅作为 /usr/share 下的参考保留。

实测首启:生成 sk-gw-838d66a1... 并打印,与旧的共享固定值不同。
2026-09-28 23:27:42 +08:00
cd82835f25 chore(version): 1.7.2 -> 1.7.3
启动重复播种 admin key 与配置封存非幂等的修复,走 patch 发布。
2026-09-28 22:53:19 +08:00
26ea782350 fix(core): 修复启动重复播种 admin key + 配置封存非幂等
根因是 unseal 时序:NewFromConfig 把解密放在最后,而之前几步已经在读凭据。

1. seedKeys 重复播种(生产已累积 4 个同名 admin key)
   seedKeys 用 cfg.Keys[i].Key 与明文 gateway_keys 比对去重,但此时内存里的
   key 还是密文 enc:v1:…,比对永不命中 ⇒ 每次重启追加一个同值 admin key。
   实测:core.New(path) 连续重启,seeded key 数 2→3→4 递增。
   (旧测试用 NewFromConfig 构造全新内存对象,没有「盘上已有密文」这个前提,
    复现不出 —— 必须走 core.New 这条读盘的生产路径。)

2. 启动恒重写 config.yaml
   migratePlaintextSecrets 按内存状态判断,而 Save() 末尾会把内存恢复为明文,
   于是每次调用都判定「还有明文」并重写;注释却自称幂等。
   改为 UnsealSecrets 在解密前记录「盘上是否明文」,SealIfNeeded 据此决定
   是否写回 ⇒ 已封存的配置启动不再落盘。

原测试 TestMigratePlaintextSecretsIsIdempotent 用 ModTime 比较,两次写落在同一
时间戳刻度内就看不出来,所以表现为 ~1/6 概率的 flake 而非稳定失败。已改为比较
文件内容并走真实启动路径(UnsealSecrets + SealIfNeeded),并顺带消除该 flake。

附带更正:先前判断「rebuildRegistry 也会拿到密文 API key」不成立 ——
mergedSources → resolveSourceKey 对每个 source 独立解密(belt-and-braces),
provider 始终拿到明文。unseal 前置仍予保留,以消除对该兜底路径的隐性依赖、
并让 seedKeys 在明文下比较。

判据:
- TestRestartDoesNotDuplicateSeededKeys(敏感:回退顺序必红)
- TestSealingIsIdempotentAcrossStarts(12/12 稳定,原先 1/6 flake)
- TestProvidersGetPlaintextCredentials(钉 provider 必须拿到明文这一不变量)
2026-09-28 22:52:51 +08:00
5c58244781 chore(version): 1.7.1 -> 1.7.2
token 统计单位修复(流式改用上游真实 usage、图片不再记 token),
影响 per-model 配额计费口径,走 patch 发布。
2026-09-28 22:19:24 +08:00
0121d23f91 fix(tokens): 流式统计改用上游真实 usage,图片不再记 token
两处 token 单位错误,均影响 per-model 配额计费:

1. 流式路径的 prompt/completion 只是「字节÷3」估算。
   pumpStream 明明收到了上游最后一帧的真实 usage,却只发给客户端、
   从不回写审计记录,于是配额按估算值扣。生产实测同一请求:
   上游 prompt=37/completion=179 → 记账 27/262,prompt 低估 1.4x、
   completion 高估 1.5x(双向失真)。同模型流式 completion 中位数
   是非流式的 4-27 倍。非流式路径本就用真实值,两路不一致。
   修法:lastUsage 非零时写回 rec.Prompt/rec.Compl,估算降为兜底
   (上游不报 usage 时仍保留原估算行为)。

2. 图片请求把「图片张数」记成 completion_tokens。
   rec.Compl = int64(len(resp.ImageData)),len 是切片长度即张数
   (生产 38 条 image 记录全是 1),且被计入 token 总量。
   图片生成无 token 概念 ⇒ 新增 Req.ImageCount 独立字段,
   Prompt/Compl 归 0;UI 记录表 image 行改显示张数(新增 i18n thImgs)。

顺带补 TestUILocaleKeyParity:此前无人校验 zh/en 键集合一致,
单边加键不会报错,只会显示原始键名。

新增 token_units_test.go(定值上游 6 项),做过变异验证:
回退修复实测复现 stream=16/173 vs chat=44/100、image completion=3。
2026-09-28 22:19:12 +08:00
de7c372ad2 chore(version): 1.7.0 -> 1.7.1
v1.7.0 的配额语义(整钥总额)与最终设计不符,本 patch 版把配额改为
按模型独立计费。已在生产部署过的 v1.7.0 保留不动,语义修正走 patch。
2026-09-27 19:07:48 +08:00
c51066f0b6 refactor(quota): 配额改为按模型,删除整钥总配额
用户明确要求:配额应当是密钥对应的**每个模型的单独配额**,而非整体配额。

## 语义变更

删除 GWKey.TokenQuota / ReqQuota / Period / Hours(整钥总额)。
ModelScope 新增 ReqQuota —— 请求数配额下沉到每条模型范围。

现在:每条 models[] 各自带 token 配额 + 请求数配额 + 重置周期,
彼此独立。一个模型用满只影响该模型。

★ 为什么不保留整钥总额:它会让「把 A 模型的额度挪给 B」变成一次全局
重分配;按模型独立计费则每个模型各自可控,运维能直接看出哪个模型在吃预算。

## 连带改动

- checkQuota 合并 key 级与 scope 级判定;checkKeyQuotaRetry 整体删除
  (顺带修掉上轮遗留的双重判定:入口不再先判空再重算)
- core:CreateKeyWithQuota / UpdateKeyWithQuota / ApplyQuota 全部删除,
  改由 ValidateScopeQuotas 校验每条 scope 的配额
- admin key:scope 上的配额不强制(admin 的 scope 仍限制模型范围,
  但不强制配额)—— 否则管理员会把自己锁在门外
- /api/v1/keys 不再回显 key 级配额字段(scope 里已含)
- WebUI:删除整钥配额徽标 / 「配额」按钮 / 创建表单的配额组 /
  putScope 的整钥回传;模型砖块与范围编辑器新增「请求数配额」输入,
  徽标显示 `1.0K 77×·1h`(未设配额显示 ∞)

## 判据

- TestOneModelsQuotaDoesNotBlockAnother 是本次核心保证。
  ★ 它第一版是**假判据**:m2 从不消耗,key-wide 计数器与 m1 自己的计数器
  读数恰好相同,退回 key-wide 仍通过。变异测试抓到后改为「先用 m2 花掉
  远超 m1 配额的量,再验证 m1 仍可用」—— 这样两种设计才可区分。
- TestUncappedModelNeverBlocked / TestAdminKeyScopesAreNotEnforced 新增
- UI 契约判据重写:整钥配额界面必须彻底消失(13 个符号)、
  scope 编辑器必须往返 req_quota、putScope 只发 scope 列表
- 错误消息点名具体模型(TestKeyAPIRejectionNamesTheModel)
- 3/3 变异全被抓

实测(真实进程 + 浏览器):m2 配额 500000 连打 25 次全成功,
m1 配额 1000 立即 429「token quota exceeded for "m1" (4315/1000)」,
此后 m2/m3 仍 200。UI:整钥配额元素全为 0,砖块各显配额,
编辑器预填/保存正确,零 JS 异常。
2026-09-27 19:02:13 +08:00
5530912d32 chore(version): 1.6.0 -> 1.7.0
中版本跃迁:新开 release/v1.7.x 承载 1.7.x 全部 patch。
v1.5.x 已发到 v1.6.0(tag),不再追加。
2026-09-27 18:46:59 +08:00
cc5e725226 docs: 补齐 per-key 配额的运维视角文档
代码回流 main 时配额小节已随行,但本轮新增的三项认知此前只存在于
commit message 与判据注释里,运维查不到:

1. **配额拒绝 vs 容量拒绝是两种东西**。配额在入口检查、不占上游槽位,
   是廉价拒绝(实测 19–21ms,429 + Retry-After);容量不足要等满
   busyWait 才 503(约 2.6s)。客户端据此可以区分「等窗口重置」与
   「等上游腾容量」——前者只需耐心,后者通常该降并发或换源。
   附实测对照表(容量 8、0.6s/请求、100 并发)。

2. **配额的开销**。每请求检查 149ns(配了配额)/ 42.6ns(未配配额,
   不碰桶)/ 37ns(admin);窗口查询按窗口长度扫描而非扫全量保留
   (1h 49ns、24h 55ns、30d 3.9us);生产形态内存 3.7MB。
   明确写出「未配配额的 key 几乎不付代价」,运维可放心多建 key。

3. **按源 pin 的桶是惰性创建的,以及它的代价**。无条件维护会让
   20 密钥 × 8 模型 × 3 源多占 18MB,所以只有真被 pin 查询过才维护;
   代价是配置 pinned 配额之前的历史用量无法事后按源拆分,首个窗口
   可能少算 —— 这是个会让排障困惑的行为,必须写出来。

同时更新 WebUI 密钥页说明:卡片配额徽标、「配额」编辑按钮、创建表单
可配预算(admin 自动禁用)、「我的密钥」页展示本 key 预算。

README.md / README_EN.md 同步。
2026-09-27 18:46:39 +08:00
652842783f test(gateway): 补配额桶的真实并发竞态判据
per-key 配额桶是共享 map:每个被记录的请求写它,每个配额检查读它。
单线程单测完全看不到这里的竞态,只有让多个 goroutine 同时读写才有效。

key_quota_concurrency_test.go:64 goroutine 跑 2 秒,并发 Record +
KeyWindowTokens + KeyWindowReqs + KeyWindowModelTokens,其中一条路径
在中途 opt in 惰性创建的 pinned 桶(那条路径一次改两个桶 map)。

go test -race 结果:零 DATA RACE,5444 万 token 全部入账。
全仓 -race(./...)亦全绿。

(cherry picked from commit 18cfd6d32b)
2026-09-27 18:44:41 +08:00
a21ae84cbe perf(gateway): 拒绝路径只判定一次 + 补配额交互判据
复查后修掉一个自己引入的缺陷,并补上此前缺失的交叉场景验证。

## 修复:拒绝路径重复判定

4 个入口原本先 checkModelScope(判是否为空)再 writeScopeReject
(内部又 checkQuota 一次)。即每个【被拒】的请求要跑两遍配额统计,
且两次之间用量可能变化 —— 判定与响应存在理论竞态。

改为 checkQuota 一次判定直接把 *quotaRejection 交给 writeReject,
消息与 Retry-After 都来自同一次读,不再有二次求值。
checkModelScope 保留(只需知道放行与否的调用方仍可用)。

## 补判据:此前完全没验证过的交叉场景

1. TestKeyQuotaWinsOverSlotQuota —— key 配额与 AUTO 槽位配额是两种
   不同作用域的限额(槽位是全网关共享的上游预算,key 配额属于单个
   调用方)。两者同时耗尽时必须报【key 配额】:报槽位配额会被表述成
   「无可用容量」,读起来像上游故障,而调用方能处理的恰恰是 key 配额。
2. TestUncappedKeyNeverBlockedByEmptyScope —— 只配模型范围、不配配额的
   key(生产上 5 把 user key 全是这种)绝不能被槽位检查误伤。

## 复查补测的实测数据

配额检查的真实开销(每请求一次,走完整 checkQuota 路径):
  配了配额    149 ns  0 allocs
  未配配额     42.6 ns 0 allocs   <- 生产上 5/7 把 key 是这种
  admin key    37 ns  0 allocs

未配配额的 key 只付 FindKey 的开销、根本不碰桶。相对一次 LLM 请求
(秒级)可忽略。

生产配置副本(7 key / 16 源 / 真加密凭据 / 真上游)实测:
- 100 并发 -> 50 成功 / 50 容量拒绝,RSS 19.9 -> 25.8 MB
- 生产形态桶内存(7 key x 8 model x 2 源 x 40 天满 retention)
  = 3.73 MB,占 ~32MB 预算的 11%
- 配额记账与 stats 一致:配 63000 配额后报 64062/63000
- **跨重启存活**:重启后从审计日志回放,仍报 64062/63000 并拦截;
  未配配额的 key 仍 200

(cherry picked from commit 9811654b3e)
2026-09-27 18:44:41 +08:00
d072a03c9a perf(gateway): 配额桶扫描改为窗口化 + pinned 桶惰性创建
审查本特性线的性能时发现两个问题,均有实测数据。

## 1. 窗口查询是全扫,代价落在每个请求上

sumBuckets 原来遍历整个 map(最多 960 个小时桶),实测 5.9us/op。
配额检查在每个请求上跑 2-3 次(key 总 token、key 请求数、scope token),
于是单请求多付约 18us。

注意这**不是本改动引入的成本**:main 上既有的 WindowTokens 同样是
5907ns/op(全扫)。是本改动让它在请求路径上被调用得更多。

改为只遍历窗口可能覆盖的桶(键是整点小时,范围是精确的,不是采样):
- 24h 窗口 200ns -> 55ns
- 1h  窗口  80ns -> 49ns
- 30d 窗口 5.9us -> 3.9us(720 次查找,只有配 month 配额时才走到)

等价性由 TestSumBucketsMatchesFullScan 保证(400 组随机桶位置 x 6 种
窗口,对全扫逐项比对)。★ 第一次写错成 floor,判据立刻抓到:
30 天窗口报 8878 而全扫是 8649 —— 正确是 ceil。

## 2. pinned 桶无条件创建,内存最坏 26.7MB

每条记录写两个桶:裸 model 与 "source::model"。但 pinned 桶只有
「配额里显式 pin 了 source」时才会被查。

实测最坏情况(20 key x 8 model x 3 source x 40 天全 retention):
HeapAlloc 26.67MB —— 而 README 宣传「16 源生产实例 ~32-35MB」,
等于吃掉 80% 内存预算。

改为惰性:只有 KeyWindowModelTokens 带 source 查询过某个 (key, model)
之后,才开始维护它的 pinned 桶。

  20key x 8model x 3src   26.67MB -> 8.87MB  (-67%)
  5key x 6model(真实)     3.36MB -> 2.26MB  (-33%)
  5key x 12model            5.82MB -> 3.59MB  (-38%)

代价:配 pinned 配额之前发生的用量无法事后按源拆分(记录里虽然有
Source,但桶只存了裸 model),所以 pinned 配额的首个窗口可能少算。
已在代码注释与判据中写明。

## 其余实测

  Record        main 基线 275ns/429B/3allocs -> 283ns/429B/3allocs
                (+8ns,分配数不变;3 allocs 来自 ring buffer)
  配额检查全路径  115ns / 0 allocs(每请求新增)
  纯读路径      6.5ns / 0 allocs

## 100 并发调度/拒绝压测(真实进程 + 可报并发峰值的假上游)

  容量 100(4+96),0.15s/请求,100 并发  ok=100 fail=0   上游峰值 42
  容量 100,0.15s/请求,200 并发          ok=200 fail=0   上游峰值 97
  容量 8,3s/请求,100 并发               ok=8   fail=92  上游峰值 8
  容量 8,0.6s/请求,100 并发             ok=32  fail=68  上游峰值 8
  容量 8,0.6s/请求,40 并发              ok=32  fail=8   上游峰值 8

上游峰值恒定不超过 max_concurrent,容量拒绝返回 503 + busyWait 有界
等待(约 2.6s)。main 基线在同条件下 ok=32 fail=68、上游峰值 8、
延迟分布相同 —— 配额改动没有触碰调度/拒绝路径。

配额拒绝单独验证(低并发避开容量拒绝):req_quota=50 用尽后
100 并发全部 429 rate_limit_exceeded + Retry-After: 1661,
**延迟仅 19-21ms**、上游 total 未增加 —— 配额在入口廉价拒绝,
不占用任何上游槽位,与容量不足的昂贵等待形成明确分工。

## 判据

key_quota_perf_test.go:4 个基准 + 2 个判据(sumBuckets 等价性、
pinned 桶惰性)。3/3 变异全被抓(无条件建 pinned 桶、firstHour 用
floor、keyHour 不再写)。

(cherry picked from commit be11a06a46)
2026-09-27 18:44:41 +08:00
f814ff7468 fix(webui): 修 7 处弹窗关闭错对象 + 模板管理器变量遮蔽
上一提交只修了自己新加的两处弹窗,全站其余 7 处是同一缺陷:所有对话框
共用 id="modal-wrap"(CSS `#modal-wrap:not(:empty){display:flex}`)且可以叠
加(seed-key 提示就盖在密钥页上),而
`const w = $("#modal-wrap"); w.remove()` 移除的是**文档里第一个**,不是用户
刚提交的那一个。

逐处改为两种安全写法:
- 能拿到按钮的(saveSource / saveTemplate / sortScopeSave / scopeSave /
  keyQuotaSave / createKey / downloadStatsCsv / downloadKeysCsv /
  scrAddFromForm):`btn.closest("#modal-wrap")`
- 拿不到按钮的:新增 `closeTopModal()` 取**最后一个**(用户看到的那个),
  并作为所有 `if (w) w.remove()` 之后的兜底
- 顺带把 sortScopeSave / downloadStatsCsv / downloadKeysCsv / scrAddFromForm
  的签名补上 btn / this 参数 —— 否则 .closest 恒为 null,表单永远不关

同时修一个相邻的既有 bug:`openTemplateModal` 的 `.map((t) => ...)` 用 t 做
循环变量,模板字面量里又调 t("srcEdit"),t 被遮蔽成对象 ⇒ 打开模板管理器
直接抛 `t is not a function`,整个弹窗渲染失败(main 上就有,git show 确认)。
参数改名 tpl。修后模板管理器完整渲染(浏览器实测:DeepSeek / 智谱 / Kimi /
SiliconFlow 各行 + Edit/Delete 按钮文案全部正常,零异常)。

判据从 2 处扩到全量(internal/gateway/ui_quota_contract_test.go,+3 例):
- 全文档扫描:任何 `.remove()` 配裸 `$("#modal-wrap")` 即失败
- 9 个关闭对话框的处理器必须走 .closest 或 closeTopModal
- closeTopModal 必须取 all.length - 1(首尾颠倒就是原 bug)
- 用 .closest 的处理器,其签名必须真的有 btn 参数 —— 否则查找恒为 null,
  表单永远不关
5/5 变异全被抓:sortScopeSave 去 btn 参数、closeTopModal 取第一个、scopeSave
退回裸选择器、saveSource 退回裸选择器、downloadKeysCsv 删兜底。

浏览器实测(共享 Chromium CDP,真实进程,每处都插入一个「decoy」弹窗
占据文档首位,复现原 bug 的触发条件):
- scopeSave / keyQuotaSave / scrAddFromForm / saveSource / saveTemplate
  五个处理器:自己的表单关、decoy 保留 ✓
- 零 JS 异常
★ 测试自身踩了两个坑,都不是代码问题:① `document.querySelector(sel) && .click()`
  在 CDP 里求值为 undefined,改成箭头函数;② 编辑模板时没填名字就点保存,
  saveTemplate 因 `if (!nm)` 早退、fetch 零调用 —— 一开始我把这个误读成
  「修复失效」,加 fetch 拦截 + 读 #s-name 的值才定位到是测试数据缺失。
  ⇒ 「点按钮没反应」要先分清是「事件没触发」「请求失败」还是「早退」。

(cherry picked from commit b5c3fea0bb)
2026-09-27 18:44:41 +08:00
ef631b43dd feat(webui): 密钥配额表单 + 修复弹窗关闭错对象
功能:让 per-key 配额在 WebUI 里可配置可见,之前的实现只有 API 与
config.yaml 能配。

- 密钥卡片头部显示配额徽标(token / 请求数 + 重置窗口),admin key
  不显示编辑入口(服务端本就永不受限,给入口只会让人以为配了会生效)。
- 新增「配额」编辑弹窗:token 配额、请求数配额、重置周期(复用既有
  的 period 词表与 n-hour 联动),预填从 canvas 的 data-* 读。
- 创建密钥弹窗同步加配额字段;选 admin 角色时自动禁用(同样因为服务端
  忽略 admin 的配额)。
- 「我的密钥」页新增 KEY-WIDE QUOTA 列,用户能看到自己这把 key 的预算。

修一个真 bug:保存弹窗用 $("#modal-wrap") 关闭自己,而全站弹窗共用这个
id、且可以叠加(seed key 提示就盖在密钥页上)。实测(共享 Chromium
CDP,seed 提示与配额弹窗共存)确认:保存后被移除的是 seed 提示,配额表单
反而留在屏幕上 —— 症状是「保存了但弹窗没关」,指向的方向完全错。改为用
点击的按钮 btn.closest("#modal-wrap") 解析自己的弹窗。createKey 有同样
问题,一并修。既有文件里另有 7 处同样写法,未动(不在本次范围,且新判据
只对本次改的两处断言,避免误伤)。

判据新增 internal/gateway/ui_quota_contract_test.go(6 例):
- 两个表单必须用 .closest 解析自己的弹窗
- putScope 必须带上 4 个配额字段(API 视其为指针,省略=清空预算)
- 创建请求必须真的发出配额字段
- **数据流判据**:徽标要真读 k.token_quota 等、编辑表单要真读
  canvas 写的 data-kquota 等。只查字面量存在会漏 —— 字段躺在死分支里
  判据照样通过(这是本轮实际踩到的:keyCapBadges 经 keyPeriodSuffix
  间接读 k.period,被判据抓到后我把读取显式化而不是放宽判据)
- 弹窗扫描先剥注释,否则修复说明里引用的字面量会被当成违规
- 复用既有 ui_contract_test.go 的 jsFunctionBody(大括号配平);
  自己第一版用 2000 字符固定窗口,被长注释顶开后仍在窗口外命中后面
  函数的同名字段,读起来像通过 —— 窗口法在这里是假判据

7 个变异全部被抓(unsafe 关闭、putScope 丢字段、createKey 丢字段、
徽标不读字段、canvas 不写 data-*、kq-hours 改名、周期词表缺项)。

浏览器实测(共享 Chromium CDP,真实进程 + 加密配置):
- 徽标渲染 1.0K·1h / 5×·1h;编辑框预填 1000/5/hour,hours 框按周期联动
- 保存后回读 250000/77/nhour/6,徽标更新为 250.0K·6h,toast Saved
- 零 JS 异常
- **关键回归**:编辑模型 scope 后配额仍是 250000/77/nhour,未被清空
- 创建带配额的 key,服务端确认 {t:50000,r:300,p:week,role:user}
- user 视角「我的密钥」显示 777·1h 与 9×·1h

文档:README.md / README_EN.md 补「密钥用量配额」小节(配置示例、
周期词表、429 语义、admin 豁免、整点分桶最晚晚 1 小时释放、PUT 的
省略 vs 0 语义、429 响应样例),特性列表各加一条。

(cherry picked from commit ce66c7f6c2)
2026-09-27 18:44:41 +08:00
9c3aabb7f9 feat(gateway): per-key 用量配额(token + 请求数)与重置周期
问题:密钥控制只能限制模型范围。实测发现三个缺陷,其中前两个让
per-model token_quota 在真实链路上从未生效:

1. 桶键不含 key。scopeTokens 调 WindowTokens(model, source, win),
   桶键是 model / source::model,与调用方无关。实测两把 key 各用
   1000 token,窗口报 2000 —— A key 的额度被 B key 消耗。
2. 无 source pin 的桶永远是空的。真实记录 Source 总被填上,桶键存成
   "deepseek::m1",而无 pin 的查询找 "m1" —— 读到 0,永远 < quota,
   配额形同虚设。实测 WindowTokens("m1","",1h)=0 而 pinned=2000。
3. AUTO scope 走 KeyTokens(key),是全时段累计、永不重置。实测 30 天
   前的 200 token 仍计入 1 小时配额(报 210 而非 10)。配了
   period: hour 也不会每小时归零。

生产 5 把 user key 全是 token_quota: 0,所以前两条一直没暴露。

改动:
- Stats 新增 per-key 小时桶 keyModelHour(key → model → hour)与
  keyHour(key 总量)、keyReqHour(请求数),retention 40 天,与既有
  modelHour 对齐以覆盖最长的 month 窗口;LoadAudit 走 aggregateLocked,
  所以窗口用量跨重启存活。modelHour 保持 key-blind:它服务的是 AUTO
  槽位配额(限制整个网关对某槽位的消耗),语义不同,不应被 per-key
  改造污染。
- 每个请求写两份模型桶:裸 model 与 source::model。无 pin 的 scope
  条目读前者,有 pin 的读后者。
- GWKey 新增 TokenQuota / ReqQuota / Period / Hours:整钥配额,
  跨该 key 所有模型共享一份预算;ReqQuota 覆盖持续请求量(源上的
  RPM 只管突发)。
- 配额耗尽返回 429 + Retry-After(rate_limit_exceeded),而不是 403:
  403 让客户端以为这把 key 永远不能用该模型,直接放弃;429 + 等待
  才能在窗口重置后自动恢复。模型越权仍是 403。
- admin key 永不受配额限制 —— 否则操作者会把自己锁在门外。
- 周期词表在写入时校验,拼错的 period 被拒绝而不是静默当成永不过期
  (那与操作者输入的意图正好相反)。
- PUT /api/keys 的配额字段是指针:省略=保留原值,显式 0=解除限制。
  否则只改模型范围就会悄悄清空预算。

判据 3 个文件 24 例,9 个变异全部被抓:key 隔离、pin 桶缺失、
AUTO 周期、key-blind 退化、429→403、admin 被限、PUT 清空配额、
Validate 失效、pinned 桶缺失。前三个变异最初漏网 —— 判据只测了
Stats 层没测接线,补了走真实 HTTP 的接线层与 API 层判据后抓住。
端到端验证:真实进程 + 加密配置往返,配额字段与 enc:v1 密钥均正常。

(cherry picked from commit 5306251840)
2026-09-27 18:44:41 +08:00
a7355debed feat(deploy): 部署前校验 master.key 可解封配置 + llmsproxy -show-secrets
密钥校验(deploy.sh)
- 新增 verify_master_key,在 build/替换任何文件之前执行。失败则二进制与
  配置分毫未动、服务不受影响(已负向验证:缺钥匙、错钥匙两种情况都挡住)
- 走真实的 -show-secrets 解密路径,而不是只检查钥匙文件格式——格式合法
  但内容不匹配(重新生成、恢复了错的备份、换机器)同样会被拒
- 钥匙来源与 config 包一致:LLMS_PROXY_MASTER_KEY 优先,否则
  dirname(runtime_file)/master.key
- 配置里没有密文时跳过并提示(首次加密场景)

-show-secrets
- llmsproxy -show-secrets -config <path>:把凭据打到 stdout 后退出
- 不启动任何东西、不写任何文件(已验证 mtime 不变)
- 加密往返无损:封存前后输出逐字节一致

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
(cherry picked from commit a8cff57e24)
2026-09-27 17:12:03 +08:00
f6baa13583 feat: WebUI 改为依赖 /api/v1,UI 与 agent 共用一套 API 契约
- sources / sort / keys 三个页面的数据源从 /api/sources 切到 /api/v1/sources
  (写操作仍走 /api/sources:v1 是只读门面,不做变更)
- 编辑弹窗改用 /api/v1/sources/{name}?reveal=credentials(admin-only)取明文 key。
  这是必须的:表单要整体回传源,若不回填 key,改个端口就会把 key 清空。
- 遮蔽视图仍是默认,只有显式 reveal 才返回明文

端到端验证(真浏览器 + 临时实例,非仅 API 测试):
- sources/sort/keys 三页实际发出 GET /api/v1/sources,0 console error
- editSource('demo') → reveal=credentials,#s-key 与 #s-url 正确回填
- 写入往返:改 base_url /v1→/v2 后重开,key 仍在(未被清空)
- 落盘 api_key 明文残留 0、密文 1

测试:+1(reveal 必须 admin,否则任意 user key 可读全部凭据)
变异验证:reveal 去掉 admin 校验 → 403 断言变红

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
(cherry picked from commit 926b9f6565)
2026-09-27 17:12:03 +08:00
fd03ef6e4a feat: 密钥静态加密 + /api/v1 agent 管理 API
密钥加密(写侧封存 / 读侧解封)
- config.yaml 的 sources[].api_key、sources[].headers、keys[].key 落盘即
  AES-256-GCM 密文(enc:v1: 前缀),master.key 复用 runtime store 那把
- 内存里永远是明文:鉴权比对、API 返回新建 key、WebUI 编辑回填都不受影响
- 启动时一次性封存现存明文(幂等,已封存则不写盘);-check 不写文件
- UpsertSourceInYAML 增加 box 参数,新加的源不再以明文落盘
- 解密失败改为硬错误:原先 MustDecrypt 返回密文会被下次 Save 二次封存
  (实测:源 key 18→20、静默损坏),现在启动即失败且配置分毫不动

/api/v1:面向 agent 的管理 API(WebUI 零影响)
- GET /api/v1            机器可读索引,列出每个端点的方法/权限/用途
- GET /api/v1/overview   一次调用看全貌:源 + AUTO 链 + 密钥数 + 健康度
- GET /api/v1/health     仅健康快照
- GET /api/v1/models     按源分组的可路由模型清单
- GET /api/v1/sources[/{name}]  凭据遮蔽后的源
- GET /api/v1/auto       调度链与实时槽位状态
- GET /api/v1/keys       admin only,密钥元数据,绝不回显密钥本身
- 沿用同一套网关 key 鉴权;读端点任意角色,写仍需 admin

测试:15 个新用例(含负向:泄密、越权、写操作必须被拒)
变异验证:maskKey 不遮蔽→红、去掉 admin 校验→红

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
(cherry picked from commit ad28a924a5)
2026-09-27 17:12:03 +08:00
97bb9c6ef4 feat(opencode): 透传 completion_tokens_details.reasoning_tokens 与上游 cost
回答「opencodego 的用量与费用透传呢」时逐字段核对上游产出,发现 usage 漏了
一项、费用整项丢失。

## 上游实际发什么(实测 opencode.ai/zen/go/v1)

  {
    "choices": [...],
    "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"
  }

cost 在**顶层**且是**字符串**。流式时还会单独发一帧:
{"choices":[],"cost":"0"}

## 此前丢了两样

1. completion_tokens_details.reasoning_tokens —— 输出里有多少是思考 token。
   没有它,客户端无法判断 completion_tokens 里多少是可见回答、多少是思考,
   而两者都按输出计费。
2. cost —— 唯一的费用信号,网关整个丢弃。Go 订阅是包月制恒为 "0",
   但 Zen 按量付费模型(以及未来的其它源)有信息量。

顺带修掉一处流式/非流式不一致:命中缓存时上游同时给
prompt_tokens_details.cached_tokens 和独立的 hit/miss,流式路径写成了 elseif,
只留 details,与非流式产出不同(只认独立字段的老客户端会看不到缓存)。

## 实现

- types.TokenUsage += CompletionTokensDetails;UnifiedResponse / UnifiedChunk += Cost
- opencodego/opencodezen 适配器映射两个字段;空 choices 帧改成 usage 与 cost
  都可带(早退只带 usage 会把同帧的 cost 丢干净 —— 新测试先抓到的就是这个)
- Gateway ChatCompletion / ChatChunk += cost,随终帧发(对齐上游的
  {"choices":[],"cost":"0"} 形态)
- Go 兜底 standardSSEChunk 同步支持(openai 系适配器不再漏 reasoning_tokens;
  纯 cost 帧不再被整体丢弃),新增 rawCostString 兼容字符串/数字两种形态

费用只做**搬运**:不解析、不换算、不汇总 —— 它是上游事实,且只有部分上游提供。

## 验证

经网关实测 gozen:deepseek-v4.1-flash,流式与非流式产出逐字段一致:
  prompt_tokens_details.cached_tokens=6784
  prompt_cache_hit_tokens=6784 / miss=148
  completion_tokens_details.reasoning_tokens=16
  cost="0"

测试:TestOpenCodeCostAndReasoningPassthrough(含「无数据不得凭空造字段」反例)、
TestOpenCodeStreamCacheFieldsMatchNonStream、TestTokenUsageMarshalsCompletionTokensDetails。

(cherry picked from commit c744ee151e)
2026-09-27 17:12:03 +08:00
d81074621a fix(opencode): 采纳客户端真实会话 id + 超窗消息不再被限流措辞封杀
两处都源于同一次排查:pi 到底有没有带会话标识、超窗为什么触发不了压缩。

## 1) 客户端会话 id:pi 一直在发,只是被配置关掉了

之前结论是「通用客户端不发会话 id」——只对了一半。pi 有会话 id,且能发:
pi-ai 的 createClient 在 compat.sendSessionAffinityHeaders 为真时,会把
平台会话 id(uuidv7,整个会话恒定)放到 x-session-affinity /
x-client-request-id / session_id 上。该开关默认 false,而 llmsproxy 的
provider 配置里没开,所以此前一直收不到。

现在网关按优先级采纳:x-session-affinity → x-session-id → session_id →
body 的 prompt_cache_key,并把值经 types.ChatRequest.ClientSession 传到
适配器 meta.client_session。适配器的会号种子优先级变为:
客户端会话 id > 首条 user 消息指纹 > 按源固定。

刻意不采纳 x-client-request-id:名字含 request,部分客户端每请求都换,
拿它当会话会让上游前缀缓存永不命中(pi 总会同时发 x-session-affinity,够用)。

实测:抓 127.0.0.1:8081 的真实 pi 请求,配置打开后收到
x-session-affinity = session_id = x-client-request-id = <子会话 uuid>。
上游缓存确为会话级隔离(同前缀、不同会号:A 冷→命中,B 首次仍为 0),
两个不同 header 值互不命中,反证网关确实采纳了客户端会话 id。

## 2) 超窗消息必须「干净」,否则被同链的限流措辞反向封杀

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 明细原样拼在归一化标记后面,
等于让一条限流 tier 的措辞反过来封杀超窗识别。

现在超窗走独立的干净消息:
  context_length_exceeded: context window is full; reduce the length of
  the messages (gozen/deepseek-v4.1-flash)
只留超窗措辞 + 超窗源名,不带任何其它 tier 的文本。

测试:TestOverflowMessageSurvivesRateLimitedSiblingTier 用 pi 的完整判定
顺序(先 NON_OVERFLOW 后 OVERFLOW)断言同链限流 tier 不再封杀超窗识别;
TestClientSessionFromRequestHeaders / TestClientRequestIDIsNotUsedAsSession /
TestOpenCodePrefersClientSessionID 覆盖会话采纳与优先级。

(cherry picked from commit c9c09b2ba2)
2026-09-27 17:12:03 +08:00
4dd3431c26 feat(opencode): per-conversation session via first-user-message fingerprint
Follow-on to the session-stability fix. "Per source" already made the
prefix cache hit, but it puts every conversation into one upstream session.

Using the client's own session id is not possible: capturing real agent
traffic (tcpdump on 127.0.0.1:8081) shows generic OpenAI clients send NO
session identifier at all — no user / session_id / conversation_id /
metadata in the body, and no session header (only X-Stainless-* plus
User-Agent: pi). The x-opencode-session the Go endpoint asks for is an
OpenCode native-client concept that a generic client cannot forward.

Since history is replayed every turn, the FIRST user message is invariant
for the life of a conversation, so it is used as the conversation
fingerprint. The session becomes stable within a conversation and distinct
across conversations; requests with no user message fall back to per-source
stability.

Measured through the gateway (same 5.7k-token prompt): 2nd call
cached_tokens=5504, and an unrelated conversation gets its own session.

Test: TestOpenCodeSessionIsStableForCache covers same-conversation
stability, cross-conversation separation, per-request request ids and the
sessionless fallback.

(cherry picked from commit 39b48e556e)
2026-09-27 17:12:03 +08:00
3d1407e6fa fix(opencode): make x-opencode-session stable so the upstream prefix cache can hit
The opencode adapters derived x-opencode-session from meta.timestamp, i.e. a
brand new session on every request. The upstream prefix cache is
session-scoped, so no request could ever hit it, and the cache fields the
endpoint does report (prompt_tokens_details.cached_tokens,
prompt_cache_hit_tokens/prompt_cache_miss_tokens) always came back 0/absent.

Measured against the live endpoint, same 6032-token prompt:

  fixed session id   -> 2nd call: hit 5888, miss 144
  rotating session id -> every call: hit 0, miss 6032

Fix: derive the session from the source name (stable), matching how
x-opencode-project is already derived. x-opencode-request stays unique per
request — it is only a request identifier, not part of the cache key.
Applied to both opencodego and opencodezen.

Through the gateway the same prompt now reports, on the 2nd call:
  details={'cached_tokens': 5888} hit=5888 miss=144      (non-streaming)
  prompt_tokens_details={'cached_tokens': 5888}          (streaming)

Test: TestOpenCodeSessionIsStableForCache asserts the session is stable
across requests for one source while the request id differs.

(cherry picked from commit 791d198f47)
2026-09-27 17:12:03 +08:00
db320d2f6c fix(opencodego): inject empty reasoning_content on tool-calling turns
v1.5.4 stopped stripping reasoning_content, which fixes clients that send
it — but most agent clients (pi included) never store or replay their
reasoning, keeping only the tool call. OpenCode Go validates the field on
any assistant turn that carries tool_calls and rejects the whole request:

  400 invalid_request_error: The `reasoning_content` in the thinking mode
  must be passed back to the API.

Verified against the live endpoint that an EMPTY string satisfies the
check, so the adapter now fills in "" when a tool-calling assistant turn
has no reasoning_content. Nothing is fabricated: the reasoning shown to
the client is still exactly what the upstream returned for that turn.

Measured: with a tool_call + tool_result history and no reasoning_content,
all 25 configured Go models returned 400 before and all 25 answer
correctly now.

Test: TestOpenCodeGoVsZenReasoning also pins that a plain assistant turn
(no tool calls) must NOT gain the field.

(cherry picked from commit d1a72cd23a)
2026-09-27 17:12:03 +08:00
b544671f5e feat(adapters): split opencode into opencodezen and opencodego
Zen (https://opencode.ai/zen/v1) and Go (https://opencode.ai/zen/go/v1)
are different services with different requirements, and one shared adapter
could not satisfy both.

The decisive difference is reasoning_content:

  * OpenCode Go runs thinking models and REQUIRES the assistant turn's
    reasoning_content to be echoed back. The shared adapter stripped it
    (msg.reasoning_content = nil), so every replay of a thinking turn
    failed with:
      400 invalid_request_error: The `reasoning_content` in the thinking
      mode must be passed back to the API.
    Reproduced directly: the same request with reasoning_content -> 200,
    without -> 400. That is why the Go tier never worked in an agent loop.

  * The Zen free pool must not receive it, so it keeps stripping.

Both adapters keep the earlier fixes they share (never drop an assistant
turn carrying tool_calls; send stream_options only when streaming; role
whitelist; multimodal strip) and the opencode client fingerprint headers —
the Go endpoint additionally REQUIRES x-opencode-session, which the
adapter already sends.

config: localzen -> opencodezen, gozen -> opencodego.
Verified: all 25 gozen models answer correctly through the gateway with a
thinking + tool_call + tool_result history (was 0/25 before), streaming
included; the Zen free models still pass.

Test: TestOpenCodeGoVsZenReasoning pins the Go-keeps / Zen-strips split.
(cherry picked from commit 71e9040a45)
2026-09-27 17:12:03 +08:00
0881a49286 fix(opencode): only send stream_options with stream:true
OpenCode Go (and other strict OpenAI-compatible upstreams) reject a
non-streaming request that carries stream_options with
"stream_options should be set along with stream". The adapter attached it
unconditionally, so every non-stream call through the opencode adapter
failed on those upstreams.

Verified against OpenCode Go: 25/25 configured models now pass a real
completion through the gateway (they previously 400'd).

Test: TestOpenCodeStreamOptionsOnlyWhenStreaming (absent when
non-streaming, include_usage present when streaming).

(cherry picked from commit 0528941e24)
2026-09-27 17:12:03 +08:00
4e2b5bfda9 fix: empty-array content becomes invalid {} on every pass-through adapter; gemini/ollama drop tool calls
Three related forwarding defects found by auditing every adapter with a
tool-calling replay (assistant turn with content:[] + tool_calls).

1) content:[] -> content:{} (all 12 openai-adapter sources, plus
   deepseek/trae/sensenova/agentrouter/github/groq/kimicode/mistral)

   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 a tool-calling assistant turn with
   no text as content:[], so every pass-through adapter rewrote it to
   content:{} — not valid OpenAI (content is string|array|null). Verified
   against a live upstream: content:[] produced "400 invalid arguments"
   while content:"" was accepted.

   Fixed once at the decode boundary (types.ChatMessage.UnmarshalJSON):
   empty-array content normalises to "" and an empty tool_calls array is
   dropped, so every adapter — including future ones — sees a valid shape.

2) gemini dropped tool_calls and never emitted functionCall /
   functionResponse; the tool role also stayed as an invalid role inside
   contents and system was not moved to systemInstruction.

3) ollama copied only role/content, dropping tool_calls and the call
   attribution entirely (it needs tool_name, not tool_call_id).

Test: TestAdaptersPreserveToolCalls asserts, for every adapter, that the
call id (or function name where the wire format has no id), the function
name, the tool result and the trailing user turn all survive, plus a
negative control for plain text.

(cherry picked from commit 9114468753)
2026-09-27 17:12:03 +08:00
61b847da86 fix(opencode): never drop an assistant turn that carries tool_calls
An agent client (pi) serialises an assistant turn whose content is only
[thinking, toolCall] as content:[] with tool_calls. The multimodal-strip
pass treated an empty content array as 'nothing left, drop the whole
message' and discarded the tool_calls with it.

The next message is that call's tool result, so it arrived orphaned: the
model saw a result for a call it had never made and re-issued the same
call on every turn — an endless repeated-tool-call loop. Reproduced
against a capture sink: content:[] lost tool_calls, while content:"" and
content:null kept them.

Only messages with neither usable content NOR a tool call now get
dropped. Content that collapses to empty but still has tool_calls or a
tool_call_id is emitted as "" instead.

Test: TestOpenCodeKeepsToolCallWithEmptyContent (plus a negative control
that an image-only message without tool calls is still dropped).

(cherry picked from commit 499f0cac2f)
2026-09-27 17:12:03 +08:00
6432b61eb6 fix(gateway): AUTO scope grants all models — restrict to routing mode only
A key with scope=[AUTO] could previously:
1. request ANY concrete model id directly (hasScopeModel/checkModelScope
   treated AUTO as a wildcard)
2. see the full 56-model list on /v1/models (intersectModels considered
   AUTO as grant-everything)

AUTO now only authorizes the AUTO routing mode. Direct requests to a
specific model require an explicit scope entry.

Also carries agentrouter.lua WAF fingerprint headers (Origin/Referer/
X-Requested-With) already staged on this branch.

Tests: TestHasScopeModelWithSourcePrefix updated; full suite green.
(cherry picked from commit 7fb8f96b82)
2026-09-27 17:12:03 +08:00
dd708c9680 fix(gateway): 上下文超窗错误归一化,客户端才能压缩重试
问题:上游返回上下文超窗时,客户端(pi)既不压缩也不重试,只看到一条
普通上游错误。链路两处叠加:

1. 措辞不在客户端识别列表里。pi 靠 @earendil-works/pi-ai 的
   OVERFLOW_PATTERNS(25 条正则)判断超窗,而 justworker 返回的是
   「请精简对话历史…(Context window is full…)」——与那 25 条一条都不匹配
   (最接近的 /context window exceeds limit/i 也不命中,因为措辞是 "is full"
   而非 "exceeds limit")。
2. AUTO 链路的每条 tier 错误按 80 字节 OneLine 截断,而 "Context window
   is full" 这类短诊断词常出现在尾部,正好被按字节切掉(直连路径是 160
   才保住)。

修法:
- 新增 overflowMarkers 识别超窗措辞(含中文写法),命中时把客户端可见
  错误归一化为 context_length_exceeded 前缀——它命中 pi 的
  /context[_ ]length[_ ]exceeded/i,超窗因此可被发现并触发压缩重试。
- AUTO 链路每条 tier 错误宽度 80→160,短诊断词不再被截断。

归一化只加前缀,原始诊断信息保留,便于定位是哪一层超窗。

测试:internal/gateway/overflow_err_test.go(5 例,含「标记必须命中 pi
正则」、非超窗不得误标、80 vs 160 宽度的回归对比)。
2026-09-10 20:38:39 +08:00
46 changed files with 6319 additions and 175 deletions

View File

@ -24,6 +24,7 @@
### 强大的多租户调度能力
- **多密钥多租户**:支持无限密钥,每个密钥独立角色、模型范围、Token 配额、重置周期
- **按模型配额**:每把 key 的每个模型单独配 token + 请求数配额与重置周期(小时/周/月/自定义 N 小时);一个模型用满只影响该模型,同 key 其它模型照常;耗尽返 429 + `Retry-After` 并点名模型,admin key 永不受限
- **AUTO 智能调度**:基于优先级档位的分级调度,同优先级源自动轮询负载均衡,故障自动毫秒级故障转移
- **自愈冷却**:冷却上限 5 分钟,过半后放行 1 个探测请求,上游/额度恢复即刻回归轮询,无需等满冷却窗口
- **Token 配额管理**:精确到模型级别的 Token 配额控制,支持小时/周/月/自定义小时周期自动重置
@ -178,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` 解析顺序:
@ -344,7 +420,10 @@ Environment=MALLOC_ARENA_MAX=2
支持按时间范围导出 CSV;点击模型可生成 pin 到该模型的连接配置。
- **对话页**:流式/非流式调试。
- **密钥页**:创建/编辑网关 key,为每个 key 配模型范围(模型 + 源 + token 配额 + 周期),
管理员管理全部 key,用户只看到自己的 key。
管理员管理全部 key,用户只看到自己的 key。每个模型砖块显示自己的配额徽标
(如 `1.0K 77×·1h`,未设配额显示 `∞`),点开可编辑该模型的 token 配额、
请求数配额与重置周期 —— 配额按模型独立生效,一个用满不影响同一 key 的其它模型。
「我的密钥」页对用户展示本 key 的模型范围与各自配额。
- **优先级页**:拖拽积木配置 AUTO 链档位。
- **源页**:在线增删改上游源(API key 等敏感字段加密落盘)。
- **适配器页**:上传 / 删除 Lua 适配器脚本。

View File

@ -38,6 +38,11 @@ Extracted and independently evolved from the multi-source LLM adapter layer of
`reasoning_content`, `tool_calls`, `usage`).
- **Image generation**: `POST /v1/images/generations`, routed to models with
`kind: image`.
- **Per-model quota**: each key gives every model its own token and request caps
plus a reset period (hour/week/month/custom N hours). One model running out
affects only that model — the key's other models keep working. Exhaustion
answers 429 + `Retry-After` naming the model, so a client resumes when the
window rolls over; admin keys are never capped.
- **Multimodal**: `content` arrays (`image_url` etc.) pass through losslessly;
Anthropic/Gemini/Ollama are translated automatically.
- **LuaJIT VM**: golua-binding LuaJIT; each adapter has its own VM + worker
@ -171,11 +176,103 @@ under the `keys` field of the runtime file (encrypted at rest):
delete the seed key.
- The WebUI **Keys page** creates/deletes keys. Each key has a role (`admin`
manages everything, `user` sees only its own key) and an optional **model
scope** (model + source + token quota + reset period).
scope** (model + source + token quota + reset period). Each model brick shows
its own budget badge (e.g. `1.0K 77×·1h`, `∞` when uncapped); clicking it
edits that model's token quota, request quota and reset period. Quotas apply
per model, so one model running out never blocks the key's others. The
"My key" view shows a user its model scopes and their budgets.
- Clients authenticate with any authorized key's plaintext as
`Authorization: Bearer <key>`.
- Deleting a key removes it from the store immediately.
#### Per-key usage quota
**Quotas are per model.** Each key's `models[]` list gives every model its own
token budget and request budget. One model running out affects only that model
— the key's other models keep working.
```yaml
keys:
- key: sk-gw-<hex>
role: user
name: agent-alice
models:
- model: deepseek-v4-flash
token_quota: 1000000 # this key's token budget for this model
req_quota: 20000 # requests within the window
period: nhour # "" | hour | week | month | nhour
hours: 6 # n-hour only
- model: AUTO
token_quota: 5000000 # AUTO is a quota entry like any other
period: hour
- model: kimi-k3 # no quota listed = unlimited
```
- **There is deliberately no key-wide total.** A key-wide cap would make
"move A's budget to B" a global reallocation; per-model budgets keep each
model independently controllable, so it stays visible which model is
actually consuming the spend.
- `period`: empty = never resets (lifetime total); `hour` / `week` / `month` =
fixed windows; `nhour` + `hours` = a custom hour count. **A misspelled
period is rejected at write time** rather than silently becoming a
never-resetting quota.
- Quotas are isolated per key *and* per model within a key: one key exhausting
`m1` never draws on another key's budget, and never blocks the same key's
`m2`.
- Usage counts chat, streaming and image requests, and survives a restart by
replaying the audit log (40 days retained, covering the longest `month`
window).
- An exhausted quota returns **429 + `Retry-After`**
(`rate_limit_exceeded`) and the message names the model that ran out, so a
client resumes when the window rolls over; a model the key may not use stays
403. **Admin keys are never capped** (quotas on their scopes are not
enforced either), so a cap can never lock the operator out.
- Buckets are whole unix hours, so a window frees up at most an hour late
(deliberately freeing late rather than overspending).
- `PUT /api/keys/{key}` submits quotas by submitting `models` — the caps are
part of the scope, so there is no second budget that can drift out of sync
with the model list. An explicit `0` lifts that model's cap.
##### Quota rejection vs capacity rejection
A quota is checked at the entry point and occupies no upstream slot, so it is a
cheap rejection. Running out of capacity instead waits out the bounded
`busyWait` before answering 503. Measured (capacity 8, 0.6 s/request, 100
concurrent clients):
| Case | Status | Latency | Reaches upstream |
|---|---|---|---|
| Quota spent (`req_quota` 50, 100 concurrent) | `429 rate_limit_exceeded` + `Retry-After` | **19–21 ms** | no |
| Capacity exhausted (same topology, no quota) | `503 upstream_error` | ~2.6 s (`busyWait`) | overflow: no |
So a client can tell "wait for the window to roll over" apart from "wait for
upstream capacity" — the first only needs patience, the second usually means
back off concurrency or switch sources.
##### What the quota costs
Buckets are kept per (key, model, whole unix hour) for 40 days. Measured on an
AMD 7840HS:
- A window query scans the window, not the whole retention: 49 ns for 1 h,
55 ns for 24 h, 3.9 µs for 30 d.
- Recording one request: 283 ns, 3 allocations (unchanged from before this
feature; the allocations come from the record ring buffer).
- Memory: the production shape (7 keys x 8 models x 2 sources at full 40-day
retention) costs about **3.7 MB**. The `source::model` bucket used by a
source-pinned quota is created **lazily** — it is maintained only once some
quota actually pins that source, because writing it on every record would
cost an extra 18 MB at 20 keys x 8 models x 3 sources.
The trade-off: usage recorded before a pinned quota was configured cannot be
split by source afterwards, so **a pinned quota may under-count its first
window**.
```
HTTP/1.1 429 Too Many Requests
Retry-After: 2100
{"error":{"type":"rate_limit_exceeded","message":"token quota exceeded for \"deepseek-v4-flash\" (5000000/5000000)"}}
```
### Model routing
`/v1/chat/completions` `model` resolution order:

View File

@ -1,12 +1,12 @@
{
"name": "modelrouter-gui",
"version": "1.4.2",
"version": "1.7.5",
"lockfileVersion": 3,
"requires": true,
"packages": {
"": {
"name": "modelrouter-gui",
"version": "1.4.2",
"version": "1.7.5",
"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,

View File

@ -1,6 +1,6 @@
{
"name": "modelrouter-gui",
"version": "1.4.2",
"version": "1.7.5",
"description": "ModelRouter Desktop — embedded ModelRouter core with tray",
"author": "ModelRouter",
"main": "main.js",

View File

@ -24,6 +24,7 @@ import (
func main() {
cfgPath := flag.String("config", "config.yaml", "path to gateway config file")
checkOnly := flag.Bool("check", false, "validate the config file and exit (0 = valid, 1 = invalid); nothing is started and no file is written")
showSecrets := flag.Bool("show-secrets", false, "print the config's credentials in the clear and exit; nothing is started and no file is written")
flag.Parse()
// -check is the deploy-time preflight: parse and validate the config
@ -41,6 +42,20 @@ func main() {
return
}
// -show-secrets is the operator escape hatch for a config whose
// credentials are sealed at rest: it prints them to stdout and exits
// without starting anything or writing anything. Nothing else in the
// program prints a credential, and this path never logs to a file.
if *showSecrets {
if _, err := os.Stat(*cfgPath); err != nil {
log.Fatalf("[llmsproxy] show-secrets: %v", err)
}
if err := config.PrintSecrets(*cfgPath); err != nil {
log.Fatalf("[llmsproxy] show-secrets: %v", err)
}
return
}
created, err := config.EnsureDefault(*cfgPath)
if err != nil {
log.Fatalf("[llmsproxy] config: %v", err)
@ -58,19 +73,26 @@ func main() {
}
}
gw, err := gateway.New(c, c.GatewayKeys())
gw, err := gateway.New(c)
if err != nil {
log.Fatalf("[llmsproxy] gateway: %v", err)
}
// Ops hygiene: surface the two most common footguns instead of silently
// running with them.
if keys := c.GatewayKeys(); len(keys) == 0 {
log.Printf("[llmsproxy] WARNING: gateway_keys is EMPTY — without a key every request is rejected")
//
// Read the AUTHORITATIVE source. Auth uses core.ListKeys() (cfg.Keys), not
// the legacy cfg.GatewayKeys list: seedKeys copies gateway_keys into keys[]
// on first start and the YAML list is ignored for auth afterwards. Checking
// GatewayKeys() here made the warning lie — once the legacy list was empty
// (e.g. after rotating the starter key away) it reported "every request is
// rejected" while seven working keys, one of them admin, were in service.
if keys := c.ListKeys(); len(keys) == 0 {
log.Printf("[llmsproxy] WARNING: no gateway key configured — every request will be rejected. Add one in the WebUI or set gateway_keys")
} else {
for _, k := range keys {
if k == "sk-gw-local-0001" || k == "sk-local-0001" {
log.Printf("[llmsproxy] WARNING: gateway key %q looks like the starter/example key — rotate it before exposing the gateway", k)
if k.Key == "sk-gw-local-0001" || k.Key == "sk-local-0001" {
log.Printf("[llmsproxy] WARNING: gateway key %q looks like the starter/example key — rotate it before exposing the gateway", k.Key)
}
}
}

View File

@ -244,6 +244,63 @@ verify_config() {
return 1
}
# ---------- 密钥可解封校验 ----------
# config.yaml 里的凭据以 enc:v1: 密文落盘,解密钥匙是 master.key(或
# LLMS_PROXY_MASTER_KEY)。一旦钥匙丢失/错位,配置在**启动时**才会报错,
# 那时服务已经处于「二进制已替换 + 配置已替换」的半死状态,只能靠回滚。
# 所以在替换任何东西之前先验:能不能用现有的钥匙解开现有的密文。
#
# 两种情形要分开:
# 1. 配置里**有**密文 ⇒ 钥匙必须能解开(否则拒绝部署)
# 2. 配置里**没有**密文 ⇒ 全新安装/首次加密,无需钥匙(但会顺手记一条,
# 因为它意味着这次部署会把明文写成密文,而钥匙必须同时就位)
verify_master_key() {
log "密钥可解封校验"
[[ -f "$TARGET_CONFIG" ]] || return 0 # 配置都不在,交给 verify_config 报错
local sealed
sealed=$(grep -c 'enc:v1:' "$TARGET_CONFIG" 2>/dev/null || echo 0)
if [[ "$sealed" -eq 0 ]]; then
warn "配置内无密文(尚未加密或全为明文)——首次加密需要 master.key 可写"
return 0
fi
# 钥匙来源与 config 包一致:env 优先,否则 runtime_file 同目录的 master.key
local runtime_file key_file
runtime_file=$(awk '/^runtime_file:/{print $2; exit}' "$TARGET_CONFIG")
[[ -n "$runtime_file" ]] || runtime_file="$(dirname "$TARGET_CONFIG")/runtime.json"
key_file="$(dirname "$runtime_file")/master.key"
if [[ -n "${LLMS_PROXY_MASTER_KEY:-}" ]]; then
log " 钥匙来源: LLMS_PROXY_MASTER_KEY(环境变量)"
# 走真实解密路径验证,而不是只检查格式
if out=$("$TARGET_BIN" -show-secrets -config "$TARGET_CONFIG" 2>&1); then
log "✓ 凭据可解封($(echo "$out" | grep -cE '^(source|key) ') 项)"
return 0
fi
err "凭据无法用 LLMS_PROXY_MASTER_KEY 解封,未替换任何文件:"
printf '%s\n' "$out" >&2
return 1
fi
if [[ ! -f "$key_file" ]]; then
err "配置含 $sealed 处密文,但钥匙文件不存在: $key_file"
err " 没有它,部署后服务将无法启动。恢复 master.key 或设置 LLMS_PROXY_MASTER_KEY 后重试。"
return 1
fi
log " 钥匙文件: $key_file ($(stat -c '%A' "$key_file" 2>/dev/null || echo '?'),$(wc -c < "$key_file") 字节)"
if out=$("$TARGET_BIN" -show-secrets -config "$TARGET_CONFIG" 2>&1); then
log "✓ 凭据可解封($(echo "$out" | grep -cE '^(source|key) ') 项)"
return 0
fi
err "凭据无法解密(钥匙与配置不匹配),未替换任何文件:"
printf '%s\n' "$out" | head -5 >&2
err " 密钥来源: $key_file"
err " 常见原因:master.key 被重新生成、恢复了不匹配的备份,或换了机器。"
return 1
}
# ---------- 同步适配器 ----------
sync_adapters() {
log "同步适配器"
@ -393,6 +450,10 @@ main() {
parse_args "$@"
log "═══ llmsproxy 原子部署开始 ═══"
precheck
# 密钥校验放在 build/替换之前:它的失败后果最重(配置里的密文永久解不开),
# 必须在动任何文件之前就挡住。此阶段用的是**当前线上二进制**,它已经
# 支持 -show-secrets(首次引入该子命令的版本除外,此时跳过并提示)。
verify_master_key
build
sync_adapters
backup_old_state

View File

@ -326,3 +326,144 @@ func statusOf(resp *http.Response) int {
}
return resp.StatusCode
}
// The startup warning about key configuration must describe the key set that
// AUTH actually uses (cfg.Keys, via core.ListKeys), not the legacy
// gateway_keys YAML list.
//
// It used to check GatewayKeys() only. seedKeys copies gateway_keys into keys[]
// on first start and the YAML list stops being consulted, so once that list was
// emptied — e.g. after rotating the starter key away — the process logged
// "every request will be rejected" while seven working keys, one of them admin,
// were in service. Observed in production: the warning appeared and that same
// key returned HTTP 200 on /v1/models.
//
// The fixture reproduces exactly that shape: gateway_keys EMPTY, with the
// working key present only in keys[].
func TestStartupWarningReflectsRealKeysNotLegacyList(t *testing.T) {
up := newMockUpstream(t)
dir := t.TempDir()
// Seed a real install once so keys[] gets populated from gateway_keys, then
// rewrite the file to look like a rotated install: legacy list cleared, the
// seeded key still present in keys[].
path := writeConfig(t, dir, "127.0.0.1:0", map[string]*mockUpstream{"good": up})
bin := buildBinary(t)
// First run creates the sealed config with keys[] populated.
seedPort := freePort(t)
cfg1 := writeConfig(t, dir, "127.0.0.1:"+seedPort, map[string]*mockUpstream{"good": up})
g1 := startGateway(t, bin, "127.0.0.1:"+seedPort, cfg1)
g1.stop()
data, err := os.ReadFile(cfg1)
if err != nil {
t.Fatal(err)
}
if !strings.Contains(string(data), "keys:") {
t.Fatalf("first run did not populate keys[]:\n%s", data)
}
// Now clear the legacy gateway_keys list, keeping keys[] intact. The
// rewritten file may use different indentation, so match loosely.
cleared, n := clearGatewayKeys(string(data))
if n == 0 {
t.Fatalf("could not find a gateway_keys entry to clear in:\n%s", data)
}
if err := os.WriteFile(cfg1, []byte(cleared), 0o644); err != nil {
t.Fatal(err)
}
// Start again and inspect the log. The config pins its listen address, so
// rewrite that field to the port we will actually poll.
port2 := freePort(t)
rehosted, ok := replaceListen(string(cleared), "127.0.0.1:"+port2)
if !ok {
t.Fatalf("could not rewrite listen in:\n%s", cleared)
}
cfg2 := filepath.Join(dir, "config2.yaml")
if err := os.WriteFile(cfg2, []byte(rehosted), 0o644); err != nil {
t.Fatal(err)
}
g2 := startGateway(t, bin, "127.0.0.1:"+port2, cfg2)
defer g2.stop()
// Sanity: the key really does authenticate, so the warning would be a lie.
resp, _ := g2.do("GET", "/v1/models", "", true)
if resp == nil || resp.StatusCode != http.StatusOK {
t.Fatalf("key from keys[] did not authenticate (status %v); test fixture is wrong",
resp)
}
logs := g2.logString()
if strings.Contains(logs, "gateway_keys is EMPTY") {
t.Errorf("startup logged 'gateway_keys is EMPTY ... every request is rejected' while "+
"a working admin key was in service — the warning must read the authoritative "+
"key set (core.ListKeys), not the legacy YAML list.\nlog:\n%s", logs)
}
_ = path
_ = g1
}
func freePort(t *testing.T) string {
t.Helper()
l, err := net.Listen("tcp", "127.0.0.1:0")
if err != nil {
t.Fatal(err)
}
defer l.Close()
_, port, _ := net.SplitHostPort(l.Addr().String())
return port
}
// stop terminates the gateway and waits, so its log buffer is complete.
func (g *gatewayUnderTest) stop() {
if g == nil || g.cmd == nil || g.cmd.Process == nil {
return
}
_ = g.cmd.Process.Kill()
_, _ = g.cmd.Process.Wait()
}
// logString returns everything the gateway has written so far.
func (g *gatewayUnderTest) logString() string { return g.log.String() }
// clearGatewayKeys empties the gateway_keys list in a config, tolerating any
// indentation the YAML writer chose. Returns the new text and how many entries
// it removed.
func clearGatewayKeys(src string) (string, int) {
lines := strings.Split(src, "\n")
var out []string
removed := 0
inList := false
for _, l := range lines {
trimmed := strings.TrimSpace(l)
if strings.HasPrefix(l, "gateway_keys:") {
out = append(out, "gateway_keys: []")
inList = true
removed++
continue
}
if inList {
if strings.HasPrefix(trimmed, "- ") {
removed++
continue
}
inList = false
}
out = append(out, l)
}
return strings.Join(out, "\n"), removed
}
// replaceListen rewrites the top-level listen: value.
func replaceListen(src, addr string) (string, bool) {
lines := strings.Split(src, "\n")
for i, l := range lines {
if strings.HasPrefix(l, "listen:") {
lines[i] = "listen: " + addr
return strings.Join(lines, "\n"), true
}
}
return src, false
}

View File

@ -27,6 +27,11 @@ type Config struct {
Auto []ModelScope `yaml:"auto,omitempty"` // AUTO 调度链规则(WebUI 优先级页编辑,chat)
AutoImage []ModelScope `yaml:"auto_image,omitempty"` // AUTO 生图调度链规则(WebUI 优先级页·生图)
Keys []GWKey `yaml:"keys,omitempty"` // 网关密钥(WebUI 密钥页管理)
// 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.
@ -192,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
@ -253,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)
}
@ -361,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"`
}
@ -388,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"`
}
@ -398,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

View 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
}

View 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)
}
}

View File

@ -7,6 +7,7 @@ import (
"crypto/rand"
"encoding/hex"
"fmt"
"log"
"os"
"path/filepath"
"sort"
@ -53,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)
}
@ -74,6 +97,12 @@ func NewFromConfig(cfg *config.Config) (*Core, error) {
if err := c.seedPresetTemplates(); err != nil {
return nil, fmt.Errorf("seed preset templates: %w", err)
}
// Seal any credential still in the clear in config.yaml. Startup writes
// nothing when the file was already sealed (hadPlaintextSecrets is decided
// from the on-disk state, before the unseal above).
if err := cfg.SealIfNeeded(hadPlaintextSecrets); err != nil {
return nil, fmt.Errorf("normalize secrets: %w", err)
}
return c, nil
}
@ -238,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
@ -254,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
@ -264,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 != "" {
@ -439,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 != "" {
@ -447,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
}
@ -655,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.
@ -756,3 +803,20 @@ func (c *Core) Close() {
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
}

View 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)
}
}
}

417
internal/gateway/apiv1.go Normal file
View File

@ -0,0 +1,417 @@
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"},
{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",
"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")

View 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")
}
}

View File

@ -7,6 +7,8 @@ import (
"fmt"
"log"
"net/http"
"regexp"
"strconv"
"strings"
"sync/atomic"
"time"
@ -28,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.
@ -38,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 {
@ -62,6 +71,10 @@ type ChatChunk struct {
// Usage is sent in the final chunk of a stream (empty choices) so
// OpenAI-compatible clients can read token usage.
Usage *types.TokenUsage `json:"usage,omitempty"`
// Cost is the upstream-reported charge for this request, passed through
// verbatim (OpenCode reports it as a decimal string). Emitted on the
// terminal chunk, mirroring OpenCode's own {"choices":[],"cost":"0"}.
Cost string `json:"cost,omitempty"`
}
type ChunkChoice struct {
@ -115,13 +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
}
}
@ -152,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{}
@ -179,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 &quotaRejection{
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 &quotaRejection{
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 &quotaRejection{msg: fmt.Sprintf("model %q is not allowed for this key", model)}
}
// scopeTokens returns the tokens a scope entry has consumed within its reset
// window (total for AUTO / per model otherwise).
// quotaWindowSuffix describes a quota's reset window for an error message, so
// a rejected caller can tell a permanent block from one that clears in an hour.
func quotaWindowSuffix(period string, hours int64) string {
switch {
case period == "hour":
return ", resets hourly"
case period == "week":
return ", resets weekly"
case period == "month":
return ", resets monthly"
case period == "nhour" && hours > 1:
return fmt.Sprintf(", resets every %dh", hours)
}
return ""
} // scopeTokens returns the tokens this key consumed on the scope entry's model
// within its reset window, isolated per key. For an AUTO entry the cap covers
// everything the key routed through AUTO; for a model entry it covers that
// model only.
//
// It reads the per-key hourly buckets rather than the key-blind model
// buckets, so one key's usage can never exhaust another's quota.
func (g *Gateway) scopeTokens(ctx context.Context, sc config.ModelScope) int64 {
k := keyID(reqKey(ctx))
if sc.Model != "" && strings.EqualFold(sc.Model, "AUTO") {
return g.stats.KeyTokens(k)
}
win := AutoPeriodSeconds(sc.Period, sc.Hours)
return g.stats.WindowTokens(sc.Model, sc.Source, win)
if sc.Model != "" && strings.EqualFold(sc.Model, "AUTO") {
return g.stats.KeyWindowTokens(k, win)
}
return g.stats.KeyWindowModelTokens(k, sc.Model, sc.Source, win)
}
// hasScopeModel reports whether a model (possibly with a "source-model" /
@ -229,18 +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"), ""
@ -287,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
@ -301,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()
@ -317,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)),
@ -354,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()
@ -372,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)),
@ -517,21 +609,115 @@ func upstreamErrStatus(err error) int {
return http.StatusBadGateway
}
// overflowMarkers 匹配上游「上下文超窗」类措辞。
//
// 上游写法五花八门,且**不在 pi 客户端的识别列表里**。pi 靠
// @earendil-works/pi-ai 的 OVERFLOW_PATTERNS 判断超窗并据此触发压缩重试,
// 而 justworker 返回的是「请精简对话历史…(Context window is full…)」——
// 与那 25 条正则一条都不匹配,于是 pi 既不压缩也不重试,只把它当成一条
// 普通上游错误。这里把可识别的超窗措辞归一化成 pi 一定认得的标记。
var overflowMarkers = []*regexp.Regexp{
regexp.MustCompile(`(?i)context[ _-]?window is full`),
regexp.MustCompile(`(?i)context[_ ]length[_ ]exceeded`),
regexp.MustCompile(`(?i)exceeds? the context window`),
regexp.MustCompile(`(?i)maximum context length`),
regexp.MustCompile(`(?i)reduce the length of the messages`),
regexp.MustCompile(`(?i)too many tokens`),
regexp.MustCompile(`(?i)token limit exceeded`),
regexp.MustCompile(`请精简对话历史`),
regexp.MustCompile(`上下文(长度)?超(出|限)`),
regexp.MustCompile(`对话历史过长`),
}
// overflowCanonical 命中 pi 的 /context[_ ]length[_ ]exceeded/i。
//
// 注意:这个前缀只是必要条件,不是充分条件——pi 还会先用
// NON_OVERFLOW_PATTERNS 排除整条消息,见 overflowClientMessage。
const overflowCanonical = "context_length_exceeded"
// looksLikeOverflow 判断错误文本是否属于上下文超窗。
func looksLikeOverflow(s string) bool {
for _, re := range overflowMarkers {
if re.MatchString(s) {
return true
}
}
return false
}
// clientUpstreamErr collapses upstream failure details into a short message
// for the client: per-tier bodies (WAF HTML pages, quota payloads, ...) stay
// in rec.Err / the stats API and the server log instead of the response.
// Per-tier one-line reasons are kept (quota/cooling skips carry no error body
// and are the actionable part); each is capped so HTML dumps can't leak.
//
// 超窗会被归一化成「干净」的 overflowCanonical 消息,见 overflowClientMessage。
func clientUpstreamErr(err error) string {
log.Printf("[gateway] upstream failure surfaced to client: %v", err)
if looksLikeOverflow(err.Error()) {
return overflowClientMessage(err)
}
return upstreamErrSummary(err)
}
// overflowHint 是给客户端的超窗说明,刻意只用 pi 认得的措辞。
const overflowHint = "context window is full; reduce the length of the messages"
// clientSessionFromRequest 取客户端自带的会话标识,拿不到时返回 ""。
//
// 通用 OpenAI 客户端默认不带会话 id,但 pi 支持:只要 provider 的 compat 里打开
// sendSessionAffinityHeaders,pi 就会把**平台真实的会话 id**(uuidv7,整个会话
// 恒定)放到 x-session-affinity / x-client-request-id 上(sessionAffinityFormat
// 为 openrouter 时是 x-session-id,为 openai 时是 session_id)。
// 有了它就不必再从首条 user 消息推导会话指纹(那只在客户端不发会话 id 时才作为
// 退路,且历史压缩后会漂移)。prompt_cache_key 是 OpenAI 原生的缓存键,
// 客户端若带也一并当会话用。
//
// 刻意不采纳 x-client-request-id:名字含 request,部分客户端每请求都换,
// 拿它当会话会让上游缓存永不命中。pi 总会同时发 x-session-affinity,够用。
func clientSessionFromRequest(r *http.Request, bodyKey string) string {
for _, h := range []string{"X-Session-Affinity", "X-Session-Id", "Session-Id"} {
if v := strings.TrimSpace(r.Header.Get(h)); v != "" {
return v
}
}
return strings.TrimSpace(bodyKey)
}
// overflowClientMessage 为超窗失败生成「干净」的客户端消息。
//
// 为什么不能沿用 upstreamErrSummary:pi 的 isContextOverflow 先查
// NON_OVERFLOW_PATTERNS(/rate limit/、/too many requests/、Bedrock 前缀),
// 一旦命中就直接判为「非超窗」——**哪怕消息里已经有 context_length_exceeded**,
// pi 也不会压缩重试。而 AUTO 链的失败消息天生是多 tier 原因的拼接:
// 超窗 tier(gozen 400 maximum context length)常与配额/限流 tier
// (429 token plan exhausted、cooling、no free slot)同时出现。
// 把明细原样带出去,等于让一条限流 tier 的措辞反过来封杀超窗识别。
// 所以这里只保留「超窗」措辞 + 是哪个源超的窗,其余一律不带。
func overflowClientMessage(err error) string {
var ce *scheduler.ChainErr
if errors.As(err, &ce) {
for _, t := range ce.Tiers {
if looksLikeOverflow(t.Err.Error()) {
return fmt.Sprintf("%s: %s (%s/%s)", overflowCanonical, overflowHint, t.Source, t.Model)
}
}
}
return fmt.Sprintf("%s: %s", overflowCanonical, overflowHint)
}
// upstreamErrSummary 把上游失败压成一行短消息。
func upstreamErrSummary(err error) string {
var ce *scheduler.ChainErr
if errors.As(err, &ce) {
parts := make([]string, 0, len(ce.Tiers)+len(ce.Skipped))
for _, t := range ce.Tiers {
parts = append(parts, types.OneLine(fmt.Sprintf("%s/%s: %v", t.Source, t.Model, t.Err), 80))
// 160 而不是 80:短诊断词("Context window is full")常落在尾部,
// 80 字节按字节截断正好会把它切掉,超窗就再也认不出来。
parts = append(parts, types.OneLine(fmt.Sprintf("%s/%s: %v", t.Source, t.Model, t.Err), 160))
}
for _, sk := range ce.Skipped {
parts = append(parts, types.OneLine(sk, 80))
parts = append(parts, types.OneLine(sk, 160))
}
msg := strings.Join(parts, "; ")
if len(msg) > 300 {
@ -606,6 +792,7 @@ func writeChatCompletion(w http.ResponseWriter, resp *types.UnifiedResponse, mod
Created: time.Now().Unix(),
Model: modelName,
Choices: []ChatChoice{{Index: 0, Message: msg, FinishReason: resp.FinishReason}},
Cost: resp.Cost,
}
if resp.TokenUsage.Total > 0 || resp.TokenUsage.Prompt > 0 || resp.TokenUsage.Completion > 0 {
out.Usage = &resp.TokenUsage
@ -725,11 +912,15 @@ func (g *Gateway) pumpStream(w http.ResponseWriter, rec *Req, chunks <-chan type
// source's status-page latency average.
rec.FirstByteMs = time.Since(t0).Milliseconds()
var lastUsage *types.TokenUsage
var lastCost string
lastFinish := ""
for ck := range chunks {
if ck.Usage != nil {
lastUsage = mergeUsage(lastUsage, ck.Usage)
}
if ck.Cost != "" {
lastCost = ck.Cost
}
chunk := ChatChunk{
ID: id, Object: "chat.completion.chunk", Created: created, Model: modelName,
}
@ -766,6 +957,19 @@ func (g *Gateway) pumpStream(w http.ResponseWriter, rec *Req, chunks <-chan type
// Final usage chunk (OpenAI standard: empty choices + usage before [DONE]).
// Prefer the upstream's exact usage if the stream carried it; fall back to
// the gateway's estimate otherwise.
// Write the upstream's exact numbers back onto the audit record. Without
// this the streamed path kept only the per-chunk byte estimate, so the same
// request recorded ~1.4x its real completion tokens (measured: upstream 100,
// audit 145) while the non-streaming path recorded 100. Two paths, two
// different numbers for one request is a reporting bug, not a rounding one.
if lastUsage != nil {
if lastUsage.Prompt > 0 {
rec.Prompt = int64(lastUsage.Prompt)
}
if lastUsage.Completion > 0 {
rec.Compl = int64(lastUsage.Completion)
}
}
var tut *types.TokenUsage
if lastUsage != nil {
tut = lastUsage
@ -796,6 +1000,14 @@ func (g *Gateway) pumpStream(w http.ResponseWriter, rec *Req, chunks <-chan type
ID: id, Object: "chat.completion.chunk", Created: created, Model: modelName,
Choices: []ChunkChoice{},
Usage: tut,
Cost: lastCost,
})
} else if lastCost != "" {
// Cost arrived without any usage (OpenCode sends it on its own frame).
send(ChatChunk{
ID: id, Object: "chat.completion.chunk", Created: created, Model: modelName,
Choices: []ChunkChoice{},
Cost: lastCost,
})
}
fmt.Fprintf(w, "data: [DONE]\n\n")
@ -903,8 +1115,8 @@ func (g *Gateway) handleImage(w http.ResponseWriter, r *http.Request) {
}
if isAuto(model) {
if chain := g.core.AutoImageChain(); chain != nil && len(chain.Tiers) > 0 {
if msg := g.checkModelScope(r.Context(), "AUTO"); msg != "" {
writeError(w, http.StatusForbidden, "model_not_allowed", msg)
if q := g.checkQuota(r.Context(), "AUTO"); q != nil {
g.writeReject(w, q)
return
}
done := g.stats.Begin()
@ -926,7 +1138,10 @@ func (g *Gateway) handleImage(w http.ResponseWriter, r *http.Request) {
}
rec.OK = true
rec.Status = http.StatusOK
rec.Compl = int64(len(resp.ImageData))
// Image generation has no token concept. Recording len(ImageData)
// (the image COUNT) in completion_tokens mislabels image count as
// tokens and feeds it into the token totals; leave it 0.
rec.ImageCount = len(resp.ImageData)
g.writeRec(rec)
writeJSON(w, http.StatusOK, types.ImageGenResponse{
Created: time.Now().Unix(),
@ -950,8 +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()
@ -975,7 +1190,8 @@ func (g *Gateway) handleImage(w http.ResponseWriter, r *http.Request) {
}
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(),

View 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)
}
}

View File

@ -58,7 +58,7 @@ func newTestGateway(t *testing.T, srcs ...config.Source) *Gateway {
t.Fatalf("core: %v", err)
}
t.Cleanup(c.Close)
g, err := New(c, []string{"sk-test"})
g, err := New(c)
if err != nil {
t.Fatalf("gateway: %v", err)
}
@ -670,8 +670,11 @@ func TestHasScopeModelWithSourcePrefix(t *testing.T) {
if g.hasScopeModel(scope, "deepseek-v4-flash-free-extra") {
t.Error("hasScopeModel returned true for unrelated model")
}
if !g.hasScopeModel([]config.ModelScope{{Model: "AUTO"}}, "zen:anything") {
t.Error("AUTO scope should allow any prefixed model")
if g.hasScopeModel([]config.ModelScope{{Model: "AUTO"}}, "zen:deepseek-v4-flash-free") {
t.Error("AUTO scope must NOT grant specific model ids even with prefix")
}
if g.hasScopeModel([]config.ModelScope{{Model: "AUTO"}}, "deepseek-v4-pro") {
t.Error("AUTO scope must NOT grant specific model ids")
}
}

View 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

View 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")
}
}

View 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")
}
}

View 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)
}
}

View 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)
}
}
}

View File

@ -48,13 +48,7 @@ func (g *Gateway) handleKeysAPI(w http.ResponseWriter, r *http.Request) {
writeError(w, http.StatusBadRequest, "invalid_request", "invalid json: "+err.Error())
return
}
if body.Role == "" {
body.Role = "user"
}
if body.Role != "admin" && body.Role != "user" {
writeError(w, http.StatusBadRequest, "invalid_request", "role must be admin or user")
return
}
body.Role = config.NormalizeRole(body.Role)
rec, err := g.core.CreateKey(body.Name, body.Role, body.Models, body.Note)
if err != nil {
writeError(w, http.StatusBadRequest, "key_error", err.Error())
@ -76,6 +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())

View 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)
}
}

View File

@ -115,7 +115,7 @@ func (g *Gateway) loginRecord(ip string, ok bool, now int64) {
}
}
func New(c *core.Core, gatewayKeys []string) (*Gateway, error) {
func New(c *core.Core) (*Gateway, error) {
sub, err := fs.Sub(uiFS, "ui")
if err != nil {
return nil, err
@ -222,6 +222,8 @@ 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/"):

View File

@ -48,6 +48,10 @@ type Req struct {
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.
@ -84,6 +88,21 @@ type Stats struct {
auditPath string
replayPartial bool // aggregates built from a bounded audit tail
modelHour map[string]map[int64]int64 // model -> unix-hour bucket -> tokens
// keyModelHour buckets the same tokens as modelHour but keyed by
// (gateway key id, model) so per-key model quotas are isolated from each
// other. modelHour stays key-blind on purpose: it backs the AUTO slot
// quota, which is a gateway-wide limit on a slot, not a per-key one.
keyModelHour map[string]map[string]map[int64]int64 // key -> model -> unix-hour -> tokens
// keyHour buckets a key's total tokens per unix hour, backing the
// whole-key quota (all models of one key share one budget).
keyHour map[string]map[int64]int64
// keyReqHour buckets a key's request count per unix hour, backing the
// key-wide request quota. Counted for every request type (chat, stream,
// image) including failed ones, so a failing client cannot loop for free.
keyReqHour map[string]map[int64]int64
// pinnedWanted records which (key, model) pairs have a quota that pins a
// source, so the pinned bucket is maintained only where it is read.
pinnedWanted map[string]bool
}
const hourSec = 3600
@ -116,14 +135,18 @@ func NewStats(maxRecords int) *Stats {
maxRecords = defaultRingSize
}
return &Stats{
byKey: map[string]*Stat{},
byModel: map[string]*Stat{},
bySrc: map[string]*Stat{},
byKeyModel: map[string]map[string]*Stat{},
byKeySrc: map[string]map[string]*Stat{},
byStatus: map[int]*Stat{},
modelHour: map[string]map[int64]int64{},
maxRecs: maxRecords,
byKey: map[string]*Stat{},
byModel: map[string]*Stat{},
bySrc: map[string]*Stat{},
byKeyModel: map[string]map[string]*Stat{},
byKeySrc: map[string]map[string]*Stat{},
byStatus: map[int]*Stat{},
modelHour: map[string]map[int64]int64{},
keyModelHour: map[string]map[string]map[int64]int64{},
keyHour: map[string]map[int64]int64{},
keyReqHour: map[string]map[int64]int64{},
pinnedWanted: map[string]bool{},
maxRecs: maxRecords,
}
}
@ -318,6 +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)
@ -362,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 {
@ -374,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.
@ -464,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

View 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)
}
}

View File

@ -818,6 +818,10 @@
kAnySrc: "任意源",
kQuotaB: "Token 配额",
kQuotaHintB: "0 / 留空 = 无限",
kReqQuota: "请求数配额",
kReqQuotaHint: "限制周期内的请求次数。0 / 留空 = 无限。",
kQuotaPerModelHint:
"配额按模型单独设置:创建后点「+ 添加模型」,逐个模型配 token 配额与周期。一个模型用满只影响该模型,同一密钥的其它模型照常。",
kPeriodB: "重置周期",
kPerNothing: "不限",
kPerHour: "每 小时",
@ -980,6 +984,7 @@
thErr: "失败",
thPrompt: "输入 Tokens",
thCompl: "输出 Tokens",
thImgs: "张",
thAvgLat: "平均延迟",
thMaxLat: "最长延迟",
thTime: "时间",
@ -1043,6 +1048,10 @@
kAnySrc: "any source",
kQuotaB: "Token quota",
kQuotaHintB: "0 / empty = unlimited",
kReqQuota: "Request quota",
kReqQuotaHint: "Caps requests per window. 0 / empty = unlimited.",
kQuotaPerModelHint:
"Quotas are per model: after creating the key, use \"Add model\" to give each model its own token budget and reset period. One model running out affects only that model; the key's other models keep working.",
kPeriodB: "Reset period",
kPerNothing: "Never",
kPerHour: "Every hour",
@ -1206,6 +1215,7 @@
thErr: "Err",
thPrompt: "Prompt Tokens",
thCompl: "Completion Tokens",
thImgs: "imgs",
thAvgLat: "Avg latency",
thMaxLat: "Max latency",
thTime: "Time",
@ -1692,7 +1702,20 @@
</div>`;
document.body.appendChild(wrap);
}
function downloadStatsCsv(from, to) {
// closeTopModal removes the topmost dialog. Several dialogs share the id
// "modal-wrap" and more than one can be open at once (the seed-key notice
// sits on top of the keys page), so `$("#modal-wrap")` resolves to
// whichever comes FIRST in the document — not necessarily the one the
// user is looking at. Handlers that receive the clicked button resolve
// their own dialog with .closest(); handlers with no button reference
// use this, which matches what the user sees.
function closeTopModal() {
const all = [...document.querySelectorAll("#modal-wrap")];
const w = all[all.length - 1];
if (w) w.remove();
return w;
}
function downloadStatsCsv(from, to, btn) {
const q = new URLSearchParams({
export: "csv",
from: String(Math.floor(from)),
@ -1700,8 +1723,7 @@
});
if (statsKeyF) q.set("key", statsKeyF);
location.href = "/api/stats?" + q.toString();
const w = $("#modal-wrap");
if (w) w.remove();
(btn ? btn.closest("#modal-wrap") : null) || closeTopModal();
}
// downloadStatsCsvFromForm reads the custom date range from the export
// modal's date inputs and forwards to downloadStatsCsv. (Previously this
@ -1719,12 +1741,11 @@
if (!isFinite(from) || from <= 0) return toast(t("expRange"));
downloadStatsCsv(from, to);
}
function downloadKeysCsv() {
function downloadKeysCsv(btn) {
location.href =
"/api/stats?export=keys-csv" +
(statsKeyF ? "&key=" + encodeURIComponent(statsKeyF) : "");
const w = $("#modal-wrap");
if (w) w.remove();
(btn ? btn.closest("#modal-wrap") : null) || closeTopModal();
}
async function paintStats() {
try {
@ -2174,7 +2195,7 @@
<td class="t-tag">${fmtTime(r.time)}</td>
<td class="num">${r.ok ? `<span class="tag tag-green">${r.status || 200}</span>` : `<span class="tag tag-red" title="${esc(r.error || "")}">${r.status || 500}</span>`}</td>
<td>${esc(keyNames[r.key] ? keyNames[r.key] + " · " + r.key : r.key)}</td><td class="t-tag">${esc(r.type)}</td><td>${esc(r.model)}</td><td>${esc(r.source || "")}</td>
<td class="num">${fmtTok(r.prompt_tokens)}</td><td class="num">${fmtTok(r.completion_tokens)}</td>
<td class="num">${fmtTok(r.prompt_tokens)}</td><td class="num">${r.type === "image" ? fmtN(r.image_count || 0) + " " + t("thImgs") : fmtTok(r.completion_tokens)}</td>
<td class="num">${cacheCell(r)}</td>
<td class="num">${fmtMs(r.latency_ms)}</td></tr>`;
}
@ -2772,7 +2793,10 @@
/* ---------- sources tab ---------- */
async function renderSources() {
const j = await api("/api/sources");
// Read through /api/v1: the UI is a client of the same public API an
// agent uses, so there is one contract rather than a private path.
// The list view needs no credentials, so the masked view is enough.
const j = await api("/api/v1/sources");
const rows = j.sources
.map(
(
@ -2794,13 +2818,23 @@
function editSource(name) {
const wrap = document.createElement("div");
wrap.id = "modal-wrap";
Promise.all([api("/api/sources"), api("/api/status")])
.then(([src, st]) => {
const s = src.sources.find((x) => x.name === name) || {
name: name,
models: [{ id: "", priority: 0, kind: "chat" }],
};
const cur = s.adapter || "openai";
// The edit dialog round-trips the whole source, so it must load the
// real api_key — otherwise saving an unrelated field would blank
// it. That is what ?reveal=credentials is for (admin-only).
const srcPath = name
? "/api/v1/sources/" + encodeURIComponent(name) + "?reveal=credentials"
: "/api/v1/sources";
Promise.all([api(srcPath), api("/api/status")])
.then(([src, st]) => {
// Editing: the single-source shape. Creating: an empty form.
const s = name
? src.source
: { name: "", models: [{ id: "", priority: 0, kind: "chat" }] };
if (!s) {
toast(t("toastSave"));
return;
}
const cur = s.adapter || "openai";
const apps = ["", ...(st.adapters || []).map((a) => a.name)];
if (cur && !apps.includes(cur)) apps.push(cur);
const adSel =
@ -2905,9 +2939,10 @@
body: JSON.stringify(payload),
});
toast(t("toastSaved"));
const w = $("#modal-wrap");
const w = btn && btn.closest("#modal-wrap");
if (w) w.remove();
else if (window._modal) window._modal.remove();
else closeTopModal();
renderSources();
} catch (e) {
toast(tFmt("toastSaveFail", e.message));
@ -3064,11 +3099,11 @@
? `<div class="tbl-wrap"><table><tr><th>${t("tName")}</th><th>${t("tURL")}</th><th>${t("tAdapter")}</th><th>${t("tModels")}</th><th></th></tr>` +
tpls
.map(
(t) =>
`<tr><td><b>${esc(t.name)}</b></td><td>${esc(t.base_url)}</td><td>${esc(t.adapter)}</td>
<td><div class="src-models">${(t.models || []).map((m) => `<span class="tag tag-blue">${esc(m.id)}</span>`).join("")}</div></td>
<td><button class="ghost small" onclick="editTemplate(${JSON.stringify(t.name).replace(/["<>]/g, (c) => ({ '"': "&quot;", "<": "&lt;", ">": "&gt;" }[c]))})">${t("srcEdit")}</button>
<button class="danger small" onclick="delTemplate('${escAttr(t.name)}')">${t("srcDel")}</button></td></tr>`,
(tpl) =>
`<tr><td><b>${esc(tpl.name)}</b></td><td>${esc(tpl.base_url)}</td><td>${esc(tpl.adapter)}</td>
<td><div class="src-models">${(tpl.models || []).map((m) => `<span class="tag tag-blue">${esc(m.id)}</span>`).join("")}</div></td>
<td><button class="ghost small" onclick="editTemplate(${JSON.stringify(tpl.name).replace(/["<>]/g, (c) => ({ '"': "&quot;", "<": "&lt;", ">": "&gt;" }[c]))})">${t("srcEdit")}</button>
<button class="danger small" onclick="delTemplate('${escAttr(tpl.name)}')">${t("srcDel")}</button></td></tr>`,
)
.join("") +
"</table></div>"
@ -3168,8 +3203,9 @@
body: JSON.stringify({ name: nm, ...body }),
});
toast(t("toastSaved"));
const w = $("#modal-wrap");
const w = btn && btn.closest("#modal-wrap");
if (w) w.remove();
else closeTopModal();
openTemplateModal();
} catch (e) {
toast(e.message);
@ -3215,7 +3251,7 @@
const sortLanes = { chat: [], image: [] };
let sortStateMap = new Map();
async function renderSort() {
const j = await api("/api/sources");
const j = await api("/api/v1/sources");
let autoR = [];
let autoImg = [];
try {
@ -3875,7 +3911,7 @@
</select>
<div id="a-hours-box" style="display:none"><label>${t("kPerNHint")}</label>
<input id="a-hours" type="number" min="1" step="1" value="24"></div>
<p><button onclick="scrAddFromForm()">${t("kSaveScope")}</button>
<p><button onclick="scrAddFromForm(this)">${t("kSaveScope")}</button>
<button class="ghost" onclick="this.closest('#modal-wrap').remove()">${t("mCancel")}</button></p>
</div>`;
document.body.appendChild(wrap);
@ -3885,7 +3921,7 @@
});
$("#a-model").focus();
}
function scrAddFromForm() {
function scrAddFromForm(btn) {
const raw = $("#a-model").value.trim();
if (!raw) {
toast(t("kName"));
@ -3908,14 +3944,14 @@
],
});
const li = sortState.lanes.length - 1;
const w = $("#modal-wrap");
if (w) w.remove();
if (btn) btn.closest("#modal-wrap")?.remove();
else closeTopModal();
paintSort([li]);
persistAuto()
.then(() => toast(t("sortHintSave")))
.catch((e) => toast(e.message));
}
function sortScopeSave(li, ji) {
function sortScopeSave(li, ji, btn) {
const it = sortState.lanes[li].models[ji];
if (!it) return;
let q = parseInt($("#q-quota").value);
@ -3924,8 +3960,7 @@
let h = parseInt($("#q-hours").value);
if (isNaN(h) || h < 1) h = 1;
it.meta = { quota: q, period: p, hours: p || q ? h : 0 };
const w = $("#modal-wrap");
if (w) w.remove();
(btn ? btn.closest("#modal-wrap") : null) || closeTopModal();
paintSort([li]);
persistAuto()
.then(() => toast(t("sortHintSave")))
@ -4155,7 +4190,7 @@
async function loadModelPairs() {
let srcs = [];
try {
srcs = (await api("/api/sources")).sources || [];
srcs = (await api("/api/v1/sources")).sources || [];
} catch (e) {}
const pairs = [];
(srcs || []).forEach((s) =>
@ -4284,13 +4319,14 @@
const comb = scopeComb(m);
const src = normSrc(m.source);
const attrs = `data-key="${escAttr(key)}" data-model="${escAttr(comb)}"
data-quota="${m.token_quota || 0}" data-period="${escAttr(m.period || "")}" data-hours="${m.hours || 0}"`;
data-quota="${m.token_quota || 0}" data-reqquota="${m.req_quota || 0}"
data-period="${escAttr(m.period || "")}" data-hours="${m.hours || 0}"`;
return `<span class="mb" draggable="true" ${attrs} title="${escAttr(t("kBrickH"))}"
onclick="scopeEdit('${escAttr(key)}','${escAttr(comb)}')"
oncontextmenu="scopeCtx(event,'${escAttr(key)}','${escAttr(comb)}')">
<span class="mb-ico"><svg viewBox="0 0 24 24" fill="none" stroke="currentColor" stroke-width="2" stroke-linecap="round" stroke-linejoin="round" width="10" height="10"><circle cx="12" cy="12" r="10"/></svg></span>
<span class="mb-name">${esc(m.model)}${src ? `<em class="mb-src">${esc(src)}</em>` : ""}</span>
<span class="mb-quota">${esc(quantBadge(m.token_quota, m.period, m.hours))}</span>
<span class="mb-quota">${esc(scopeQuotaBadge(m))}</span>
<span class="copy-b" title="${escAttr(t("kCopyB"))}" onclick="event.stopPropagation();scopeDup('${escAttr(key)}','${escAttr(comb)}')"><svg viewBox="0 0 24 24" fill="none" stroke="currentColor" stroke-width="2" stroke-linecap="round" stroke-linejoin="round" width="12" height="12"><rect x="9" y="9" width="13" height="13" rx="2" ry="2"/><path d="M5 15H4a2 2 0 0 1-2-2V4a2 2 0 0 1 2-2h9a2 2 0 0 1 2 2v1"/></svg></span>
<span class="mb-x" title="${escAttr(t("kDelB"))}" onclick="event.stopPropagation();scopeRm('${escAttr(key)}','${escAttr(comb)}')"><svg viewBox="0 0 24 24" fill="none" stroke="currentColor" stroke-width="2" stroke-linecap="round" stroke-linejoin="round" width="12" height="12"><line x1="18" y1="6" x2="6" y2="18"/><line x1="6" y1="6" x2="18" y2="18"/></svg></span>
</span>`;
@ -4322,6 +4358,16 @@
if (p === "nhour") return "·" + Math.max(1, h) + "h";
return "";
}
// scopeQuotaBadge shows a scope entry's budgets: token quota and, when
// set, the request count. They are per model — a spent budget blocks
// only that model, not the whole key.
function scopeQuotaBadge(m) {
const parts = [];
if (+m.token_quota > 0) parts.push(fmtQuota(m.token_quota));
if (+m.req_quota > 0) parts.push(fmtQuota(m.req_quota) + "\u00d7");
if (!parts.length) return "\u221e";
return parts.join(" ") + periodText(m.period, m.hours);
}
function quantBadge(quota, period, hours) {
quota = +quota || 0;
period = period || "";
@ -4336,16 +4382,21 @@
model: comb[0],
source: src || undefined,
token_quota: parseInt(b.dataset.quota) || 0,
req_quota: parseInt(b.dataset.reqquota) || 0,
period: b.dataset.period || "",
hours: parseInt(b.dataset.hours) || 0,
};
});
}
async function putScope(key, scopes) {
// Each scope entry carries its own quotas, so the whole budget travels
// with the models it applies to. There is no separate key-wide total
// that could drift out of sync with the model list.
const body = { models: scopes };
await api("/api/keys/" + encodeURIComponent(key), {
method: "PUT",
headers: { "Content-Type": "application/json" },
body: JSON.stringify({ models: scopes }),
body: JSON.stringify(body),
});
}
async function scopePush(key) {
@ -4448,6 +4499,9 @@
<label>${t("kQuotaB")} <span class="muted">${t("kQuotaHintB")}</span></label>
<input id="sc-quota" type="number" min="0" step="1"
placeholder="${escAttr(t("kQuotaHintB"))}" value="${sc.token_quota ? sc.token_quota : ""}">
<label>${t("kReqQuota")} <span class="muted">${t("kReqQuotaHint")}</span></label>
<input id="sc-reqs" type="number" min="0" step="1"
placeholder="${escAttr(t("kReqQuotaHint"))}" value="${sc.req_quota || ""}">
<label>${t("kPeriodB")}</label>
<select id="sc-period">
<option value="" ${!sc.period ? "selected" : ""}>${t("kPerNothing")}</option>
@ -4481,6 +4535,8 @@
const parts = splitCombKey(comb);
let q = parseInt($("#sc-quota").value);
if (isNaN(q) || q < 0) q = 0;
let rq = parseInt($("#sc-reqs").value);
if (isNaN(rq) || rq < 0) rq = 0;
let hours = parseInt($("#sc-hours").value);
if (isNaN(hours) || hours < 1) hours = 1;
const period = $("#sc-period").value;
@ -4489,6 +4545,7 @@
model: parts.model,
source: parts.source || undefined,
token_quota: q,
req_quota: rq,
period,
hours,
};
@ -4498,8 +4555,9 @@
if (btn) btn.disabled = true;
try {
await putScope(key, scopes);
const w = $("#modal-wrap");
const w = btn && btn.closest("#modal-wrap");
if (w) w.remove();
else closeTopModal();
toast(t("kSaved"));
await loadKeys();
} catch (e) {
@ -4648,6 +4706,7 @@
<option value="user">${t("kRoleUser")}</option>
<option value="admin">${t("kRoleAdmin")}</option>
</select>
<p class="muted" style="font-size:12px">${t("kQuotaPerModelHint")}</p>
<label>${t("kNote")}</label>
<input id="kc-note">
<p><button onclick="createKey(this)">${t("kCreateBtn")}</button>
@ -4679,8 +4738,12 @@
if (btn) btn.disabled = false;
return;
}
const w = $("#modal-wrap");
// Close THIS dialog (resolved from the clicked button), not whichever
// #modal-wrap comes first: the seed-key notice may already be open and
// would otherwise be the one that gets removed.
const w = btn ? btn.closest("#modal-wrap") : null;
if (w) w.remove();
else closeTopModal();
$("#k-newbox").innerHTML = `
<div class="key-canvas" style="background:var(--card2)">
<div class="kc-head">

View File

@ -347,3 +347,60 @@ func cssBlock(src, selector string) (string, bool) {
}
return src[i : i+j+1], true
}
// The WebUI ships two locale objects (`zh` and `en`) that every rendered string
// goes through. A key added to only one of them does not error: t() falls back
// to `t("key") || "literal"` in some call sites and to the bare key string in
// others, so the user sees either a hardcoded local string or a raw key name.
// Nothing in the Go test suite noticed — this was found by hand while adding
// thImgs. Pin the key sets so a one-sided edit fails here instead of shipping.
func TestUILocaleKeyParity(t *testing.T) {
src := uiSource(t)
iZh := strings.Index(src, "zh: {")
iEn := strings.Index(src, "en: {")
if iZh < 0 || iEn < 0 || iEn < iZh {
t.Fatalf("locale blocks not found (zh=%d en=%d)", iZh, iEn)
}
zh := src[iZh:iEn]
en := src[iEn:]
// End the EN block at its closing brace, so keys from later objects
// (config templates etc.) do not pollute the comparison.
if end := strings.Index(en, "\n },"); end > 0 {
en = en[:end]
}
keyRe := regexp.MustCompile(`(?m)^\s*([A-Za-z_][A-Za-z0-9_]*)\s*:`)
collect := func(seg string) map[string]bool {
m := map[string]bool{}
for _, g := range keyRe.FindAllStringSubmatch(seg, -1) {
m[g[1]] = true
}
return m
}
zk, ek := collect(zh), collect(en)
// The locale markers themselves are the two block headers.
delete(zk, "zh")
delete(ek, "en")
var missingEN, missingZH []string
for k := range zk {
if !ek[k] {
missingEN = append(missingEN, k)
}
}
for k := range ek {
if !zk[k] {
missingZH = append(missingZH, k)
}
}
if len(missingEN) > 0 {
t.Errorf("translation keys present in zh but missing in en: %v", missingEN)
}
if len(missingZH) > 0 {
t.Errorf("translation keys present in en but missing in zh: %v", missingZH)
}
if len(zk) < 100 {
t.Errorf("only %d zh keys parsed — the block boundary regexp has drifted "+
"and this test is no longer checking anything", len(zk))
}
}

View 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)
}
}
}

View File

@ -5,10 +5,12 @@ adapter.version = "1.0.0"
adapter.endpoint = "/chat/completions"
adapter.headers = {}
-- AgentRouter 的 WAF 会按客户端指纹白名单校验,只放行官方客户端。
-- 这里通过 build_headers 注入官方客户端 User-Agent 以通过校验。
-- 默认 QwenCode 指纹(实测可通);可在 source.meta.user_agent 覆盖成 Claude Code 指纹。
-- AgentRouter 的阿里云 WAF 按三重维度校验:出口 IP(必须海外)、
-- TLS 指纹(JS/Go HTTP 层与 curl 不同)、浏览器指纹 header 集。
-- 纯 UA 已不够:WAF 拦 12 种 agent CLI UA 中的 10 种(全 405 HTML),
-- 只有同时带 Origin: qwen.ai + Referer + X-Requested-With 才放行,
-- 且同 IP 短时间连续 3+ 次请求会被速率封禁,靠 adapter 侧低频调用 +
-- llmsproxy 的 prefFailStep 冷却自然限制,来源无法控制速率。
local default_ua = "QwenCode/0.2.0 (linux; x64)"
-- AgentRouter fronts Claude models (claude-opus-4-8), and Claude upstreams
@ -65,6 +67,13 @@ function adapter.build_headers(meta)
["Content-Type"] = "application/json",
["Authorization"] = "Bearer " .. meta.api_key,
["User-Agent"] = ua,
-- QwenCode client fingerprint: WAF rejects bare UA without these.
["Accept"] = "application/json, text/plain, */*",
["Accept-Language"] = "zh-CN,zh;q=0.9",
["Origin"] = "https://qwen.ai",
["Referer"] = "https://qwen.ai/",
["X-Requested-With"] = "XMLHttpRequest",
["Connection"] = "keep-alive",
}
end

View File

@ -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

View File

@ -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

View File

@ -34,8 +34,10 @@ end
-- OpenAI /chat/completions format (pass-through, strip provider-specific fields)
-- zen 上游 schema 只接受 text content part(无视觉/音频能力):多模态 part
-- (image_url / input_audio / file 等)一律剥离;因此失去全部 content 的
-- 消息整条丢弃,避免上游 "unknown variant `image_url`, expected `text`"。
-- (image_url / input_audio / file 等)一律剥离。剥离后 content 变空的消息
-- 若不再携带 tool_calls / tool_call_id 才整条丢弃(避免上游
-- "unknown variant `image_url`, expected `text`");带工具调用的必须保留,
-- 否则会把紧随其后的 tool 结果变成孤儿,模型会反复重发同一个调用。
-- zen 上游角色白名单只有 system / user / assistant / tool / latest_reminder:
-- OpenAI 的 developer(及 function 等)不在其中,直接透传会触发上游
-- "unknown variant `developer`, expected one of ..." 错误;统一归一化为 system。
@ -52,8 +54,14 @@ function adapter.transform_request(raw_body)
if not ok then return raw_body end
req.disable_thinking = nil
req.extra_body = nil
if type(req.stream_options) ~= "table" then req.stream_options = {} end
req.stream_options.include_usage = true
-- stream_options is only valid alongside stream:true. Sending it on a
-- non-streaming request is rejected by stricter upstreams (OpenCode Go:
-- "stream_options should be set along with stream"), which failed every
-- non-stream call. Only attach it when the request actually streams.
if req.stream then
if type(req.stream_options) ~= "table" then req.stream_options = {} end
req.stream_options.include_usage = true
end
if req.messages then
local kept = {}
for _, msg in ipairs(req.messages) do
@ -72,7 +80,24 @@ function adapter.transform_request(raw_body)
end
end
if #parts == 0 then
drop = true
-- Content collapsed to nothing after stripping unsupported
-- parts. A message that still carries a tool call must
-- NEVER be dropped: the very next message is its tool
-- result, and dropping the call orphans that result. The
-- model then sees a result for a call it never made and
-- re-issues the same tool call on every turn (observed as
-- an infinite "repeated tool call" loop).
--
-- This is the common shape for agent clients: an assistant
-- turn whose content is only [thinking, toolCall] serialises
-- to content:[] with tool_calls — exactly the case that used
-- to vanish here. Emit "" instead, which zen accepts.
if (type(msg.tool_calls) == "table" and #msg.tool_calls > 0)
or msg.tool_call_id ~= nil then
msg.content = ""
else
drop = true
end
else
msg.content = parts
end

View 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

View 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

View 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()
}
}

View 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()
}
}

View File

@ -221,6 +221,70 @@ func TestOpenCodeStripsMultiModalParts(t *testing.T) {
}
}
// TestOpenCodeKeepsToolCallWithEmptyContent pins the fix for an infinite
// "repeated tool call" loop: an assistant turn whose content serialises to an
// empty ARRAY while carrying tool_calls must survive the multimodal strip.
// Dropping it orphans the following tool result, so the model never sees the
// call it already made and re-issues it on every turn.
func TestOpenCodeKeepsToolCallWithEmptyContent(t *testing.T) {
vm := NewVM(freshAdapterDir(t))
if err := vm.Start(); err != nil {
t.Fatal(err)
}
defer vm.Stop()
body := `{"model":"x","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"},
{"role":"user","content":"now summarize"}
]}`
out, err := vm.Transform("opencode", "transform_request", body)
if err != nil {
t.Fatalf("transform_request: %v", err)
}
var req struct {
Messages []struct {
Role string `json:"role"`
Content json.RawMessage `json:"content"`
ToolCalls []struct {
ID string `json:"id"`
} `json:"tool_calls"`
ToolCallID string `json:"tool_call_id"`
} `json:"messages"`
}
if err := json.Unmarshal([]byte(out), &req); err != nil {
t.Fatalf("unmarshal: %v (%s)", err, out)
}
if len(req.Messages) != 4 {
t.Fatalf("expected 4 messages (tool call must not be dropped), got %d: %s", len(req.Messages), out)
}
asst := req.Messages[1]
if asst.Role != "assistant" || len(asst.ToolCalls) != 1 || asst.ToolCalls[0].ID != "call_abc" {
t.Fatalf("assistant tool_calls lost: %s", out)
}
if string(asst.Content) != `""` {
t.Fatalf("empty content must be rewritten to \"\", got %s", asst.Content)
}
if req.Messages[2].Role != "tool" || req.Messages[2].ToolCallID != "call_abc" {
t.Fatalf("tool result must stay paired with its call: %s", out)
}
// Negative control: a message with neither text nor tool calls is still dropped.
only := `{"model":"x","messages":[
{"role":"user","content":[{"type":"image_url","image_url":{"url":"https://ex.com/a.png"}}]},
{"role":"user","content":"keep me"}
]}`
out2, err := vm.Transform("opencode", "transform_request", only)
if err != nil {
t.Fatal(err)
}
if got := strings.Count(out2, `"role"`); got != 1 {
t.Fatalf("image-only message without tool calls must still be dropped, got %d: %s", got, out2)
}
}
func TestBuildHeadersCustomHook(t *testing.T) {
vm := NewVM(freshAdapterDir(t))
if err := vm.Start(); err != nil {

View File

@ -707,7 +707,7 @@ func (p *Provider) probeModels(ctx context.Context) (bool, string) {
if err != nil {
return false, ""
}
if hdrs, herr := p.buildHeaders("{}", u); herr == nil {
if hdrs, herr := p.buildHeaders("{}", u, ""); herr == nil {
req.Header = hdrs
}
resp, err := p.client.Do(req)
@ -750,7 +750,7 @@ func (p *Provider) probeChat(ctx context.Context) (bool, string) {
body, err := json.Marshal(probe)
if err == nil {
var hdr http.Header
if hdrs, herr := p.buildHeaders(string(body), p.URL()); herr == nil {
if hdrs, herr := p.buildHeaders(string(body), p.URL(), ""); herr == nil {
hdr = hdrs
}
raw, status, derr := p.do(ctx, p.URL(), string(body), hdr)
@ -1031,13 +1031,15 @@ func (p *Provider) Throttle(ctx context.Context) error {
// ---- request construction ----
func (p *Provider) buildHeaders(body, url string) (http.Header, error) {
func (p *Provider) buildHeaders(body, url, clientSession string) (http.Header, error) {
meta := map[string]interface{}{
"url": url,
"method": http.MethodPost,
"body": body,
"api_key": p.cfg.APIKey,
"timestamp": types.Now(),
// client_session 是客户端自带的会话标识(请求头),为空时适配器自行推导。
"client_session": clientSession,
"source": map[string]interface{}{
"name": p.cfg.Name,
"meta": p.cfg.Meta,
@ -1082,7 +1084,7 @@ func (p *Provider) Chat(ctx context.Context, req *types.ChatRequest) (*types.Uni
if err != nil {
return nil, err
}
hdrs, err := p.buildHeaders(body, p.URL())
hdrs, err := p.buildHeaders(body, p.URL(), req.ClientSession)
if err != nil {
return nil, err
}
@ -1142,7 +1144,7 @@ func (p *Provider) ChatStream(ctx context.Context, req *types.ChatRequest) (<-ch
p.Release()
return nil, err
}
hdrs, err := p.buildHeaders(body, p.URL())
hdrs, err := p.buildHeaders(body, p.URL(), req.ClientSession)
if err != nil {
p.Release()
return nil, err
@ -1411,7 +1413,7 @@ func (p *Provider) Image(ctx context.Context, req *types.ImageGenRequest) (*type
// fall back to passthrough adapter (openai-style)
transformed = string(b)
}
hdrs, err := p.buildHeaders(transformed, p.ImageURL())
hdrs, err := p.buildHeaders(transformed, p.ImageURL(), "")
if err != nil {
return nil, err
}
@ -1516,12 +1518,24 @@ func standardSSEChunk(data string) string {
PromptTokensDetails *struct {
CachedTokens int `json:"cached_tokens"`
} `json:"prompt_tokens_details"`
CompletionTokensDetails *struct {
ReasoningTokens int `json:"reasoning_tokens"`
} `json:"completion_tokens_details"`
} `json:"usage"`
// Cost is a top-level decimal string on OpenCode Zen/Go responses.
Cost json.RawMessage `json:"cost"`
}
if err := json.Unmarshal([]byte(data), &raw); err != nil {
return ""
}
if len(raw.Choices) == 0 && raw.UpstreamUsage.Total == 0 && raw.UpstreamUsage.TotalTokens == 0 {
// Cost-only chunk: OpenCode sends the charge on its own frame
// {"choices":[],"cost":"0"} with no usage at all. Dropping it here
// would lose the only cost signal the client can see.
if c := rawCostString(raw.Cost); c != "" {
out, _ := json.Marshal(types.UnifiedChunk{Cost: c})
return string(out)
}
return ""
}
var usage *types.TokenUsage
@ -1539,6 +1553,11 @@ func standardSSEChunk(data string) string {
CachedTokens: pu.PromptTokensDetails.CachedTokens,
}
}
if pu.CompletionTokensDetails != nil {
usage.CompletionTokensDetails = &types.CompletionTokensDetails{
ReasoningTokens: pu.CompletionTokensDetails.ReasoningTokens,
}
}
}
finish := ""
done := false
@ -1563,6 +1582,24 @@ func standardSSEChunk(data string) string {
return string(out)
}
// rawCostString normalizes an upstream cost field into a plain string. The
// value is a decimal **string** on OpenCode Zen/Go ("0", "0.0012"), but other
// upstreams have been seen to send a bare number, so both are accepted.
// Returns "" when absent, null, or empty.
func rawCostString(raw json.RawMessage) string {
s := strings.TrimSpace(string(raw))
if s == "" || s == "null" {
return ""
}
if s[0] == '"' {
var unq string
if json.Unmarshal(raw, &unq) == nil {
return strings.TrimSpace(unq)
}
}
return s
}
// pickFirst returns a if non-zero, else b (for usage keys that may appear in
// either standard *_tokens or legacy short form).
func pickFirst(a, b int) int {

View File

@ -3,6 +3,7 @@
package types
import (
"bytes"
"encoding/json"
"errors"
"strings"
@ -28,6 +29,10 @@ type ChatRequest struct {
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) {
@ -61,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"`
@ -79,6 +117,11 @@ 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 {
@ -95,6 +138,16 @@ type TokenUsage struct {
// 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
@ -119,15 +172,16 @@ func (t TokenUsage) MarshalJSON() ([]byte, error) {
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"`
PromptTokensDetails *PromptTokensDetails `json:"prompt_tokens_details,omitempty"`
PromptCacheHitTokens int `json:"prompt_cache_hit_tokens,omitempty"`
PromptCacheMissTokens int `json:"prompt_cache_miss_tokens,omitempty"`
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,
@ -138,6 +192,7 @@ func (t TokenUsage) MarshalJSON() ([]byte, error) {
PromptTokensDetails: pdetails,
PromptCacheHitTokens: t.PromptCacheHit,
PromptCacheMissTokens: t.PromptCacheMiss,
CompletionDetails: t.CompletionTokensDetails,
})
}
@ -181,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

View File

@ -2,6 +2,7 @@ package types
import (
"encoding/json"
"strings"
"testing"
)
@ -68,3 +69,86 @@ func TestTokenUsageMarshalCacheFields(t *testing.T) {
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)
}
}

View 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: []

View File

@ -68,7 +68,12 @@ 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"
cp "$ROOT/config.yaml" "$TAR_DIR/config.example.yaml"
# 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/"

View File

@ -28,8 +28,11 @@ contents:
dst: /etc/systemd/system/llmsproxy.service
file_info:
mode: 0644
# example config (seeded to /etc/llmsproxy on first install by postinst)
- src: ./config.yaml
# 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

View File

@ -9,9 +9,19 @@ mkdir -p /etc/llmsproxy/adapters
if [ -f /etc/llmsproxy/config.yaml ]; then
echo "llmsproxy: keeping existing /etc/llmsproxy/config.yaml"
else
# seed a starter config from the packaged example
cp -f /usr/share/llmsproxy/config.example.yaml /etc/llmsproxy/config.yaml 2>/dev/null || true
chmod 0644 /etc/llmsproxy/config.yaml
# 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