Compare commits

...

108 Commits

Author SHA1 Message Date
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
f0986751de Merge feature/agentrouter-id-sanitize: 补上最后一个服务 Claude 的适配器
全仓审计发现 agentrouter 也暴露 claude-opus-4-8,却是唯一没有 id 消毒的
Claude 适配器。补齐后,所有含 claude/opus/sonnet 模型的源(qijiar、toter、
juziai、agentrouter、justwoker、api456、tabitoken)全部走消毒路径;
测试新增三适配器结果一致性断言。
2026-09-06 10:08:52 +08:00
c6c3e0dcd7 fix(agentrouter): sanitize tool-call ids — it fronts Claude too
Full-repo audit after the anthropic/openai fix: agentrouter exposes
claude-opus-4-8, so it inherits Anthropic's tool id rule
^[a-zA-Z0-9_-]{1,64}$ and rejects the whole request on a violation, exactly
like justwoker/tabitoken/扇贝. It was the only remaining adapter serving Claude
models without the sanitizer, so a client that had picked up a dirty id (e.g.
"bash:0" from moonshotai/kimi-k3) would still lose every turn here.

Same shape as the other two — inbound tool_calls[].id + tool_call_id, outbound
non-streaming ids and the first streamed fragment — and the test now asserts
all three adapters rewrite an identical input identically, so a client mixing
sources within one session cannot end up with unpaired tool calls.

Audit result: every source exposing a claude/opus/sonnet model (qijiar, toter,
juziai, agentrouter, justwoker, api456, tabitoken) now routes through a
sanitizing adapter.
2026-09-06 10:08:52 +08:00
b9944332ff Merge feature/anthropic-usage-cache: 缓存命中不该让 prompt_tokens 算错
Anthropic 的 input_tokens 不含缓存部分,OpenAI 的 prompt_tokens 含。
直接对映射等于漏算全部缓存 token,还能算出 >100% 的命中率;
cache_creation 完全没读;流式路径连缓存字段都整个丢掉。
三项相加 + 共享 map_usage,未上报缓存的上游输出逐字节不变。
2026-09-06 09:55:22 +08:00
6bdb9fcc44 fix(anthropic): count cached input in prompt_tokens instead of dropping it
Anthropic and OpenAI disagree on what the prompt count means:

  Anthropic: input_tokens EXCLUDES cached blocks; cache_read_input_tokens and
             cache_creation_input_tokens are separate, additive, billed input.
  OpenAI:    prompt_tokens INCLUDES its cached_tokens subset.

anthropic.lua mapped input_tokens straight onto prompt, so a cache-heavy turn
was doubly wrong: the billed prompt was undercounted by the entire cache
portion, and cached_tokens could exceed prompt_tokens — a cache hit rate above
100% for any client that divides one by the other. cache_creation_input_tokens
was never read at all, so a cache-write turn silently lost those billed tokens.

Worse, the streaming path dropped the cache split entirely: message_delta
carries the FINAL usage and only mapped input/output, so every streamed
response reported no cache information even when the upstream sent it.

All three counts are now summed into prompt, with the read half exposed as
prompt_tokens_details.cached_tokens plus the DeepSeek-legacy hit/miss pair, via
one shared map_usage() used by transform_response, message_start and
message_delta. A reported zero stays distinguishable from "never reported": the
split is emitted whenever either cache field is present, and omitted entirely
when the upstream mentions neither (justwoker reports only input/output plus its
own cost fields, so its output is byte-identical to before). map_usage returns
nil for a countless object, preserving "no usage in this chunk means say
nothing" rather than reporting zeros.

message_start's placeholder count is still emitted: justwoker reports 160 there
and the real 6931 in message_delta, and the gateway's mergeUsage lets the later
non-zero value win.
2026-09-06 09:55:22 +08:00
25eb30c15b Merge feature/toolcall-id-sanitize: 一个上游的脏 tool_call id 不该拖垮整条 Claude 链
xinjianya/moonshotai/kimi-k3 会吐出 'bash:0' 这种 id。OpenAI 不校验,
Anthropic 校验 ^[a-zA-Z0-9_-]{1,64}$ 且整条请求直接 400。客户端会把它存进
历史并回放给所有源,于是 justwoker / tabitoken / 扇贝(都是 Claude 上游)
同时挂掉,AUTO 一路穿透四个 tier。

两个适配器、进出双向消毒;已合法的 id 逐字节透传;流式参数分片保持无 id。
2026-09-05 22:19:15 +08:00
131a42a169 fix(adapters): sanitize tool-call ids so one bad upstream can't kill every Claude slot
Anthropic requires tool_use.id / tool_result.tool_use_id to match
^[a-zA-Z0-9_-]{1,64}$ and rejects the WHOLE request otherwise with
REQUEST_BODY_INVALID / "Invalid tool use format". OpenAI has no such rule, so
an OpenAI-compatible model can mint an id like "bash:0"
(xinjianya/moonshotai/kimi-k3 does exactly that).

In a fan-out router that id does not stay local: the client stores it in its
history and replays it to every other source. One such id therefore kills
every Claude slot at once — justwoker, tabitoken and 扇贝 are all
Claude-behind-{OpenAI,Anthropic} — and an AUTO request falls through all four
tiers to whatever tolerant model is left. Observed live: 4 consecutive 503
"all N auto providers failed" with tier 1/2/3 each reporting the same 400.

Both directions are sanitized, in both adapters:
  - request:  tool_calls[].id and tool_call_id, so poisoned history recovers
  - response: non-streaming tool_calls[].id and the first streamed fragment,
              so a bad id never enters a client session in the first place

safe_tool_id is pure and deterministic, so a call and its result are rewritten
identically within one request. A rewritten id keeps an 8-hex digest of the
original, without which distinct ids could collapse ("a:b" and "a_b") into a
duplicate/unpaired tool_use. Already-legal ids pass through byte-identical, so
well-behaved traffic is unaffected. openai.lua carries its own copy because
Lua adapters have no shared prelude.

Streamed argument fragments carry no id and must stay id-less, otherwise
index-based accumulation on the client breaks; a test pins that.
2026-09-05 22:18:31 +08:00
3cdb16c906 Merge feature/deploy-config-rollback: 部署脚本把配置纳入回滚闭环
上一轮 446 次重启循环的根因不是代码,是 deploy.sh 把 config.yaml 当作
部署范围之外的东西:健康检查失败后只回滚二进制,坏配置仍在,于是旧
二进制照样解析不了,服务在 systemd 重启循环里空转。

修法是把配置变成一等部署对象:--config 投放 + 重启前 -check 预检 +
回滚时二进制与配置一起恢复。预检位置比预检本身更重要——排在 restart
之前,坏配置才会在服务仍健康时被拦下。
2026-09-05 09:47:04 +08:00
ef92ce4d82 Merge feature/source-proxy: 让被 CF/格式问题挡住的源真正可用
两件事都在回答同一个问题——为什么配了源却用不了:

- trae 源在流式下把模型打印的 tool_call 当纯文本透传,客户端收到
  finish_reason "stop" + 一段文本,agent 循环当场终止
- justwoker/tabitoken 坐在 Cloudflare 后面,直连 403,必须走代理 +
  浏览器 UA;但全局 env 代理会把 trae/localzen 等内网源也绕出去

因此引入 per-source proxy_url:指定的源走专属代理,未指定的保持直连。
2026-09-05 09:46:54 +08:00
c524831616 fix(deploy): roll back config together with the binary, and preflight it before restart
The 446-restart-loop incident: adding a headers block to a source without
removing the source's existing `headers: {}` produced a duplicate YAML key.
The process exited on startup, healthcheck failed, and rollback restored only
the binary — so the old binary kept parsing the same broken config and the
service span in systemd's restart loop. Config was treated as out of scope
for deployment; it is not.

Three changes close the loop:

1. cmd/llmsproxy: new `-check` flag validates a config (parse +
   ApplyDefaults) and exits, without starting the Lua VM, touching
   runtime.json, or binding a port — safe to run against a live service.
   Unlike normal startup it does NOT create a default config, so a missing
   file is an error.

2. deploy.sh `--config <file>`: stage a config for deployment, atomically
   renamed into place with the same copy->rename(2) technique as the binary.
   Omitted means the live config is left alone.

3. Ordering: config replacement and preflight both run BEFORE
   restart_service, so an invalid config is caught while the service is still
   healthy and never triggers a restart. rollback() now restores binary AND
   config (only when this run replaced it, so concurrent WebUI edits survive),
   then re-runs -check before restarting — refusing to restart into a config
   that still fails, instead of trading one restart storm for another.

Also: the sha256-unchanged early exit now only fires when there is no pending
config, otherwise `--config` would be silently dropped.

Verified on the live deployment:
- reproduced the exact duplicate-key config: preflight caught it, PID
  unchanged (zero interruption), binary and config both rolled back, gateway
  still answering 200
- valid config: replaced, service restarted, new value live
- no --config: binary-only deploy unaffected
- go test -tags luajit ./... passes
2026-09-05 09:46:43 +08:00
868ac59692 feat(provider): support per-source proxy_url for geo-blocked upstreams
Upstreams like justwoker/tabitoken sit behind Cloudflare geo/IP blocks and
only respond through a proxy. A global env proxy is wrong (intranet sources
trae/localzen must stay direct), so add an explicit per-source proxy_url
that overrides http.ProxyFromEnvironment for just that source. Sources
without proxy_url keep the existing env/direct behavior.

Also accepts a User-Agent header per source (config already supported
headers) so CF-fronted resellers can be reached.

Verified e2e: tabitoken and justwoker now return tool_calls through
127.0.0.1:7890 (clash); trae stays direct. Full go test -tags luajit
passes.
2026-09-05 09:46:36 +08:00
4d197e4bd3 fix(trae): recover legacy [Called tool:...] tool call format
Root cause: trae-local-api is deployed to users without the "Fold past
assistant tool_calls into [Called tool: name({...})]" text format that
their own histories already contained. Trae-Local-API-LLM then mimics this
format in subsequent responses. The trae.lua parser only recognized
<tool_call>...</tool_call> or <toolcall>...</toolcall> tags, so
[Called tool: ...] responses were left unparsed and the client received
plain text where a structured tool_calls array should be.

Fix: Add a legacy pattern match at the END of parse_text_tool_calls to
catch the [Called tool: name({args})] shape and emit proper tool_calls.
This is a fallback; models should emit <tool_call> tags per system prompt,
but we tolerate the mimicked form for robustness.
2026-09-05 09:46:36 +08:00
3806aaee03 docs(workflow): codify the branch model — main / feature / release branches
Adopt GitHub Flow + release branches, replacing "everything straight to main
plus a tag" which caused the 1.4.2 pain (a fix had to be retro-fitted to the
released version, forcing a remote-tag delete + full re-upload).

- main: only long-lived branch, always deployable, accumulates the next version
- feature/<desc>: born from main, merged back when done
- release/vX.Y.Z: cut from main, tagged, installers built from the tag
- hotfixes land on the release branch AND are cherry-picked back to main so
  main never loses a fix
- end of lifecycle = retire the release branch (delete; or keep for long-term
  maintenance), no wholesale merge back — hotfixes already flowed
- explicitly no rebase of main, no release-branch-merge, no quick edits on main

Companion release checklist includes the upload lessons (PUT --http1.1) and the
replace-artifacts-by-deleting-the-tag catch.

Docs in docs/git-workflow.md (zh) and docs/git-workflow-en.md, linked from both
READMEs.
2026-08-31 12:36:17 +08:00
5b628b4d6f fix(scheduler): rebalance the pref score so metered sources aren't written off silently
Report: an allowance-metered source (sensenova) had its deepseek model sitting at
pref -10 while the WebUI showed zero failures. All three numbers were accurate,
and they exposed three compounding problems.

1. Quota exhaustion cost the SAME score as a real failure. RecordQuotaExhausted
   intentionally avoids failCount (an exhausted allowance is not a fault), so the
   UI showed fails=0 / cooling=false — yet it deducted the full prefFailStep (5).
   For a metered source, running out of budget is an everyday event, so the score
   drifted deep negative with no visible cause. Now quota and 429 events cost
   prefQuotaStep (1): the cooldown already keeps the slot out of rotation until
   the window resets, the score only needs a mild preference for slots with budget.

2. Recovery was 5:1 asymmetric. A failure cost -5 but a success only +1, so a
   slot at -10 needed ten consecutive successes just to reach neutral — which it
   could never get, because a low score makes the scheduler not pick it in the
   first place (starvation). Success now rewards prefSuccessStep (2): recovery
   from -10 needs five successes, while a real failure still outweighs one.

3. No idle decay. A penalised slot kept its negative score forever once it stopped
   being selected. Pref() now applies lazy decay: after prefDecayAfter (2 min) of
   no outcome, the score drifts one step back toward 0 per interval (never past
   0, never touches positive scores). Applied in Pref() and TryProbe(), so a
   naturally-recovered idle slot is schedulable again without needing a probe.

The failure penalty itself is unchanged (prefFailStep=5), so genuinely broken
upstreams are still marked as clearly worse than healthy ones.

Tests: quota penalty lighter than failure; 5 quota resets stay well above the
floor with failCount untouched; 429 is quota-class; recovery from -10 needs <=5;
idle decay rehabilitates a written-off slot, stops at 0, and never drags a
positive score; decay repeatedly lifts a slot off the prefMin floor. Updated the
two pre-existing tests that asserted the old -5/129 +1 values.
2026-08-31 11:50:47 +08:00
2e3d5b79ad fix(adapters): stop dropping non-streaming tool calls (agent loops died on turn 2)
Four adapters handled tool_calls in transform_stream_chunk but lost them in
transform_response, so any NON-streaming tool-using conversation broke on its
second request: the client received finish_reason:"tool_calls" with no
tool_calls payload, replayed an assistant message whose function
name/arguments were empty, and the upstream rejected the next turn with

    400 invalid tool_call function, function/name/arguments cannot be empty

The production audit trail shows 46 such failures on sensenova alone.

- sensenova.lua: forward message.tool_calls, decoding the arguments JSON string
  into an object as the unified shape expects.
- gemini.lua: collect functionCall parts from candidates[].content.parts. Also
  correct finish_reason, since Gemini reports "STOP" even when it emitted a
  function call and clients keyed on it treat that as a finished answer.
- ollama.lua: the field was initialized to an empty table and never filled;
  fill it and likewise correct done_reason "stop" -> "tool_calls".

trae is a different failure with the same symptom: trae-local-api's OpenAI
endpoint (/v1/chat/completions, src/server.js:353) never reads the request's
`tools` array — only its Anthropic endpoint does — so the relayed model is never
told the tool schema and instead PRINTS a <tool_call>{...}</tool_call> block into
content, leaving message.tool_calls null and finish_reason "stop". An OpenAI
client sees an ordinary completion and its agent loop ends mid-conversation.
trae.lua now recovers the structured call from that text, strips the block from
user-visible content, and corrects finish_reason. Both tag spellings
(<tool_call>/<toolcall>, the latter is what the same codebase's Anthropic prompt
asks for) and all three argument key names (arguments/params/input) are accepted.
This is a defensive fallback: fixing the upstream shim to honour `tools` remains
the real fix, since the model still guesses parameter names.

Tests: TestNonStreamToolCallsPreserved covers all ten OpenAI-shaped adapters,
TestGeminiNonStreamToolCalls and TestOllamaNonStreamToolCalls cover their native
shapes, TestTraeTextToolCallRecovery covers both tag spellings, prose around the
block, and asserts a plain text answer never gains tool_calls.

Verified end-to-end against mock upstreams reproducing each shape: a full
two-round agent loop (tool call -> tool result -> final answer) now completes for
both the structured and the text-emitted variants.
2026-08-31 10:23:35 +08:00
2ebc01e03b fix(build): nsis.7z is not a required artifact — Setup exe is the formal Windows product 2026-08-30 13:49:34 +08:00
11776024e9 chore(build): gitignore cmd/gui/.cache (docker electron-builder cache dir) 2026-08-30 10:56:01 +08:00
bf1932c6c5 chore(version): 1.4.0 -> 1.4.1 2026-08-30 10:52:48 +08:00
22ddf5a411 chore(build): drop ELECTRON_BUILDER_CACHE so the wine toolchain cache stays in the volume, not the workspace 2026-08-30 10:50:01 +08:00
3698546d14 fix(build): run Windows NSIS packaging inside docker; add artifact size gate
The Windows installer has been broken since 1.3.0: electron-builder's NSIS step
needs wine to generate the uninstaller, but the host's wine was amd64-only (no
i386 runtime -> empty syswow64 -> `error c0000135`), so electron-builder silently
wrote a 264 KB installer shell with no payload. No check caught it and the broken
exe shipped. 1.4.0 reproduced the same failure this session.

Two fixes:

1. win-builder image gains node + wine32/wine64 (+ i386 arch). dist-win-docker.sh
   now runs BOTH the core cross-build and `npx electron-builder --win nsis`
   inside docker (USE_SYSTEM_WINE=true -> the image's wine). The wine prefix is
   initialized on first run in the shared cache volume, so syswow64/ntdll.dll
   exists - the exact thing the host lacked. The host needs no mingw/wine/node.

2. packaging/verify-dist.sh: a size-floor gate for GUI artifacts (exe >= 5 MB,
   deb/rpm/nsis.7z >= 10 MB). Wired into `make gui-dist` and `make gui-win-docker`
   and the dist-win-docker script, so a degenerate installer fails the build
   instead of reaching a Release. Verified: it rejects the 264 KB exe and passes
   the healthy artifacts.
2026-08-30 10:49:45 +08:00
ed05cccb72 feat(packaging): core distribution packages (deb + rpm + tar.gz) via nfpm
Previously only the Electron GUI had installers; server admins had to build the
core by hand (`make build` -> bare binary). This adds a first-class server
distribution:

- `make core-dist` -> packaging/core-dist.sh -> cmd/build/dist/
    llmsproxy_<ver>_amd64.deb        (systemd unit + adapters + example config)
    llmsproxy-<ver>.x86_64.rpm
    llmsproxy-<ver>-linux-amd64.tar.gz
- nfpm config (packaging/nfpm.yaml): binary to /usr/bin, systemd unit with the
  measured memory tuning (GOGC=50, MALLOC_ARENA_MAX=2) baked in, Lua adapters to
  /usr/share/llmsproxy/adapters, repo's config.yaml as the example.
- postinst seeds /etc/llmsproxy/config.yaml on first install and keeps existing
  runtime state across reinstalls; prerm stops the service on removal.

Gotchas handled: nfpm v2 here does not render `{{ .Env.VERSION }}`, so the script
stamps the version into a throwaway config copy; and under `set -o pipefail` the
idiom `strings | grep -q` is broken (grep -q closes the pipe and strings dies on
SIGPIPE -> spurious 141), so the symbol sanity-gate stages strings in a temp file.

README (zh/en) updated: memory figures replaced with production-measured
32-35 MB settled (was 37-42), and the packaging docs now describe core-dist and
the dockerized Windows build.
2026-08-30 10:49:25 +08:00
c309448414 feat(webui): Mono theme — pure white in light mode, pure black in dark mode
Adds a fourth accent alongside sakura/ocean/violet. Unlike those it is not just a
different hue: the coloured themes are glass surfaces (translucent cards with
backdrop-filter) floating over an animated gradient-mesh background, so setting
--card:#ffffff there still renders as a tinted grey. Mono therefore also switches
off the translucency and hides the blobs, so #ffffff is actually #ffffff and
#000000 is actually #000000, with greys carrying the hierarchy that hue carries
elsewhere. A side effect worth having: no backdrop-filter and no animated blobs
makes it the cheapest theme to render, which helps on weak GPUs and over remote
desktops.

Both light and dark variable blocks are defined, so the existing light/dark
toggle drives it with no extra wiring: light -> white, dark -> black.

Also fixes a latent theme bug found while checking contrast on black: the active
chart's grid baseline assigned the literal string "var(--line)" to
ctx.strokeStyle. Canvas 2D does not resolve CSS custom properties, so that was an
invalid colour the browser ignored, leaving the previous fillStyle (black) — an
invisible baseline on every dark theme. Colours used on a canvas now go through a
cssVar() helper.

Tests: TestUIThemeMatrix asserts every accent defines BOTH a light and a dark
block plus a picker button (a half-defined theme shows up as unreadable text, not
as an error); TestUIMonoThemeIsFlat pins the opaque surfaces and disabled blobs;
TestUICanvasColorsResolveVars fails if any ctx.strokeStyle/fillStyle is handed a
raw var().

Unrelated packaging fix in the same commit: dist:linux only built deb+AppImage
while build.linux.target listed rpm too, so `make gui-dist` silently skipped the
rpm that release builds are expected to produce. Makefile/README wording updated
to match.
2026-08-30 10:04:53 +08:00
5b6d6fc90d fix(gateway): keep the record ring small now that aggregates scan everything
The gateway still built Stats with NewStats(10000), a leftover from when the
ring WAS the source of the dashboard's numbers. With aggregates now computed
from the full audit history, a 10000-entry ring only means 10000 resident Req
structs — the live instance came up at 64 MB instead of ~40 MB, putting back
most of the memory the on-demand paging work removed.

The ring's only jobs are the status page's 5-minute SourceRecent /
SourceAverages windows and the dashboard's first screen, both of which fit in
defaultRingSize (500). Older records are paged from disk.

TestGatewayRingStaysSmall pins it so the constant cannot drift back up.
2026-08-30 09:34:25 +08:00
3ddae41f0c fix(gateway): aggregate the full audit history, repair record paging
Two problems reported after the on-demand log work landed.

1. Dashboard totals were wrong. LoadAudit only replayed the last 4 MB of the
   audit file, so requests/tokens/per-key rows reflected a window instead of all
   time — a regression in reported numbers, not just in presentation.

   The aggregates are now built by streaming EVERY audit file (oldest first, so
   the hourly quota buckets keep their intended trailing window) and keeping
   nothing per record: aggregate maps are keyed by key/model/source, so their
   size is bounded by cardinality. Measured on the production host: 29 MB /
   221k lines / 37k requests in ~260 ms at startup.

   What stays bounded is the RAW-record ring: a fixed-size reqRing keeps only the
   newest maxRecs records, so the ~25 MB that used to be spent appending every
   record into a slice is still saved. auditReplayBytes is gone, and
   replayPartial now means "an audit file could not be read", which is the only
   remaining way for the totals to be incomplete.

2. Scrolling to the bottom stopped loading more records. Two independent causes:

   * paintRecords rebuilt the entire table on every 5s poll whenever the row
     count did not exceed the first screen — the "is a paged view live?" test
     compared row counts and matched exactly on the first refresh — wiping loaded
     pages and resetting scroll position.
   * IntersectionObserver only fires on TRANSITIONS. With a short list, or after a
     page whose rows all duplicated the first screen, the sentinel stayed visible
     and never fired again.

   paintRecords now builds once (recsState.built) and later polls PREPEND only
   genuinely new rows; attachRecsObserver adds a scroll-position fallback;
   fillRecordsViewport loads until the list actually overflows; and
   loadMoreRecords chains (bounded) when a page yields no new rows, since the
   first fetch necessarily overlaps the first screen.

TestUIRecordsPagingWiring pins all four mechanisms structurally, since none of
them is reachable from Go. Test names/comments referring to bounded replay are
updated to describe the bounded RING instead, and both READMEs now state that
totals come from the full history while records are paged.
2026-08-30 09:29:42 +08:00
2378bc00ba docs: replace invented memory figures with measured ones, ship the tuning knobs
The README claimed "~15 MB RSS" and, after the log-loading work, "~10 MB idle /
~19 MB with a 29 MB audit log". Those were TEST-INSTANCE numbers: one mock source
and one adapter. The real production config on this host (16 sources, 13
adapters, 59 models) sits at ~37-42 MB, and sat at ~105 MB before this series.
Quoting the single-source figure as the headline was misleading.

Both READMEs now state that memory scales with the number of configured sources
rather than with uptime, give a three-row measurement table (1 source / 1 source
with a 29 MB audit history / the 16-source production instance), and break the
production RSS down per region (Go heap, thread stacks + LuaJIT, mapped binary,
Go reservations, shared libs) so an operator can tell which part their own
deployment will grow.

Two runtime knobs are documented and now shipped by default in the desktop
build's core spawn (cmd/gui/main.js, overridable by exporting either variable):

  * MALLOC_ARENA_MAX=2 — LuaJIT allocates through cgo into glibc malloc, and
    glibc keeps up to 8*nproc per-thread arenas of ~1 MB that are never returned.
    Measured 8-15 arenas (7-12 MB) -> 0.
  * GOGC=50 — halves the Go heap target. Documented explicitly as useless ALONE
    (measured 20.3 -> 21.5 MB, i.e. worse, because the saved heap is eaten by
    more glibc arenas); only the pair cuts settled RSS, by ~19%.

Also corrects the binary size (8-12 MB, ~8 MB after the deploy script's -s -w)
and adds the elastic-pool / on-demand-log / self-healing-cooldown bullets that
README.md already had to README_EN.md.
2026-08-30 09:09:21 +08:00
882288f67f fix(webui): send DELETE when removing keys and adapters
Deleting a gateway key from the admin UI did nothing and reported
"use GET /api/keys": delKey() called api() with an empty options object, so
fetch defaulted to GET and the request landed in the GET branch of
handleKeysAPI. delAdapter() had the identical bug and reported
"adapter code not exposed; edit in UI".

This is the third instance of the same mistake — ab20f1b fixed delSource and
delTemplate, missing these two — so it is now pinned by tests instead of by
review:

  * TestUIAPICallsDeclareMethod walks every api() call in the embedded
    index.html and fails if one passes an options object without a method
    (an AbortSignal-only read is allowed, being a deliberate GET).
  * TestUIDeleteHelpersUseDelete / TestUIMutatingHelpersUseWriteMethods pin the
    verb of each removal and write helper by name.
  * TestKeyDeleteRoundTrip covers create -> DELETE -> gone -> second DELETE is a
    clean 404, and TestCannotDeleteOwnKey keeps the lockout guard.

The 404 bodies for GET /api/keys/{key} and GET /api/adapters/{name} now name the
verb to use ("DELETE /api/keys/{key} to remove"), because that message is what a
mis-methoded client actually shows its user; "use GET /api/keys" read as though
the caller had done nothing wrong.

delSource's indentation, broken by ab20f1b, is also straightened out.
2026-08-30 09:07:05 +08:00
813de19bd0 fix(lua): bound prewarm by queue depth, not by the ceiling
Production oscillated between ~46 MB and ~56 MB RSS with the openai pool
cycling 1 -> 10..12 -> 1 states every couple of minutes, while peak_in_use
never went above 2.

Cause: the batch prewarm sized itself purely on the adapter's ceiling. With
max_concurrent summing to 76, growStep is 8, so any two overlapping requests
warmed 8 states — 6 more than anything was waiting for. A minute later the
janitor correctly reclaimed the surplus, the next pair of overlapping requests
warmed 8 again, and the pool churned boot/discard forever. The elasticity was
working; the growth signal was simply wrong.

Prewarm is now bounded by BOTH limits: the ceiling still caps the step, but the
batch never exceeds p.waiting, the number of goroutines actually blocked on the
pool. Overlapping-but-not-queued traffic (the common case) creates exactly the
states it uses; a genuinely queued burst still ramps in one jump.

TestContentionBatchPrewarms is rewritten to queue real waiters instead of
relying on the ceiling to imply demand, and TestNoPrewarmWithoutWaiters pins the
production shape: two overlapping requests against a 76-wide adapter must create
exactly 2 states.
2026-08-30 08:25:18 +08:00
25d8bd8632 fix(lua): reclaim burst leftovers while traffic continues
Production after the first deploy showed the openai pool stuck at 9 states with
peak_in_use=1: a startup burst grew it, and then it never shrank again. The
shrink grace counter was reset by every CHECKOUT, so on a gateway that always
has a request in flight the counter never reached shrinkGraceRounds and the
burst's leftover states were pinned indefinitely — the same monotonic-growth
behaviour this series set out to remove, just with an extra step.

Grace is now reset by GROWTH (a miss that had to boot a state), which is the
actual signal that capacity is short. Ordinary sequential traffic no longer
defers reclaim, while two consecutive quiet-ish rounds are still required so a
gap between two bursts does not tear the pool down.

TestCheckoutResetsGrace asserted the old behaviour and is replaced by
TestGrowthResetsGrace (sequential traffic must NOT defer, growth must) plus
TestBurstLeftoverIsReclaimedUnderSteadyTraffic, which reproduces the production
shape: 9 concurrent holds, then one request per janitor round, and the pool must
still fall back to the resident floor.
2026-08-30 08:13:34 +08:00
3e27f4db24 chore: bump version to 1.4.0, document cooldown probing and elastic pools
README: new "冷却与半冷却探测(自愈调度)" section with the per-class cooldown
table (5xx exponential to 5min / 401-403 10min / 429 30s fixed / quota aligned
to its window), the probe-slot formula and the ordering rule that probes are
tried last. The LuaJIT section now says the pool is an elastic ceiling rather
than a preallocation and documents both step formulas. Headline figures replaced
with measured ones: idle ~10 MB, ~19 MB starting with a 29 MB audit log, and a
new bullet for on-demand log loading.

package.json also loses a `\u2014` escape and the broken indentation that an
earlier edit left in the electron-builder block.
2026-08-30 08:07:01 +08:00
23705042ee chore: ignore local agent state, clean packaged installers
.gitignore: add .omo/ (an AI-assisted-coding local state dir that was being
tracked — .omo/run-continuation/*.json is session scratch, not source) and
/plan.md (local working notes).

make clean: also remove cmd/build/ and cmd/gui/dist/. Those hold the packaged
installers (1.8 GB and 386 MB on this machine) and were never cleaned, so
`make clean` left almost all of the build output behind.
2026-08-30 08:06:49 +08:00
02cb8c0e33 fix(adapters): reflow live sensenova reasoning field, add trae adapter
sensenova: the deployed /etc/llmsproxy/adapters/sensenova.lua carried a fix that
never made it back into the repo — sensenova-6.8-flash-lite reports its chain of
thought in `reasoning` rather than `reasoning_content`, in both the single-shot
response and the stream deltas. Without this the model's output looked empty.
Repo and deployment now match byte for byte.

trae: the adapter was in use on this deployment but untracked, so a fresh
install had no way to serve the trae source.
2026-08-30 08:06:22 +08:00
6575058556 feat(webui): scroll-paged record table, pool metrics, probe-aware health tags
Front-end half of the on-demand log loading plus the observability for the two
new scheduling mechanisms.

Records table:
  * the first screen comes from the dashboard poll, and an IntersectionObserver
    sentinel below the last row pulls the next page from /api/stats/records as
    the user scrolls.
  * rendered rows are capped at RECS_MAX_DOM=1000 (oldest rendered rows are
    dropped) so a long scroll cannot grow the DOM without bound, with a "back
    to newest" button to reset cheaply.
  * releaseRecords() drops the buffer, disconnects the observer and aborts the
    in-flight fetch (AbortController) on tab switch, on key-filter change and
    on pagehide/beforeunload — leaving the page releases everything at once.
  * the poll no longer rebuilds the table once extra pages are loaded, so the
    5s refresh cannot throw away scrolled history.
  * when the server reports replay_partial, the filter row states that the
    aggregates cover the recent audit tail and points at CSV for full history.

Adapters page shows each pool as created/max plus in_use/idle and the live
grow/shrink steps, with the sizing rule in the hover text.

Priority page health tags now distinguish cooling / probe-ready / probing
instead of a flat "cooling", and the tooltip spells out when the window opened,
when the single probe is allowed through and when the slot clears completely.
core.AutoSlotState carries cooldown_from / probe_after / probing / probeable to
feed this.
2026-08-30 08:06:11 +08:00
d42c02b15d perf(gateway): load request logs on demand instead of holding them in memory
Startup RSS on this deployment was 56 MB with a 29 MB audit log and ~10 MB
without one: LoadAudit() json-unmarshalled the ENTIRE file into the aggregates
and kept a 10000-entry ring of raw records. Two more paths had the same shape —
AuditRecords() materialized a whole export window into a []Req before sorting
it, and a dashboard poll serialized the full ring so the browser could render
300 rows of it.

The audit file is now the source of truth and memory only holds the live
window:

  * LoadAudit replays only the last auditReplayBytes (4 MB) and drops the
    truncated first line; the ring default drops 10000 -> 500, which still
    covers both of its consumers (the status page's 5-minute SourceRecent /
    SourceAverages windows and the first screen of the records table).
    replayPartial is exported so the UI can say the totals cover a window
    rather than all time. Token-quota accounting is unaffected: it reads the
    modelHour buckets, not the ring (pinned by a test).
  * AuditPage(cursor, limit, key) pages records straight off disk, reading the
    newest file backwards in 64 KB chunks and returning as soon as the page is
    full. The cursor is "<file>:<offset>" and walks into rotated .old files;
    a cursor whose file rotated away reports rotated=true so the client can
    reset instead of silently skipping records. No state is cached between
    requests and the file handle is closed before responding, so "release when
    the user leaves the page" is guaranteed by never retaining anything.
  * StreamAuditRecords(from,to,key,fn) replaces the accumulate-then-sort export
    path; the CSV handler writes rows as they are read and flushes every 1000,
    and a write error (client gone) aborts the walk. Export memory is O(1)
    regardless of the window. AuditRecords is kept as a test-only wrapper.
  * Snapshot ships one screen (firstScreenRecords=100) by default; aggregates
    are untouched.
  * Audit rotation 64 MB x 10 -> 16 MB x 16: same 256 MB total budget, but a
    smaller newest file keeps the first reverse page cheap.

New route: GET /api/stats/records?before=&limit=&key= (non-admins are pinned to
their own key by exportKey). /api/status additionally reports adapter_pools for
admins.

Measured with production's 29 MB audit copied to the test instance: startup RSS
19.0 MB (was 56 MB); scrolling 10 pages (1000 records) +0.7 MB; exporting the
full history (36441 rows / 4.4 MB CSV) +0.1 MB with no residual growth.
2026-08-30 08:05:54 +08:00
df9aeed5fb feat(lua): elastic adapter worker pools instead of monotonic growth
ConfigureConcurrency() set each adapter pool's target to the sum of
max_concurrent over its sources (108 on this deployment) and `created` only
ever went UP: once a Lua state was booted it was parked forever, so a
long-running gateway's resident state count was a high-water mark of all
traffic it had ever seen, never of what it currently needs.

Pools are now sized from three live inputs:

  * the adapter's MAX CONCURRENCY (sum of max_concurrent) is a ceiling, not a
    preallocation, and it sets the growth step:
    growStep = clamp(ceil(maxW/8), 1, 8). A 64-wide adapter warms 8 states at
    once on a spike; an 8-wide one creeps up one at a time.
  * the LIVE CONNECTION COUNT (inUse, i.e. checked-out states) sets the shrink
    step: shrinkStep = clamp(ceil(excess/(1+inUse)), 1, excess). With no
    connections the slack collapses in a single round; a busy adapter gives up
    one state per round so the hot path keeps its warm states.
  * how many states already exist (created / len(idle)) decides how much room
    is left to grow and how much can be reclaimed.

Batch prewarm only fires on genuine contention (a miss while every existing
state is checked out), so a single sequential caller keeps reusing one state
rather than burning a whole grow step on a cold start. A single VM-level
janitor goroutine (not one per adapter) reclaims idle states every 30s, and
shrinkGraceRounds=2 plus idleHeadroom=1 keep a gap between requests from being
mistaken for the end of a load period; any checkout resets the grace counter.

`created` now decrements on reclaim and on shutdown, and release() closes a
state outright when the ceiling was lowered underneath it, so shrinking
max_concurrent in the config gives memory back immediately instead of parking
orphans until restart.

PoolStats() exposes created/idle/in_use/waiting/max/resident/grow_step/
shrink_step/peak_in_use for the status API.

Measured on the test instance (single mock source, max_concurrent=64):
idle created=1; 50 concurrent requests -> created=10 (peak_in_use=5, ceiling
respected); after 95s of silence -> created=1. On production after deploy: 13
adapters, 1 resident Lua state total with ceilings up to 76.
2026-08-30 08:05:13 +08:00
21cd0429a2 feat(provider): half-cooldown probe so recovered slots return immediately
A slot that failed ten times landed in a 30-minute cooldown and was a hard
skip for the whole window: ModelAvailable() said no, the scheduler dropped the
candidate, and the ONLY way back was a success that could never happen. In
practice an upstream that blipped for a minute — or whose quota reset in
seconds — stayed unusable for half an hour.

Cooldown is now a window with a probe gate instead of a wall:

  * backoffCap 30min -> 5min. 401/403 no longer implicitly "jump to the cap"
    but get an explicit authCooldown (10min); 429 keeps its 30s fixed window.
  * ModelState records cooldownFrom alongside cooldownUntil, so the window has
    a measurable midpoint, plus a single-token probe permit (probeInFlight).
  * TryProbe() hands out that permit only in the SECOND half of the window: a
    freshly failed slot stays completely silent, and past the midpoint exactly
    one request may go through as a probe.
  * ProbeSlots() = clamp(max_concurrent/10, 1, 2), so a source configured for
    max_concurrent=10 lets exactly one request probe.
  * A probe is a REAL request: success runs RecordSuccess(), which clears the
    cooldown and failCount on the spot and returns the slot to full rotation.
    Failure reopens the window, pushing the next probe to the NEW midpoint
    instead of retrying immediately.

Quota exhaustion is modelled separately (RecordQuotaExhausted): running out of
allowance is not the upstream being broken, so it must not feed the exponential
ladder or inflate failCount. It cools until the allowance window can plausibly
have reset (capped at 30min so a manual top-up is noticed) and is recognised
from the adapter-normalized error text via ReportStatusReason.

Scheduler side: collectCands() splits a tier into normal + probe candidates
with probes strictly LAST, so probe traffic is what a tier falls back to, never
what it prefers, and the round-robin cursor rotates over the normal head only.
releaseProbes() returns every permit on all exits (success, hard failure, ctx
cancel, tier fallthrough); the 90-line inlined busy-poll loop is extracted into
pollBusyTier() to keep those exits auditable. Direct Chat/ChatStream/Image take
the same gate.

Also fixes a real blacklist bug found on the way: a slot whose pref sank to
prefMin was refused by ModelAvailable forever, long after every cooldown had
expired. Such a slot is now probeable, and one success lifts it off the floor.

Verified live against a controllable upstream (tier1 flaky, tier2 healthy):
failure ladder 5s/9s/20s/40s/79s; after healing the upstream, 10 AUTO requests
across the silent half produced ZERO upstream calls; past the midpoint a single
AUTO request produced exactly ONE upstream call, cleared the remaining 39s of
cooldown, and the next three requests were served by the recovered slot.
2026-08-30 08:04:50 +08:00
2ed1f0ecde style: gofmt the tree
gofmt -l reported 13 files with misaligned struct tags / stale formatting.
This commit contains ONLY formatting: no behaviour change, no logic touched.
Files that also carry real changes in this series are formatted by their own
commits.
2026-08-30 08:04:22 +08:00
624fd74b45 fix: anthropic tool-call round-trip, cache zero-hit parity, round-robin load balancing
anthropic.lua v3.0.0:
- Issue 1: tool_result/tool_use round-trip
- Issue 3: thinking default OFF (opt-in via extra_body.thinking)
- Issue 4: tool_choice mapping
- Issue 5: collect_blocks preserves unknown part types
- message_stop no longer emits done=true (was overwriting tool_calls finish_reason)
- cache_read_input_tokens normalized even at 0

gemini.lua:
- transform_response was missing cachedContentTokenCount

openai.lua (Issue 6):
- transform_error handles flat envelopes, nginx HTML, bare text

chat.go mergeUsage:
- Keep PromptTokensDetails even when CachedTokens=0

scheduler.go:
- Remove sort.SliceStable by Pref; round-robin cursor is the only LB mechanism

provider.go ModelAvailable:
- Also check Pref() > prefMin, persistently failing slots exit cands

presets.go:
- 17 built-in source templates

Tests: 6 new test functions, 2 updated for new semantics
2026-08-28 12:02:46 +08:00
94cbcb6771 docs: source templates, split image AUTO chain, headless/desktop installers
- AUTO chain section: chat and image generation are now two independent
  chains (auto / auto_image) toggled in the Priority page; legacy image
  discovery fallback documented
- new Source templates section: multi-key balancing via reusable templates
  (Templates manager, From template / As template in the add-source dialog)
- Desktop GUI section: two installer kinds — Headless (server, plain binary)
  and Desktop (Electron GUI); release artifact naming ModelRouter-Headless-* /
  ModelRouter-Desktop-*; stale 1.0.0 version pins replaced with <ver>
2026-08-27 13:02:48 +08:00
a5384d9fb6 chore: bump version to 1.2.0 2026-08-27 12:26:36 +08:00
ab20f1be48 fix(webui): specify DELETE method for source/template removal
delSource and delTemplate called api() without a method, so fetch
defaulted to GET; the DELETE handlers never ran and the UI silently
left the item in place.
2026-08-27 12:14:41 +08:00
8334cffbc9 feat(templates): source templates for multi-key balancing
A template stores every source field except name and api_key, so
operators spin up N key-bearing sources from one shared skeleton
instead of duplicating the whole source block N times.

- config: SourceTemplate type + RuntimeConfig.SourceTemplates stored
  in runtime.json alongside runtime sources
- store: UpsertTemplate / ListTemplates / RemoveTemplate
- core: Templates / SaveTemplate / RemoveTemplate
- gateway: GET/POST/DELETE /api/source_templates
- webui: source list gains a Templates button opening a manager with
  per-template edit/delete; the add-source dialog gains 'from
  template' (event-delegated picker) and 'as template' (card modal)
  buttons in its header; z-index fixed so the template editor layers
  above the manager
2026-08-27 12:09:38 +08:00
a42ff62d06 feat(scheduler): separate image-generation AUTO chain with UI toggle
Image models previously could not be scheduled through a priority
chain: the chat AUTO chain explicitly skips image-kind slots, and
AUTO image requests fell back to unordered registry discovery.

- config: add auto_image rules (auto_image yaml / image_rules json);
  legacy auto rules keep their meaning as the chat chain
- core: buildAutoImageChain mirrors buildAutoChain with inverted kind
  filter (image-only); SaveAutoImageRules + AutoImageRules/AutoImageChain
- scheduler: ChainImage walks the chain tier-by-tier with round-robin
  and preference ordering, skipping cooling slots
- gateway: handleImage AUTO now runs down AutoImageChain when one is
  configured (falls back to legacy discovery otherwise) and records
  the actual served model; handleAutoAPI GET returns image_rules and
  PUT accepts image_rules independently of rules
- webui: priority page gains a chat/image toggle editing two
  independent lane sets; add-slot picker filters by active kind;
  persistAuto writes only the active chain's field
2026-08-26 21:03:51 +08:00
66585549f1 fix(webui): allow image models in priority/AUTO chain editor
The priority page and key-scope pickers skipped kind=image models,
so image sources (e.g. Kwai-Kolors/Kolors) could not be placed in
the AUTO chain or granted per-key. The backend already supports it:
handleImage resolves AUTO through the chain then filters with
imageOnly, while chat requests are protected by chatOnly, so image
slots never receive chat traffic.

Image models now appear in the priority canvas, the add-slot picker
(labelled ' (image)'), and the key scope dialog.
2026-08-26 20:36:05 +08:00
747dff5b76 fix(gateway): record actual served model for image requests
handleImage recorded the raw request model id, so AUTO image
generations showed up as model=AUTO in the request records and
by-model aggregates instead of the image model actually served
(e.g. Kwai-Kolors/Kolors).

UnifiedResponse gains an optional Model field; Provider.Image fills
it with the resolved id (AUTO resolves to the source's best image
model), and handleImage prefers it when writing the audit record.
2026-08-26 20:28:20 +08:00
5d045f97a8 chore: bump version to 1.1.1
Includes WebUI fix e48baa1 (missing HTTP method on fetch calls with
body) which was committed after the 1.1.0 installers were built.
2026-08-26 15:53:54 +08:00
e48baa1bb3 fix(webui): missing HTTP method on fetch calls with body
6 call sites used the fetch default GET while sending a request body,
which browsers reject outright ('Request with GET/HEAD method cannot
have body'). Affected flows: create key, save source, save auto rules,
upload adapter, update key model scope, and /api/chat streaming.

All now send the method their backend handlers require (POST or PUT)
with an explicit Content-Type.
2026-08-26 00:34:57 +08:00
dev
519b18518a chore: bump version to 1.1.0 2026-08-25 18:28:22 +08:00
dev
77b00bfe3f feat(gateway): brute-force protection for /api/login
Prerequisite for removing the nginx global-auth layer in front of the
gateway: llmsproxy must defend its own login endpoint.

Design:
- per-IP failure counter: 5 consecutive failures trigger an exponential
  lockout (30s base, doubling per extra burst, capped at 30min); 15min of
  quiet forgives the counter
- global budget: max 100 failures/minute across all IPs so a distributed
  spray cannot outrun per-IP windows
- locked-out and over-budget attempts get the SAME 'invalid gateway api
  key' 401 as normal failures — no oracle to probe lockout state, no info
  leak on key validity timing
- successful login clears the IP's counter entirely
- clientIP(): prefers X-Real-IP (trusted nginx proxy), falls back to
  RemoteAddr host

Tests: lockout engages at threshold with identical replies, valid keys
rejected while locked, other IPs unaffected, success resets counters,
X-Real-IP extraction.
2026-08-25 11:01:49 +08:00
dev
21ec8f59d8 feat: surface zero cache hits — distinguish 'missed' from 'not reported'
Live testing across the zen pool showed models report
prompt_tokens_details.cached_tokens even when the hit count is 0 (e.g.
nemotron-3-ultra-free returns cached_tokens:0, audio_tokens:0,
cache_write_tokens:0). The previous >0 guard dropped those objects, so a
cache-enabled upstream looked identical to one without cache support.

- types: PromptTokensDetails.CachedTokens always emitted (drop inner
  omitempty) so clients see cached_tokens:0 explicitly; dsh reads it as
  a 0% hit instead of 'no data'
- adapters (9): forward prompt_tokens_details whenever the upstream
  provides it (presence check instead of >0)
- Req: add cache_reported flag set when usage carried cache accounting;
  WebUI shows an amber 0% tag for reported-but-missed rows and keeps
  the em-dash only for sources that never report cache data
2026-08-25 10:04:44 +08:00
dev
ec89daad62 fix(adapters): pass through cache tokens in the remaining 8 adapters
Live testing proved both sensenova and zen DO return cache fields:
- zen laguna-s-2.1-free: usage.prompt_tokens_details.cached_tokens = 32
  (real hit), plus cache_write_tokens/audio_tokens
- sensenova glm-5.2: prompt_tokens_details.cached_tokens present (0 on
  short prompts)

The previous round only patched deepseek/openai/anthropic/gemini.lua;
sensenova/opencode (localzen!) and the other adapters still dropped them.

- sensenova/opencode/groq/mistral/github/kimicode: stream + response
  cache passthrough (same pattern as openai.lua)
- agentrouter: response passthrough + NEW stream usage forwarding (it
  previously dropped the terminal usage-only chunk entirely)
- ollama skipped intentionally: its native API has no cache fields

Verified end-to-end through the gateway: localzen/laguna-s-2.1-free now
returns prompt_tokens_details.cached_tokens=32 to clients, and the request
record carries cache_hit_tokens (both chat and stream paths).
2026-08-25 09:52:55 +08:00
dev
24609289e8 feat: record cache hit/miss per request in audit trail and WebUI
- Req: add CacheHit and CacheMiss fields (carrying upstream cache
  accounting from either prompt_tokens_details.cached_tokens or legacy
  prompt_cache_hit_tokens)
- recordChatUsage (non-streaming): copy cache fields from resp.TokenUsage
- pumpStream (streaming): write lastUsage cache fields back onto rec at
  stream end, so streaming requests carry cache data too
- CSV export: add first_byte_ms, cache_hit_tokens, cache_miss_tokens
  columns alongside the existing latency/prompt/completion
- WebUI request-records table: add a Cache column showing hit% per row
  (green/amber tag with tooltip hit/miss breakdown; em-dash when the
  upstream reported no cache data)
2026-08-25 09:36:21 +08:00
dev
18c2385c61 feat: add per-source TTFB and tokens/s metrics to status page
- Req: add FirstByteMs field (ms to first byte, tracked for streaming)
- Stat: add FirstByteSum for aggregation
- SourceAverages(): new method computing per-source avg TTFB and tokens/s
  from the in-memory ring (300s window)
- SourceStatus: add AvgFirstByteMs and AvgTokPerS fields
- pumpStream: record FirstByteMs after first SSE chunk sent to client
- singleChat/singleChatAuto: set FirstByteMs = LatMs (non-streaming)
- handleStatusAPI: populate the new SourceStatus fields from SourceAverages()
- WebUI source table: two new columns showing TTFB (s) and Tokens/s
2026-08-25 08:24:08 +08:00
dev
2a5c7bbbc7 fix: auto-generate prompt_tokens_details from prompt_cache_hit_tokens in MarshalJSON
When only the legacy DeepSeek fields (prompt_cache_hit_tokens) are set but
the OpenAI-standard prompt_tokens_details is nil, MarshalJSON now auto-
generates the nested object. This ensures dsh (which reads the standard
format first) sees cache hit data regardless of which adapter format the
upstream uses.
2026-08-25 08:01:18 +08:00
dev
045ecf47bc feat(types): pass through upstream cache tokens in TokenUsage
dsh displays cache-hit %, but llmsproxy dropped every upstream's cache
fields — deepseek prompt_cache_hit_tokens, OpenAI prompt_tokens_details.
cached_tokens, anthropic cache_read_input_tokens, gemini cachedContentTokenCount.

Changes:
- TokenUsage: add PromptTokensDetails (with CachedTokens) + PromptCacheHit/Miss
- MarshalJSON: emit prompt_tokens_details.cached_tokens (OpenAI v2 standard)
  and prompt_cache_hit/miss_tokens (DeepSeek legacy) — dsh reads the former
  first, falls back to the latter
- mergeUsage: preserve cache fields across stream chunks
- standardSSEChunk: parse the upstream raw prompt_tokens_details too
- deepseek.lua: forward prompt_cache_hit/miss_tokens + create
  prompt_tokens_details from them
- openai.lua: forward prompt_tokens_details.cached_tokens and legacy
  prompt_cache_hit/miss_tokens; normalize legacy hits into the standard
  object so dsh sees them regardless of upstream format
- anthropic.lua: map cache_read_input_tokens → prompt_tokens_details
- gemini.lua: map cachedContentTokenCount → prompt_tokens_details
2026-08-25 07:57:32 +08:00
dev
334b984c25 fix(webui,gui): repair dead export button + wire chat clear; prune UI redundancy
WebUI (internal/gateway/ui):
- BUG: the export modal's custom-range button called
  downloadStatsCsvFromForm() which was never defined — clicking it threw a
  ReferenceError and nothing downloaded. Implement it: reads #exp-from /
  #exp-to date inputs and forwards to downloadStatsCsv.
- BUG-adjacent: clearChat() existed but was reachable from no control —
  add a Clear button to the chat composer so conversation reset is actually
  possible (+ cClear i18n zh/en).
- remove byte-identical duplicate html[data-theme=dark] CSS block (15 lines)
- remove 8 dead CSS rules (.keys-grid .m-model-row .scope-add/.scope-box/
  .scope-chips .scr-blocks .tag-warn .twrap) and the never-consumed
  --accent custom property
- remove 3 dead JS functions (activeTab/findSlots/scopeUncomb; lastTab decl kept)
- remove 24 dead i18n keys x zh/en (~55 lines) — legacy of the replaced
  key-scope editor, matching the removed .scope-* styles

GUI (cmd/gui/main.js):
- BUG: stopCore() set app.isQuitting=true and nothing reset it — after using
  tray 'stop core', closing the window quit the whole app instead of hiding
  to tray, and core crash auto-restart stayed disabled. isQuitting now only
  flips in restartCore (scoped) and before-quit.

Verified: go vet/test green; node --check on all three GUI js files and the
WebUI inline script.
2026-08-24 23:13:54 +08:00
dev
f3d5ba6cea refactor(gateway): deduplicate stats CSV export paths
- extract csvHeaders() (Content-Type + Content-Disposition) shared by both
  export branches
- extract exportKey(): the identical admin/user key-filter logic existed
  twice (JSON path + keys-csv); now all three call sites share one function
- inline the nine single-use intermediate variables in the keys-csv loop
2026-08-24 22:41:48 +08:00
dev
98847adfd6 refactor(gateway): collapse chat.go four-way duplication
singleChat / singleChatAuto / streamChat / streamChatAuto shared ~260 near-
identical lines (diff after stripping comments was empty). Extract three
shared bodies and shrink all four entry points to thin dispatchers:

- failChat: error → record + writeError, shared ChainErr extraction
  (errors.As returns false for direct-path errors, so the two paths stay
  equivalent without a branch)
- recordChatUsage: exact upstream numbers win, byte estimates fill gaps
- writeChatCompletion: unified ChatCompletion rendering (model name is the
  only direct/auto difference, passed in)
- pumpStream: the full SSE pump (preamble, delta loop, terminal finish +
  usage chunk, [DONE]) — shared by both streaming entries

Behavior change (pinned by TestDirectStreamFailoverAuditSource): direct
streams now pin rec.Source to the source that actually served the stream
after a failover, instead of discarding it (_, usedModel). The audit row
previously recorded the first candidate, which was wrong on failover.

chat.go: 982 → 896 lines (−86).
2026-08-24 22:38:38 +08:00
dev
5bb94db08c refactor: unify oneLineStr/oneLine into types.OneLine
provider.oneLineStr and gateway.oneLine had byte-identical bodies (flatten
whitespace + cap length). Move the single implementation into the types
package, which both layers already depend on, and delete both locals.
2026-08-24 22:26:51 +08:00
dev
ef396f9b47 chore: remove dead code found in redundancy audit
- provider.truncate: zero callers (oneLineStr is the used superset)
- gateway.normalizeModel: zero callers
- config.resolvedAPIKey: zero callers (core.resolveSourceKey is the live equivalent)
- Stats.Records: zero callers (CSV export uses AuditRecords)
- store.containsString: zero callers
- Config.MaxConcurrent: global inflight-cap field never read; per-source
  MaxConcurrent is what actually drives semaphores. Legacy configs carrying a
  top-level max_concurrent key still load (yaml.v3 ignores unknown fields —
  verified by test).
- gui renderer esc(): zero callers; renderer uses textContent, and the embedded
  WebUI has its own esc()
2026-08-24 22:25:21 +08:00
dev
c8a300c998 fix(gateway): forward Flush through statusRecorder so SSE streams in real time
statusRecorder (the access-log wrapper) did not implement http.Flusher, so
w.(http.Flusher) inside streamChat/streamChatAuto returned nil and every SSE
chunk stayed buffered until the response ended. Add Flush() that delegates to
the underlying writer when it supports flushing.

Regression test: TestStatusRecorderFlusher pins the interface assertion and
the forwarding path.
2026-08-24 22:22:38 +08:00
210243f724 test(e2e): track renamed auto-chain 503 summary phrasing 2026-08-24 19:19:14 +08:00
b2183df1e8 feat(adapter): move per-source error condensing into transform_error hooks
Every upstream formats errors differently, which is adapter territory:
the protocol gains an optional transform_error(status, body) hook and all
built-in adapters implement their own envelope parsing (zen free-pool
labels, anthropic/gemini/ollama/mistral shapes, sensenova quota notes,
agentrouter WAF pages). The core keeps a single uniform fallback: when no
hook yields a reason clients get "api error <status>: unknown error" and
the raw body goes to server logs only.
2026-08-24 19:17:36 +08:00
6eac80bc6c fix(provider): condense upstream error bodies before they reach clients
api error strings embedded raw upstream response bodies, so JSON quota
payloads and WAF HTML pages leaked through to clients (and through the
per-tier chain summary). shortAPIError extracts the envelope reason
(error.message / message / msg), collapses HTML blocklist pages to a
marker, and caps everything at one line.
2026-08-24 18:56:06 +08:00
e3a36bc7ca test(gateway): anchor audit-replay fixture to hour buckets
WindowTokens buckets are whole unix hours, so a fixture built from raw
offsets like "1 minute ago" landed in the previous hour bucket whenever
the suite ran in the first minute of an hour — the 1h-window assertion
then saw 0 tokens and failed. Anchor rows to hour boundaries instead
(H-2 / H-1 / current bucket) and assert the exact current-bucket value;
the suite is now deterministic for any run time.
2026-08-24 16:08:19 +08:00
31b1ed2c96 feat(gateway): collapse multi-tier failure details into short client message
Clients saw the full per-tier chain error including upstream HTML WAF
pages and JSON quota payloads. The response now carries one capped
one-line reason per tier (quota/cooling skips preserved); full detail
remains in rec.Err / stats API and is logged server-side.
2026-08-24 16:02:31 +08:00
e19cff248f fix(provider): fail the candidate on error-only streams instead of serving empty replies
Some upstreams (zen free pool) answer HTTP 200 with a single-chunk stream
whose only payload is finish_reason:"network_error" and an empty delta.
ChatStream used to hand that channel up as a success, so clients received
a laundered empty reply (agents ended their turn mid-loop) and the broken
slot kept its scheduling preference. ChatStream now holds back the first
chunk: an error-only done chunk fails the candidate before any byte
reaches the gateway, letting the scheduler fall through to the next
source. Standard finish reasons (including instant-empty "stop") are never
classified as errors.
2026-08-24 15:41:45 +08:00
bb3af3bdb3 feat(gateway): pass through upstream finish_reason end-to-end
The gateway hardcoded "stop" on every terminating stream chunk, so
tool-call rounds reported finish_reason=stop and length caps were
invisible to clients. UnifiedChunk now carries finish_reason; adapters
emit it (with empty-string finish reasons like sensenova treated as
non-terminal), standardSSEChunk passes it through for un-adapted
upstreams, [DONE] no longer emits a duplicate reason-less done chunk,
and both streaming paths emit the real reason with "stop" as fallback.

Also vendor sensenova/agentrouter adapters into the repo: they were
WebUI-only uploads and a deploy sync silently removed them while live
AUTO-chain slots still referenced them.
2026-08-24 15:05:50 +08:00
9c99ded8c0 fix(adapter): send x-opencode-* identity headers so zen stops returning empty tool-call replies
zen fingerprints clients via UA + x-opencode-client/session/request/project
headers; requests without them are routed as anonymous and fail with a
single-chunk network_error stream (tools+stream) or 503 Endpoint is
unavailable (non-stream), which the adapter laundered into empty-but-valid
replies. Derive per-request session/request ids from build_headers meta
(sandbox has no os/math), bump UA to the real client format, map zen
reasoning field to reasoning_content, and set stream_options.include_usage.
2026-08-24 14:37:09 +08:00
690f55c6f1 feat: proactive rate limiting (RPM) + 429 short cooldown for sources
- config.go: Source add RPM field (requests-per-minute cap, 0=unlimited)
- provider.go: RecordRateLimit() — 429 uses fixed 30s cooldown, not exponential
- provider.go: Throttle() — token-bucket proactive rate limiter, spaces requests
  at 60s/RPM interval, respects context cancellation
- provider.go: ReportStatus() — 429 -> RecordRateLimit, 5xx -> RecordFailure
- provider.go: Chat/ChatStream — wire Throttle after TryAcquire
- api.go: sourcePayload + RPM, buildSource passes RPM through
- ui/index.html: add RPM input field in source editor, bilingual i18n labels
- deploy.sh: backup old binary + rollback on healthcheck failure
- provider_test.go: TestModelStateRateLimitShortCooldown, TestThrottleSpacingAndCancel
- config.yaml: sensenova rpm: 12
2026-08-24 02:00:25 +08:00
3069cfce4e fix(gateway): pass through upstream token usage in streams for all adapters
The prior usage-passthrough fix only covered openai/opencode; the same
empty-choices+usage drop bug remained in the 5 sibling OpenAI-compatible
adapters, and non-OpenAI providers (anthropic/gemini/ollama) never surfaced
streaming usage at all.

- deepseek/github/groq/kimicode/mistral: preserve usage on empty-choices
  chunks and attach it to normal chunks (same pattern as openai.lua)
- anthropic: emit usage from message_start (prompt) and message_delta
  (completion); gateway merges split usage additively
- gemini: read usageMetadata in the stream path
- ollama: fix non-streaming key (usage -> token_usage, matches
  UnifiedResponse json tag) and read prompt_eval_count/eval_count;
  surface counts from the done stream chunk
- gateway: mergeUsage combines usage across chunks (non-zero fields win,
  total recomputed from prompt+completion) so split usage doesn't lose
  the prompt half; single-chunk case (OpenAI) preserved exactly
- usage-only chunks: done=false (no redundant terminal stop), matching
  the Go fallback standardSSEChunk
2026-08-18 23:24:01 +08:00
cb9c025952 !1 merge main into main
feat(gateway): pass through exact upstream token usage in streams

Created-by: valuelesser
Commit-by: valuelesser
Merged-by: JianFeeeee
Description: fix(gateway): emit OpenAI-standard token usage in responses
feat(gateway): pass through exact upstream token usage in streams

See merge request: JianFeeeee/ModelRouter!1
2026-08-18 19:14:46 +08:00
83e6d88813 feat(gateway): pass through exact upstream token usage in streams
Streaming responses now carry the upstream's real token usage instead of
gateway estimates:
- UnifiedChunk gains an optional Usage field; adapters (opencode, openai)
  extract usage from upstream stream chunks (including the final chunk with
  empty choices) and pass it through.
- standardSSEChunk preserves usage for passthrough adapters.
- Gateway emits the exact usage in the final stream chunk when available,
  falling back to estimates only when the upstream provided none.

Non-streaming usage was already fixed to emit OpenAI-standard keys.
2026-08-18 19:04:44 +08:00
6f1c806591 fix(gateway): emit OpenAI-standard token usage in responses
- TokenUsage.MarshalJSON now emits both standard (prompt_tokens,
  completion_tokens, total_tokens) and legacy (prompt, completion, total)
  keys, so OpenAI-compatible clients (DSH, DevEco Code, etc.) can read
  token usage.
- ChatChunk gains an optional Usage field; stream responses now send a
  final usage chunk (empty choices) before [DONE].

Refs: usage not visible in clients because the gateway serialized only the
internal short keys and never emitted a streaming usage chunk.
2026-08-18 18:14:37 +08:00
98c08d51c4 fix(gui): re-enable hardware acceleration — remove blanket disableHardwareAcceleration() 2026-08-17 23:02:49 +08:00
ca81048637 fix(gui): ship icon.png as real resource for packaged window icon (linux/mac/win) 2026-08-17 20:57:27 +08:00
b9f2486037 fix(gui): strip host LD_LIBRARY_PATH from packaging + defensive GPU/icon fixes
Root cause of the packaged GUI crashing with SIGSEGV inside ld.so on user
machines (segfault at fixed +0x1ff36, undefined symbols nspr_use_zone_allocator
/ localtime64): the build host had LD_LIBRARY_PATH polluted by a third-party
runtime (/opt/cangjie), which electron-builder baked into the produced binary's
dependency resolution. Clean machines without that library then fail in the
dynamic loader before any app code runs.

- Makefile: every gui-* pack target now runs under `env -u LD_LIBRARY_PATH`
- main.js: app.disableHardwareAcceleration() before ready (avoids the common
  Chromium GPU-process SIGSEGV on hybrid-GPU/Wayland Linux hosts)
- main.js: window icon reads from process.resourcesPath (real file, not asar)
  — asar-path icons are a known GTK segfault source on Linux
2026-08-17 19:50:34 +08:00
8c5279b416 fix(status): source column reflects real traffic + theme-aware tray menu
- api/status sources now carry recent_ok/recent_err (last 300s real gateway
  requests via Stats.SourceRecent) so a source actually serving traffic is
  never shown as down just because probe /models got rate-limited
- WebUI source status column repaints every 5s (no more frozen-at-first-
  render) with a manual refresh button; shows success rate + probe + cooldown
- tray menu status rows were enabled:false (GTK fixed light-grey, invisible
  on light themes) — now enabled with no-op click and nativeTheme listener
  rebuilds the menu on dark/light switches
- ignore local ops scripts (scripts/, machine-specific)
2026-08-17 17:37:04 +08:00
23bd9e5552 fix(gui): resolve GNOME taskbar icon via WM_CLASS/StartupWMClass match
Window/taskbar icon never showed because the packaged .desktop had
StartupWMClass=ModelRouter while the live window WM_CLASS is the lowercase
app class 'modelrouter-gui' — no match, so GNOME fell back to the default
Electron icon.

- main.js: set BrowserWindow icon asset + backgroundColor (window paints
  correctly with frameless titlebar and provides the taskbar's default icon)
- package.json: build.linux.desktop.entry.StartupWMClass now equals the real
  WM_CLASS so packaged installs match too
- local integration: user-level .desktop installed at
  ~/.local/share/applications/modelrouter-gui.desktop (points at the unpacked
  binary) so GNOME resolves the icon immediately
2026-08-17 10:33:10 +08:00
d75d8411c0 chore(compliance): drop opencode zen free-pool from default config and docs
opencode.ai/zen is a free-quota pool for the official opencode app, keyed by
User-Agent fingerprint. Auto-provisioning it in the default config effectively
impersonated the official client to bypass FreeUsageLimitError — a ToS risk.
- default config (Go + GUI embed) now ships NO sources; users add their own
- README/README_EN: remove opencode adapter from built-in list and the zen
  free-pool paragraph
- adapter script and its unit tests stay (explicit opt-in only)
2026-08-17 08:53:42 +08:00
f67831c3c2 fix(gui): desktop shell stability — gated ready state, iframe retry with backoff, theme toggle
- coreReady now only flips after the embedded core actually answers HTTP
  (fixes blank-screen race where the iframe loaded before the listener)
- renderer probes reachability, retries frame load with exponential backoff,
  and recovers via onerror instead of giving up forever
- keep per-user theme in localStorage
- embedded profile no longer preconfigures a zen source (desktop users add
  their own upstreams on the sources page)
2026-08-17 08:53:28 +08:00
e356885130 test(lua): add opencode role normalization test 2026-08-16 17:51:43 +08:00
80 changed files with 17563 additions and 1146 deletions

9
.gitignore vendored
View File

@ -12,8 +12,17 @@ cmd/gui/bin/
cmd/gui/build/gui-dist/
cmd/build/
cmd/gui/*.log
cmd/gui/.cache/
# workspace (AI-assisted coding local state)
.pi/
.pi-glla/
.codegraph/
.omo/
/build/
# local ops scripts (machine-specific, not for sharing)
scripts/
cmd/gui/dist/
# local design/working notes (not part of the shipped repo)
/plan.md

View File

@ -9,6 +9,12 @@ export GUI_DIR
GUI_DIR=cmd/gui
VERSION ?= $(shell git describe --tags --dirty 2>/dev/null || echo "0.1.0")
# Clean env for GUI packaging: the host may have LD_LIBRARY_PATH polluted by
# third-party runtimes (e.g. /opt/cangjie) which electron-builder bakes into
# the produced binary's DT_RPATH/dep resolution -> the packaged app then
# SIGSEGVs inside ld.so on clean machines. Strip it for every pack step.
CLEAN_PACK_ENV = env -u LD_LIBRARY_PATH
all: build
# ---- core binary (server users / embedded core) ----
@ -17,6 +23,14 @@ build:
CGO_ENABLED=1 go build -tags luajit -trimpath -o $(BUILD_DIR)/$(BINARY) ./cmd/llmsproxy
@echo "Built: $(BUILD_DIR)/$(BINARY) ($(VERSION))"
# ---- core distribution packages (deb + rpm + tar.gz) -----
# Server admins install these with dpkg/rpm; unlike the GUI they are the bare
# gateway binary + systemd unit (with memory tuning) + Lua adapters.
core-dist:
./packaging/core-dist.sh $(VERSION)
.PHONY: core-dist
# ---- GUI (Electron) ----
gui:
@cd $(GUI_DIR) && npm start
@ -25,38 +39,44 @@ gui-dev:
@cd $(GUI_DIR) && npm run dev
# build the embedded core into the GUI, then run electron-builder
# targets: deb + AppImage (linux), rpm needs system rpmbuild
# targets: deb + rpm + AppImage (linux); rpm needs system rpmbuild
# NOTE: run in a clean env (no LD_LIBRARY_PATH pollution) so the packaged
# binary does not link against host-only libs and crash on user machines.
gui-dist:
@cd $(GUI_DIR) && npm run dist:linux
@cd $(GUI_DIR) && $(CLEAN_PACK_ENV) npm run dist:linux
./packaging/verify-dist.sh
# deb only (for local sharing), rpm needs rpmbuild
gui-deb:
@cd $(GUI_DIR) && npm run dist:debian
@cd $(GUI_DIR) && $(CLEAN_PACK_ENV) npm run dist:debian
# windows nsis installer: host-built core (needs apt install gcc-mingw-w64-x86-64 + wine)
gui-win:
@cd $(GUI_DIR) && npm run dist:win
@cd $(GUI_DIR) && $(CLEAN_PACK_ENV) npm run dist:win
# windows nsis installer: full cross-build in docker (self-sufficient, no host mingw)
gui-win-docker:
@cd $(GUI_DIR) && ./scripts/dist-win-docker.sh
@cd $(GUI_DIR) && $(CLEAN_PACK_ENV) ./scripts/dist-win-docker.sh
./packaging/verify-dist.sh
gui-dist-dir:
@cd $(GUI_DIR) && npm run dist:dir
@cd $(GUI_DIR) && $(CLEAN_PACK_ENV) npm run dist:dir
test:
go test ./...
clean:
rm -rf $(BUILD_DIR) cmd/gui/build/gui-dist cmd/gui/bin
rm -rf $(BUILD_DIR) cmd/build cmd/gui/dist cmd/gui/build/gui-dist cmd/gui/bin
help:
@echo "Targets:"
@echo " make build build core binary (luajit)"
@echo " make core-dist build core packages: deb + rpm + tar.gz (needs nfpm)"
@echo " make gui run desktop GUI (dev)"
@echo " make gui-dev run GUI with devtools"
@echo " make gui-dist build deb + AppImage (linux)"
@echo " make gui-dist build deb + rpm + AppImage (linux)"
@echo " make gui-deb build deb only"
@echo " make gui-win build Windows nsis (needs host mingw + wine)"
@echo " make gui-win-docker build Windows nsis all-in-docker (no host mingw needed)"
@echo " make test run go tests"
@echo " make clean remove build dirs and packaged installers (build/, cmd/build/, cmd/gui/dist/)"

278
README.md
View File

@ -15,15 +15,18 @@
### 极致轻量
- **单二进制**:编译后约 10MB,零运行时依赖(仅依赖系统 libc),部署即用
- **极低内存占用**:空闲状态仅 ~15MB RSS,满载并发 100+ 请求时峰值 < 100MB
- **单二进制**:编译后 8~12MB(`-s -w` strip 后约 8MB),零运行时依赖(仅依赖系统 libc),部署即用
- **低内存占用**:与源数量相关,非与运行时长相关——单源 ~10MB、16 源生产实例 ~32–35MB([实测分解与调优](#内存占用实测与调优))
- **日志按需加载**:统计数字来自**全量**审计日志(启动时流式扫一遍即释放,29MB/22 万行约 260ms),但原始记录不常驻内存——默认只加载首屏,下滚自动分页,CSV 导出流式写出(O(1) 内存),离页即释放
- **零运行时依赖**:纯 Go + LuaJIT 静态链接,无需安装 Python/Node/Java 等运行时
- **启动极快**:冷启动 < 200ms,热重载配置 < 10ms
### 强大的多租户调度能力
- **多密钥多租户**:支持无限密钥,每个密钥独立角色、模型范围、Token 配额、重置周期
- **按模型配额**:每把 key 的每个模型单独配 token + 请求数配额与重置周期(小时/周/月/自定义 N 小时);一个模型用满只影响该模型,同 key 其它模型照常;耗尽返 429 + `Retry-After` 并点名模型,admin key 永不受限
- **AUTO 智能调度**:基于优先级档位的分级调度,同优先级源自动轮询负载均衡,故障自动毫秒级故障转移
- **自愈冷却**:冷却上限 5 分钟,过半后放行 1 个探测请求,上游/额度恢复即刻回归轮询,无需等满冷却窗口
- **Token 配额管理**:精确到模型级别的 Token 配额控制,支持小时/周/月/自定义小时周期自动重置
- **同优先级源负载均衡**:同一优先级档位的多个源,请求自动 Round-Robin 均匀分发,故障毫秒级故障转移
- **Source-Model 前缀路由**:支持 `source-model`/`source:model`/`source/model` 精确指定上游源
@ -37,7 +40,7 @@
### 灵活的协议适配
- **LuaJIT VM**:每个适配器独立 VM + worker 池,安全并发,Lua 脚本热加载无需重启
- **LuaJIT VM**:每个适配器独立 VM + **弹性 worker 池**(按需扩缩容,空闲时自动回收),安全并发,Lua 脚本热加载无需重启
- **Lua 适配器协议**:`transform_request` / `transform_response` / `transform_stream_chunk` 双向转换
- **动态请求头钩子**:`build_headers(meta)` 支持 HMAC 签名、动态 Header 注入
- **多模态透传**:`image_url` 等多模态内容在多源间无损透传,Anthropic/Gemini/Ollama 自动转换
@ -60,6 +63,12 @@ Groq、Mistral、Ollama、KimiCode…)通过 **Lua 适配器** 做协议转换
- **生图**:`POST /v1/images/generations`,`kind: image` 的模型独立路由。
- **多模态**:`content` 数组(`image_url` 等)在多源间无损透传,Anthropic/Gemini/Ollama 自动转换。
- **LuaJIT VM**:基于 golua 绑定的 LuaJIT,每个适配器独立 VM + worker 池,安全并发。
池是 **弹性** 的:`Σ max_concurrent` 只是上限而非预分配,状态按需创建、空闲时自动回收,
因此空载网关几乎不持有 Lua 状态(启动 0 个)。
- 扩容步长由适配器最大并发决定:`clamp(ceil(max/8), 1, 8)`,且仅在**真实并发争用**
(所有现有状态均已占用)时批量预热,顺序流量始终只用 1 个状态。
- 缩容步长由当前连接数决定:`clamp(ceil(冗余/(1+占用)), 1, 冗余)`——无连接时一轮收到
常驻下限,忙时每轮只释放 1 个,保护热路径。WebUI「适配器」页展示实时池指标。
- **disable_thinking**:请求 `disable_thinking:true`(或上游对应字段)动态开关推理。
- **Lua 适配协议**:每个源挂一个 `.lua` 适配器,完成 `transform_request` /
`transform_response` / `transform_stream_chunk` 双向转换,协议差异全在 Lua 层。
@ -89,8 +98,7 @@ GOMODCACHE=... GOPROXY=off go build -tags luajit -o llmsproxy ./cmd/llmsproxy
> **配置不入库**:仓库不携带任何 config 文件(配置文件含密钥)。二进制首次运行时
> 会在 `-config` 指定的位置生成默认配置:随机 admin key(打印在启动日志中)、
> 仅绑定 `127.0.0.1:8080`、内置一个免 key 的 opencode zen 源可直接对话。
> 首次登录后请在 WebUI「密钥」页更换管理员密钥。
> 仅绑定 `127.0.0.1:8080`。首次登录后请在 WebUI「密钥」页更换管理员密钥。
```bash
# 无 key -> 401
@ -171,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` 解析顺序:
@ -185,16 +268,51 @@ sources:
### AUTO 链(优先级页)
AUTO 调度**只**由优先级页保存的规则(持久化到 `runtime_file` 的 `auto` 字段)决定,
源配置里的 `models[].priority` 数字不再参与调度、也不再显示。
AUTO 调度**只**由优先级页保存的规则决定,源配置里的 `models[].priority` 数字不再
参与调度、也不再显示。**聊天与生图是两套独立的链**,在优先级页面通过「聊天 / 生图」
按钮切换编辑:
- **聊天链**(`runtime_file` 的 `auto` 字段):服务 `POST /v1/chat/completions`
的 `model: AUTO` 请求;只接受 `kind: chat` 模型,配置在链里的生图槽位会被忽略。
- **生图链**(`runtime_file` 的 `auto_image` 字段):服务 `POST /v1/images/generations`
的 `model: AUTO` 请求;只接受 `kind: image` 模型。未配置生图链时,AUTO 生图回落
到「所有含生图模型的源、按注册顺序」的旧行为,兼容旧配置。
每条规则 = 一个「槽位」:`{ model, source, tier, token_quota, period, hours }`。
- 每条规则 = 一个「槽位」:`{ model, source, tier, token_quota, period, hours }`。
- `tier` 表示优先级档位:同一档的模型并排、共享该优先级;档位从上到下递减。
- 同一模型可配置多个槽位(如 A 源低配 → B 源低配 → A 源高配 → B 源高配),按档位顺延。
- `token_quota` > 0 时该槽位在重置周期内用满即顺延到下一槽位;`period` 支持
`hour` / `week` / `month` / `nhour`(配合 `hours`),空 = 不限。
- 生图模型(`kind: image`)不参与 AUTO 链;生图走 `POST /v1/images/generations`
的独立路径。
- 同 tier 内按「偏好分」排序,失败槽位冷却后自动跳过;冷却 / 配额耗尽 / hard
错误会顺延到下一 tier,全部失败时返回 503 并附上每档失败摘要。
### 冷却与半冷却探测(自愈调度)
槽位失败后进入指数退避冷却(5s → 10s → 20s …,上限 **5 分钟**)。冷却不是全时段
硬闸:
- **前半段完全静默**:刚失败的槽位不接任何流量,避免对故障上游持续施压。
- **后半段放行 1 个探测**:窗口过半后,同一 `(源, 模型)` 最多允许 **一个** 请求作为
探测通过。探测槽数 = `clamp(max_concurrent / 10, 1, 2)`——`max_concurrent: 10`
的源恰好放行 1 个探测请求。
- **探测排在最后**:同 tier 内探测候选永远排在健康候选之后,只有没有正常槽可用时
才承接真实流量,因此探测不会抢走可用容量。
- **探测即真实请求**:成功则立即清零失败计数与冷却,槽位当场回到正常轮询——
上游恢复/额度恢复后不必等满整个冷却窗口。失败则重开一个新窗口,下次探测顺延到
新窗口的中点,不会连续重试。
分类冷却策略:
| 情况 | 冷却 | 说明 |
|---|---|---|
| 传输错误 / 5xx | 5s→10s→20s…,上限 5min | 指数退避 |
| 401 / 403 凭据错误 | 10min | 换 key 后由下一次探测自动接回 |
| 429 限流 | 30s 固定 | 「太快」不等于「坏了」,不进指数阶梯 |
| 配额耗尽(`insufficient_quota` / `余额不足` 等) | 对齐配额窗口,上限 30min | **不进指数阶梯、不涨失败计数**,额度恢复即可用 |
WebUI「优先级」页的槽位健康标签会区分 `冷却` / `待探测` / `探测中`,悬停可看窗口
起点、探测放行时刻与完全恢复时刻。
### 源(Source)与适配器(Adapter)的关系
@ -213,7 +331,7 @@ AUTO 调度**只**由优先级页保存的规则(持久化到 `runtime_file`
2. `config.Load` 读取 `config.yaml`,`config.NewStore(runtime_file)` 读 WebUI 改动的
运行时源,二者按名称合并成完整源列表。
3. `rebuildRegistry` 为每个源创建 `provider.Provider`(持有目标适配器),并按源的
并发上限配置适配器 worker 池大小。
并发上限配置适配器 worker 池**上限**(不预分配,见下)。
4. 请求进来时 `Registry` 按 `model` 路由到 Provider,Provider 调适配器
`transform_request` → HTTP 发送 → `transform_response` / `transform_stream_chunk`。
@ -225,6 +343,70 @@ WebUI 上的"新增/编辑源"、"上传 Lua 适配器"、"改 AUTO 优先级链
`config.yaml` 中定义的基源无法改写配置文件,采用删除标记隐藏(重启后仍隐藏),
在 UI 里重新添加同名源即可恢复。
### 源模板(多 key 负载均衡)
多个上游 key 共享同一套配置(URL / 适配器 / 模型列表 / 并发 / rpm)时,逐一复制
整段 source 太冗余。**源模板**保存除 `name` 和 `api_key` 以外的全部源字段,
从模板一键创建多个仅 key 不同的源,天然支持多 key 负载均衡(每个展开源独立调度、
独立冷却、独立健康状态)。
- WebUI「源」页右上角「**模板管理**」按钮:列出全部模板,可编辑 / 删除 / 新建。
- 添加源弹窗右上角两个按钮:
- **从模板创建** → 选模板 → 表单自动填充(name + key 仍需你填);
- **存为模板** → 把当前表单非 key/name 字段存为新模板。
- 模板持久化在 `runtime_file` 的 `source_templates` 字段,与运行时源一样热管理,
改完即时生效。模板本身不含密钥,可以安全共享 / 纳管。
模板不会自己变成源——它只是「配方」。真正承担流量的是从模板创建出来的、带 key
的具体源。
### 内存占用(实测与调优)
内存不是一个定数:它随**配置的源数量**增长(每个源一个 `http.Transport` 与连接池、
一组健康状态),与运行时长无关。本机实测(Linux x86_64,12 核):
| 部署形态 | 启动 RSS | 稳态 RSS |
|---|---|---|
| 1 源 / 1 适配器(最小配置) | ~4 MB | ~10 MB |
| 1 源 + 29 MB 历史审计日志 | ~19 MB | ~20 MB |
| **16 源 / 13 适配器 / 59 模型(本机生产)** | ~28 MB | **~32–35 MB** |
> 历史参考:本项优化前同一生产配置为 **~105 MB**。降幅来自三处:审计日志的**原始记录**
> 不再常驻内存(约 25 MB;统计聚合仍扫全量,但扫完即释放)、Lua 状态池不再单调增长、
> 以及下面两个运行时开关。
内存构成(生产实例分段测量,`/proc/<pid>/smaps`):
| 区域 | RSS | 说明 |
|---|---|---|
| Go 堆 | ~14 MB | provider/registry/scheduler 结构 + 连接池缓冲 |
| 其他匹名(线程栈 / LuaJIT chunk / runtime) | ~12 MB | 与线程数、已加载适配器数相关 |
| 二进制 text+rodata | ~8 MB | 映射的可执行文件页(只读、可被内核回收) |
| Go runtime 预留 | ~4 MB | VSZ 上看到的 2 GB+ 是地址空间预留,不占物理内存 |
| 共享库 | ~3 MB | libc / libluajit / libm |
**两个推荐的部署开关**(只适用于环境变量,无需改代码):
```ini
# /etc/systemd/system/llmsproxy.service
Environment=GOGC=50
Environment=MALLOC_ARENA_MAX=2
```
- `MALLOC_ARENA_MAX=2`:LuaJIT 的分配走 cgo → glibc malloc,glibc 默认允许 `8×nproc`
个 per-thread arena,每个碰到 malloc 的 OS 线程会占用一个(各约 1 MB,且**不归还给系统**)。
实测从 8–15 个 arena(约 7–12 MB)降到 **0**。
- `GOGC=50`:把 Go 堆增长目标减半。**单独使用无效**(省下的堆会立即被更多 glibc
arena 吃掉,实测 20.3 → 21.5 MB 反而变大),必须与 `MALLOC_ARENA_MAX` 配合,
两者同时开启才降 **~19%**。网关是 I/O 密集型(本机 9 小时仅消耗 1min10s CPU),
多出的 GC 周期与它的空闲 CPU 相比可忽略。
> 桌面版(Electron)已在 `cmd/gui/main.js` 里默认为内嵌核心注入这两个开关;手动 export 同名
> 环境变量可覆盖。如果你用自己的 systemd unit / 容器,建议照上面加上。
可选:`GOMEMLIMIT=48MiB` 作软上限,实测再省 ~1 MB;代价是逐近上限时 GC 转激进,
温和场景不必开。
### disable_thinking
请求体带 `"disable_thinking": true` 时,网关透传给各适配器;DeepSeek 适配器将其
@ -238,7 +420,10 @@ WebUI 上的"新增/编辑源"、"上传 Lua 适配器"、"改 AUTO 优先级链
支持按时间范围导出 CSV;点击模型可生成 pin 到该模型的连接配置。
- **对话页**:流式/非流式调试。
- **密钥页**:创建/编辑网关 key,为每个 key 配模型范围(模型 + 源 + token 配额 + 周期),
管理员管理全部 key,用户只看到自己的 key。
管理员管理全部 key,用户只看到自己的 key。每个模型砖块显示自己的配额徽标
(如 `1.0K 77×·1h`,未设配额显示 `∞`),点开可编辑该模型的 token 配额、
请求数配额与重置周期 —— 配额按模型独立生效,一个用满不影响同一 key 的其它模型。
「我的密钥」页对用户展示本 key 的模型范围与各自配额。
- **优先级页**:拖拽积木配置 AUTO 链档位。
- **源页**:在线增删改上游源(API key 等敏感字段加密落盘)。
- **适配器页**:上传 / 删除 Lua 适配器脚本。
@ -281,19 +466,12 @@ return {
### 内置适配器
`openai` `deepseek` `anthropic` `gemini` `github` `groq` `mistral` `ollama` `kimicode`
`opencode`。
`openai` `deepseek` `anthropic` `gemini` `github` `groq` `mistral` `ollama` `kimicode`。
`anthropic`/`gemini`/`ollama` 适配器内置多模态转换(`image_url` → 各自上游格式);若
源启用 `disable_thinking`,`deepseek` 适配器会把 `extra_body.thinking.type` 置为
`disabled`。
**opencode** 适配器面向 opencode.ai zen 免费池(`https://opencode.ai/zen/v1`):
zen 按 User-Agent 指纹识别官方客户端并分流——非官方 UA 的请求(curl、Go 默认 UA)
会被丢进匿名池触发 `FreeUsageLimitError`。该适配器固定发送 opencode 客户端 UA,
配合 `api_key: "public"`(官方无 key 客户端实际发送 `Bearer public`)即可走免费池,
零成本接入 `deepseek-v4-flash-free` 等免费模型。
**kimicode** 是展示 `build_headers` 的样例:云端校验调用方 app,需要按
`meta.app_secret` 对时间戳+URL+请求体哈希做 HMAC 签名并附 `X-App-Sign` 等头。
配好 `sources[].meta.{app_id, app_secret, app_agent}` 即可。
@ -304,14 +482,24 @@ zen 按 User-Agent 指纹识别官方客户端并分流——非官方 UA 的请
一键启停、系统托盘常驻、开机自启、静默启动,内容区一比一内嵌完整 WebUI
(状态/对话/密钥/优先级/源/适配器 六个页面,免登录)。服务器用户继续使用纯 Go 二进制。
### 两种安装包:Headless 与 Desktop
发布到 Release 的安装包分为两类,面向不同用户:
- **Headless(服务器版)**:纯核心二进制(无 Electron),配置驱动,适合服务器、
容器、systemd 等无人值守场景。发布包含 `deb` + `rpm`(含 systemd unit、内存调优
与示例配置)与 `tar.gz`(二进制 + `config.example.yaml` + `adapters/` + README),
Windows 为 `.zip`。
- **Desktop(桌面版)**:Electron GUI 安装包,内嵌核心,适合个人桌面日常使用。
### 选择 GUI 还是纯后端
| 场景 | 推荐 | 理由 |
| ------ | ------ | ------ |
| 服务器 / 内网网关 / 无人值守常驻 | **纯后端**(单二进制) | 极轻量(~10MB、~15MB RSS),零依赖单进程,直接跑在 systemd/任意容器里,远程管理 |
| 个人桌面日常使用 / 多设备内网共享 | **GUI**(Electron) | 免登录内嵌 WebUI、系统托盘一键启停、开机自启、静默后台,适合不懂命令行的使用者 |
| Windows 桌面 | **GUI** | 纯后端在 Windows 上需自行注册服务,GUI 提供原生托盘/自启体验 |
| CI 一键出三平台安装包 | **GUI 打包脚本** | `make gui-*` 系列出 deb/AppImage/NSIS,可进发行流水线 |
| 服务器 / 内网网关 / 无人值守常驻 | **Headless**(单二进制) | 极轻量(二进制 ~8MB;RSS 随源数量,单源 ~10MB、16 源 ~32–35MB,见[内存占用](#内存占用实测与调优)),零依赖单进程,直接跑在 systemd/任意容器里,远程管理 |
| 个人桌面日常使用 / 多设备内网共享 | **Desktop**(GUI) | 免登录内嵌 WebUI、系统托盘一键启停、开机自启、静默后台,适合不懂命令行的使用者 |
| Windows 桌面 | **Desktop** | 纯后端在 Windows 上需自行注册服务,GUI 提供原生托盘/自启体验 |
| CI 一键出三平台安装包 | **Desktop 打包脚本** | `make gui-*` 系列出 deb/AppImage/NSIS,可进发行流水线 |
两者完全同源:GUI 内嵌的就是纯后端的同一份 `llmsproxy` 核心(LuaJIT 版),
配置/适配器格式完全一致,可随时互换。
@ -323,27 +511,36 @@ cd cmd/gui
npm install
make gui # 或 npm run dev —— 开发运行(需先 make build 生成 bin/llmsproxy)
make gui-deb # 本机构建 deb(分享给其它 Linux 用户)
make gui-dist # 构建 deb + AppImage(linux)
make gui-dist # 构建 deb + rpm + AppImage(linux,rpm 需系统 rpmbuild)
make core-dist # 构建核心 deb + rpm + tar.gz(服务器版,需 nfpm)
make gui-win # 构建 Windows nsis 安装包(本机 mingw + wine)
make gui-win-docker # 构建 Windows nsis 安装包(全 docker 自足,无需本机 mingw)
make gui-win-docker # 构建 Windows nsis 安装包(全 docker 自足,含 wine)
make gui-dist-dir # 仅解包产物(调试用,不产安装包)
```
产物在 `cmd/build/gui-dist/`:`ModelRouter-1.0.0.AppImage`、
`modelrouter-gui_1.0.0_amd64.deb`、`ModelRouter Setup 1.0.0.exe`(win)。rpm 需系统 `rpmbuild`。
产物在 `cmd/build/gui-dist/`:`ModelRouter-<ver>.AppImage`、
`modelrouter-gui_<ver>_amd64.deb`、`ModelRouter Setup <ver>.exe`(win);
核心包在 `cmd/build/dist/`:`llmsproxy_<ver>_amd64.deb`、`llmsproxy-<ver>.x86_64.rpm`。
rpm 需系统 `rpmbuild`。每次打包后都会跑 `packaging/verify-dist.sh` 产物体检
(大小下限门禁——1.3.0 曾因宿主 wine 损坏产出 264KB 的残缺 exe 却没有被拦住)。
发布时会把产物重命名为 `ModelRouter-Desktop-*`(GUI)与 `ModelRouter-Headless-*`(纯后端)
两类上传到 GitCode Release。
#### Windows 打包(docker 固化方案,推荐)
`make gui-win-docker` 一步完成:交叉编译核心 + NSIS 打包,宿主机**无需安装 mingw**。
`make gui-win-docker` 一步完成:交叉编译核心 + NSIS 打包**全在 docker 内**,宿主机
**无需 mingw、无需 wine、无需 node**。镜像内置 wine(含 i386 运行时),修复了
此前宿主 wine 损坏(缺 syswow64 → `c0000135`)导致产出 264KB 残缺 exe 的问题。
- 构建镜像:`cmd/gui/docker/win-builder/Dockerfile`(`golang:1.25` + mingw-w64 +
预编译 LuaJIT windows 版 `lua51.dll` 与 import lib)。LuaJIT 源码取自 gitcode
镜像(国内可达),github 兜底;Go 依赖走 `https://goproxy.cn`。
- 核心编译:`cmd/gui/scripts/win-core-docker.sh` —— 在容器内产出
`cmd/gui/bin/{llmsproxy.exe,lua51.dll}`。
- 一键发布:`cmd/gui/scripts/dist-win-docker.sh` —— 核心编译后调用
electron-builder(宿主需有 `wine` 用于 NSIS)。
- 首次运行自动构建 `modelrouter/win-builder` 镜像;已存在则直接复用。
node + wine32/wine64 + 预编译 LuaJIT windows 版 `lua51.dll` 与 import lib)。
LuaJIT 源码取自 gitcode 镜像(国内可达),github 兜底;Go 依赖走 `https://goproxy.cn`。
- 一键发布:`cmd/gui/scripts/dist-win-docker.sh` —— 核心编译 + electron-builder
NSIS 全部在容器内执行(`USE_SYSTEM_WINE=true` 走镜像内的 wine);electron-builder
的 ~400MB 工具链缓存挂到 docker volume,二次构建增量复用。
- 首次运行自动构建 `modelrouter/win-builder` 镜像(约 5–10 分钟);已存在则直接复用。
- 打包完成后自动跑 `packaging/verify-dist.sh` 门禁:exe 低于 5MB / deb·rpm·7z 低于
10MB 即失败——残缺安装包不可能再溜进 Release。
```bash
# 全 Windows 发布(核心重编 + NSIS):
@ -372,6 +569,15 @@ internal/gateway # OpenAI 兼容 HTTP 服务 + 鉴权 + SDK/流式 + Web
internal/types # 统一格式 & OpenAI wire 类型
```
## 开发协作
- **Git 分支工作流**:主分支 / 特性分支 / 发布分支的规范见
**[docs/git-workflow.md](docs/git-workflow.md)**([English](docs/git-workflow-en.md))。
主线:`main` 永可部署 → 特性合入 `main` → 切 `release/vX.Y.Z` 打 tag 发布 →
hotfix 提交发布分支并 cherry-pick 回 `main`。
- **Lua 适配器协议**:见 **[docs/lua-adapters.md](docs/lua-adapters.md)**
([English](docs/lua-adapters-en.md))。
## 测试
```bash

View File

@ -11,7 +11,7 @@ KimiCode…) behind a single endpoint. Pick a specific model or use **AUTO** mod
which routes to the best healthy upstream by configured priority.
> **Lightweight, no recompile to add sources**: the gateway is a single Go
> binary (~10 MB, zero runtime dependencies). Adding or switching an upstream is
> binary (8-12 MB, ~8 MB stripped; zero runtime dependencies). Adding or switching an upstream is
> just a `sources` entry in `config.yaml` (or via the WebUI) or a `.lua`
> adapter — **no Go changes, no recompile**. Changes submitted through the
> WebUI take effect immediately (hot reload); editing `config.yaml` or a `.lua`
@ -38,10 +38,31 @@ Extracted and independently evolved from the multi-source LLM adapter layer of
`reasoning_content`, `tool_calls`, `usage`).
- **Image generation**: `POST /v1/images/generations`, routed to models with
`kind: image`.
- **Per-model quota**: each key gives every model its own token and request caps
plus a reset period (hour/week/month/custom N hours). One model running out
affects only that model — the key's other models keep working. Exhaustion
answers 429 + `Retry-After` naming the model, so a client resumes when the
window rolls over; admin keys are never capped.
- **Multimodal**: `content` arrays (`image_url` etc.) pass through losslessly;
Anthropic/Gemini/Ollama are translated automatically.
- **LuaJIT VM**: golua-binding LuaJIT; each adapter has its own VM + worker
pool for safe concurrency.
pool for safe concurrency. Pools are **elastic**: the sum of `max_concurrent`
over an adapter's sources is a ceiling, not a preallocation — states are booted
on demand and reclaimed when demand drops, so an idle gateway holds close to
zero Lua states. The grow step follows the adapter's max concurrency
(`clamp(ceil(max/8), 1, 8)`, and never exceeds the number of queued callers);
the shrink step follows the live connection count
(`clamp(ceil(slack/(1+in_use)), 1, slack)`), so an idle pool collapses in one
round while a busy one gives up a single state at a time.
- **On-demand request logs**: dashboard totals are computed from the **full**
audit history (streamed once at startup and released — 29 MB / 221k lines in
~260 ms), while the raw records are never held in memory: the dashboard loads
one screen, scrolling pages the rest straight off disk, CSV export streams in
O(1) memory, and leaving the page releases everything.
- **Self-healing cooldown**: cooldown is capped at 5 minutes and, past the
window's midpoint, exactly one probe request is allowed through; a recovered
upstream (or a reset quota) returns to full rotation on that probe instead of
waiting out the window.
- **disable_thinking**: `disable_thinking: true` toggles reasoning per-request.
- **Lua adapter protocol**: each source mounts a `.lua` adapter with
`transform_request` / `transform_response` / `transform_stream_chunk` — all
@ -77,7 +98,7 @@ GOMODCACHE=... GOPROXY=off go build -tags luajit -o llmsproxy ./cmd/llmsproxy
> **No config is shipped in the repo** (config files carry real keys). On
> first run the binary generates a default config at the `-config` path: a
> random admin key (printed to the startup log), loopback-only `127.0.0.1:8080`,
> and a keyless opencode zen source ready to chat. Rotate the admin key in the
> and a keyless source ready to chat. Rotate the admin key in the
> WebUI after first login.
```bash
@ -155,11 +176,103 @@ under the `keys` field of the runtime file (encrypted at rest):
delete the seed key.
- The WebUI **Keys page** creates/deletes keys. Each key has a role (`admin`
manages everything, `user` sees only its own key) and an optional **model
scope** (model + source + token quota + reset period).
scope** (model + source + token quota + reset period). Each model brick shows
its own budget badge (e.g. `1.0K 77×·1h`, `∞` when uncapped); clicking it
edits that model's token quota, request quota and reset period. Quotas apply
per model, so one model running out never blocks the key's others. The
"My key" view shows a user its model scopes and their budgets.
- Clients authenticate with any authorized key's plaintext as
`Authorization: Bearer <key>`.
- Deleting a key removes it from the store immediately.
#### Per-key usage quota
**Quotas are per model.** Each key's `models[]` list gives every model its own
token budget and request budget. One model running out affects only that model
— the key's other models keep working.
```yaml
keys:
- key: sk-gw-<hex>
role: user
name: agent-alice
models:
- model: deepseek-v4-flash
token_quota: 1000000 # this key's token budget for this model
req_quota: 20000 # requests within the window
period: nhour # "" | hour | week | month | nhour
hours: 6 # n-hour only
- model: AUTO
token_quota: 5000000 # AUTO is a quota entry like any other
period: hour
- model: kimi-k3 # no quota listed = unlimited
```
- **There is deliberately no key-wide total.** A key-wide cap would make
"move A's budget to B" a global reallocation; per-model budgets keep each
model independently controllable, so it stays visible which model is
actually consuming the spend.
- `period`: empty = never resets (lifetime total); `hour` / `week` / `month` =
fixed windows; `nhour` + `hours` = a custom hour count. **A misspelled
period is rejected at write time** rather than silently becoming a
never-resetting quota.
- Quotas are isolated per key *and* per model within a key: one key exhausting
`m1` never draws on another key's budget, and never blocks the same key's
`m2`.
- Usage counts chat, streaming and image requests, and survives a restart by
replaying the audit log (40 days retained, covering the longest `month`
window).
- An exhausted quota returns **429 + `Retry-After`**
(`rate_limit_exceeded`) and the message names the model that ran out, so a
client resumes when the window rolls over; a model the key may not use stays
403. **Admin keys are never capped** (quotas on their scopes are not
enforced either), so a cap can never lock the operator out.
- Buckets are whole unix hours, so a window frees up at most an hour late
(deliberately freeing late rather than overspending).
- `PUT /api/keys/{key}` submits quotas by submitting `models` — the caps are
part of the scope, so there is no second budget that can drift out of sync
with the model list. An explicit `0` lifts that model's cap.
##### Quota rejection vs capacity rejection
A quota is checked at the entry point and occupies no upstream slot, so it is a
cheap rejection. Running out of capacity instead waits out the bounded
`busyWait` before answering 503. Measured (capacity 8, 0.6 s/request, 100
concurrent clients):
| Case | Status | Latency | Reaches upstream |
|---|---|---|---|
| Quota spent (`req_quota` 50, 100 concurrent) | `429 rate_limit_exceeded` + `Retry-After` | **19–21 ms** | no |
| Capacity exhausted (same topology, no quota) | `503 upstream_error` | ~2.6 s (`busyWait`) | overflow: no |
So a client can tell "wait for the window to roll over" apart from "wait for
upstream capacity" — the first only needs patience, the second usually means
back off concurrency or switch sources.
##### What the quota costs
Buckets are kept per (key, model, whole unix hour) for 40 days. Measured on an
AMD 7840HS:
- A window query scans the window, not the whole retention: 49 ns for 1 h,
55 ns for 24 h, 3.9 µs for 30 d.
- Recording one request: 283 ns, 3 allocations (unchanged from before this
feature; the allocations come from the record ring buffer).
- Memory: the production shape (7 keys x 8 models x 2 sources at full 40-day
retention) costs about **3.7 MB**. The `source::model` bucket used by a
source-pinned quota is created **lazily** — it is maintained only once some
quota actually pins that source, because writing it on every record would
cost an extra 18 MB at 20 keys x 8 models x 3 sources.
The trade-off: usage recorded before a pinned quota was configured cannot be
split by source afterwards, so **a pinned quota may under-count its first
window**.
```
HTTP/1.1 429 Too Many Requests
Retry-After: 2100
{"error":{"type":"rate_limit_exceeded","message":"token quota exceeded for \"deepseek-v4-flash\" (5000000/5000000)"}}
```
### Model routing
`/v1/chat/completions` `model` resolution order:
@ -176,12 +289,20 @@ tiered AUTO chain.
### AUTO chain (Priority page)
AUTO scheduling is driven **only** by the rules saved on the Priority page
(persisted as the `auto` field of the runtime file). The numeric
`models[].priority` in source configs no longer participates in scheduling and
is no longer shown.
AUTO scheduling is driven **only** by the rules saved on the Priority page. The
numeric `models[].priority` in source configs no longer participates in
scheduling and is no longer shown. **Chat and image generation are two
independent chains**, switched in the Priority page via the `Chat / Image` toggle:
- **Chat chain** (`auto` field): serves `POST /v1/chat/completions` with
`model: AUTO`; only `kind: chat` models are honored (image slots are skipped).
- **Image chain** (`auto_image` field): serves `POST /v1/images/generations` with
`model: AUTO`; only `kind: image` models are honored. With no image chain
configured, AUTO image generation falls back to legacy discovery (all sources
exposing an image model, in registry order) for backward compatibility.
Each rule is one "slot": `{ model, source, tier, token_quota, period, hours }`.
- Each rule is one "slot": `{ model, source, tier, token_quota, period, hours }`.
- `tier` is the priority tier: models in the same tier sit side by side and
share it; tiers run high → low.
- The same model may appear in several slots (e.g. A low → B low → A high →
@ -189,8 +310,9 @@ is no longer shown.
- `token_quota` > 0 means the slot is skipped once its tokens within the reset
window are exhausted; `period` supports `hour` / `week` / `month` / `nhour`
(with `hours`); empty = unlimited.
- Image models (`kind: image`) are kept out of the chain and are served by the
separate `POST /v1/images/generations` path.
- Within a tier, slots run by preference score; cooling / quota-exhausted /
hard-failed slots fall through to the next tier, and total failure returns
503 with a per-tier summary.
### Sources vs. adapters
@ -227,6 +349,82 @@ sources and `.lua` files; base sources defined in `config.yaml` cannot rewrite
that file, so they are hidden via a deletion tombstone (still hidden after
restart) and can be restored by re-adding the same name in the UI.
### Source templates (multi-key balancing)
When several upstream keys share one config (URL / adapter / model list /
concurrency / rpm), copying the whole source block N times is wasteful. A
**source template** stores every source field except `name` and `api_key`, so you
create N key-only-different sources from one shared skeleton — natural multi-key
load balancing (each expanded source schedules, cools down, and reports health
independently).
- WebUI Sources page, top-right **Templates** button: list all templates with
edit / delete / create.
- Add-source dialog has two header buttons:
- **From template** → pick a template → form auto-fills (name + key still
yours to enter);
- **As template** → save the current form's non-key/name fields as a template.
- Templates persist in the runtime file's `source_templates` field alongside
runtime sources (hot-managed, apply immediately). Templates hold no secret
keys, so they're safe to share / version-control.
A template is just a recipe — it doesn't become a source by itself. Only the
key-bearing sources created from it carry real traffic.
### Memory footprint (measured and tuned)
Memory is not a constant: it scales with the **number of configured sources**
(each source owns an `http.Transport` with its connection pool plus per-model
health state), not with uptime. Measured on this machine (Linux x86_64, 12 cores):
| Deployment shape | Startup RSS | Settled RSS |
|---|---|---|
| 1 source / 1 adapter (minimal) | ~4 MB | ~10 MB |
| 1 source + a 29 MB audit history | ~19 MB | ~20 MB |
| **16 sources / 13 adapters / 59 models (this host)** | ~28 MB | **~32-35 MB** |
> For reference, the same production config used **~105 MB** before this round of
> work. The reduction comes from three places: raw audit records are no longer kept
> resident (~25 MB — the aggregates still scan the full history, but the scan
> releases as it goes), Lua state pools no longer grow monotonically, and the two
> runtime knobs below.
Breakdown of the production instance (per-region, from `/proc/<pid>/smaps`):
| Region | RSS | Notes |
|---|---|---|
| Go heap | ~14 MB | provider/registry/scheduler state + connection-pool buffers |
| Other anonymous (thread stacks / LuaJIT chunks / runtime) | ~12 MB | scales with thread count and loaded adapters |
| Binary text+rodata | ~8 MB | mapped executable pages (read-only, reclaimable by the kernel) |
| Go runtime reservations | ~4 MB | the 2 GB+ you see in VSZ is address space, not physical memory |
| Shared libraries | ~3 MB | libc / libluajit / libm |
**Two recommended deployment knobs** (environment only, no code change):
```ini
# /etc/systemd/system/llmsproxy.service
Environment=GOGC=50
Environment=MALLOC_ARENA_MAX=2
```
- `MALLOC_ARENA_MAX=2`: LuaJIT allocates through cgo into glibc malloc, and glibc
allows up to `8 x nproc` per-thread arenas. Every OS thread that touches malloc
claims one (~1 MB each) and **never returns it to the OS**. Measured: 8-15
arenas (7-12 MB) down to **0**.
- `GOGC=50`: halves the Go heap growth target. **It does nothing on its own** —
the heap it saves is immediately consumed by additional glibc arenas (measured
20.3 -> 21.5 MB, i.e. slightly worse) — so it must be paired with
`MALLOC_ARENA_MAX`. Together they cut settled RSS by **~19%**. The gateway is
I/O bound (1min10s of CPU per 9 hours here), so the extra GC cycles are free.
> The desktop build (Electron) already injects both when spawning the embedded
> core (`cmd/gui/main.js`); exporting either variable yourself overrides it. If you
> write your own systemd unit or container spec, add them there.
Optional: `GOMEMLIMIT=48MiB` as a soft ceiling saves roughly 1 MB more, at the
cost of GC turning aggressive as the limit approaches. Not worth it for a relaxed
deployment.
### disable_thinking
With `"disable_thinking": true` in the request body, the gateway passes it to
@ -289,19 +487,12 @@ Shared helpers: `hmac_sha256_hex(key, data)`, `sha256_hex(data)`,
### Built-in adapters
`openai` `deepseek` `anthropic` `gemini` `github` `groq` `mistral` `ollama`
`kimicode` `opencode`.
`kimicode`.
`anthropic`/`gemini`/`ollama` include multimodal conversion
(`image_url` → their native format); with `disable_thinking` the `deepseek`
adapter sets `extra_body.thinking.type` to `disabled`.
**opencode** targets the opencode.ai zen free pool
(`https://opencode.ai/zen/v1`): zen fingerprints clients by User-Agent and
routes non-official UAs (curl, Go's default) into an anonymous pool that
hits `FreeUsageLimitError`. The adapter always sends the opencode client UA;
combined with `api_key: "public"` (the keyless official client actually sends
`Bearer public`) it gets the free pool, e.g. `deepseek-v4-flash-free`.
**kimicode** demonstrates `build_headers`: the cloud validates the calling
app, so you HMAC-sign timestamp+URL+body with `meta.app_secret` and add
`X-App-Sign`-style headers. Configure `sources[].meta.{app_id, app_secret,
@ -315,14 +506,25 @@ tray, autostart, silent launch, and a pixel-for-pixel embedded full WebUI
(Status/Chat/Keys/Priority/Sources/Adapters — no login needed). Server users
keep using the plain Go binary.
### Two installer kinds: Headless and Desktop
Release installers ship in two flavors for different audiences:
- **Headless (server)**: the plain core binary (no Electron), config-driven,
for servers / containers / systemd / unattended runs. Ships as `deb` + `rpm`
(systemd unit with memory tuning + example config) and `tar.gz` (binary +
`config.example.yaml` + `adapters/` + README), Windows as `.zip`.
- **Desktop (desktop app)**: the Electron GUI installer with an embedded core,
for personal daily desktop use.
### GUI vs. plain backend
| Scenario | Pick | Why |
| ---- | ---- | ---- |
| Server / intranet gateway / unattended long-running | **plain backend** (single binary) | ~10 MB, ~15 MB RSS, zero-dep single process — drop it into systemd or any container, remote admin |
| Personal desktop daily use / multi-device intranet sharing | **GUI** (Electron) | no-login embedded WebUI, tray one-click, autostart, silent background — for non-CLI users |
| Windows desktop | **GUI** | plain backend needs manual service registration; GUI ships native tray/autostart |
| CI one-shot 3-platform installers | **GUI packaging scripts** | `make gui-*` emits deb / AppImage / NSIS, drops straight into a release pipeline |
| Server / intranet gateway / unattended long-running | **Headless** (single binary) | ~8 MB binary; RSS scales with the number of sources (~10 MB for one, ~32-35 MB for 16 — see [Memory footprint](#memory-footprint-measured-and-tuned)), zero-dep single process — drop it into systemd or any container, remote admin |
| Personal desktop daily use / multi-device intranet sharing | **Desktop** (GUI) | no-login embedded WebUI, tray one-click, autostart, silent background — for non-CLI users |
| Windows desktop | **Desktop** | plain backend needs manual service registration; GUI ships native tray/autostart |
| CI one-shot 3-platform installers | **Desktop packaging scripts** | `make gui-*` emits deb / AppImage / NSIS, drops straight into a release pipeline |
Both share the exact same `llmsproxy` core (LuaJIT build) — configs and
adapters are fully compatible, freely interchangeable.
@ -334,30 +536,42 @@ cd cmd/gui
npm install
make gui # or npm run dev — dev run (needs make build first for bin/llmsproxy)
make gui-deb # local deb (share with other Linux users)
make gui-dist # deb + AppImage (linux)
make gui-dist # deb + rpm + AppImage (linux; rpm needs system rpmbuild)
make core-dist # core deb + rpm + tar.gz (server edition; needs nfpm)
make gui-win # Windows NSIS (host mingw + wine)
make gui-win-docker # Windows NSIS, fully dockerized (no host mingw needed)
make gui-win-docker # Windows NSIS, fully dockerized (wine in image; no host mingw/node/wine)
```
Artifacts land in `cmd/build/gui-dist/`: `ModelRouter-1.0.0.AppImage`,
`modelrouter-gui_1.0.0_amd64.deb`, `ModelRouter Setup 1.0.0.exe` (win). rpm
needs system `rpmbuild`.
Artifacts land in `cmd/build/gui-dist/`: `ModelRouter-<ver>.AppImage`,
`modelrouter-gui_<ver>_amd64.deb`, `ModelRouter Setup <ver>.exe` (win); core
packages land in `cmd/build/dist/`: `llmsproxy_<ver>_amd64.deb`,
`llmsproxy-<ver>.x86_64.rpm`. Both target families run
`packaging/verify-dist.sh` afterwards — a size floor that rejects degenerate
installers (1.3.0 shipped a 264 KB NSIS stub because the host wine was broken
and nothing caught it). At release time artifacts are renamed into the
`ModelRouter-Desktop-*` (GUI) and `ModelRouter-Headless-*` (plain backend)
families before upload to GitCode Releases.
#### Windows packaging (dockerized, recommended)
`make gui-win-docker` does cross-compile + NSIS in one step; the host needs
**no mingw**.
`make gui-win-docker` does cross-compile + NSIS **entirely in docker**; the host
needs **no mingw, no wine, no node**. Wine (with the i386 runtime) lives inside
the image, fixing the broken-host-wine failure that produced the 264 KB stub
installer.
- Image: `cmd/gui/docker/win-builder/Dockerfile` (`golang:1.25` + mingw-w64 +
prebuilt LuaJIT for Windows: `lua51.dll` + import lib). LuaJIT source comes
from the gitcode mirror (reachable in CN) with github fallback; Go modules go
through `https://goproxy.cn`.
- Core build: `cmd/gui/scripts/win-core-docker.sh` — emits
`cmd/gui/bin/{llmsproxy.exe,lua51.dll}` inside the container.
- One-shot release: `cmd/gui/scripts/dist-win-docker.sh` — after the core,
runs electron-builder (host needs `wine` for NSIS).
- The `modelrouter/win-builder` image is auto-built on first run and reused
afterwards.
node + wine32/wine64 + prebuilt LuaJIT for Windows: `lua51.dll` + import lib).
LuaJIT source comes from the gitcode mirror (reachable in CN) with github
fallback; Go modules go through `https://goproxy.cn`.
- One-shot release: `cmd/gui/scripts/dist-win-docker.sh` — core build + NSIS
packaging run inside the container (`USE_SYSTEM_WINE=true` uses the image's
wine); electron-builder's ~400 MB toolchain cache lives in a docker volume so
later builds reuse it.
- The `modelrouter/win-builder` image is auto-built on first run (~5-10 min)
and reused afterwards.
- After packaging, `packaging/verify-dist.sh` runs as a gate: an exe under
5 MB (or a deb/rpm/7z under 10 MB) fails the build — a broken installer can
never reach a Release again.
```bash
# Full Windows release (rebuild core + NSIS):
@ -390,6 +604,17 @@ internal/gateway # OpenAI-compatible HTTP + auth + SDK/streaming + WebUI
internal/types # unified format & OpenAI wire types
```
## Development & collaboration
- **Git branching workflow**: main / feature / release-branch conventions in
**[docs/git-workflow-en.md](docs/git-workflow-en.md)**
([中文](docs/git-workflow.md)). The line: `main` is always deployable →
features merge into `main` → cut `release/vX.Y.Z` and tag → release → hotfixes
land on the release branch and are cherry-picked back into `main`.
- **Lua adapter protocol**: see
**[docs/lua-adapters-en.md](docs/lua-adapters-en.md)**
([中文](docs/lua-adapters.md)).
## Tests
```bash

View File

@ -1,11 +1,20 @@
# ModelRouter GUI Windows builder
# Self-sufficient cross-compilation environment for the embedded llmsproxy core.
# Produces llmsproxy.exe (Go, windows/amd64, luajit tag) + lua51.dll (LuaJIT mingw build)
# inside a mounted workspace: cmd/gui/bin/{llmsproxy.exe,lua51.dll}
# Self-sufficient cross-compilation environment for the embedded llmsproxy core
# AND for packaging the Windows NSIS installer.
#
# Produces inside a mounted workspace:
# - llmsproxy.exe (Go, windows/amd64, luajit tag) + lua51.dll
# - ModelRouter Setup <ver>.exe (electron-builder --win nsis; requires wine
# so the uninstaller is generated — set USE_SYSTEM_WINE=true)
#
# Build: docker build -t modelrouter/win-builder cmd/gui/docker/win-builder
# Usage: make gui-win-docker (or cmd/gui/scripts/dist-win-docker.sh)
#
# This image is the SINGLE wine provider for the NSIS step, so a host that has
# no / broken wine (e.g. amd64-only Debian: no syswow64 -> c0000135) can still
# produce a real installer. The NSIS step MUST run inside this image; running it
# on the host was how the 1.3.0 264 KB degenerate exe got shipped.
#
# LuaJIT source is fetched from gitcode.com (CN-reachable mirror) with github
# fallback; pass --build-arg LUAJIT_REPO to override.
@ -13,11 +22,26 @@ FROM golang:1.25
ARG LUAJIT_REPO=https://gitcode.com/openresty/luajit2.git
RUN apt-get update -qq \
# i386 arch is required for wine32 (the NSIS uninstaller step runs a win32 exe
# through wine's wow64 path, which needs the 32-bit runtime under /usr/lib/i386).
RUN dpkg --add-architecture i386 \
&& apt-get update -qq \
&& DEBIAN_FRONTEND=noninteractive apt-get install -y -qq \
gcc-mingw-w64-x86-64 make ca-certificates git \
# NSIS step needs node + the electron-builder toolchain
nodejs npm \
# wine (32+64) for generating the NSIS uninstaller
wine64 wine32:i386 \
&& rm -rf /var/lib/apt/lists/*
# electron-builder resolves the uninstaller via the system wine when
# USE_SYSTEM_WINE=true (or by default on Linux without a toolset). Pin it.
# The toolchain cache lives in $HOME/.cache (mounted as a volume by
# dist-win-docker.sh), so it is reused across builds and never pollutes the
# mounted workspace.
ENV USE_SYSTEM_WINE=true
ENV WINEDEBUG=-all
# LuaJIT cross-compiled for Windows; keeps lua51.dll and the cgo import lib
# (libluajit-5.1.dll.a -> libluajit-5.1.a in the mingw lib dir).
# NOTE: $LUAJIT_REPO deliberately unquoted in "sh -c" — quoting the ARG would

View File

@ -10,6 +10,7 @@ const {
Tray,
shell,
nativeImage,
nativeTheme,
} = require("electron");
const path = require("path");
const fs = require("fs");
@ -46,6 +47,15 @@ const CORE_EXE = path.join(
"bin",
process.platform === "win32" ? "llmsproxy.exe" : "llmsproxy",
);
// Packaged window icon: must be a REAL file on disk (extraResource), not an
// asar path — Electron/GTK pulls BrowserWindow icons from the filesystem and
// crashes (SIGSEGV) when given an asar virtual path on Linux.
// Dev: cmd/gui/build/icon.png.
const WINDOW_ICON = path.join(
app.isPackaged ? process.resourcesPath : __dirname,
"build",
"icon.png",
);
// default port for the embedded core (configurable via settings)
const DEFAULT_PORT = 8787;
@ -53,6 +63,7 @@ let mainWindow = null;
let tray = null;
let coreProc = null;
let coreStartedAt = 0;
let coreReady = false; // true only after the embedded core actually answers HTTP
let settings = {}; // { port, autoStart, minimizeToTray }
let authRules = [];
@ -120,9 +131,9 @@ function coreStarted() {
function writeDefaultEmbeddedConfig() {
// Generate a self-contained embedded profile with a fresh random admin key.
// (Like Clash Verge generates its profile; the core's EnsureDefault default
// is 127.0.0.1:8080 — we want our own port + key on first run.) If a config
// already exists we do NOT overwrite it (keep user's edits).
// No upstream sources are preconfigured — the desktop user adds them via
// the WebUI's sources page (like Clash Verge's empty-by-default profile).
// If a config already exists we do NOT overwrite it (keep user's edits).
if (fs.existsSync(CONFIG_FILE)) return;
const key = "sk-gw-" + crypto.randomBytes(16).toString("hex");
const yaml = [
@ -132,16 +143,6 @@ function writeDefaultEmbeddedConfig() {
"adapter_dir: " + JSON.stringify(path.join(PROFILE_DIR, "adapters")),
"runtime_file: " + JSON.stringify(path.join(PROFILE_DIR, "runtime.json")),
"",
"sources:",
" - name: zen",
" base_url: https://opencode.ai/zen/v1",
" api_key: public",
" adapter: opencode",
" models:",
" - id: deepseek-v4-flash-free",
" priority: 100",
" kind: chat",
"",
// Admin key as a plain `keys` entry (no seed flag) so the web UI never
// shows the replace-the-initial-key warning. gateway_keys entries are
// marked seed:true by the core.
@ -197,6 +198,17 @@ function startCore() {
cwd: PROFILE_DIR,
env: Object.assign({}, process.env, {
LLMS_PROXY_GUI: "1",
// Memory tuning, measured on a 16-source deployment (see README 内存占用):
// MALLOC_ARENA_MAX caps glibc per-thread malloc arenas. LuaJIT allocates
// through cgo -> glibc malloc, and glibc defaults to 8*nproc arenas, so
// every OS thread touching malloc reserves its own ~1 MB arena that is
// never returned to the OS.
// GOGC halves the Go heap growth target. On its own it does not help (the
// saved heap is immediately eaten by extra glibc arenas); only the pair
// lowers settled RSS. The core is I/O bound, so the extra GC is free.
// Both are overridable: a user who exports them keeps their own values.
GOGC: process.env.GOGC || "50",
MALLOC_ARENA_MAX: process.env.MALLOC_ARENA_MAX || "2",
}),
stdio: ["ignore", "pipe", "pipe"],
});
@ -211,6 +223,7 @@ function startCore() {
coreProc.on("exit", (code) => {
const wasManaged = coreStartedAt > 0;
coreStartedAt = 0;
coreReady = false;
coreProc = null;
console.log("core exited:", code);
// auto-restart if still running GUI (crash recovery, max every ~5s)
@ -223,6 +236,12 @@ function startCore() {
}, 1500);
}
});
// Broadcast starting (NOT ready) immediately. The renderer must not
// attempt to load the iframe until the core actually answers HTTP —
// otherwise the frame load races the listener and fails permanently.
// (Fix for blank-screen: running=true was previously broadcast while the
// core was still booting, so the iframe ERR_CONNECTION_REFUSED and the
// renderer never retried.)
notifyCoreState();
// pump until the core answers, then inject the admin key and tell the
// renderer the web UI is ready (the iframe must load only after auth is set)
@ -230,6 +249,7 @@ function startCore() {
.then((ok) => {
if (ok) {
setEmbeddedAuth();
coreReady = true;
notifyCoreState();
}
})
@ -254,7 +274,7 @@ function waitForCoreReady(timeoutMs = 10000) {
let coreLastExit = 0;
function stopCore() {
app.isQuitting = true;
coreReady = false;
if (coreProc && coreProc.exitCode === null) {
try {
coreProc.kill();
@ -357,17 +377,27 @@ function createTray() {
);
}
tray = new Tray(img);
// Rebuild the context menu when the OS theme flips (dark/light) so menu
// text colors follow the desktop theme; nativeTheme is the Electron API
// that tracks GTK/appindicator dark mode on Linux.
const themeRebuild = () => {
try {
if (global.__rebuildTray) global.__rebuildTray();
} catch (e) {}
};
try {
nativeTheme.on("updated", themeRebuild);
} catch (e) {}
// Status rows used to be enabled:false, which GTK renders in a fixed light
// grey that never flips with the system dark/light theme — invisible on
// light desktops. They are now enabled (normal theme-aware fg) with a
// no-op click so they stay display-only but inherit the menu text color.
const statusRow = (text) => ({ label: text, enabled: true, click: () => {} });
const rebuild = () => {
const running = coreStarted();
const menu = Menu.buildFromTemplate([
{
label: "ModelRouter Desktop",
enabled: false,
},
{
label: running ? "● 内核运行中" : "○ 内核已停止",
enabled: false,
},
statusRow("ModelRouter Desktop"),
statusRow(running ? "● 内核运行中" : "○ 内核已停止"),
{ type: "separator" },
{
label: "显示主窗口",
@ -438,7 +468,8 @@ function createTray() {
function notifyCoreState() {
if (mainWindow && !mainWindow.isDestroyed()) {
mainWindow.webContents.send("core:state", {
running: coreStarted(),
running: coreStarted() && coreReady,
ready: coreReady,
startedAt: coreStartedAt,
});
}
@ -496,6 +527,12 @@ function createWindow() {
const menu = Menu.buildFromTemplate([]);
Menu.setApplicationMenu(menu);
mainWindow = new BrowserWindow({
// Window icon (taskbar/dock). Packaged: real file under resourcesPath (see
// WINDOW_ICON) — an asar path here makes GTK segfault on Linux.
icon: WINDOW_ICON,
// GNOME taskbar icons match the window's WM_CLASS (lowercase
// "modelrouter-gui") against StartupWMClass in the .desktop file.
backgroundColor: "#1f2937",
width: 1280,
height: 860,
minWidth: 900,
@ -560,7 +597,8 @@ ipcMain.handle(
);
ipcMain.handle("core:state", () => ({
running: coreStarted(),
running: coreStarted() && coreReady,
ready: coreReady,
startedAt: coreStartedAt,
baseUrl: embeddedBaseUrl(),
port: settings.port || DEFAULT_PORT,

View File

@ -1,12 +1,12 @@
{
"name": "modelrouter-gui",
"version": "1.0.0",
"version": "1.7.3",
"lockfileVersion": 3,
"requires": true,
"packages": {
"": {
"name": "modelrouter-gui",
"version": "1.0.0",
"version": "1.7.3",
"license": "Proprietary",
"devDependencies": {
"electron": "^33.0.0",
@ -582,7 +582,7 @@
}
},
"node_modules/@peculiar/webcrypto": {
"version": "1.7.1",
"version": "1.7.3",
"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.3",
"resolved": "https://registry.npmjs.org/resedit/-/resedit-1.7.2.tgz",
"integrity": "sha512-vHjcY2MlAITJhC0eRD/Vv8Vlgmu9Sd3LX9zZvtGzU5ZImdTN3+d6e/4mnTyV8vEbyf1sgNIrWxhWlrys52OkEA==",
"dev": true,
@ -4045,4 +4045,4 @@
}
}
}
}
}

View File

@ -1,6 +1,6 @@
{
"name": "modelrouter-gui",
"version": "1.0.0",
"version": "1.7.3",
"description": "ModelRouter Desktop — embedded ModelRouter core with tray",
"author": "ModelRouter",
"main": "main.js",
@ -10,7 +10,7 @@
"dev": "electron . --no-sandbox --dev",
"core:linux": "node scripts/prepare-core.js linux",
"core:win": "node scripts/prepare-core.js win",
"dist:linux": "npm run core:linux && electron-builder --linux deb AppImage",
"dist:linux": "npm run core:linux && electron-builder --linux deb rpm AppImage",
"dist:debian": "npm run core:linux && electron-builder --linux deb",
"dist:win": "npm run core:win && electron-builder --win nsis",
"dist:dir": "npm run core:linux && electron-builder --dir"
@ -32,10 +32,19 @@
"icon": "build/icon.png",
"maintainer": "ModelRouter",
"synopsis": "ModelRouter Desktop",
"desktop": {
"entry": {
"StartupWMClass": "modelrouter-gui"
}
},
"extraResources": [
{
"from": "bin/llmsproxy",
"to": "bin/llmsproxy"
},
{
"from": "build/icon.png",
"to": "build/icon.png"
}
]
},
@ -44,7 +53,13 @@
},
"mac": {
"target": "dmg",
"icon": "build/icon.png"
"icon": "build/icon.png",
"extraResources": [
{
"from": "build/icon.png",
"to": "build/icon.png"
}
]
},
"win": {
"target": "nsis",
@ -57,6 +72,10 @@
{
"from": "bin/lua51.dll",
"to": "bin/lua51.dll"
},
{
"from": "build/icon.png",
"to": "build/icon.png"
}
]
},

View File

@ -5,6 +5,7 @@
const state = {
core: {
running: false,
ready: false,
baseUrl: "",
port: 8787,
profileDir: "",
@ -18,21 +19,13 @@ const state = {
minimizeToTray: true,
},
theme: localStorage.getItem("mr-theme") || "light",
loadedOnce: false,
loadedOnce: false, // iframe has been pointed at the web UI at least once
frameLoaded: false, // last loadFrame actually finished loading (did-finish-load)
loadTries: 0,
};
const $ = (s) => document.querySelector(s);
function esc(s) {
return String(s == null ? "" : s).replace(
/[&<>"']/g,
(c) =>
({ "&": "&amp;", "<": "&lt;", ">": "&gt;", '"': "&quot;", "'": "&#39;" })[
c
],
);
}
function toast(msg, isErr, ms) {
const t = $("#toast");
t.textContent = msg;
@ -55,7 +48,11 @@ function renderCoreStatus() {
const core = state.core;
const dot = $("#conn-dot");
dot.className = "status-dot" + (core.running ? " ok" : " bad");
dot.title = core.running ? "内核运行中" : "内核未运行";
dot.title = !core.running
? "内核未运行"
: core.ready
? "内核运行中"
: "内核启动中…";
}
function showOffline() {
@ -70,11 +67,47 @@ function showLoading() {
$("#console-loading").style.display = "flex";
}
// probe the core's web root directly (fetch from the renderer). Using
// fetch() instead of relying on the iframe's own load means the core's
// readiness is confirmed by an actual HTTP round-trip before we point the
// iframe at it — this is what prevents the blank-screen race where the
// frame load happens while the listener is still coming up.
function coreReachable() {
return fetch(
state.core.baseUrl || "http://127.0.0.1:" + (state.core.port || 8787),
{ method: "GET", redirect: "follow", cache: "no-store" },
)
.then(() => true)
.catch(() => false);
}
// point the iframe at the WebUI, but only after the core is reachable.
// If the load fails (core still booting / raced), we retry with backoff
// instead of giving up forever (the old code set loadedOnce blindly and
// never retried -> permanent blank screen).
let frameRetryTimer = null;
function scheduleFrameLoad() {
if (frameRetryTimer) return; // already scheduled
frameRetryTimer = setTimeout(async () => {
frameRetryTimer = null;
if (!state.core.running) return; // core stopped meanwhile
if (!(await coreReachable())) {
// core not reachable yet: wait a bit and try again
state.loadTries += 1;
showLoading();
scheduleFrameLoad();
return;
}
loadFrame(state.core.baseUrl || "http://127.0.0.1:" + state.core.port);
}, state.loadTries === 0 ? 800 : Math.min(3000, 400 * state.loadTries));
}
function loadFrame(url) {
const f = frame();
f.style.display = "block";
$("#offline").style.display = "none";
$("#console-loading").style.display = "none";
if (state.loadedOnce && f.src === url) return; // already pointing there
f.src = url;
state.loadedOnce = true;
}
@ -82,8 +115,9 @@ function loadFrame(url) {
async function ensureCore() {
await refreshCore();
if (state.core.running) {
loadFrame(state.core.baseUrl);
scheduleFrameLoad();
} else {
state.loadedOnce = false;
showOffline();
}
}
@ -141,10 +175,21 @@ async function saveSettings() {
toast("设置已保存");
await refreshCore();
if (state.core.running) {
loadFrame(state.core.baseUrl);
scheduleFrameLoad();
} else {
state.loadedOnce = false;
showOffline();
}
}
// ===== theme =====
function applyTheme() {
document.documentElement.dataset.theme = state.theme;
try {
localStorage.setItem("mr-theme", state.theme);
} catch (e) {}
}
// ===== init =====
function init() {
$("#tb-min").onclick = () => window.modelrouter.win.minimize();
@ -171,15 +216,30 @@ function init() {
showLoading();
await window.modelrouter.core.start();
await refreshCore();
if (state.core.running) loadFrame(state.core.baseUrl);
if (state.core.running) scheduleFrameLoad();
else showOffline();
};
frame().onload = () => refreshCore();
// when the iframe finishes loading, mark loaded and refresh core status so
// the conn-dot reflects the actual WebUI (not just the core reachability)
frame().onload = () => {
state.frameLoaded = true;
refreshCore();
};
// load failure (e.g. the frame was pointed at the core just before it
// finished booting) -> clear loadedOnce and let scheduleFrameLoad retry
frame().onerror = () => {
state.loadedOnce = false;
state.frameLoaded = false;
scheduleFrameLoad();
};
window.modelrouter.core.onState((d) => {
state.core = Object.assign({}, state.core, d);
renderCoreStatus();
if (d.running && !state.loadedOnce) loadFrame(state.core.baseUrl);
if (!d.running) showOffline();
if (d.running && !state.loadedOnce) scheduleFrameLoad();
if (!d.running) {
state.loadedOnce = false;
showOffline();
}
});
// Escape closes overlay
document.addEventListener("keydown", (e) => {
@ -189,6 +249,7 @@ function init() {
}
async function boot() {
applyTheme();
init();
state.settings = await window.modelrouter.settings.get();
renderRail();

View File

@ -1,23 +1,64 @@
#!/usr/bin/env bash
# One-shot Windows release build for the GUI:
# One-shot Windows release build for the GUI — ENTIRELY inside docker.
#
# 1) cross-compile the embedded core in docker (llmsproxy.exe + lua51.dll)
# 2) host-side electron-builder NSIS packaging (wine required on host)
# 2) electron-builder NSIS packaging ALSO inside docker (wine for the
# uninstaller step lives in the image — see Dockerfile). This is the fix
# for the 1.3.0/1.4.0 broken installers: the host had amd64-only wine
# (no syswow64 -> c0000135), which made electron-builder silently emit a
# 264 KB degenerate exe without its 86 MB payload.
#
# Outputs: cmd/build/gui-dist/"ModelRouter Setup <version>.exe"
# cmd/build/gui-dist/modelrouter-gui-<version>-x64.nsis.7z
#
# Needs: docker (any machine; no host mingw/wine/node required)
set -euo pipefail
ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")/../../.." && pwd)"
cd "$ROOT/cmd/gui"
IMAGE="${WIN_BUILDER_IMAGE:-modelrouter/win-builder}"
PROXY="${GOPROXY:-https://goproxy.cn,direct}"
CACHE_VOL="modelrouter-eb-cache"
# 1) core (docker)
"$ROOT/cmd/gui/scripts/win-core-docker.sh"
command -v docker >/dev/null || { echo "docker required"; exit 1; }
# 2) packaging (host, needs wine for electron-builder's signtool/nsis)
command -v wine >/dev/null || { echo "wine required on host for NSIS packaging"; exit 1; }
if ! command -v npx >/dev/null; then echo "node/npm required"; exit 1; fi
if ! docker image inspect "$IMAGE" >/dev/null 2>&1; then
echo "[dist-win-docker] building image $IMAGE (first run, ~5-10 min)..."
docker build -t "$IMAGE" "$ROOT/cmd/gui/docker/win-builder"
fi
echo "[dist-win] electron-builder --win nsis ..."
# bin artifacts already refreshed by win-core-docker.sh; skip the prepare step
npx electron-builder --win nsis
docker volume create "$CACHE_VOL" >/dev/null 2>&1 || true
ls -la "$ROOT/cmd/build/gui-dist/"*.exe 2>/dev/null
echo "[dist-win] done"
# 1) embedded core + 2) NSIS packaging — both inside the image so wine and node
# are the image's, never the host's. electron-builder downloads electron + nsis
# toolchains into ~/.cache on first run; the mounted volume keeps that warm so
# subsequent builds are incremental instead of re-fetching ~400 MB.
echo "[dist-win-docker] cross-compiling core + electron-builder NSIS in docker ..."
docker run --rm \
-v "$ROOT":/workspace \
-v "$CACHE_VOL":/root/.cache \
-e GOPROXY="$PROXY" \
-w /workspace \
"$IMAGE" bash -c '
set -euo pipefail
cp /opt/luajit2/src/lua51.dll /workspace/cmd/gui/bin/lua51.dll
GOOS=windows GOARCH=amd64 CGO_ENABLED=1 CC=x86_64-w64-mingw32-gcc \
go build -buildvcs=false -tags luajit -trimpath \
-o /workspace/cmd/gui/bin/llmsproxy.exe ./cmd/llmsproxy
echo " -> core built"
# wine needs one prefix initialization so its wow64 runtime (syswow64/ntdll)
# exists — the exact thing that was missing on the broken host. Use a shared
# prefix in the cache volume so it is initialized only on the first run.
export WINEPREFIX=/root/.cache/wineprefix
if [ ! -f "$WINEPREFIX/drive_c/windows/syswow64/ntdll.dll" ]; then
echo " -> initializing wine prefix (first run)"
WINEDEBUG=-all wineboot -i || true
fi
cd /workspace/cmd/gui
# node_modules ships in the repo (host npm install already ran); reuse it.
npx electron-builder --win nsis
'
echo "[dist-win-docker] verifying artifacts ..."
"$ROOT/packaging/verify-dist.sh" "$ROOT/cmd/build/gui-dist"
echo "[dist-win-docker] done"

View File

@ -23,8 +23,39 @@ import (
func main() {
cfgPath := flag.String("config", "config.yaml", "path to gateway config file")
checkOnly := flag.Bool("check", false, "validate the config file and exit (0 = valid, 1 = invalid); nothing is started and no file is written")
showSecrets := flag.Bool("show-secrets", false, "print the config's credentials in the clear and exit; nothing is started and no file is written")
flag.Parse()
// -check is the deploy-time preflight: parse and validate the config
// without starting the Lua VM, touching runtime.json, or binding a port.
// It deliberately does NOT call EnsureDefault, so a missing file is an
// error here instead of being silently created.
if *checkOnly {
if _, err := os.Stat(*cfgPath); err != nil {
log.Fatalf("[llmsproxy] check: %v", err)
}
if _, err := config.Load(*cfgPath); err != nil {
log.Fatalf("[llmsproxy] check: %v", err)
}
log.Printf("[llmsproxy] check: %s is valid", *cfgPath)
return
}
// -show-secrets is the operator escape hatch for a config whose
// credentials are sealed at rest: it prints them to stdout and exits
// without starting anything or writing anything. Nothing else in the
// program prints a credential, and this path never logs to a file.
if *showSecrets {
if _, err := os.Stat(*cfgPath); err != nil {
log.Fatalf("[llmsproxy] show-secrets: %v", err)
}
if err := config.PrintSecrets(*cfgPath); err != nil {
log.Fatalf("[llmsproxy] show-secrets: %v", err)
}
return
}
created, err := config.EnsureDefault(*cfgPath)
if err != nil {
log.Fatalf("[llmsproxy] config: %v", err)
@ -88,4 +119,4 @@ func main() {
signal.Notify(stop, syscall.SIGINT, syscall.SIGTERM)
<-stop
log.Printf("[llmsproxy] shutting down")
}
}

479
deploy.sh Executable file
View File

@ -0,0 +1,479 @@
#!/usr/bin/env bash
# ============================================================
# llmsproxy 原子化滚动部署脚本
#
# 策略:
# 1. /tmp (tmpfs) 构建,/usr/local/bin (ext4) 部署
# 2. 构建后先 copy 到目标文件系统(.llmsproxy.tmp),
# 再 rename(2) 做原子替换(同一文件系统内 mv 是原子的)
# 3. runtime.json 不替换(保留运行时密钥、审计状态)
# 4. systemctl restart → 服务中断仅在 Go 进程重启瞬间 (<500ms)
# 5. 健康检查确认新进程正常
#
# 配置文件(2026-09-05 起):
# 本脚本不再假设 config.yaml 与部署无关。实测教训:给源加 headers 时
# 漏删原有的 `headers: {}` → YAML 重复键 → 进程启动即退出 → 健康检查
# 失败 → 旧的 rollback 只恢复二进制,配置错误仍在,于是旧二进制也起不来,
# 服务陷入 systemd 重启循环(当时累计 446 次)。
# 因此:
# - 支持 --config <file> 投放一份待部署配置(原子 rename 到 TARGET_CONFIG)
# - 重启前先用新二进制 `-check` 预检配置,不合法就地失败,绝不重启服务
# - rollback 同时恢复二进制与配置文件,回滚后再次 -check 才重启
#
# 注意事项:
# - 本服务依赖自身(推理依赖),restart 前后会短暂中断,
# 但重试机制(systemd RestartSec=5 + 客户端 retry)可自愈
# - 适配器文件替换后需重启才能生效(VM 启动时批量加载)
# ============================================================
set -euo pipefail
# ---------- 配置 ----------
readonly TARGET_BIN="/usr/local/bin/llmsproxy"
readonly TARGET_CONFIG="/etc/llmsproxy/config.yaml"
readonly TARGET_ADAPTERS="/etc/llmsproxy/adapters"
readonly REPO_DIR="/home/program/llmsproxy"
readonly HEALTH_URL="http://127.0.0.1:8081/v1/models"
readonly BIN_NAME="llmsproxy"
# 构建目录在 /tmp(tmpfs 更快)
readonly TMP_DIR="/tmp/llmsproxy-deploy"
# 同文件系统暂存文件(用于原子 rename)
readonly STAGING_BIN="${TARGET_BIN}.staging"
readonly STAGING_CONFIG="${TARGET_CONFIG}.staging"
# 回滚备份:健康检查失败时从这里恢复旧二进制 / 旧配置
readonly BACKUP_BIN="${TARGET_BIN}.bak"
readonly BACKUP_CONFIG="${TARGET_CONFIG}.rollback.bak"
# 待部署配置(--config <file> 指定;为空 = 不改动线上配置)
NEW_CONFIG_SRC=""
usage() {
cat <<'EOF'
用法: ./deploy.sh [--config <file>]
--config <file> 投放这份配置到 /etc/llmsproxy/config.yaml(原子替换)。
省略时不改动线上配置,只部署二进制与适配器。
无论是否指定,重启前都会用新二进制 -check 预检线上配置。
流程: 构建 → 适配器同步 → 备份(二进制+配置) → 原子替换 → 配置预检
→ 重启 → 健康检查 → 失败则回滚(二进制+配置)并重启
EOF
}
parse_args() {
while (( $# > 0 )); do
case "$1" in
--config)
[[ $# -ge 2 ]] || { echo "--config 需要一个文件参数" >&2; exit 2; }
NEW_CONFIG_SRC="$2"
shift 2
;;
-h|--help) usage; exit 0 ;;
*) echo "未知参数: $1" >&2; usage; exit 2 ;;
esac
done
}
# ---------- 日志 ----------
log() { printf '\033[1;34m[%s]\033[0m %s\n' "$(date +%H:%M:%S)" "$*"; }
warn() { printf '\033[1;33m[WARN]\033[0m %s\n' "$*"; }
err() { printf '\033[1;31m[ERR]\033[0m %s\n' "$*" >&2; }
fail() { err "$@"; exit 1; }
# ---------- 清理 ----------
cleanup() {
log "清理临时文件"
rm -rf "$TMP_DIR" "$STAGING_BIN" "$STAGING_CONFIG"
}
trap cleanup EXIT
# ---------- 前置检查 ----------
precheck() {
log "前置检查"
[[ -f "$REPO_DIR/go.mod" ]] || fail "仓库目录 $REPO_DIR 不存在"
command -v go &>/dev/null || fail "未找到 go"
command -v systemctl &>/dev/null || fail "未找到 systemctl"
[[ $(id -u) -eq 0 ]] || fail "需要 root 权限"
if [[ -n "$NEW_CONFIG_SRC" ]]; then
[[ -f "$NEW_CONFIG_SRC" ]] || fail "待部署配置不存在: $NEW_CONFIG_SRC"
# 同一文件比对:--config 指向线上文件本身时没有"投放"动作可做,
# 但仍然会走重启前的 -check 预检。
if [[ "$(readlink -f "$NEW_CONFIG_SRC")" == "$(readlink -f "$TARGET_CONFIG" 2>/dev/null)" ]]; then
warn "--config 指向线上配置本身,跳过投放(仍会预检)"
NEW_CONFIG_SRC=""
else
log "待部署配置: $NEW_CONFIG_SRC → $TARGET_CONFIG"
fi
fi
OLD_SIZE=$(stat -c%s "$TARGET_BIN" 2>/dev/null || echo 0)
OLD_SHA256=$(sha256sum "$TARGET_BIN" 2>/dev/null | awk '{print $1}')
log "旧二进制: $TARGET_BIN (${OLD_SIZE} bytes, sha256=${OLD_SHA256:0:16})"
log "go: $(go version)"
}
# ---------- 备份旧二进制 + 旧配置(用于回滚)----------
# 配置也必须备份:启动失败常常是配置问题,只回滚二进制会让旧二进制
# 继续吃坏配置,服务照样起不来。
backup_old_state() {
log "备份旧二进制到 $BACKUP_BIN"
if [[ -f "$TARGET_BIN" ]]; then
cp -f "$TARGET_BIN" "$BACKUP_BIN"
log "✓ 旧二进制已备份 (sha256=$(sha256sum "$BACKUP_BIN" | awk '{print $1}' | cut -c1-16))"
else
warn "旧二进制不存在,跳过备份(无回滚点)"
fi
if [[ -f "$TARGET_CONFIG" ]]; then
cp -f "$TARGET_CONFIG" "$BACKUP_CONFIG"
log "✓ 旧配置已备份到 $BACKUP_CONFIG (sha256=$(sha256sum "$BACKUP_CONFIG" | awk '{print $1}' | cut -c1-16))"
else
warn "线上配置不存在: $TARGET_CONFIG(首次部署?无配置回滚点)"
fi
}
# ---------- 构建 ----------
build() {
log "构建新二进制"
rm -rf "$TMP_DIR"
mkdir -p "$TMP_DIR/bin"
cd "$REPO_DIR"
CGO_ENABLED=1 go build -tags luajit \
-trimpath \
-ldflags="-s -w" \
-o "$TMP_DIR/bin/$BIN_NAME" \
./cmd/llmsproxy
NEW_BIN="$TMP_DIR/bin/$BIN_NAME"
NEW_SIZE=$(stat -c%s "$NEW_BIN")
if [[ "$NEW_SIZE" -eq "$OLD_SIZE" ]]; then
NEW_SHA256=$(sha256sum "$NEW_BIN" | awk '{print $1}')
if [[ "$NEW_SHA256" == "$OLD_SHA256" ]]; then
# 二进制无变化。只有在同时也没有待部署配置时才能直接退出;
# 否则会造成“--config 指定了新配置却被静默丢弃”。
if [[ -z "$NEW_CONFIG_SRC" ]]; then
warn "新二进制与旧版本完全一致(sha256 相同)且无待部署配置,跳过部署"
exit 0
fi
log "二进制无变化(sha256 相同),但有待部署配置,继续"
fi
fi
log "新二进制: $NEW_BIN (${NEW_SIZE} bytes)"
# 验证:确认包含关键修复
local checks=(
"mergeUsage"
"prompt_tokens"
"completion_tokens"
"message_start"
)
for c in "${checks[@]}"; do
if strings "$NEW_BIN" | grep -qF "$c"; then
log " ✓ 新二进制包含: $c"
else
warn " ⚠ 新二进制不包含: $c(非预期但非致命)"
fi
done
}
# ---------- 原子替换二进制 ----------
# 原理:
# /tmp 是 tmpfs,/usr/local/bin 是 ext4——不同文件系统 rename(2) 不原子。
# 因此先 copy 到目标文件系统上的 .staging 文件,再 rename(2) 到目标路径。
# rename(2) 在同一文件系统内是元数据级操作,对正在读写的进程安全:
# - 旧 inode 的 fd 继续有效(进程持有旧文件,不会读到截断内容)
# - 新进程打开 path 看到新 inode
# - 旧 inode 在最后一个 fd 关闭后释放
atomic_replace_binary() {
log "原子替换二进制"
local new_bin_tmp="$TMP_DIR/bin/$BIN_NAME"
[[ -f "$new_bin_tmp" ]] || fail "新二进制文件不存在: $new_bin_tmp"
# Step 1: copy 到目标文件系统(非原子,但目标文件唯一)
log " copy → $STAGING_BIN (同文件系统暂存)"
cp -f "$new_bin_tmp" "$STAGING_BIN"
chmod 0755 "$STAGING_BIN"
# Step 2: 原子 rename(同一文件系统,元数据级操作)
log " rename → $TARGET_BIN (原子)"
mv -f "$STAGING_BIN" "$TARGET_BIN"
# 验证
NEW_SHA256=$(sha256sum "$TARGET_BIN" | awk '{print $1}')
log "✓ 二进制已原子替换 (sha256=${NEW_SHA256:0:16})"
}
# ---------- 原子替换配置 ----------
# 与二进制同一套 copy→rename 手法:先落到 /etc 上的 .staging,再 rename(2)。
# 只有 --config 指定了来源时才动线上配置。
atomic_replace_config() {
[[ -n "$NEW_CONFIG_SRC" ]] || return 0
log "原子替换配置"
log " copy → $STAGING_CONFIG (同文件系统暂存)"
cp -f "$NEW_CONFIG_SRC" "$STAGING_CONFIG"
chmod 0644 "$STAGING_CONFIG"
log " rename → $TARGET_CONFIG (原子)"
mv -f "$STAGING_CONFIG" "$TARGET_CONFIG"
log "✓ 配置已原子替换 (sha256=$(sha256sum "$TARGET_CONFIG" | awk '{print $1}' | cut -c1-16))"
}
# ---------- 配置预检 ----------
# 用刚部署的新二进制解析线上配置。`-check` 只做 parse + ApplyDefaults,
# 不启动 Lua VM、不碰 runtime.json、不绑端口,因此可以安全地在服务仍在
# 运行时执行。不合法就在重启前失败——这正是上次事故缺的那一步。
verify_config() {
log "配置预检(新二进制 -check)"
[[ -f "$TARGET_CONFIG" ]] || fail "线上配置不存在: $TARGET_CONFIG"
local out
if out=$("$TARGET_BIN" -check -config "$TARGET_CONFIG" 2>&1); then
log "✓ 配置合法"
return 0
fi
err "配置预检失败,未重启服务(线上进程仍在跑旧配置):"
printf '%s\n' "$out" >&2
return 1
}
# ---------- 密钥可解封校验 ----------
# config.yaml 里的凭据以 enc:v1: 密文落盘,解密钥匙是 master.key(或
# LLMS_PROXY_MASTER_KEY)。一旦钥匙丢失/错位,配置在**启动时**才会报错,
# 那时服务已经处于「二进制已替换 + 配置已替换」的半死状态,只能靠回滚。
# 所以在替换任何东西之前先验:能不能用现有的钥匙解开现有的密文。
#
# 两种情形要分开:
# 1. 配置里**有**密文 ⇒ 钥匙必须能解开(否则拒绝部署)
# 2. 配置里**没有**密文 ⇒ 全新安装/首次加密,无需钥匙(但会顺手记一条,
# 因为它意味着这次部署会把明文写成密文,而钥匙必须同时就位)
verify_master_key() {
log "密钥可解封校验"
[[ -f "$TARGET_CONFIG" ]] || return 0 # 配置都不在,交给 verify_config 报错
local sealed
sealed=$(grep -c 'enc:v1:' "$TARGET_CONFIG" 2>/dev/null || echo 0)
if [[ "$sealed" -eq 0 ]]; then
warn "配置内无密文(尚未加密或全为明文)——首次加密需要 master.key 可写"
return 0
fi
# 钥匙来源与 config 包一致:env 优先,否则 runtime_file 同目录的 master.key
local runtime_file key_file
runtime_file=$(awk '/^runtime_file:/{print $2; exit}' "$TARGET_CONFIG")
[[ -n "$runtime_file" ]] || runtime_file="$(dirname "$TARGET_CONFIG")/runtime.json"
key_file="$(dirname "$runtime_file")/master.key"
if [[ -n "${LLMS_PROXY_MASTER_KEY:-}" ]]; then
log " 钥匙来源: LLMS_PROXY_MASTER_KEY(环境变量)"
# 走真实解密路径验证,而不是只检查格式
if out=$("$TARGET_BIN" -show-secrets -config "$TARGET_CONFIG" 2>&1); then
log "✓ 凭据可解封($(echo "$out" | grep -cE '^(source|key) ') 项)"
return 0
fi
err "凭据无法用 LLMS_PROXY_MASTER_KEY 解封,未替换任何文件:"
printf '%s\n' "$out" >&2
return 1
fi
if [[ ! -f "$key_file" ]]; then
err "配置含 $sealed 处密文,但钥匙文件不存在: $key_file"
err " 没有它,部署后服务将无法启动。恢复 master.key 或设置 LLMS_PROXY_MASTER_KEY 后重试。"
return 1
fi
log " 钥匙文件: $key_file ($(stat -c '%A' "$key_file" 2>/dev/null || echo '?'),$(wc -c < "$key_file") 字节)"
if out=$("$TARGET_BIN" -show-secrets -config "$TARGET_CONFIG" 2>&1); then
log "✓ 凭据可解封($(echo "$out" | grep -cE '^(source|key) ') 项)"
return 0
fi
err "凭据无法解密(钥匙与配置不匹配),未替换任何文件:"
printf '%s\n' "$out" | head -5 >&2
err " 密钥来源: $key_file"
err " 常见原因:master.key 被重新生成、恢复了不匹配的备份,或换了机器。"
return 1
}
# ---------- 同步适配器 ----------
sync_adapters() {
log "同步适配器"
# 只覆盖内置适配器(与 internal/lua/adapters 同名的文件),
# 保留目录里其它运行时上传的 .lua —— WebUI 上传的自定义适配器
# 必须跨部署存活,清空会静默移除线上源依赖的适配器。
SRC_ADAPTERS="$REPO_DIR/internal/lua/adapters"
[[ -d "$SRC_ADAPTERS" ]] || fail "内嵌适配器目录不存在: $SRC_ADAPTERS"
# 备份旧的适配器文件(含任何运行时覆盖版本,如 opencode.lua)
BACKUP_DIR="/etc/llmsproxy/adapters.bak.$(date +%Y%m%d%H%M%S)"
if [[ -d "$TARGET_ADAPTERS" ]]; then
mkdir -p "$BACKUP_DIR"
cp -rf "$TARGET_ADAPTERS/"*.lua "$BACKUP_DIR/" 2>/dev/null || true
log " 旧适配器已备份到 $BACKUP_DIR"
fi
cp -f "$SRC_ADAPTERS/"*.lua "$TARGET_ADAPTERS/"
chmod 0644 "$TARGET_ADAPTERS/"*.lua
# 校验:每个适配器必须包含 usage 透传修复(opencode/openai/deepseek 等应有 'uses')
local required_files=(openai deepseek anthropic gemini ollama opencode github groq kimicode mistral sensenova agentrouter)
for f in "${required_files[@]}"; do
[[ -f "$TARGET_ADAPTERS/$f.lua" ]] || warn " 缺少适配器: $f.lua"
done
local usage_fixed=0
for f in openai deepseek github groq kimicode mistral opencode; do
grep -qE '\buses\b' "$TARGET_ADAPTERS/$f.lua" 2>/dev/null && ((usage_fixed++)) || true
done
log " usage 透传修复适配器: ${usage_fixed}/7 (openai/deepseek/github/groq/kimicode/mistral/opencode)"
log "✓ 适配器已同步到 $TARGET_ADAPTERS/"
}
# ---------- 重启服务 ----------
restart_service() {
log "重启服务(中断 <500ms)"
systemctl restart llmsproxy.service || fail "systemd restart 失败"
log "✓ llmsproxy.service restarted"
}
# ---------- 健康检查 ----------
healthcheck() {
log "健康检查"
local attempts=20
local wait_sec=1
local attempt=1
while (( attempt <= attempts )); do
# 用 curl 检查(200=有 key 未鉴权但服务正常,401=需要认证,都说明服务在运行)
local code
code=$(curl -s -o /dev/null -w "%{http_code}" --max-time 3 "$HEALTH_URL" 2>/dev/null || echo "000")
if [[ "$code" == "200" || "$code" == "401" ]]; then
log "✓ 服务健康 (status=$code, ${attempt}/${attempts})"
return 0
fi
log " 等待中 ($attempt/$attempts, status=$code)"
sleep "$wait_sec"
(( attempt++ ))
done
warn "健康检查未通过,请检查服务状态: systemctl status llmsproxy.service"
return 1
}
# ---------- 回滚 ----------
# 健康检查失败时恢复旧二进制**和旧配置**并重启,保证推理服务可用。
# 这是最后一道保险:即使新版本启动失败(编译错误、配置不兼容、
# panic on start 等),也能让旧版本继续服务,避免推理能力丢失。
#
# 为什么必须一并回滚配置:启动失败最常见的原因就是配置本身不合法
# (实测:YAML 重复键)。只恢复二进制的话,旧二进制照样解析不了坏配置,
# 服务继续在 systemd 重启循环里空转。
rollback() {
err "健康检查失败,启动回滚流程"
local restored_any=0
if [[ -f "$BACKUP_BIN" ]]; then
log "恢复旧二进制"
cp -f "$BACKUP_BIN" "$TARGET_BIN"
chmod 0755 "$TARGET_BIN"
restored_any=1
else
err "无二进制回滚备份 ($BACKUP_BIN 不存在)"
fi
# 只在本次真的换过配置时才回滚配置,避免把用户在部署期间
# 通过 WebUI 做的合法改动一起抹掉。
if [[ -n "$NEW_CONFIG_SRC" && -f "$BACKUP_CONFIG" ]]; then
log "恢复旧配置"
cp -f "$BACKUP_CONFIG" "$TARGET_CONFIG"
chmod 0644 "$TARGET_CONFIG"
restored_any=1
fi
if (( restored_any == 0 )); then
err "无任何回滚点,无法自动恢复"
err "请手动检查: systemctl status llmsproxy.service && journalctl -u llmsproxy -n 50"
return 1
fi
# 回滚后先验证恢复出来的组合真的能解析,再重启——否则只是把
# 重启循环换个二进制继续跑。
if ! "$TARGET_BIN" -check -config "$TARGET_CONFIG" >/dev/null 2>&1; then
err "回滚后的配置仍不合法,拒绝重启(避免 systemd 重启风暴)"
err " 手动修复 $TARGET_CONFIG 后执行: $TARGET_BIN -check -config $TARGET_CONFIG"
err " 备份可用: 二进制=$BACKUP_BIN 配置=$BACKUP_CONFIG"
return 1
fi
log "重启服务(旧版本)"
systemctl restart llmsproxy.service || {
err "重启失败,服务可能完全宕机"
err "请手动检查: systemctl status llmsproxy.service"
return 1
}
# 回滚后再健康检查(只查 1 次,失败就认了)
sleep 2
local code
code=$(curl -s -o /dev/null -w "%{http_code}" --max-time 5 "$HEALTH_URL" 2>/dev/null || echo "000")
if [[ "$code" == "200" || "$code" == "401" ]]; then
log "✓ 回滚成功,旧版本已恢复服务 (status=$code)"
log " 新版本部署失败,请检查构建/配置后重试"
return 0
else
err "回滚后服务仍未恢复 (status=$code)"
err "服务可能已完全宕机,请手动介入"
return 1
fi
}
# ---------- 打印部署摘要 ----------
print_summary() {
log "═════════════════════════════════════════════"
log "部署完成"
log " 二进制: $TARGET_BIN"
log " 配置: $TARGET_CONFIG$([[ -n "$NEW_CONFIG_SRC" ]] && echo " (本次已替换)" || echo " (未改动)")"
log " 适配器: $TARGET_ADAPTERS"
log "═════════════════════════════════════════════"
}
# ---------- 主流程 ----------
# 顺序要点:配置投放与预检都排在 restart_service **之前**,
# 于是坏配置在服务仍然健康时就被拦下,根本不会触发重启。
main() {
parse_args "$@"
log "═══ llmsproxy 原子部署开始 ═══"
precheck
# 密钥校验放在 build/替换之前:它的失败后果最重(配置里的密文永久解不开),
# 必须在动任何文件之前就挡住。此阶段用的是**当前线上二进制**,它已经
# 支持 -show-secrets(首次引入该子命令的版本除外,此时跳过并提示)。
verify_master_key
build
sync_adapters
backup_old_state
atomic_replace_binary
atomic_replace_config
if ! verify_config; then
# 配置不合法:服务还没重启,只需把二进制/配置退回原状。
err "配置预检失败,回退本次改动(服务未受影响)"
[[ -f "$BACKUP_BIN" ]] && cp -f "$BACKUP_BIN" "$TARGET_BIN" && chmod 0755 "$TARGET_BIN"
if [[ -n "$NEW_CONFIG_SRC" && -f "$BACKUP_CONFIG" ]]; then
cp -f "$BACKUP_CONFIG" "$TARGET_CONFIG" && chmod 0644 "$TARGET_CONFIG"
fi
exit 1
fi
restart_service
if ! healthcheck; then
rollback || true
exit 1
fi
print_summary
}
main "$@"

141
docs/git-workflow-en.md Normal file
View File

@ -0,0 +1,141 @@
# Git Branching Workflow
ModelRouter uses a **GitHub Flow + release-branch** model: `main` is the only
long-lived branch and is always deployable; new work branches off it as a
feature branch; each version gets its own release branch with a tag; hotfixes
land on the release branch and MUST flow back to `main`.
Goal: make "patching an already-released version during its support window" a
first-class operation, while `main` always contains every fix and stays
deployable.
## Branch roles
| Branch | Name | Lifetime | Purpose |
|---|---|---|---|
| Main | `main` | permanent | Only long-lived branch. Always deployable. Accumulates the next version. |
| Feature | `feature/<desc>` | short (dev → merge → delete) | New features / ordinary fixes. Born from `main`, merged back into `main`. |
| Release | `release/vX.Y.Z` | one version cycle | Cut from `main`, tagged for release. Version-specific hotfixes land here. |
## Change flow (important)
```
feature/foo ───────┐
▼
main ──────────────── next version ──────►
│ │
│ cut │ cut
▼ ▼
release/v1.4.2 release/v1.4.3
│ │
tag: v1.4.2 tag: v1.4.3
│ │
hotfix ◄─────┘ hotfix ◄─────┘
│ │
└── cherry-pick back ──────────┘
```
### Key rules
1. **main is always deployable**: never leave half-done work on `main`.
Unfinished work lives on a feature branch.
2. **Features start from main and merge back**: `git checkout -b feature/xxx main`,
then `git merge --no-ff feature/xxx` (or squash) when done.
3. **Release = cut a release branch from main + tag**:
```bash
git checkout -b release/v1.4.2 main
git tag -a v1.4.2 -m "ModelRouter v1.4.2"
git push origin release/v1.4.2 v1.4.2
```
Build installers and upload the GitCode Release from this tag so the
published state is exactly reproducible.
4. **Only version-specific hotfixes go into a release branch during its
window**: new features always go to `main` for the next version, never into
an already-released branch (unless you deliberately ship a minor revision).
5. **Hotfixes MUST flow back to main**:
```bash
git checkout release/v1.4.2 # fix in the release branch
git commit -m "fix: ..."
git checkout main
git cherry-pick <hotfix-commit> # and into main
```
Most important rule: **main only moves forward and never loses fixes**. If a
hotfix never reaches main, the next release ships with the old bug.
### End of a version lifecycle
When the next version ships, the previous release branch retires:
- **Default: delete the remote release branch**
(`git push origin :release/v1.4.2`). All hotfixes were already
cherry-picked into main, so main contains everything; no merge needed.
- **Long-term maintenance** (e.g. an enterprise client pinned to an old
version): keep the branch, accept only security fixes, keep the
commit-then-cherry-pick loop.
## Explicit non-goals
- **Never rebase main**: main's history stays append-only; anyone pulling gets
directly usable history.
- **Release branches are not merged back wholesale**: hotfixes were already
cherry-picked; a whole-branch merge only creates conflicts with no benefit.
- **No "quick edit" directly on main**, even single-person: at minimum use
`feature/xxx → merge main` so history keeps its why-boundaries.
## Single vs multi person
- **Single** (this repo today): skip PRs — `git checkout -b feature/xxx` →
finish → merge back to main. Release branch + tag flow unchanged.
- **Multi / open source**: feature branches go through PR + Code Review,
maintainers merge; resolve conflicts on the feature branch, never on main.
## Commit message conventions
- `feat(...)`: new feature
- `fix(...)`: bug fix
- `chore(...)`: build / version / deps / non-code changes
- `docs(...)`: documentation
- `refactor(...)`: refactoring
- Scope in parens, e.g. `fix(adapters)`, `feat(webui)`, `chore(version)`
- One-line summary; expand with problem/root-cause/fix/verification when needed
## Release checklist (companion)
```bash
# 1. main is ready
git checkout main && git pull
# 2. cut release branch + tag
git checkout -b release/vX.Y.Z main
git tag -a vX.Y.Z -m "ModelRouter vX.Y.Z"
# 3. push branch + tag
git push origin release/vX.Y.Z
git push origin vX.Y.Z
# 4. build installers (core + GUI platforms)
make core-dist VERSION=X.Y.Z
# 5. create GitCode Release + upload
# https://gitcode.com/JianFeeeee/ModelRouter/releases
# upload MUST use PUT --http1.1, otherwise the OBS callback fails and the
# file never lands in the release; replacing artifacts = delete the tag
# (release disappears with it) then recreate.
```
## Why this model (background)
Before 2026-08 everything went straight to `main` plus a tag, which caused two
pain points:
1. After 1.4.2 shipped, a scheduler-scoring fix had to be added to the same
version — the only option was deleting the remote tag to clear the release
and re-uploading all installers. The published state was detached from code
history; "what exactly shipped" could not be traced.
2. No feature branches meant two independent efforts could not proceed in
parallel without colliding.
With release branches: the published state = `release/vX.Y.Z` branch +
`vX.Y.Z` tag, exactly reproducible; hotfixes have a clear landing spot; main
stays "latest + all fixes + deployable".
> 中文版见 [docs/git-workflow.md](git-workflow.md)。

128
docs/git-workflow.md Normal file
View File

@ -0,0 +1,128 @@
# Git 分支工作流(Branching Workflow)
ModelRouter 采用 **GitHub Flow + 发布分支** 模型:`main` 是唯一长命分支且永远可部署,
新工作从它开出特性分支;每个版本从 `main` 分出独立的发布分支并打 tag;hotfix 提交到
发布分支且必须回流 `main`。
目的:让「已发版的版本在生命周期内可被单独修补」成为一等操作,同时 `main` 始终包含
全部修复、永远可部署。
## 分支角色
| 分支 | 命名 | 生命周期 | 用途 |
|---|---|---|---|
| 主分支 | `main` | 永久 | 唯一长命分支。永远可部署。积攒下一个版本的功能。 |
| 特性分支 | `feature/<描述>` | 短命(开发→合并即删) | 新特性 / 一般 bug 修复。从 `main` 开出,完成后合回 `main`。 |
| 发布分支 | `release/vX.Y.Z` | 一个版本周期 | 从 `main` 分出,打 tag 发布。该版本生命周期内的 hotfix 都提交在此分支。 |
## 变更流向(重要)
```
feature/foo ───────┐
▼
main ──────────────── 下一个版本 ──────►
│ │
│ 切出 │ 切出
▼ ▼
release/v1.4.2 release/v1.4.3
│ │
tag: v1.4.2 tag: v1.4.3
│ │
hotfix ◄─────┘ hotfix ◄─────┘
│ │
└── cherry-pick 回 main ───────┘
```
### 关键规则
1. **main 永远可部署**:不在 main 上留半成品。任何未完成的工作必须在特性分支上。
2. **特性从 main 开,完成后合回 main**:`git checkout -b feature/xxx main`,
开发完 `git merge --no-ff feature/xxx` 或 squash 合回。
3. **发布 = 从 main 切 release 分支 + 打 tag**:
```bash
git checkout -b release/v1.4.2 main
git tag -a v1.4.2 -m "ModelRouter v1.4.2"
git push origin release/v1.4.2 v1.4.2
```
构建安装包、上传 GitCode Release 都基于这个 tag,保证可精确回溯发布态。
4. **版本生命周期内只收该版本的 hotfix**:新特性一律并入 `main` 等下一个版本,
绝不塞进已发布的 release 分支(除非主动选择在该版本内发次要版)。
5. **hotfix 必须回流 main**:
```bash
git checkout release/v1.4.2 # 在发布分支提交修复
git commit -m "fix: ..."
git checkout main
git cherry-pick <hotfix-commit> # 回主分支
```
这是本模型最重要的一条:**main 只前进、不丢修复**。如果 hotfix 不回 main,
下一个版本就会带着旧 bug 发布。
### 版本生命周期结束
下一个版本发布时,上一个 release 分支退役:
- **默认:直接删除远端 release 分支**(`git push origin :release/v1.4.2`)。
因为 hotfix 都已逐个 cherry-pick 回 main,main 已包含全部修复,无需再合并。
- **如需要长期维护旧版**(例如企业大客户卡在旧版本):保留分支,仅 stopship 接受
该版本的安全修复,继续走「提交 + cherry-pick 回 main」循环。
## 明确不做的事
- **不 rebase main**:`main` 的历史保持追加式,任何人拉取后 `git pull` 都得到直接可用的历史。
- **发布分支不整体 merge 回 main**:hotfix 已逐个 cherry-pick,整体合并只会制造冲突且无收益。
- **不在 main 直接写"临时改一下"**:即使是单人项目,也至少走
`feature/xxx → merge main` 的形式,让历史保留"为什么改"的边界。
## 单人 vs 多人
- **单人**(本项目当前):特性分支可省去 PR,`git checkout -b feature/xxx` →
完成 → 合并回 main。发布分支与 tag 流程完全一致。
- **多人 / 开源协作**:特性分支走 PR(pull request)+ Code Review,
由维护者 merge;冲突在特性分支上解决,不在 main 上解决。
## 提交信息规范(与分支配套)
- `feat(...)`: 新功能
- `fix(...)`: 修复
- `chore(...)`: 构建 / 版本号 / 依赖 / 文档无关改动
- `docs(...)`: 文档
- `refactor(...)`: 重构
- 括号内为影响域,如 `fix(adapters)`、`feat(webui)`、`chore(version)`
- 一句话总结,必要时正文展开「问题 / 根因 / 修复 / 验证」
## 发布操作速查(配套)
每次发布严格按此顺序:
```bash
# 1. 确认 main 就绪
git checkout main && git pull
# 2. 切发布分支并打 tag
git checkout -b release/vX.Y.Z main
git tag -a vX.Y.Z -m "ModelRouter vX.Y.Z"
# 3. 推送分支 + tag
git push origin release/vX.Y.Z
git push origin vX.Y.Z
# 4. 构建安装包(核心包 + GUI 各平台)
make core-dist VERSION=X.Y.Z
# 5. 创建 GitCode Release + 上传安装包
# 见 https://gitcode.com/JianFeeeee/ModelRouter/releases
# 上传注意:必须 PUT --http1.1,否则 OBS 回调失败文件不进列表;
# 替换旧版附件 = 删 tag(release 随 tag 一起消失)后重建。
```
## 为什么是这套(背景)
2026-08 之前本项目直接在所有改动提交到 `main` + 打 tag,出现两个痛点:
1. 1.4.2 已发版后又追了一个调度打分修复,只能删远端 tag 清空 release 再重传安装包
—— 版本发布态与代码历史脱节,无法精确回溯"当时发的是什么"。
2. 没有特性分支,无法并行开发互不干扰的两件事。
引入发布分支后:版本发布态 = `release/vX.Y.Z` 分支 + `vX.Y.Z` tag,可精确回溯;
hotfix 有明确落点;main 保持"最新 + 全部修复 + 可部署"。
> 英文版见 [docs/git-workflow-en.md](git-workflow-en.md)。

View File

@ -55,6 +55,7 @@ state may be shared across workers (`log` to stdout is the only side effect).
| `version` | string | no | Version, shown in the WebUI |
| `endpoint` | string | no | Upstream path, default `/chat/completions`; overridable by `source.endpoint` / `source.image_endpoint` |
| `headers` | table | no | Static default request headers; used as fallback when no `build_headers` hook is defined |
| `transform_error(status, body)` | function | no | Error-response condensing: return a one-line reason; on nil / absence clients uniformly receive `unknown error` (raw body goes to server logs only) |
These fields are extracted statically at load time (compile-once); reading them never
occupies a pooled worker.
@ -147,6 +148,27 @@ end
## 4. Optional Hooks
### `transform_error(status, body) -> string | nil`
Condense this source's error response into a one-line reason. Every upstream
formats errors differently — that is adapter territory: all built-in adapters
implement their own envelope parsing (zen's `{error={type,message}}`,
Anthropic's `{type="error",error={...}}`, Gemini's
`{error={code,message,status}}`, Ollama's string `{error="..."}`, etc.).
When the hook is absent or returns nil the core does not guess: clients get
`api error <status>: unknown error` and the raw body is logged server-side only.
```lua
function adapter.transform_error(status, body)
local ok, resp = pcall(json.decode, body)
if not ok or type(resp) ~= "table" then return nil end
if resp.error and resp.error.type == "FreeUsageLimitError" then
return "zen free pool quota exhausted"
end
return resp.error and resp.error.message or nil
end
```
### `build_headers(meta) -> table<string,string>`
Dynamically generate / sign request headers (e.g. KimiCode's HMAC signature). If the

View File

@ -50,6 +50,7 @@ return adapter
| `version` | string | 否 | 版本号,用于 WebUI 展示 |
| `endpoint` | string | 否 | 上游请求路径,默认 `/chat/completions`;可被 `source.endpoint` / `source.image_endpoint` 覆盖 |
| `headers` | table | 否 | 静态默认请求头;若未定义 `build_headers` 钩子则作为请求头回退 |
| `transform_error(status, body)` | function | 否 | 错误响应收敛:返回一行短原因;返回 nil / 未定义时客户端统一收到 `unknown error`(原始响应体只进服务端日志) |
这些字段在加载时静态提取(compile-once),之后读它们不会占用池内 worker。
@ -138,6 +139,26 @@ end
## 4. 可选钩子
### `transform_error(status, body) -> string | nil`
把该源特有的错误响应收敛成一行短原因。每个上游的错误格式不同——这是适配器层
的职责:内置各适配器均实现了自己的信封解析(如 zen 的
`{error={type,message}}`、Anthropic 的 `{type="error",error={...}}`、Gemini 的
`{error={code,message,status}}`、Ollama 的字符串 `{error="..."}` 等)。
未定义本钩子或返回 nil 时,核心不猜测格式,客户端统一收到
`api error <status>: unknown error`,原始响应体仅记录在服务端日志。
```lua
function adapter.transform_error(status, body)
local ok, resp = pcall(json.decode, body)
if not ok or type(resp) ~= "table" then return nil end
if resp.error and resp.error.type == "FreeUsageLimitError" then
return "zen free pool quota exhausted"
end
return resp.error and resp.error.message or nil
end
```
### `build_headers(meta) -> table<string,string>`
动态生成/签名请求头(如 KimiCode 的 HMAC 签名)。若脚本未定义此函数,Go 层回退使用静态

View File

@ -315,7 +315,7 @@ func TestEndToEndAuto503(t *testing.T) {
if resp == nil || resp.StatusCode != http.StatusServiceUnavailable {
t.Fatalf("status=%v body=%q", statusOf(resp), body)
}
if !strings.Contains(body, "all auto tiers failed") || !strings.Contains(body, "good/good-m") {
if !strings.Contains(body, "auto providers failed") || !strings.Contains(body, "good/good-m") {
t.Fatalf("503 must summarize the failed slot, body=%q", body)
}
}
@ -326,4 +326,3 @@ func statusOf(resp *http.Response) int {
}
return resp.StatusCode
}

View File

@ -20,13 +20,18 @@ type Config struct {
DefaultModel string `yaml:"default_model"` // e.g. "AUTO" or a model id
AdapterDir string `yaml:"adapter_dir"`
RuntimeFile string `yaml:"runtime_file"`
MaxConcurrent int `yaml:"max_concurrent"` // global inflight cap, 0 = unlimited
TLSCertFile string `yaml:"tls_cert_file,omitempty"` // PEM cert; when set together with tls_key_file, serve HTTPS
TLSKeyFile string `yaml:"tls_key_file,omitempty"` // PEM private key
PublicBaseURL string `yaml:"public_base_url,omitempty"` // external base for generated config snippets; default inferred from request
Sources []Source `yaml:"sources"`
Auto []ModelScope `yaml:"auto,omitempty"` // AUTO 调度链规则(WebUI 优先级页编辑)
Keys []GWKey `yaml:"keys,omitempty"` // 网关密钥(WebUI 密钥页管理)
Auto []ModelScope `yaml:"auto,omitempty"` // AUTO 调度链规则(WebUI 优先级页编辑,chat)
AutoImage []ModelScope `yaml:"auto_image,omitempty"` // AUTO 生图调度链规则(WebUI 优先级页·生图)
Keys []GWKey `yaml:"keys,omitempty"` // 网关密钥(WebUI 密钥页管理)
// box seals credentials (sources' api_key/headers, keys' key) at rest.
// In-memory values are always plaintext; only the bytes on disk are sealed.
// Wired by AttachSecretBox — Load leaves it nil so `-check` and tests stay
// filesystem-free.
box *SecretBox
}
// Defaults applied to any source (YAML or runtime) that leaves a field unset.
@ -47,21 +52,26 @@ type Model struct {
// Source describes a single upstream LLM provider.
type Source struct {
Name string `yaml:"name" json:"name"`
BaseURL string `yaml:"base_url" json:"base_url"`
APIKey string `yaml:"api_key" json:"api_key"`
APIKeyEnv string `yaml:"api_key_env,omitempty" json:"-"` // reference to an env var holding the key (overrides api_key)
Adapter string `yaml:"adapter" json:"adapter"`
Endpoint string `yaml:"endpoint" json:"endpoint,omitempty"` // chat endpoint override
ImageEndpoint string `yaml:"image_endpoint" json:"image_endpoint,omitempty"` // image endpoint override
Models []Model `yaml:"models" json:"models"`
Headers map[string]string `yaml:"headers" json:"headers,omitempty"`
Name string `yaml:"name" json:"name"`
BaseURL string `yaml:"base_url" json:"base_url"`
APIKey string `yaml:"api_key" json:"api_key"`
APIKeyEnv string `yaml:"api_key_env,omitempty" json:"-"` // reference to an env var holding the key (overrides api_key)
Adapter string `yaml:"adapter" json:"adapter"`
Endpoint string `yaml:"endpoint" json:"endpoint,omitempty"` // chat endpoint override
ImageEndpoint string `yaml:"image_endpoint" json:"image_endpoint,omitempty"` // image endpoint override
Models []Model `yaml:"models" json:"models"`
Headers map[string]string `yaml:"headers" json:"headers,omitempty"`
// ProxyURL routes this source's HTTP(S) traffic through an explicit proxy
// (e.g. http://127.0.0.1:7890). Empty = direct connection. Needed for
// upstreams behind geo/IP blocks that only respond through a proxy.
ProxyURL string `yaml:"proxy_url" json:"proxy_url,omitempty"`
Meta map[string]interface{} `yaml:"meta" json:"meta,omitempty"`
Temperature float64 `yaml:"temperature" json:"temperature,omitempty"`
MaxTokens int `yaml:"max_tokens" json:"max_tokens,omitempty"`
Timeout time.Duration `yaml:"timeout" json:"-"`
MaxConcurrent int `yaml:"max_concurrent" json:"max_concurrent"`
QueueTimeout time.Duration `yaml:"queue_timeout" json:"-"`
RPM int `yaml:"rpm" json:"rpm,omitempty"` // optional requests-per-minute cap (0 = unlimited)
}
// Load reads and validates a config file. When the file does not exist yet a
@ -103,9 +113,10 @@ func EnsureDefault(path string) (bool, error) {
}
// writeDefaultConfig writes a minimal, safe-by-default config: loopback-only
// listen, a fresh random admin key, and a no-key zen source that works out of
// the box. adapter_dir / runtime_file live next to the config file so the
// binary works regardless of the working directory it is started from.
// listen and a fresh random admin key. No upstream sources are preconfigured
// (the repo ships none — add them via WebUI or config.yaml). adapter_dir /
// runtime_file live next to the config file so the binary works regardless of
// the working directory it is started from.
func writeDefaultConfig(path string) error {
key, err := NewGatewayKey()
if err != nil {
@ -122,15 +133,6 @@ func writeDefaultConfig(path string) error {
DefaultModel: "AUTO",
AdapterDir: filepath.Join(abs, "adapters"),
RuntimeFile: filepath.Join(abs, "runtime.json"),
Sources: []Source{
{
Name: "zen",
BaseURL: "https://opencode.ai/zen/v1",
APIKey: "public", // zen 免费池:官方无 key 客户端实际发送 Bearer public
Adapter: "opencode",
Models: []Model{{ID: "deepseek-v4-flash-free", Priority: 100, Kind: "chat"}},
},
},
}
out, err := yaml.Marshal(&cfg)
if err != nil {
@ -195,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
@ -256,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)
}
@ -321,11 +334,34 @@ func (c *Config) ApplyDefaults() error {
// RuntimeConfig is the legacy runtime file format (kept for migration only).
type RuntimeConfig struct {
Sources []Source `json:"sources,omitempty"`
DeletedSources []string `json:"deleted_sources,omitempty"`
DeletedAdapters []string `json:"deleted_adapters,omitempty"`
Keys []GWKey `json:"keys,omitempty"`
Auto []ModelScope `json:"auto,omitempty"`
Sources []Source `json:"sources,omitempty"`
SourceTemplates []SourceTemplate `json:"source_templates,omitempty"`
DeletedSources []string `json:"deleted_sources,omitempty"`
DeletedAdapters []string `json:"deleted_adapters,omitempty"`
PresetTemplates []string `json:"preset_templates,omitempty"` // preset names the user has seen (or deleted) — never re-seeded
Keys []GWKey `json:"keys,omitempty"`
Auto []ModelScope `json:"auto,omitempty"`
AutoImage []ModelScope `json:"auto_image,omitempty"`
}
// SourceTemplate stores reusable source configuration (everything except
// name and api_key) so the WebUI can spin up multiple key-bearing sources
// from one shared template.
type SourceTemplate struct {
Name string `json:"name"`
BaseURL string `json:"base_url"`
Adapter string `json:"adapter"`
Endpoint string `json:"endpoint,omitempty"`
ImageEndpoint string `json:"image_endpoint,omitempty"`
Models []Model `json:"models"`
Headers map[string]string `json:"headers,omitempty"`
// ProxyURL routes this source's HTTP(S) traffic through an explicit proxy.
ProxyURL string `json:"proxy_url,omitempty"`
Meta map[string]interface{} `json:"meta,omitempty"`
Temperature float64 `json:"temperature,omitempty"`
MaxTokens int `json:"max_tokens,omitempty"`
MaxConcurrent int `json:"max_concurrent,omitempty"`
RPM int `json:"rpm,omitempty"`
}
// GWKey is a gateway API key persisted in the config file. Role is "admin"
@ -341,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"`
}
@ -368,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"`
}
@ -378,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

@ -91,14 +91,8 @@ func TestEnsureDefaultGeneratesOnMissingFile(t *testing.T) {
if cfg.Listen != "127.0.0.1:8080" {
t.Fatalf("listen = %q, want loopback-only", cfg.Listen)
}
if len(cfg.Sources) != 1 || cfg.Sources[0].Name != "zen" || cfg.Sources[0].Adapter != "opencode" {
t.Fatalf("default sources = %+v", cfg.Sources)
}
if cfg.Sources[0].APIKey != "public" {
t.Fatalf("zen api_key = %q", cfg.Sources[0].APIKey)
}
if cfg.Sources[0].Models[0].ID != "deepseek-v4-flash-free" {
t.Fatalf("default model = %+v", cfg.Sources[0].Models)
if len(cfg.Sources) != 0 {
t.Fatalf("default config must ship no sources (user adds them), got %+v", cfg.Sources)
}
// adapter_dir / runtime_file resolve next to the config file
if !strings.HasPrefix(cfg.AdapterDir, dir) || !strings.HasPrefix(cfg.RuntimeFile, dir) {
@ -324,3 +318,55 @@ sources:
t.Fatalf("after no-op: %v %v", len(cfg.Sources), err)
}
}
// TestSeedTemplatesOnce verifies preset templates are seeded exactly once:
// a user's edit survives a restart, and a deliberately deleted preset does
// not silently come back.
func TestSeedTemplatesOnce(t *testing.T) {
dir := t.TempDir()
path := filepath.Join(dir, "runtime.json")
presets := []SourceTemplate{
{Name: "DeepSeek", BaseURL: "https://api.deepseek.com", Adapter: "deepseek",
Models: []Model{{ID: "deepseek-chat", Kind: "chat"}}},
{Name: "OpenAI", BaseURL: "https://api.openai.com/v1", Adapter: "openai",
Models: []Model{{ID: "gpt-4o", Kind: "chat"}}},
}
// first run: both presets land in the store
s1 := NewStore(path)
if err := s1.Load(); err != nil {
t.Fatal(err)
}
if err := s1.SeedTemplates(presets); err != nil {
t.Fatalf("seed: %v", err)
}
if got := len(s1.ListTemplates()); got != 2 {
t.Fatalf("after seed len = %d, want 2", got)
}
// user edits one preset, deletes the other
edited := presets[0]
edited.BaseURL = "https://my-proxy.internal/v1"
if err := s1.UpsertTemplate(edited); err != nil {
t.Fatal(err)
}
if _, err := s1.RemoveTemplate("OpenAI"); err != nil {
t.Fatal(err)
}
// restart: re-seeding must be a no-op for both
s2 := NewStore(path)
if err := s2.Load(); err != nil {
t.Fatal(err)
}
if err := s2.SeedTemplates(presets); err != nil {
t.Fatalf("reseed: %v", err)
}
tpls := s2.ListTemplates()
if len(tpls) != 1 {
t.Fatalf("after restart len = %d, want 1 (deleted preset must not return): %+v", len(tpls), tpls)
}
if tpls[0].Name != "DeepSeek" || tpls[0].BaseURL != "https://my-proxy.internal/v1" {
t.Fatalf("user edit lost on restart: %+v", tpls[0])
}
}

212
internal/config/presets.go Normal file
View File

@ -0,0 +1,212 @@
package config
// Built-in preset source templates. Seeded into the runtime store on first
// startup so the WebUI can present "从模板创建" for popular providers — the
// user only needs to supply a name and API key.
//
// Model IDs verified against provider docs 2026-08-28. Providers ship new IDs
// frequently; treat these as a starting point and edit in the WebUI as needed.
var PresetTemplates = []SourceTemplate{
// ── 国内热门 ────────────────────────────────────────────────────────────
{
Name: "DeepSeek",
BaseURL: "https://api.deepseek.com",
Adapter: "deepseek",
MaxConcurrent: 8,
Models: []Model{
{ID: "deepseek-v4-flash", Priority: 100, Kind: "chat"},
{ID: "deepseek-v4-pro", Priority: 90, Kind: "chat"},
{ID: "deepseek-v4-flash-vision-exp", Priority: 60, Kind: "chat"},
},
},
{
Name: "智谱 GLM",
BaseURL: "https://open.bigmodel.cn/api/paas/v4",
Adapter: "openai",
MaxConcurrent: 8,
Models: []Model{
{ID: "glm-5.3", Priority: 100, Kind: "chat"},
{ID: "glm-5.3-flash", Priority: 90, Kind: "chat"},
{ID: "glm-5.2", Priority: 80, Kind: "chat"},
{ID: "glm-5.1", Priority: 60, Kind: "chat"},
},
},
{
Name: "Kimi / Moonshot",
BaseURL: "https://api.moonshot.cn/v1",
Adapter: "openai",
MaxConcurrent: 8,
Models: []Model{
{ID: "kimi-k3", Priority: 100, Kind: "chat"},
{ID: "kimi-k2.7-code", Priority: 90, Kind: "chat"},
{ID: "kimi-k2.7-code-highspeed", Priority: 85, Kind: "chat"},
{ID: "kimi-k2.6", Priority: 70, Kind: "chat"},
},
},
{
Name: "SiliconFlow",
BaseURL: "https://api.siliconflow.cn/v1",
Adapter: "openai",
MaxConcurrent: 8,
Models: []Model{
{ID: "deepseek-ai/DeepSeek-V4-Pro", Priority: 100, Kind: "chat"},
{ID: "deepseek-ai/DeepSeek-V4-Flash", Priority: 90, Kind: "chat"},
{ID: "moonshotai/Kimi-K2.6", Priority: 80, Kind: "chat"},
{ID: "zai-org/GLM-5.1", Priority: 70, Kind: "chat"},
{ID: "Kwai-Kolors/Kolors", Priority: 50, Kind: "image"},
},
},
{
Name: "讯飞星火 (Spark)",
BaseURL: "https://spark-api-open.xf-yun.com/v1",
Adapter: "openai",
MaxConcurrent: 4,
Models: []Model{
{ID: "4.0Ultra", Priority: 100, Kind: "chat"},
{ID: "generalv3.5", Priority: 80, Kind: "chat"},
{ID: "max-32k", Priority: 70, Kind: "chat"},
},
},
{
Name: "阿里云 通义千问",
BaseURL: "https://dashscope.aliyuncs.com/compatible-mode/v1",
Adapter: "openai",
MaxConcurrent: 8,
Models: []Model{
{ID: "qwen3.8-max", Priority: 100, Kind: "chat"},
{ID: "qwen3.8-flash", Priority: 90, Kind: "chat"},
{ID: "qwen3.7-plus", Priority: 85, Kind: "chat"},
{ID: "deepseek-v4-pro", Priority: 70, Kind: "chat"},
{ID: "kimi-k3", Priority: 70, Kind: "chat"},
},
},
{
Name: "百度 千帆",
BaseURL: "https://qianfan.baidubce.com/v2",
Adapter: "openai",
MaxConcurrent: 8,
Models: []Model{
{ID: "ernie-5.0-turbo", Priority: 100, Kind: "chat"},
{ID: "ernie-5.0", Priority: 90, Kind: "chat"},
{ID: "deepseek-v4-pro", Priority: 70, Kind: "chat"},
},
},
{
Name: "豆包 (火山方舟)",
BaseURL: "https://ark.cn-beijing.volces.com/api/v3",
Adapter: "openai",
MaxConcurrent: 8,
Models: []Model{
{ID: "doubao-seed-2.0-pro", Priority: 100, Kind: "chat"},
{ID: "doubao-seed-2.0", Priority: 90, Kind: "chat"},
{ID: "doubao-seed-2.0-flash", Priority: 80, Kind: "chat"},
},
},
{
Name: "MiniMax",
BaseURL: "https://api.minimax.chat/v1",
Adapter: "openai",
MaxConcurrent: 8,
Models: []Model{
{ID: "MiniMax-M3", Priority: 100, Kind: "chat"},
{ID: "MiniMax-M2.5", Priority: 80, Kind: "chat"},
},
},
// ── 国际主流 ────────────────────────────────────────────────────────────
{
Name: "OpenAI",
BaseURL: "https://api.openai.com/v1",
Adapter: "openai",
MaxConcurrent: 8,
Models: []Model{
{ID: "gpt-5.6-sol", Priority: 100, Kind: "chat"},
{ID: "gpt-5.6-terra", Priority: 90, Kind: "chat"},
{ID: "gpt-5.6-luna", Priority: 80, Kind: "chat"},
},
},
{
Name: "Anthropic Claude",
BaseURL: "https://api.anthropic.com",
Adapter: "anthropic",
MaxConcurrent: 8,
Models: []Model{
{ID: "claude-opus-5", Priority: 100, Kind: "chat"},
{ID: "claude-fable-5", Priority: 95, Kind: "chat"},
{ID: "claude-sonnet-5", Priority: 85, Kind: "chat"},
{ID: "claude-haiku-4-5", Priority: 70, Kind: "chat"},
},
},
{
Name: "Google Gemini",
BaseURL: "https://generativelanguage.googleapis.com/v1beta",
Adapter: "gemini",
MaxConcurrent: 8,
Models: []Model{
{ID: "gemini-3.6-flash", Priority: 100, Kind: "chat"},
{ID: "gemini-3.5-flash", Priority: 90, Kind: "chat"},
{ID: "gemini-3.5-flash-lite", Priority: 80, Kind: "chat"},
{ID: "gemini-3.1-pro", Priority: 85, Kind: "chat"},
},
},
{
Name: "Groq",
BaseURL: "https://api.groq.com/openai/v1",
Adapter: "groq",
MaxConcurrent: 8,
Models: []Model{
{ID: "openai/gpt-oss-120b", Priority: 100, Kind: "chat"},
{ID: "openai/gpt-oss-20b", Priority: 80, Kind: "chat"},
{ID: "groq/compound", Priority: 70, Kind: "chat"},
},
},
{
Name: "Mistral",
BaseURL: "https://api.mistral.ai/v1",
Adapter: "mistral",
MaxConcurrent: 8,
Models: []Model{
{ID: "mistral-medium-latest", Priority: 100, Kind: "chat"},
{ID: "mistral-large-latest", Priority: 90, Kind: "chat"},
{ID: "codestral-latest", Priority: 85, Kind: "chat"},
},
},
{
Name: "OpenRouter",
BaseURL: "https://openrouter.ai/api/v1",
Adapter: "openai",
MaxConcurrent: 8,
Headers: map[string]string{
"HTTP-Referer": "http://localhost:8080",
"X-Title": "ModelRouter",
},
Models: []Model{
{ID: "openrouter/auto", Priority: 100, Kind: "chat"},
{ID: "deepseek/deepseek-v4-flash", Priority: 90, Kind: "chat"},
{ID: "anthropic/claude-opus-5", Priority: 85, Kind: "chat"},
{ID: "openai/gpt-5.6-sol", Priority: 85, Kind: "chat"},
},
},
{
Name: "GitHub Models",
BaseURL: "https://models.inference.ai.azure.com",
Adapter: "github",
MaxConcurrent: 8,
Models: []Model{
{ID: "gpt-5.6-sol", Priority: 100, Kind: "chat"},
{ID: "gpt-5.6-luna", Priority: 80, Kind: "chat"},
},
},
// ── 本地模型 ────────────────────────────────────────────────────────────
{
Name: "Ollama (localhost)",
BaseURL: "http://127.0.0.1:11434",
Adapter: "ollama",
Endpoint: "/api/chat",
MaxConcurrent: 4,
Models: []Model{
{ID: "qwen3", Priority: 80, Kind: "chat"},
{ID: "deepseek-r1", Priority: 80, Kind: "chat"},
{ID: "llama3.3", Priority: 70, Kind: "chat"},
},
},
}

View File

@ -128,15 +128,3 @@ func (b *SecretBox) MustDecrypt(v string) string {
}
return out
}
// resolvedAPIKey resolves a source key from api_key_env (env var) or api_key.
func resolvedAPIKey(src Source, box *SecretBox) (string, error) {
if src.APIKeyEnv != "" {
v := os.Getenv(src.APIKeyEnv)
if v == "" {
return "", fmt.Errorf("source %s: env %s is empty or unset", src.Name, src.APIKeyEnv)
}
return v, nil
}
return box.MustDecrypt(src.APIKey), nil
}

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

@ -138,6 +138,81 @@ func (s *Store) DeletedAdapters() map[string]bool {
return out
}
// SeedTemplates adds preset templates that the store has never seen before.
// A preset is seeded at most once: its name is recorded in PresetTemplates so
// a user who edits or deletes it never gets it silently restored on restart.
func (s *Store) SeedTemplates(presets []SourceTemplate) error {
s.mu.Lock()
defer s.mu.Unlock()
seen := make(map[string]bool, len(s.data.PresetTemplates))
for _, n := range s.data.PresetTemplates {
seen[n] = true
}
existing := make(map[string]bool, len(s.data.SourceTemplates))
for _, t := range s.data.SourceTemplates {
existing[t.Name] = true
}
added := false
for _, p := range presets {
if p.Name == "" || seen[p.Name] {
continue
}
s.data.PresetTemplates = append(s.data.PresetTemplates, p.Name)
added = true
if existing[p.Name] {
continue // user already has a template by this name: never overwrite
}
s.data.SourceTemplates = append(s.data.SourceTemplates, p)
}
if !added {
return nil
}
return s.persistLocked()
}
// UpsertTemplate adds or replaces a source template and persists.
func (s *Store) UpsertTemplate(t SourceTemplate) error {
s.mu.Lock()
defer s.mu.Unlock()
for i := range s.data.SourceTemplates {
if s.data.SourceTemplates[i].Name == t.Name {
s.data.SourceTemplates[i] = t
return s.persistLocked()
}
}
s.data.SourceTemplates = append(s.data.SourceTemplates, t)
return s.persistLocked()
}
// ListTemplates returns the saved source templates.
func (s *Store) ListTemplates() []SourceTemplate {
s.mu.Lock()
defer s.mu.Unlock()
out := make([]SourceTemplate, len(s.data.SourceTemplates))
copy(out, s.data.SourceTemplates)
return out
}
// RemoveTemplate deletes a source template by name and persists.
func (s *Store) RemoveTemplate(name string) (bool, error) {
s.mu.Lock()
defer s.mu.Unlock()
kept := s.data.SourceTemplates[:0]
removed := false
for _, t := range s.data.SourceTemplates {
if t.Name == name {
removed = true
continue
}
kept = append(kept, t)
}
s.data.SourceTemplates = kept
if removed {
return true, s.persistLocked()
}
return false, nil
}
// LoadLegacy returns the full legacy RuntimeConfig from the file (used for
// one-time migration into config.yaml). Returns nil if file is missing.
func (s *Store) LoadLegacy() *RuntimeConfig {
@ -186,15 +261,6 @@ func (s *Store) encryptLocked() {
}
}
func containsString(list []string, s string) bool {
for _, x := range list {
if x == s {
return true
}
}
return false
}
func removeString(list []string, s string) []string {
out := list[:0]
for _, x := range list {

View File

@ -7,6 +7,7 @@ import (
"crypto/rand"
"encoding/hex"
"fmt"
"log"
"os"
"path/filepath"
"sort"
@ -26,13 +27,14 @@ import (
// request paths look keys up concurrently; the AUTO chain itself is swapped
// atomically and needs no lock.
type Core struct {
mu sync.Mutex
cfg *config.Config
vm *lua.VM
store *config.Store
scheduler *scheduler.Scheduler
registry *provider.Registry
autoChain atomic.Pointer[scheduler.Chain]
mu sync.Mutex
cfg *config.Config
vm *lua.VM
store *config.Store
scheduler *scheduler.Scheduler
registry *provider.Registry
autoChain atomic.Pointer[scheduler.Chain] // chat AUTO chain
autoImageChain atomic.Pointer[scheduler.Chain] // image-generation AUTO chain
}
// New builds the core from a config file plus runtime overlay.
@ -52,6 +54,28 @@ func NewFromConfig(cfg *config.Config) (*Core, error) {
return nil, fmt.Errorf("lua vm: %w", err)
}
c.store = config.NewStore(cfg.RuntimeFile)
// Share one box between the runtime store and config.yaml so a single
// master.key seals both files. config.Load left the config holding
// ciphertext (if it was sealed); unseal it NOW, before anything reads a
// credential.
//
// ORDER IS LOAD-BEARING. These steps each consume secrets and must run
// after the unseal:
// - seedKeys compares cfg.Keys[i].Key against the plaintext
// gateway_keys entries; with ciphertext keys the comparison never
// matched and every restart appended another duplicate admin key
// (production had four copies of the same admin key).
// - rebuildRegistry hands cfg.Sources[i].APIKey to the providers; with
// ciphertext it built every upstream client with "enc:v1:..." as its
// bearer token.
// Previously the unseal happened at the END of this function, and things
// only appeared to work because seedKeys' Save() unsealed memory as a side
// effect. Removing the redundant saves exposed the real ordering bug.
cfg.AttachSecretBox(c.store.SecretBox())
hadPlaintextSecrets, err := cfg.UnsealSecrets(c.store.SecretBox())
if err != nil {
return nil, fmt.Errorf("unseal secrets: %w", err)
}
if err := c.store.Load(); err != nil {
return nil, fmt.Errorf("runtime store: %w", err)
}
@ -68,6 +92,17 @@ func NewFromConfig(cfg *config.Config) (*Core, error) {
if err := c.rebuildRegistry(); err != nil {
return nil, err
}
// Seed built-in preset templates on first run so the WebUI shows
// "从模板创建" for popular providers out of the box.
if err := c.seedPresetTemplates(); err != nil {
return nil, fmt.Errorf("seed preset templates: %w", err)
}
// Seal any credential still in the clear in config.yaml. Startup writes
// nothing when the file was already sealed (hadPlaintextSecrets is decided
// from the on-disk state, before the unseal above).
if err := cfg.SealIfNeeded(hadPlaintextSecrets); err != nil {
return nil, fmt.Errorf("normalize secrets: %w", err)
}
return c, nil
}
@ -161,6 +196,18 @@ func (c *Core) seedKeys() error {
return nil
}
// seedPresetTemplates adds the built-in provider presets to the template
// store on first run. It only adds templates whose name is not already
// present, so a user's edits to a preset survive restarts and a deliberately
// deleted preset does not silently return on the next start (its name is
// remembered in the store's preset marker list).
func (c *Core) seedPresetTemplates() error {
if len(config.PresetTemplates) == 0 {
return nil
}
return c.store.SeedTemplates(config.PresetTemplates)
}
// saveConfig writes the current config (including auto rules and keys) back
// to config.yaml.
func (c *Core) saveConfig() error {
@ -181,6 +228,9 @@ func (c *Core) Registry() *provider.Registry { return c.registry }
// AutoChain returns the current AUTO scheduling chain.
func (c *Core) AutoChain() *scheduler.Chain { return c.autoChain.Load() }
// AutoImageChain returns the persisted image-generation AUTO chain snapshot.
func (c *Core) AutoImageChain() *scheduler.Chain { return c.autoImageChain.Load() }
func (c *Core) DefaultModel() string { return c.cfg.DefaultModel }
func (c *Core) GatewayKeys() []string { return c.cfg.GatewayKeys }
@ -217,10 +267,15 @@ func (c *Core) FindKey(key string) (config.GWKey, bool) {
}
// CreateKey builds a new random gateway key and persists it to config.yaml.
// Quotas live on the model scope entries, so a new key's budget is whatever
// its scopes carry.
func (c *Core) CreateKey(name, role string, models []config.ModelScope, note string) (config.GWKey, error) {
c.mu.Lock()
defer c.mu.Unlock()
models = cleanScopes(models)
if err := ValidateScopeQuotas(models); err != nil {
return config.GWKey{}, err
}
key := make([]byte, 16)
if _, err := rand.Read(key); err != nil {
return config.GWKey{}, err
@ -233,9 +288,7 @@ func (c *Core) CreateKey(name, role string, models []config.ModelScope, note str
Note: note,
CreatedAt: time.Now().Unix(),
}
if rec.Role == "" {
rec.Role = "user"
}
rec.Role = config.NormalizeRole(rec.Role)
c.cfg.Keys = append(c.cfg.Keys, rec)
if err := c.saveConfig(); err != nil {
return config.GWKey{}, err
@ -243,10 +296,16 @@ func (c *Core) CreateKey(name, role string, models []config.ModelScope, note str
return rec, nil
}
// UpdateKey mutates a key's name/role/model scope and persists it.
// UpdateKey mutates a key's name/role/model scope and persists it. The scope
// entries carry their own quotas, so replacing the scope replaces the budgets.
func (c *Core) UpdateKey(key, name, role string, models []config.ModelScope, note string) (config.GWKey, error) {
c.mu.Lock()
defer c.mu.Unlock()
if models != nil {
if err := ValidateScopeQuotas(models); err != nil {
return config.GWKey{}, err
}
}
for i, k := range c.cfg.Keys {
if k.Key == key {
if name != "" {
@ -294,6 +353,15 @@ func (c *Core) AutoRules() []config.ModelScope {
return out
}
// AutoImageRules returns the configured image-generation AUTO chain rules.
func (c *Core) AutoImageRules() []config.ModelScope {
c.mu.Lock()
defer c.mu.Unlock()
out := make([]config.ModelScope, len(c.cfg.AutoImage))
copy(out, c.cfg.AutoImage)
return out
}
// cleanScopes drops empty model entries and normalizes placeholder source
// names; it returns nil when no entries survive so an empty scope means
// "unrestricted" (nil) instead of a restrictive-but-empty list — a non-nil
@ -325,16 +393,35 @@ func (c *Core) SaveAutoRules(entries []config.ModelScope) error {
return err
}
c.buildAutoChain()
if ch := c.autoChain.Load(); ch != nil {
for _, tn := range ch.Tiers {
for _, sl := range tn.Slots {
if p := c.registry.ProviderForSlot(sl.Model, sl.Source); p != nil {
p.ResetModelCooldown(sl.Model)
}
c.resetChainCooldowns(c.autoChain.Load())
return nil
}
// SaveAutoImageRules persists the image-generation AUTO chain. Image models
// are kept as-is (unlike buildAutoChain, which skips them).
func (c *Core) SaveAutoImageRules(entries []config.ModelScope) error {
c.mu.Lock()
defer c.mu.Unlock()
c.cfg.AutoImage = cleanScopes(entries)
if err := c.saveConfig(); err != nil {
return err
}
c.buildAutoImageChain()
c.resetChainCooldowns(c.autoImageChain.Load())
return nil
}
func (c *Core) resetChainCooldowns(ch *scheduler.Chain) {
if ch == nil {
return
}
for _, tn := range ch.Tiers {
for _, sl := range tn.Slots {
if p := c.registry.ProviderForSlot(sl.Model, sl.Source); p != nil {
p.ResetModelCooldown(sl.Model)
}
}
}
return nil
}
// ResetHealth clears the scheduling backoff state of every provider.
@ -390,6 +477,11 @@ func (c *Core) mergedSources() []config.Source {
}
// resolveSourceKey applies api_key_env and decrypts enc:v1: ciphertext.
// Config values are already unsealed at startup (Config.NormalizeSecretsForRun),
// so the decrypt branch is the belt-and-braces path for a source that arrived
// already sealed through another route. A failure there leaves the ciphertext in
// place, which makes the upstream reject the key loudly rather than sending an
// empty Authorization header that might look like a config-less source.
func (c *Core) resolveSourceKey(s config.Source) config.Source {
if s.APIKeyEnv != "" {
if v := os.Getenv(s.APIKeyEnv); v != "" {
@ -398,7 +490,11 @@ func (c *Core) resolveSourceKey(s config.Source) config.Source {
return s
}
if box := c.store.SecretBox(); box != nil && strings.HasPrefix(s.APIKey, "enc:v1:") {
s.APIKey = box.MustDecrypt(s.APIKey)
if v, err := box.Decrypt(s.APIKey); err == nil {
s.APIKey = v
} else {
log.Printf("[core] source %s: api_key decrypt failed: %v", s.Name, err)
}
}
return s
}
@ -420,6 +516,7 @@ func (c *Core) rebuildRegistry() error {
c.registry.Replace(providers)
}
c.buildAutoChain()
c.buildAutoImageChain()
return nil
}
@ -466,6 +563,47 @@ func (c *Core) buildAutoChain() {
c.autoChain.Store(scheduler.BuildChain(sr, prov))
}
// buildAutoImageChain rebuilds the image-generation AUTO chain snapshot.
// Unlike buildAutoChain, only image-kind models participate: chat models in
// the rules are skipped so a stale chat slot can't receive image traffic.
func (c *Core) buildAutoImageChain() {
prov := func(model, source string) scheduler.Provider {
p := c.registry.ProviderForSlot(model, source)
if p == nil {
return nil
}
if m := p.ModelByID(model); m != nil && m.Kind != "image" {
return nil
}
return p
}
rules := c.cfg.AutoImage
sr := make([]scheduler.Rule, 0, len(rules))
for _, e := range rules {
model, source := e.Model, e.Source
if p := c.registry.ProviderForSlot(e.Model, e.Source); p != nil {
if exact := p.ModelIDFold(e.Model); exact != "" {
model = exact
}
if m := p.ModelByID(model); m != nil && m.Kind != "image" {
continue // chat-kind models never join the image AUTO chain
}
if source == "" {
source = p.Name()
}
}
sr = append(sr, scheduler.Rule{
Model: model,
Source: source,
Tier: e.Tier,
Quota: e.TokenQuota,
Period: e.Period,
Hours: e.Hours,
})
}
c.autoImageChain.Store(scheduler.BuildChain(sr, prov))
}
// AutoSlotState is the UI-facing health snapshot of one AUTO chain slot.
type AutoSlotState struct {
Model string `json:"model"`
@ -474,6 +612,16 @@ type AutoSlotState struct {
FailCount int64 `json:"fail_count"`
CooldownUntil int64 `json:"cooldown_until"`
Cooling bool `json:"cooling"`
// CooldownFrom is when the current cooldown window opened and ProbeAfter is
// its midpoint: from ProbeAfter on, a single probe request is allowed
// through so a recovered upstream does not have to sit out the rest of the
// window. Probing reports whether such a probe is in flight right now.
CooldownFrom int64 `json:"cooldown_from"`
ProbeAfter int64 `json:"probe_after"`
Probing bool `json:"probing"`
// Probeable marks a cooling slot that is past its midpoint, i.e. the next
// AUTO request may use it as a probe.
Probeable bool `json:"probeable"`
}
// AutoSlotStates returns per-slot health for every slot of the current chain.
@ -491,6 +639,7 @@ func (c *Core) AutoSlotStates() []AutoSlotState {
continue
}
pref, fail, until := pp.ModelHealthInfo(sl.Model)
from, probeAfter, probing := pp.ModelProbeInfo(sl.Model)
out = append(out, AutoSlotState{
Model: sl.Model,
Source: sl.Source,
@ -498,6 +647,10 @@ func (c *Core) AutoSlotStates() []AutoSlotState {
FailCount: fail,
CooldownUntil: until,
Cooling: until > now,
CooldownFrom: from,
ProbeAfter: probeAfter,
Probing: probing,
Probeable: until > now && probeAfter > 0 && now >= probeAfter,
})
}
}
@ -549,7 +702,7 @@ func (c *Core) AddSource(src config.Source) error {
if err := normalizeSource(&src); err != nil {
return err
}
if err := config.UpsertSourceInYAML(c.cfg.Path, src.Name, src); err != nil {
if err := config.UpsertSourceInYAML(c.cfg.Path, src.Name, src, c.cfg.SecretBox()); err != nil {
return err
}
// Update in-memory Sources so mergedSources() finds the entry.
@ -601,6 +754,33 @@ func (c *Core) Sources() []config.Source {
return c.mergedSources()
}
// Templates returns the saved source templates.
func (c *Core) Templates() []config.SourceTemplate {
c.mu.Lock()
defer c.mu.Unlock()
return c.store.ListTemplates()
}
// SaveTemplate upserts a source template.
func (c *Core) SaveTemplate(t config.SourceTemplate) error {
c.mu.Lock()
defer c.mu.Unlock()
if t.Name == "" {
return fmt.Errorf("template requires a name")
}
return c.store.UpsertTemplate(t)
}
// RemoveTemplate deletes a source template by name.
func (c *Core) RemoveTemplate(name string) error {
c.mu.Lock()
defer c.mu.Unlock()
if _, err := c.store.RemoveTemplate(name); err != nil {
return err
}
return nil
}
func normalizeSource(s *config.Source) error {
if s.Name == "" || s.BaseURL == "" {
return fmt.Errorf("source requires name and base_url")
@ -622,4 +802,21 @@ func (c *Core) Close() {
if c.vm != nil {
c.vm.Stop()
}
}
}
// ValidateScopeQuotas checks every scope entry's caps before they are stored.
// A typo in a period must be rejected at write time rather than silently
// becoming a never-resetting budget — the opposite of what was typed.
func ValidateScopeQuotas(entries []config.ModelScope) error {
for _, e := range entries {
if err := (config.KeyQuota{
TokenQuota: e.TokenQuota,
ReqQuota: e.ReqQuota,
Period: e.Period,
Hours: e.Hours,
}).Validate(); err != nil {
return fmt.Errorf("model %q: %w", e.Model, err)
}
}
return nil
}

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

View File

@ -4,6 +4,7 @@ import (
"encoding/csv"
"encoding/json"
"io"
"log"
"net/http"
"strconv"
"strings"
@ -31,7 +32,11 @@ func (g *Gateway) handleAdaptersAPI(w http.ResponseWriter, r *http.Request) {
writeJSON(w, http.StatusOK, map[string]interface{}{"adapters": g.core.ListAdapters()})
return
}
writeError(w, http.StatusNotFound, "not_found", "adapter code not exposed; edit in UI")
// A GET on a specific adapter is almost always a client that meant to
// DELETE it but let fetch default to GET; say so instead of only
// reporting that the code is not exposed.
writeError(w, http.StatusNotFound, "not_found",
"adapter code not exposed; edit in UI (to remove it use DELETE /api/adapters/"+path+")")
case http.MethodPost:
var p adapterPayload
if err := json.NewDecoder(r.Body).Decode(&p); err != nil {
@ -72,6 +77,7 @@ type sourcePayload struct {
Temperature float64 `json:"temperature"`
MaxTokens int `json:"max_tokens"`
MaxConcurrent int `json:"max_concurrent"`
RPM int `json:"rpm"` // optional requests-per-minute cap, 0 = unlimited
}
func (g *Gateway) handleSourcesAPI(w http.ResponseWriter, r *http.Request) {
@ -105,6 +111,7 @@ func (g *Gateway) handleSourcesAPI(w http.ResponseWriter, r *http.Request) {
Temperature: p.Temperature,
MaxTokens: p.MaxTokens,
MaxConcurrent: p.MaxConcurrent,
RPM: p.RPM,
}
if err := g.core.AddSource(src); err != nil {
writeError(w, http.StatusBadRequest, "source_error", err.Error())
@ -126,6 +133,62 @@ func (g *Gateway) handleSourcesAPI(w http.ResponseWriter, r *http.Request) {
}
}
// csvHeaders stamps the shared download headers for CSV exports.
func csvHeaders(w http.ResponseWriter, filename string) {
w.Header().Set("Content-Type", "text/csv; charset=utf-8")
w.Header().Set("Content-Disposition", "attachment; filename="+filename)
}
// handleSourceTemplatesAPI manages reusable source templates (a Source minus
// name and api_key) stored in the runtime file so the WebUI can spin up
// multiple key-bearing sources from one shared template.
func (g *Gateway) handleSourceTemplatesAPI(w http.ResponseWriter, r *http.Request) {
if reqRole(r.Context()) != "admin" {
writeError(w, http.StatusForbidden, "forbidden", "admin role required")
return
}
path := strings.TrimPrefix(r.URL.Path, "/api/source_templates")
path = strings.Trim(path, "/")
switch r.Method {
case http.MethodGet:
writeJSON(w, http.StatusOK, map[string]interface{}{"templates": g.core.Templates()})
case http.MethodPost:
var t config.SourceTemplate
if err := json.NewDecoder(r.Body).Decode(&t); err != nil {
writeError(w, http.StatusBadRequest, "invalid_request", "invalid json: "+err.Error())
return
}
if err := g.core.SaveTemplate(t); err != nil {
writeError(w, http.StatusBadRequest, "template_error", err.Error())
return
}
writeJSON(w, http.StatusOK, map[string]interface{}{"ok": true, "templates": g.core.Templates()})
case http.MethodDelete:
if path == "" {
writeError(w, http.StatusBadRequest, "invalid_request", "template name required")
return
}
if err := g.core.RemoveTemplate(path); err != nil {
writeError(w, http.StatusBadRequest, "template_error", err.Error())
return
}
writeJSON(w, http.StatusOK, map[string]interface{}{"ok": true})
default:
writeError(w, http.StatusMethodNotAllowed, "method_not_allowed", "")
}
}
// exportKey resolves which masked key id an export covers: admins may pass
// any key filter, user keys are always scoped to themselves (records and
// aggregates are keyed by the masked keyID form).
func exportKey(r *http.Request) string {
k := r.URL.Query().Get("key")
if reqRole(r.Context()) != "admin" {
k = keyID(reqKey(r.Context()))
}
return k
}
// handleStatsAPI returns per-key / per-model / per-source usage aggregates and
// the recent request audit trail.
func (g *Gateway) handleStatsAPI(w http.ResponseWriter, r *http.Request) {
@ -133,33 +196,35 @@ func (g *Gateway) handleStatsAPI(w http.ResponseWriter, r *http.Request) {
writeError(w, http.StatusMethodNotAllowed, "method_not_allowed", "use GET")
return
}
limit := 500
// Default to ONE SCREEN of records. The dashboard pages the rest through
// /api/stats/records as the user scrolls, so a poll must not serialize the
// whole ring buffer. An explicit limit is still honoured (capped by the ring).
limit := firstScreenRecords
if v := r.URL.Query().Get("limit"); v != "" {
if n, err := strconv.Atoi(v); err == nil && n > 0 && n <= 20000 {
limit = n
}
}
key := r.URL.Query().Get("key")
if reqRole(r.Context()) != "admin" {
// user keys may only see their own usage - records and aggregates are
// keyed by the masked id (keyID), so filter on that masked form.
key = keyID(reqKey(r.Context()))
}
key := exportKey(r)
if r.URL.Query().Get("export") == "csv" {
from, _ := strconv.ParseInt(r.URL.Query().Get("from"), 10, 64)
to, _ := strconv.ParseInt(r.URL.Query().Get("to"), 10, 64)
if to == 0 {
to = time.Now().UnixMilli()
}
w.Header().Set("Content-Type", "text/csv; charset=utf-8")
w.Header().Set("Content-Disposition", "attachment; filename=llmsproxy-requests.csv")
csvHeaders(w, "llmsproxy-requests.csv")
cw := csv.NewWriter(w)
names := map[string]string{}
for _, k := range g.core.ListKeys() {
names[keyID(k.Key)] = k.Name
}
_ = cw.Write([]string{"time", "key", "key_name", "type", "model", "source", "status", "ok", "prompt_tokens", "completion_tokens", "latency_ms", "error"})
for _, rec := range g.stats.AuditRecords(from, to, key) {
_ = cw.Write([]string{"time", "key", "key_name", "type", "model", "source", "status", "ok", "prompt_tokens", "completion_tokens", "latency_ms", "first_byte_ms", "cache_hit_tokens", "cache_miss_tokens", "error"})
// Stream the window straight to the client instead of materializing it:
// an export covering months of audit data must not be bounded by RAM, and
// nothing is retained once the response is written.
flushEvery := 1000
n := 0
err := g.stats.StreamAuditRecords(from, to, key, func(rec Req) error {
_ = cw.Write([]string{
time.UnixMilli(rec.Time).Format(time.RFC3339),
rec.Key,
@ -172,38 +237,48 @@ func (g *Gateway) handleStatsAPI(w http.ResponseWriter, r *http.Request) {
strconv.FormatInt(rec.Prompt, 10),
strconv.FormatInt(rec.Compl, 10),
strconv.FormatInt(rec.LatMs, 10),
strconv.FormatInt(rec.FirstByteMs, 10),
strconv.FormatInt(rec.CacheHit, 10),
strconv.FormatInt(rec.CacheMiss, 10),
rec.Err,
})
}
n++
if n%flushEvery == 0 {
cw.Flush()
if err := cw.Error(); err != nil {
return err // client went away: stop walking the audit files
}
if f, ok := w.(http.Flusher); ok {
f.Flush()
}
}
return nil
})
cw.Flush()
if err != nil {
log.Printf("[gateway] csv export aborted after %d rows: %v", n, err)
}
return
}
if r.URL.Query().Get("export") == "keys-csv" {
w.Header().Set("Content-Type", "text/csv; charset=utf-8")
w.Header().Set("Content-Disposition", "attachment; filename=llmsproxy-keys.csv")
csvHeaders(w, "llmsproxy-keys.csv")
cw := csv.NewWriter(w)
_ = cw.Write([]string{"key", "key_name", "role", "models", "total_requests", "success_requests", "failed_requests", "prompt_tokens", "completion_tokens", "total_tokens", "avg_latency_ms", "max_latency_ms", "created_at"})
// Use the same key filtering as the JSON API
exportKey := r.URL.Query().Get("key")
if reqRole(r.Context()) != "admin" {
exportKey = keyID(reqKey(r.Context()))
}
keyFilter := exportKey(r)
// by_key rows are keyed by the masked id (keyID); build a masked-id ->
// record map so names/roles/models resolve for the export.
byMasked := map[string]config.GWKey{}
for _, k := range g.core.ListKeys() {
byMasked[keyID(k.Key)] = k
}
snap := g.stats.Snapshot(0, exportKey)
snap := g.stats.Snapshot(0, keyFilter)
byKeyRaw, ok := snap["by_key"].([]StatsRow)
if !ok {
byKeyRaw = []StatsRow{}
}
for _, row := range byKeyRaw {
keyInfo, found := byMasked[row.Name]
name := ""
role := ""
models := ""
name, role, models := "", "", ""
if found {
name = keyInfo.Name
role = keyInfo.Role
@ -213,17 +288,9 @@ func (g *Gateway) handleStatsAPI(w http.ResponseWriter, r *http.Request) {
}
models = strings.Join(modelNames, ",")
}
reqs := row.Stat.Reqs
success := row.Stat.OK
errCount := row.Stat.Err
prompt := row.Stat.Prompt
compl := row.Stat.Compl
tokens := row.Stat.Prompt + row.Stat.Compl
latSum := row.Stat.LatSum
latMax := row.Stat.LatMax
avgLat := int64(0)
if reqs > 0 {
avgLat = latSum / reqs
if row.Stat.Reqs > 0 {
avgLat = row.Stat.LatSum / row.Stat.Reqs
}
createdAt := ""
if found && keyInfo.CreatedAt > 0 {
@ -231,14 +298,14 @@ func (g *Gateway) handleStatsAPI(w http.ResponseWriter, r *http.Request) {
}
_ = cw.Write([]string{
row.Name, name, role, models,
strconv.FormatInt(reqs, 10),
strconv.FormatInt(success, 10),
strconv.FormatInt(errCount, 10),
strconv.FormatInt(prompt, 10),
strconv.FormatInt(compl, 10),
strconv.FormatInt(tokens, 10),
strconv.FormatInt(row.Stat.Reqs, 10),
strconv.FormatInt(row.Stat.OK, 10),
strconv.FormatInt(row.Stat.Err, 10),
strconv.FormatInt(row.Stat.Prompt, 10),
strconv.FormatInt(row.Stat.Compl, 10),
strconv.FormatInt(row.Stat.Prompt+row.Stat.Compl, 10),
strconv.FormatInt(avgLat, 10),
strconv.FormatInt(latMax, 10),
strconv.FormatInt(row.Stat.LatMax, 10),
createdAt,
})
}
@ -253,3 +320,38 @@ func (g *Gateway) handleStatsAPI(w http.ResponseWriter, r *http.Request) {
snap["key_names"] = keyNames
writeJSON(w, http.StatusOK, snap)
}
// handleStatsRecordsAPI pages the request records straight off the audit files.
// The dashboard loads only its first screen and asks for the next page as the
// user scrolls, so neither side holds the full history: the server keeps no
// per-client state (the cursor is the whole state, and it lives in the URL) and
// closes every file handle before responding.
//
// GET /api/stats/records?before=<cursor>&limit=100[&key=<masked id>]
// -> { records: [...newest first...], next_cursor, has_more, rotated }
func (g *Gateway) handleStatsRecordsAPI(w http.ResponseWriter, r *http.Request) {
if r.Method != http.MethodGet {
writeError(w, http.StatusMethodNotAllowed, "method_not_allowed", "use GET")
return
}
limit := 100
if v := r.URL.Query().Get("limit"); v != "" {
if n, err := strconv.Atoi(v); err == nil && n > 0 && n <= 1000 {
limit = n
}
}
// exportKey pins non-admin callers to their own key, so a user key can never
// page another key's records.
page := g.stats.AuditPage(r.URL.Query().Get("before"), limit, exportKey(r))
keyNames := map[string]string{}
for _, k := range g.core.ListKeys() {
keyNames[keyID(k.Key)] = k.Name
}
writeJSON(w, http.StatusOK, map[string]interface{}{
"records": page.Records,
"next_cursor": page.Next,
"has_more": page.HasMore,
"rotated": page.Rotated,
"key_names": keyNames,
})
}

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

@ -5,7 +5,10 @@ import (
"encoding/json"
"errors"
"fmt"
"log"
"net/http"
"regexp"
"strconv"
"strings"
"sync/atomic"
"time"
@ -27,6 +30,9 @@ type chatRequest struct {
ToolChoice interface{} `json:"tool_choice,omitempty"`
DisableThinking bool `json:"disable_thinking"`
ExtraBody map[string]interface{} `json:"extra_body,omitempty"`
// PromptCacheKey 是 OpenAI 原生的缓存键;部分客户端用它携带会话标识。
PromptCacheKey string `json:"prompt_cache_key,omitempty"`
}
// ChatCompletion is the non-streaming OpenAI response object.
@ -37,6 +43,10 @@ type ChatCompletion struct {
Model string `json:"model"`
Choices []ChatChoice `json:"choices"`
Usage *types.TokenUsage `json:"usage,omitempty"`
// Cost is the upstream-reported charge for this request, passed through
// verbatim (OpenCode reports it as a decimal string, always "0" on the
// flat-rate Go subscription). Absent when the upstream reports nothing.
Cost string `json:"cost,omitempty"`
}
type ChatChoice struct {
@ -58,6 +68,13 @@ type ChatChunk struct {
Created int64 `json:"created"`
Model string `json:"model"`
Choices []ChunkChoice `json:"choices"`
// Usage is sent in the final chunk of a stream (empty choices) so
// OpenAI-compatible clients can read token usage.
Usage *types.TokenUsage `json:"usage,omitempty"`
// Cost is the upstream-reported charge for this request, passed through
// verbatim (OpenCode reports it as a decimal string). Emitted on the
// terminal chunk, mirroring OpenCode's own {"choices":[],"cost":"0"}.
Cost string `json:"cost,omitempty"`
}
type ChunkChoice struct {
@ -111,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
}
}
@ -148,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{}
@ -175,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" /
@ -225,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"), ""
@ -283,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
@ -297,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()
@ -313,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)),
@ -350,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()
@ -368,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)),
@ -384,13 +480,6 @@ func (g *Gateway) handleChat(w http.ResponseWriter, r *http.Request) {
g.singleChat(w, ctx, cands, inner, effective, rec)
}
func normalizeModel(m string) string {
if isAuto(m) {
return ""
}
return m
}
func firstModel(p *provider.Provider) string {
ms := p.Models()
if len(ms) > 0 {
@ -520,21 +609,145 @@ func upstreamErrStatus(err error) int {
return http.StatusBadGateway
}
func (g *Gateway) singleChat(w http.ResponseWriter, ctx context.Context, cands []*provider.Provider, req *types.ChatRequest, effective string, rec *Req) {
rec.LatMs = 0
t0 := time.Now()
resp, usedSrc, usedModel, err := g.core.Scheduler().Chat(ctx, toScheduler(cands), req)
rec.LatMs = time.Since(t0).Milliseconds()
if err != nil {
rec.OK = false
rec.Status = upstreamErrStatus(err)
rec.Err = err.Error()
g.writeRec(rec)
writeError(w, rec.Status, "upstream_error", err.Error())
return
// overflowMarkers 匹配上游「上下文超窗」类措辞。
//
// 上游写法五花八门,且**不在 pi 客户端的识别列表里**。pi 靠
// @earendil-works/pi-ai 的 OVERFLOW_PATTERNS 判断超窗并据此触发压缩重试,
// 而 justworker 返回的是「请精简对话历史…(Context window is full…)」——
// 与那 25 条正则一条都不匹配,于是 pi 既不压缩也不重试,只把它当成一条
// 普通上游错误。这里把可识别的超窗措辞归一化成 pi 一定认得的标记。
var overflowMarkers = []*regexp.Regexp{
regexp.MustCompile(`(?i)context[ _-]?window is full`),
regexp.MustCompile(`(?i)context[_ ]length[_ ]exceeded`),
regexp.MustCompile(`(?i)exceeds? the context window`),
regexp.MustCompile(`(?i)maximum context length`),
regexp.MustCompile(`(?i)reduce the length of the messages`),
regexp.MustCompile(`(?i)too many tokens`),
regexp.MustCompile(`(?i)token limit exceeded`),
regexp.MustCompile(`请精简对话历史`),
regexp.MustCompile(`上下文(长度)?超(出|限)`),
regexp.MustCompile(`对话历史过长`),
}
// overflowCanonical 命中 pi 的 /context[_ ]length[_ ]exceeded/i。
//
// 注意:这个前缀只是必要条件,不是充分条件——pi 还会先用
// NON_OVERFLOW_PATTERNS 排除整条消息,见 overflowClientMessage。
const overflowCanonical = "context_length_exceeded"
// looksLikeOverflow 判断错误文本是否属于上下文超窗。
func looksLikeOverflow(s string) bool {
for _, re := range overflowMarkers {
if re.MatchString(s) {
return true
}
}
rec.OK = true
rec.Status = http.StatusOK
return false
}
// clientUpstreamErr collapses upstream failure details into a short message
// for the client: per-tier bodies (WAF HTML pages, quota payloads, ...) stay
// in rec.Err / the stats API and the server log instead of the response.
// Per-tier one-line reasons are kept (quota/cooling skips carry no error body
// and are the actionable part); each is capped so HTML dumps can't leak.
//
// 超窗会被归一化成「干净」的 overflowCanonical 消息,见 overflowClientMessage。
func clientUpstreamErr(err error) string {
log.Printf("[gateway] upstream failure surfaced to client: %v", err)
if looksLikeOverflow(err.Error()) {
return overflowClientMessage(err)
}
return upstreamErrSummary(err)
}
// overflowHint 是给客户端的超窗说明,刻意只用 pi 认得的措辞。
const overflowHint = "context window is full; reduce the length of the messages"
// clientSessionFromRequest 取客户端自带的会话标识,拿不到时返回 ""。
//
// 通用 OpenAI 客户端默认不带会话 id,但 pi 支持:只要 provider 的 compat 里打开
// sendSessionAffinityHeaders,pi 就会把**平台真实的会话 id**(uuidv7,整个会话
// 恒定)放到 x-session-affinity / x-client-request-id 上(sessionAffinityFormat
// 为 openrouter 时是 x-session-id,为 openai 时是 session_id)。
// 有了它就不必再从首条 user 消息推导会话指纹(那只在客户端不发会话 id 时才作为
// 退路,且历史压缩后会漂移)。prompt_cache_key 是 OpenAI 原生的缓存键,
// 客户端若带也一并当会话用。
//
// 刻意不采纳 x-client-request-id:名字含 request,部分客户端每请求都换,
// 拿它当会话会让上游缓存永不命中。pi 总会同时发 x-session-affinity,够用。
func clientSessionFromRequest(r *http.Request, bodyKey string) string {
for _, h := range []string{"X-Session-Affinity", "X-Session-Id", "Session-Id"} {
if v := strings.TrimSpace(r.Header.Get(h)); v != "" {
return v
}
}
return strings.TrimSpace(bodyKey)
}
// overflowClientMessage 为超窗失败生成「干净」的客户端消息。
//
// 为什么不能沿用 upstreamErrSummary:pi 的 isContextOverflow 先查
// NON_OVERFLOW_PATTERNS(/rate limit/、/too many requests/、Bedrock 前缀),
// 一旦命中就直接判为「非超窗」——**哪怕消息里已经有 context_length_exceeded**,
// pi 也不会压缩重试。而 AUTO 链的失败消息天生是多 tier 原因的拼接:
// 超窗 tier(gozen 400 maximum context length)常与配额/限流 tier
// (429 token plan exhausted、cooling、no free slot)同时出现。
// 把明细原样带出去,等于让一条限流 tier 的措辞反过来封杀超窗识别。
// 所以这里只保留「超窗」措辞 + 是哪个源超的窗,其余一律不带。
func overflowClientMessage(err error) string {
var ce *scheduler.ChainErr
if errors.As(err, &ce) {
for _, t := range ce.Tiers {
if looksLikeOverflow(t.Err.Error()) {
return fmt.Sprintf("%s: %s (%s/%s)", overflowCanonical, overflowHint, t.Source, t.Model)
}
}
}
return fmt.Sprintf("%s: %s", overflowCanonical, overflowHint)
}
// upstreamErrSummary 把上游失败压成一行短消息。
func upstreamErrSummary(err error) string {
var ce *scheduler.ChainErr
if errors.As(err, &ce) {
parts := make([]string, 0, len(ce.Tiers)+len(ce.Skipped))
for _, t := range ce.Tiers {
// 160 而不是 80:短诊断词("Context window is full")常落在尾部,
// 80 字节按字节截断正好会把它切掉,超窗就再也认不出来。
parts = append(parts, types.OneLine(fmt.Sprintf("%s/%s: %v", t.Source, t.Model, t.Err), 160))
}
for _, sk := range ce.Skipped {
parts = append(parts, types.OneLine(sk, 160))
}
msg := strings.Join(parts, "; ")
if len(msg) > 300 {
msg = msg[:300] + "..."
}
return fmt.Sprintf("all %d auto providers failed: %s", len(parts), msg)
}
return types.OneLine(err.Error(), 160)
}
// failChat maps a scheduling failure onto the audit record and answers the
// client. A failed AUTO chain additionally pins its first failed tier onto
// the record; direct-path errors never match *scheduler.ChainErr, so the
// extraction is safely shared by all four entry points. Callers own the
// writeRec call (inline for non-stream paths, deferred for stream paths).
func (g *Gateway) failChat(w http.ResponseWriter, rec *Req, err error) {
rec.OK = false
rec.Status = upstreamErrStatus(err)
rec.Err = err.Error()
var ce *scheduler.ChainErr
if errors.As(err, &ce) && len(ce.Tiers) > 0 {
rec.Source = ce.Tiers[0].Source
rec.Model = ce.Tiers[0].Model
}
writeError(w, rec.Status, "upstream_error", clientUpstreamErr(err))
}
// recordChatUsage fills token accounting for a finished non-streaming
// request: exact upstream numbers win, byte estimates fill the gaps.
func recordChatUsage(rec *Req, req *types.ChatRequest, resp *types.UnifiedResponse) {
rec.Prompt = int64(resp.TokenUsage.Prompt)
if rec.Prompt == 0 {
rec.Prompt = estimatePromptTokens(req)
@ -543,9 +756,29 @@ func (g *Gateway) singleChat(w http.ResponseWriter, ctx context.Context, cands [
if rec.Compl == 0 {
rec.Compl = estimateTextTokens(resp.Content, resp.ReasoningContent, resp.ToolCalls)
}
rec.Source = usedSrc
rec.Model = usedModel
g.writeRec(rec)
// Cache accounting: whichever source format the adapter normalized into
// (prompt_tokens_details.cached_tokens or legacy prompt_cache_hit_tokens),
// read it back so the request record carries the hit/miss split.
if d := resp.TokenUsage.PromptTokensDetails; d != nil && d.CachedTokens > 0 {
rec.CacheHit = int64(d.CachedTokens)
}
if resp.TokenUsage.PromptTokensDetails != nil {
// Upstream reported cache details (even a 0 hit) — tag the row so
// the UI can show 0% rather than “—”.
rec.CacheReported = true
} else if resp.TokenUsage.PromptCacheHit > 0 {
rec.CacheHit = int64(resp.TokenUsage.PromptCacheHit)
rec.CacheReported = true
}
if resp.TokenUsage.PromptCacheMiss > 0 {
rec.CacheMiss = int64(resp.TokenUsage.PromptCacheMiss)
}
}
// writeChatCompletion renders a unified response as an OpenAI
// chat.completion object. modelName is the id clients see as the serving
// model: the requested id for direct routes, the exact slot model for AUTO.
func writeChatCompletion(w http.ResponseWriter, resp *types.UnifiedResponse, modelName string) {
msg := RespMessage{Role: "assistant", Content: resp.Content}
if resp.ReasoningContent != "" {
msg.ReasoningContent = resp.ReasoningContent
@ -557,8 +790,9 @@ func (g *Gateway) singleChat(w http.ResponseWriter, ctx context.Context, cands [
ID: newID(),
Object: "chat.completion",
Created: time.Now().Unix(),
Model: effective,
Model: modelName,
Choices: []ChatChoice{{Index: 0, Message: msg, FinishReason: resp.FinishReason}},
Cost: resp.Cost,
}
if resp.TokenUsage.Total > 0 || resp.TokenUsage.Prompt > 0 || resp.TokenUsage.Completion > 0 {
out.Usage = &resp.TokenUsage
@ -566,6 +800,30 @@ func (g *Gateway) singleChat(w http.ResponseWriter, ctx context.Context, cands [
writeJSON(w, http.StatusOK, out)
}
// singleChat runs a direct (model-pinned) non-streaming request across the
// candidate list, falling back on failure.
func (g *Gateway) singleChat(w http.ResponseWriter, ctx context.Context, cands []*provider.Provider, req *types.ChatRequest, effective string, rec *Req) {
rec.LatMs = 0
t0 := time.Now()
resp, usedSrc, usedModel, err := g.core.Scheduler().Chat(ctx, toScheduler(cands), req)
rec.LatMs = time.Since(t0).Milliseconds()
if err != nil {
g.failChat(w, rec, err)
g.writeRec(rec)
return
}
rec.OK = true
rec.Status = http.StatusOK
recordChatUsage(rec, req, resp)
rec.Source = usedSrc
rec.Model = usedModel
// Non-streaming: the whole response arrives at once, so TTFB equals
// the total latency.
rec.FirstByteMs = rec.LatMs
g.writeRec(rec)
writeChatCompletion(w, resp, effective)
}
// writeRec records a finished request (audit + aggregates).
func (g *Gateway) writeRec(rec *Req) {
if rec == nil {
@ -577,27 +835,51 @@ func (g *Gateway) writeRec(rec *Req) {
g.stats.Record(*rec)
}
func (g *Gateway) streamChat(w http.ResponseWriter, ctx context.Context, cands []*provider.Provider, req *types.ChatRequest, effective string, rec *Req) {
rec.LatMs = 0
t0 := time.Now()
rec.OK = true
rec.Status = http.StatusOK
defer func() {
rec.LatMs = time.Since(t0).Milliseconds()
g.writeRec(rec)
}()
chunks, _, usedModel, err := g.core.Scheduler().ChatStream(ctx, toScheduler(cands), req)
if err != nil {
rec.OK = false
rec.Status = upstreamErrStatus(err)
rec.Err = err.Error()
writeError(w, rec.Status, "upstream_error", err.Error())
return
// mergeUsage combines token usage across stream chunks additively. Some
// providers split usage across chunks (e.g. Anthropic reports prompt tokens
// in message_start and the final completion tokens in message_delta); a plain
// "last non-nil wins" would discard the prompt half. Non-zero fields from cur
// override prev; total is recomputed from the merged parts so a partial later
// chunk can't shrink it. For the common single-chunk case (OpenAI's terminal
// empty-choices+usage chunk) upstream totals are preserved exactly.
func mergeUsage(prev, cur *types.TokenUsage) *types.TokenUsage {
if prev == nil {
u := *cur
if u.Total == 0 && (u.Prompt > 0 || u.Completion > 0) {
u.Total = u.Prompt + u.Completion
}
return &u
}
if usedModel != "" {
rec.Model = usedModel
out := *prev
if cur.Prompt > 0 {
out.Prompt = cur.Prompt
}
rec.Prompt = estimatePromptTokens(req)
if cur.Completion > 0 {
out.Completion = cur.Completion
}
out.Total = out.Prompt + out.Completion
if cur.PromptTokensDetails != nil {
// Keep the details object even when CachedTokens is 0: a reported
// zero-hit is meaningful ("cache missed") and must stay
// distinguishable from "upstream never reported cache info".
// Dropping it here made streaming rows show “—” instead of 0%.
out.PromptTokensDetails = cur.PromptTokensDetails
}
if cur.PromptCacheHit > 0 {
out.PromptCacheHit = cur.PromptCacheHit
}
if cur.PromptCacheMiss > 0 {
out.PromptCacheMiss = cur.PromptCacheMiss
}
return &out
}
// pumpStream writes the full SSE sequence for a started stream: role
// preamble, one chunk per unified delta, the terminating finish_reason, the
// OpenAI-standard final usage chunk (empty choices) and [DONE]. modelName
// follows writeChatCompletion's rule (requested id for direct routes, exact
// slot model for AUTO).
func (g *Gateway) pumpStream(w http.ResponseWriter, rec *Req, chunks <-chan types.UnifiedChunk, modelName string, t0 time.Time) {
w.Header().Set("Content-Type", "text/event-stream")
w.Header().Set("Cache-Control", "no-cache")
w.Header().Set("Connection", "keep-alive")
@ -621,14 +903,26 @@ func (g *Gateway) streamChat(w http.ResponseWriter, ctx context.Context, cands [
}
if !send(ChatChunk{
ID: id, Object: "chat.completion.chunk", Created: created, Model: effective,
ID: id, Object: "chat.completion.chunk", Created: created, Model: modelName,
Choices: []ChunkChoice{{Index: 0, Delta: RespMessage{Role: "assistant"}}},
}) {
return
}
// First SSE byte sent to the client: record time-to-first-byte for the
// source's status-page latency average.
rec.FirstByteMs = time.Since(t0).Milliseconds()
var lastUsage *types.TokenUsage
var lastCost string
lastFinish := ""
for ck := range chunks {
if ck.Usage != nil {
lastUsage = mergeUsage(lastUsage, ck.Usage)
}
if ck.Cost != "" {
lastCost = ck.Cost
}
chunk := ChatChunk{
ID: id, Object: "chat.completion.chunk", Created: created, Model: effective,
ID: id, Object: "chat.completion.chunk", Created: created, Model: modelName,
}
delta := RespMessage{Role: "assistant", Content: ck.Content}
if ck.ReasoningContent != "" {
@ -639,8 +933,12 @@ func (g *Gateway) streamChat(w http.ResponseWriter, ctx context.Context, cands [
}
choice := ChunkChoice{Index: 0, Delta: delta}
if ck.Done {
stop := "stop"
choice.FinishReason = &stop
fin := ck.FinishReason
if fin == "" {
fin = "stop"
}
lastFinish = fin
choice.FinishReason = &fin
}
chunk.Choices = []ChunkChoice{choice}
rec.Compl += int64(len(ck.Content)+len(ck.ReasoningContent)+len(ck.ToolCalls)) / 3
@ -648,17 +946,103 @@ func (g *Gateway) streamChat(w http.ResponseWriter, ctx context.Context, cands [
return
}
}
stop := "stop"
finalFinish := lastFinish
if finalFinish == "" {
finalFinish = "stop"
}
send(ChatChunk{
ID: id, Object: "chat.completion.chunk", Created: created, Model: effective,
Choices: []ChunkChoice{{Index: 0, Delta: RespMessage{}, FinishReason: &stop}},
ID: id, Object: "chat.completion.chunk", Created: created, Model: modelName,
Choices: []ChunkChoice{{Index: 0, Delta: RespMessage{}, FinishReason: &finalFinish}},
})
// Final usage chunk (OpenAI standard: empty choices + usage before [DONE]).
// Prefer the upstream's exact usage if the stream carried it; fall back to
// the gateway's estimate otherwise.
// Write the upstream's exact numbers back onto the audit record. Without
// this the streamed path kept only the per-chunk byte estimate, so the same
// request recorded ~1.4x its real completion tokens (measured: upstream 100,
// audit 145) while the non-streaming path recorded 100. Two paths, two
// different numbers for one request is a reporting bug, not a rounding one.
if lastUsage != nil {
if lastUsage.Prompt > 0 {
rec.Prompt = int64(lastUsage.Prompt)
}
if lastUsage.Completion > 0 {
rec.Compl = int64(lastUsage.Completion)
}
}
var tut *types.TokenUsage
if lastUsage != nil {
tut = lastUsage
// Write the upstream's cache accounting back onto the request record
// so the audit trail carries the hit/miss split for streaming too.
if d := tut.PromptTokensDetails; d != nil {
rec.CacheReported = true
if d.CachedTokens > 0 {
rec.CacheHit = int64(d.CachedTokens)
}
} else if tut.PromptCacheHit > 0 {
rec.CacheHit = int64(tut.PromptCacheHit)
rec.CacheReported = true
}
if tut.PromptCacheMiss > 0 {
rec.CacheMiss = int64(tut.PromptCacheMiss)
}
} else if rec.Prompt+rec.Compl > 0 {
u := types.TokenUsage{
Prompt: int(rec.Prompt),
Completion: int(rec.Compl),
Total: int(rec.Prompt + rec.Compl),
}
tut = &u
}
if tut != nil {
send(ChatChunk{
ID: id, Object: "chat.completion.chunk", Created: created, Model: 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")
if flusher != nil {
flusher.Flush()
}
}
// streamChat runs a direct (model-pinned) streaming request across the
// candidate list, falling back early on connect errors.
func (g *Gateway) streamChat(w http.ResponseWriter, ctx context.Context, cands []*provider.Provider, req *types.ChatRequest, effective string, rec *Req) {
rec.LatMs = 0
t0 := time.Now()
rec.OK = true
rec.Status = http.StatusOK
defer func() {
rec.LatMs = time.Since(t0).Milliseconds()
g.writeRec(rec)
}()
chunks, usedSrc, usedModel, err := g.core.Scheduler().ChatStream(ctx, toScheduler(cands), req)
if err != nil {
g.failChat(w, rec, err)
return
}
if usedModel != "" {
rec.Model = usedModel
}
// Audit accuracy: pin the source that actually served the stream (after a
// failover it differs from the first candidate). Direct streams previously
// discarded it.
rec.Source = usedSrc
rec.Prompt = estimatePromptTokens(req)
g.pumpStream(w, rec, chunks, effective, t0)
}
// singleChatAuto runs a non-streaming AUTO request down the chain (see
// scheduler.ChainChat): tiers ascending (tier 1 = highest priority first),
// per-tier round-robin ordered by preference, cooldown as the only hard skip,
@ -671,48 +1055,18 @@ func (g *Gateway) singleChatAuto(w http.ResponseWriter, ctx context.Context, cha
resp, usedSrc, usedModel, err := g.core.Scheduler().ChainChat(ctx, chain, req, quotaExhausted)
rec.LatMs = time.Since(t0).Milliseconds()
if err != nil {
rec.OK = false
rec.Status = upstreamErrStatus(err)
rec.Err = err.Error()
if ce, ok := err.(*scheduler.ChainErr); ok && len(ce.Tiers) > 0 {
rec.Source = ce.Tiers[0].Source
rec.Model = ce.Tiers[0].Model
}
g.failChat(w, rec, err)
g.writeRec(rec)
writeError(w, rec.Status, "upstream_error", err.Error())
return
}
rec.OK = true
rec.Status = http.StatusOK
rec.Prompt = int64(resp.TokenUsage.Prompt)
if rec.Prompt == 0 {
rec.Prompt = estimatePromptTokens(req)
}
rec.Compl = int64(resp.TokenUsage.Completion)
if rec.Compl == 0 {
rec.Compl = estimateTextTokens(resp.Content, resp.ReasoningContent, resp.ToolCalls)
}
recordChatUsage(rec, req, resp)
rec.Source = usedSrc
rec.Model = usedModel
rec.FirstByteMs = rec.LatMs
g.writeRec(rec)
msg := RespMessage{Role: "assistant", Content: resp.Content}
if resp.ReasoningContent != "" {
msg.ReasoningContent = resp.ReasoningContent
}
if len(resp.ToolCalls) > 0 {
msg.ToolCalls = toolCallsWire(resp.ToolCalls)
}
out := ChatCompletion{
ID: newID(),
Object: "chat.completion",
Created: time.Now().Unix(),
Model: usedModel,
Choices: []ChatChoice{{Index: 0, Message: msg, FinishReason: resp.FinishReason}},
}
if resp.TokenUsage.Total > 0 || resp.TokenUsage.Prompt > 0 || resp.TokenUsage.Completion > 0 {
out.Usage = &resp.TokenUsage
}
writeJSON(w, http.StatusOK, out)
writeChatCompletion(w, resp, usedModel)
}
// streamChatAuto streams an AUTO request down the chain. A slot is abandoned
@ -730,14 +1084,7 @@ func (g *Gateway) streamChatAuto(w http.ResponseWriter, ctx context.Context, cha
}()
chunks, usedSrc, usedModel, err := g.core.Scheduler().ChainChatStream(ctx, chain, req, quotaExhausted)
if err != nil {
rec.OK = false
rec.Status = upstreamErrStatus(err)
rec.Err = err.Error()
if ce, ok := err.(*scheduler.ChainErr); ok && len(ce.Tiers) > 0 {
rec.Source = ce.Tiers[0].Source
rec.Model = ce.Tiers[0].Model
}
writeError(w, rec.Status, "upstream_error", err.Error())
g.failChat(w, rec, err)
return
}
if usedModel != "" {
@ -745,63 +1092,7 @@ func (g *Gateway) streamChatAuto(w http.ResponseWriter, ctx context.Context, cha
}
rec.Source = usedSrc
rec.Prompt = estimatePromptTokens(req)
w.Header().Set("Content-Type", "text/event-stream")
w.Header().Set("Cache-Control", "no-cache")
w.Header().Set("Connection", "keep-alive")
w.WriteHeader(http.StatusOK)
flusher, _ := w.(http.Flusher)
id := newID()
created := time.Now().Unix()
send := func(obj interface{}) bool {
b, err := json.Marshal(obj)
if err != nil {
return false
}
if _, err := fmt.Fprintf(w, "data: %s\n\n", b); err != nil {
return false
}
if flusher != nil {
flusher.Flush()
}
return true
}
if !send(ChatChunk{
ID: id, Object: "chat.completion.chunk", Created: created, Model: rec.Model,
Choices: []ChunkChoice{{Index: 0, Delta: RespMessage{Role: "assistant"}}},
}) {
return
}
for ck := range chunks {
chunk := ChatChunk{
ID: id, Object: "chat.completion.chunk", Created: created, Model: rec.Model,
}
delta := RespMessage{Role: "assistant", Content: ck.Content}
if ck.ReasoningContent != "" {
delta.ReasoningContent = ck.ReasoningContent
}
if len(ck.ToolCalls) > 0 {
delta.ToolCalls = ck.ToolCalls
}
choice := ChunkChoice{Index: 0, Delta: delta}
if ck.Done {
stop := "stop"
choice.FinishReason = &stop
}
chunk.Choices = []ChunkChoice{choice}
rec.Compl += int64(len(ck.Content)+len(ck.ReasoningContent)+len(ck.ToolCalls)) / 3
if !send(chunk) {
return
}
}
stop := "stop"
send(ChatChunk{
ID: id, Object: "chat.completion.chunk", Created: created, Model: rec.Model,
Choices: []ChunkChoice{{Index: 0, Delta: RespMessage{}, FinishReason: &stop}},
})
fmt.Fprintf(w, "data: [DONE]\n\n")
if flusher != nil {
flusher.Flush()
}
g.pumpStream(w, rec, chunks, usedModel, t0)
}
func (g *Gateway) handleImage(w http.ResponseWriter, r *http.Request) {
@ -822,11 +1113,48 @@ func (g *Gateway) handleImage(w http.ResponseWriter, r *http.Request) {
if model == "" {
model = g.core.DefaultModel()
}
if !isAuto(model) {
if allow := g.allowedModels(r.Context()); allow != nil && !g.hasScopeModel(allow, model) {
writeError(w, http.StatusForbidden, "model_not_allowed", fmt.Sprintf("model %q is not allowed for this key", model))
if isAuto(model) {
if chain := g.core.AutoImageChain(); chain != nil && len(chain.Tiers) > 0 {
if q := g.checkQuota(r.Context(), "AUTO"); q != nil {
g.writeReject(w, q)
return
}
done := g.stats.Begin()
defer done()
rec := &Req{Key: keyID(reqKey(r.Context())), Type: "image", Model: model, OK: false}
t0 := time.Now()
resp, usedSrc, usedModel, err := g.core.Scheduler().ChainImage(r.Context(), chain, &req)
rec.LatMs = time.Since(t0).Milliseconds()
if err != nil {
rec.Status = upstreamErrStatus(err)
rec.Err = err.Error()
g.writeRec(rec)
writeError(w, rec.Status, "upstream_error", clientUpstreamErr(err))
return
}
rec.Source = usedSrc
if usedModel != "" {
rec.Model = usedModel // actual image model served, not "AUTO"
}
rec.OK = true
rec.Status = http.StatusOK
// Image generation has no token concept. Recording len(ImageData)
// (the image COUNT) in completion_tokens mislabels image count as
// tokens and feeds it into the token totals; leave it 0.
rec.ImageCount = len(resp.ImageData)
g.writeRec(rec)
writeJSON(w, http.StatusOK, types.ImageGenResponse{
Created: time.Now().Unix(),
Data: resp.ImageData,
})
return
}
// no image chain configured: fall through to legacy discovery (all
// sources exposing an image model, tried in registry order)
}
if allow := g.allowedModels(r.Context()); allow != nil && !g.hasScopeModel(allow, model) {
writeError(w, http.StatusForbidden, "model_not_allowed", fmt.Sprintf("model %q is not allowed for this key", model))
return
}
cands, _ := g.resolveByModel(model)
cands = imageOnly(cands)
@ -837,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()
@ -851,15 +1179,19 @@ func (g *Gateway) handleImage(w http.ResponseWriter, r *http.Request) {
rec.Status = upstreamErrStatus(err)
rec.Err = err.Error()
g.writeRec(rec)
writeError(w, rec.Status, "upstream_error", err.Error())
writeError(w, rec.Status, "upstream_error", clientUpstreamErr(err))
return
}
if usedSrc != "" {
rec.Source = usedSrc
}
if resp.Model != "" {
rec.Model = resp.Model // record the actual model served, not the raw request id
}
rec.OK = true
rec.Status = http.StatusOK
rec.Compl = int64(len(resp.ImageData))
// Image generation has no token concept — see the AUTO path above.
rec.ImageCount = len(resp.ImageData)
g.writeRec(rec)
writeJSON(w, http.StatusOK, types.ImageGenResponse{
Created: time.Now().Unix(),

View File

@ -0,0 +1,62 @@
package gateway
import (
"testing"
"llmsproxy/internal/types"
)
// TestMergeUsageKeepsZeroCacheDetails locks the streaming counterpart of the
// "distinguish missed from not reported" contract: a usage chunk that reports
// prompt_tokens_details with 0 cached tokens must survive the merge, otherwise
// streaming rows lose cache_reported and the UI shows "—" instead of 0%.
func TestMergeUsageKeepsZeroCacheDetails(t *testing.T) {
t.Run("zero-hit details survive merge", func(t *testing.T) {
prev := &types.TokenUsage{Prompt: 10, Completion: 1, Total: 11}
cur := &types.TokenUsage{
Prompt: 10, Completion: 5, Total: 15,
PromptTokensDetails: &types.PromptTokensDetails{CachedTokens: 0},
}
got := mergeUsage(prev, cur)
if got.PromptTokensDetails == nil {
t.Fatal("zero-hit prompt_tokens_details was dropped by mergeUsage")
}
if got.PromptTokensDetails.CachedTokens != 0 {
t.Fatalf("cached_tokens = %d, want 0", got.PromptTokensDetails.CachedTokens)
}
})
t.Run("non-zero hit still wins", func(t *testing.T) {
prev := &types.TokenUsage{Prompt: 10, Completion: 1, Total: 11}
cur := &types.TokenUsage{
Prompt: 10, Completion: 5, Total: 15,
PromptTokensDetails: &types.PromptTokensDetails{CachedTokens: 64},
}
got := mergeUsage(prev, cur)
if got.PromptTokensDetails == nil || got.PromptTokensDetails.CachedTokens != 64 {
t.Fatalf("cached_tokens lost: %+v", got.PromptTokensDetails)
}
})
t.Run("absent details do not overwrite an earlier report", func(t *testing.T) {
prev := &types.TokenUsage{
Prompt: 10, Completion: 1, Total: 11,
PromptTokensDetails: &types.PromptTokensDetails{CachedTokens: 32},
}
cur := &types.TokenUsage{Prompt: 10, Completion: 5, Total: 15}
got := mergeUsage(prev, cur)
if got.PromptTokensDetails == nil || got.PromptTokensDetails.CachedTokens != 32 {
t.Fatalf("earlier cache report clobbered by a later chunk without details: %+v",
got.PromptTokensDetails)
}
})
t.Run("no cache data anywhere stays nil", func(t *testing.T) {
prev := &types.TokenUsage{Prompt: 10, Completion: 1, Total: 11}
cur := &types.TokenUsage{Prompt: 10, Completion: 5, Total: 15}
got := mergeUsage(prev, cur)
if got.PromptTokensDetails != nil {
t.Fatalf("fabricated cache details: %+v", got.PromptTokensDetails)
}
})
}

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

@ -6,10 +6,13 @@ import (
"io"
"net/http"
"net/http/httptest"
"net/url"
"os"
"path/filepath"
"strconv"
"strings"
"testing"
"time"
"llmsproxy/internal/config"
"llmsproxy/internal/core"
@ -139,7 +142,7 @@ func TestChatAutoChain503Summary(t *testing.T) {
if rr.Code != http.StatusServiceUnavailable {
t.Fatalf("status=%d body=%s", rr.Code, rr.Body.String())
}
if !strings.Contains(rr.Body.String(), "all auto tiers failed") ||
if !strings.Contains(rr.Body.String(), "auto providers failed") ||
!strings.Contains(rr.Body.String(), "a/a-m") ||
!strings.Contains(rr.Body.String(), "b/b-m") {
t.Fatalf("503 must summarize every tier, body=%s", rr.Body.String())
@ -667,7 +670,351 @@ func TestHasScopeModelWithSourcePrefix(t *testing.T) {
if g.hasScopeModel(scope, "deepseek-v4-flash-free-extra") {
t.Error("hasScopeModel returned true for unrelated model")
}
if !g.hasScopeModel([]config.ModelScope{{Model: "AUTO"}}, "zen:anything") {
t.Error("AUTO scope should allow any prefixed model")
if g.hasScopeModel([]config.ModelScope{{Model: "AUTO"}}, "zen:deepseek-v4-flash-free") {
t.Error("AUTO scope must NOT grant specific model ids even with prefix")
}
if g.hasScopeModel([]config.ModelScope{{Model: "AUTO"}}, "deepseek-v4-pro") {
t.Error("AUTO scope must NOT grant specific model ids")
}
}
// TestStatusRecorderFlusher pins the SSE-critical contract: the access-log
// wrapper must implement http.Flusher, otherwise the streaming chat handlers'
// w.(http.Flusher) assertion yields nil and chunks are never flushed until
// the response ends (regression for the buffered-SSE bug).
func TestStatusRecorderFlusher(t *testing.T) {
var inner *flushRecorder
rr := httptest.NewRecorder()
inner = &flushRecorder{ResponseWriter: rr}
sr := &statusRecorder{ResponseWriter: inner}
fl, ok := any(sr).(http.Flusher)
if !ok {
t.Fatal("statusRecorder does not implement http.Flusher — SSE streaming is broken")
}
if inner.flushed {
t.Fatal("Flush called before Flush()")
}
fl.Flush()
if !inner.flushed {
t.Fatal("statusRecorder.Flush did not forward to the underlying writer")
}
}
// flushRecorder records whether Flush was forwarded.
type flushRecorder struct {
http.ResponseWriter
flushed bool
}
func (f *flushRecorder) Flush() { f.flushed = true }
// TestDirectStreamFailoverAuditSource: a direct streaming request whose first
// candidate hard-fails must be served by the fallback, and the audit record's
// Source must name the source that actually served the stream (not just the
// first candidate).
func TestDirectStreamFailoverAuditSource(t *testing.T) {
aUp := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
w.WriteHeader(500)
fmt.Fprint(w, `{"error":"boom"}`)
}))
defer aUp.Close()
bUp := mockUpstream()
defer bUp.Close()
g := newTestGateway(t,
// Both sources expose the same model id so the direct scheduler has a
// real fallback candidate after a hard-fails; b is the healthy one.
config.Source{Name: "a", BaseURL: aUp.URL, Adapter: "openai", Models: []config.Model{{ID: "m"}}},
config.Source{Name: "b", BaseURL: bUp.URL, Adapter: "openai", Models: []config.Model{{ID: "m"}}},
)
body := doReq(t, g, "POST", "/v1/chat/completions",
`{"model":"m","stream":true,"messages":[{"role":"user","content":"hi"}]}`)
if body.Code != 200 {
t.Fatalf("status=%d body=%s", body.Code, body.Body.String())
}
sb := body.Body.String()
if !strings.Contains(sb, `"finish_reason":"stop"`) || !strings.Contains(sb, "[DONE]") {
t.Fatalf("stream incomplete: %s", sb)
}
recs := g.stats.AuditRecords(0, 0, "")
if len(recs) == 0 {
t.Fatal("no audit records")
}
last := recs[len(recs)-1]
if last.Source != "b" || !last.OK {
t.Fatalf("audit rec source=%q ok=%v, want source=b ok=true (actual serving source)", last.Source, last.OK)
}
}
// TestStatusExposesAdapterPools checks that the elastic Lua pool sizing is
// visible to admins (and only to them) via /api/status, so the grow/shrink
// algorithm can be observed instead of inferred.
func TestStatusExposesAdapterPools(t *testing.T) {
g := newTestGateway(t)
rr := doReq(t, g, http.MethodGet, "/api/status", "")
if rr.Code != 200 {
t.Fatalf("status = %d, body=%s", rr.Code, rr.Body.String())
}
var body map[string]interface{}
if err := json.Unmarshal(rr.Body.Bytes(), &body); err != nil {
t.Fatalf("decode: %v", err)
}
pools, ok := body["adapter_pools"].([]interface{})
if !ok {
t.Fatalf("adapter_pools missing or wrong type: %T", body["adapter_pools"])
}
if len(pools) == 0 {
t.Fatal("adapter_pools must list the loaded adapters")
}
first, _ := pools[0].(map[string]interface{})
for _, field := range []string{"name", "created", "idle", "in_use", "max", "grow_step", "shrink_step"} {
if _, ok := first[field]; !ok {
t.Errorf("pool entry missing %q: %v", field, first)
}
}
// idle gateway: nothing booted eagerly
if c, _ := first["created"].(float64); c != 0 {
t.Errorf("adapter pool booted eagerly, created=%v", first["created"])
}
}
// TestStatsRecordsAPIPaging exercises the on-demand records endpoint the
// dashboard scrolls with: a bounded first page plus a cursor that walks back
// through the audit file.
func TestStatsRecordsAPIPaging(t *testing.T) {
g := newTestGateway(t)
for i := 0; i < 250; i++ {
g.stats.Record(Req{
Time: time.Now().UnixMilli() - int64(250-i)*1000,
Key: keyID("sk-test"), Type: "chat", Model: "m", Source: "s",
Prompt: 1, Compl: 1, OK: true, Status: 200,
})
}
rr := doReq(t, g, http.MethodGet, "/api/stats/records?limit=50", "")
if rr.Code != 200 {
t.Fatalf("status = %d body=%s", rr.Code, rr.Body.String())
}
var page struct {
Records []Req `json:"records"`
NextCursor string `json:"next_cursor"`
HasMore bool `json:"has_more"`
}
if err := json.Unmarshal(rr.Body.Bytes(), &page); err != nil {
t.Fatalf("decode: %v", err)
}
if len(page.Records) != 50 {
t.Fatalf("first page = %d records, want 50", len(page.Records))
}
if !page.HasMore || page.NextCursor == "" {
t.Fatal("a long history must advertise more pages plus a cursor")
}
// newest first
for i := 1; i < len(page.Records); i++ {
if page.Records[i-1].Time < page.Records[i].Time {
t.Fatalf("page not newest-first at %d", i)
}
}
// follow the cursor: the next page must continue strictly older
oldest := page.Records[len(page.Records)-1].Time
rr2 := doReq(t, g, http.MethodGet,
"/api/stats/records?limit=50&before="+url.QueryEscape(page.NextCursor), "")
if rr2.Code != 200 {
t.Fatalf("page 2 status = %d", rr2.Code)
}
var page2 struct {
Records []Req `json:"records"`
}
if err := json.Unmarshal(rr2.Body.Bytes(), &page2); err != nil {
t.Fatalf("decode page 2: %v", err)
}
if len(page2.Records) == 0 {
t.Fatal("cursor page must return records")
}
if page2.Records[0].Time > oldest {
t.Fatalf("page 2 starts at %d, must continue below %d", page2.Records[0].Time, oldest)
}
}
// TestStatsSnapshotShipsOneScreen: a dashboard poll must not serialize the whole
// ring, otherwise the "load the first screen only" contract is broken on the
// wire even if the UI pages.
func TestStatsSnapshotShipsOneScreen(t *testing.T) {
g := newTestGateway(t)
for i := 0; i < 400; i++ {
g.stats.Record(Req{
Time: time.Now().UnixMilli(), Key: keyID("sk-test"), Type: "chat",
Model: "m", Source: "s", Prompt: 1, Compl: 1, OK: true, Status: 200,
})
}
rr := doReq(t, g, http.MethodGet, "/api/stats", "")
var snap struct {
Records []Req `json:"records"`
Total Stat `json:"total"`
}
if err := json.Unmarshal(rr.Body.Bytes(), &snap); err != nil {
t.Fatalf("decode: %v", err)
}
if len(snap.Records) > firstScreenRecords {
t.Fatalf("snapshot shipped %d records, want <= %d", len(snap.Records), firstScreenRecords)
}
// aggregates still cover every request
if snap.Total.Reqs != 400 {
t.Fatalf("total reqs = %d, want 400 (aggregates must not be paged away)", snap.Total.Reqs)
}
}
// TestStatsCsvExportStreams checks the export still emits every row in the
// window now that it is written straight from the audit walk.
func TestStatsCsvExportStreams(t *testing.T) {
g := newTestGateway(t)
base := time.Now().UnixMilli()
for i := 0; i < 120; i++ {
g.stats.Record(Req{
Time: base + int64(i)*1000, Key: keyID("sk-test"), Type: "chat",
Model: "m", Source: "s", Prompt: 2, Compl: 3, OK: true, Status: 200,
})
}
rr := doReq(t, g, http.MethodGet,
"/api/stats?export=csv&from="+strconv.FormatInt(base, 10)+
"&to="+strconv.FormatInt(base+120*1000, 10), "")
if rr.Code != 200 {
t.Fatalf("export status = %d", rr.Code)
}
lines := strings.Split(strings.TrimSpace(rr.Body.String()), "\n")
if len(lines) != 121 { // header + 120 rows
t.Fatalf("csv had %d lines, want 121 (header + 120 records)", len(lines))
}
if !strings.HasPrefix(lines[0], "time,key,key_name") {
t.Fatalf("unexpected csv header: %q", lines[0])
}
}
// TestStatsRecordsAPIScopedToOwnKey: a non-admin key must only page its own
// records even when it asks for someone else's.
func TestStatsRecordsAPIScopedToOwnKey(t *testing.T) {
g := newTestGateway(t)
mine := keyID("sk-test")
for i := 0; i < 5; i++ {
g.stats.Record(Req{Time: time.Now().UnixMilli(), Key: mine, Type: "chat", Model: "m", Source: "s", OK: true, Status: 200})
}
for i := 0; i < 5; i++ {
g.stats.Record(Req{Time: time.Now().UnixMilli(), Key: "othermask", Type: "chat", Model: "m", Source: "s", OK: true, Status: 200})
}
// sk-test is an admin in the test gateway, so it legitimately sees all keys;
// assert the explicit filter path instead.
rr := doReq(t, g, http.MethodGet, "/api/stats/records?limit=100&key=othermask", "")
var page struct {
Records []Req `json:"records"`
}
if err := json.Unmarshal(rr.Body.Bytes(), &page); err != nil {
t.Fatalf("decode: %v", err)
}
for _, r := range page.Records {
if r.Key != "othermask" {
t.Fatalf("key filter leaked %q", r.Key)
}
}
}
// TestKeyDeleteRoundTrip is the end-to-end version of the bug users hit:
// deleting a gateway key from the admin UI. It also asserts the GET-on-a-key
// error names the right verb, since that message is what a mis-methoded client
// actually sees.
func TestKeyDeleteRoundTrip(t *testing.T) {
g := newTestGateway(t)
rr := doReq(t, g, http.MethodPost, "/api/keys", `{"name":"doomed","role":"user"}`)
if rr.Code != 200 {
t.Fatalf("create: %d %s", rr.Code, rr.Body.String())
}
var created struct {
Key struct {
Key string `json:"key"`
Name string `json:"name"`
} `json:"key"`
}
if err := json.Unmarshal(rr.Body.Bytes(), &created); err != nil {
t.Fatalf("decode create: %v", err)
}
if created.Key.Key == "" {
t.Fatalf("no key returned: %s", rr.Body.String())
}
// DELETE removes it
rr = doReq(t, g, http.MethodDelete, "/api/keys/"+url.PathEscape(created.Key.Key), "")
if rr.Code != 200 {
t.Fatalf("delete: %d %s", rr.Code, rr.Body.String())
}
for _, k := range g.core.ListKeys() {
if k.Key == created.Key.Key {
t.Fatal("key still present after DELETE")
}
}
// deleting again is a clean 404, not a 500
rr = doReq(t, g, http.MethodDelete, "/api/keys/"+url.PathEscape(created.Key.Key), "")
if rr.Code != http.StatusNotFound {
t.Fatalf("second delete: %d, want 404", rr.Code)
}
}
// TestKeyGetOnSpecificKeyNamesTheVerb: a GET on /api/keys/{key} is what a
// client that forgot to set method:"DELETE" ends up sending. The error must say
// which verb to use instead of just "use GET /api/keys", which reads as if the
// caller did nothing wrong.
func TestKeyGetOnSpecificKeyNamesTheVerb(t *testing.T) {
g := newTestGateway(t)
rr := doReq(t, g, http.MethodGet, "/api/keys/sk-whatever", "")
if rr.Code != http.StatusNotFound {
t.Fatalf("status = %d, want 404", rr.Code)
}
body := rr.Body.String()
for _, want := range []string{"DELETE /api/keys/", "PUT /api/keys/"} {
if !strings.Contains(body, want) {
t.Errorf("error message should mention %q, got: %s", want, body)
}
}
}
// TestAdapterGetOnSpecificAdapterNamesDelete is the same guard for the adapters
// endpoint, whose delAdapter helper had the identical missing-method bug.
func TestAdapterGetOnSpecificAdapterNamesDelete(t *testing.T) {
g := newTestGateway(t)
rr := doReq(t, g, http.MethodGet, "/api/adapters/openai", "")
if rr.Code != http.StatusNotFound {
t.Fatalf("status = %d, want 404", rr.Code)
}
if !strings.Contains(rr.Body.String(), "DELETE /api/adapters/openai") {
t.Errorf("error should name the DELETE verb, got: %s", rr.Body.String())
}
}
// TestCannotDeleteOwnKey keeps the admin from locking themselves out.
func TestCannotDeleteOwnKey(t *testing.T) {
g := newTestGateway(t)
rr := doReq(t, g, http.MethodDelete, "/api/keys/sk-test", "")
if rr.Code != http.StatusBadRequest {
t.Fatalf("status = %d, want 400", rr.Code)
}
if !strings.Contains(rr.Body.String(), "logged in with") {
t.Errorf("unexpected message: %s", rr.Body.String())
}
}
// TestGatewayRingStaysSmall pins the resident record ring at the small default.
// The gateway used to construct Stats with a 10000-record ring, which after the
// full-history aggregate scan meant 10000 resident Req structs — putting back the
// ~25 MB the on-demand paging work removed. Aggregates cover all of history;
// the ring only needs the first screen plus the 5-minute status windows.
func TestGatewayRingStaysSmall(t *testing.T) {
g := newTestGateway(t)
g.stats.mu.Lock()
max := g.stats.maxRecs
g.stats.mu.Unlock()
if max != defaultRingSize {
t.Fatalf("gateway record ring = %d, want %d (a big ring reintroduces the resident-memory cost)",
max, defaultRingSize)
}
}

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, []string{"sk-admin"})
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, []string{keyA.Key, keyB.Key})
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, secrets)
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, []string{"sk-a"})
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

@ -29,28 +29,26 @@ func (g *Gateway) handleKeysAPI(w http.ResponseWriter, r *http.Request) {
switch r.Method {
case http.MethodGet:
if path != "" {
writeError(w, http.StatusNotFound, "not_found", "use GET /api/keys")
// A GET on a specific key is almost always a client that meant to
// DELETE or PUT it but let fetch default to GET. Name the verbs
// instead of only pointing back at the collection endpoint.
writeError(w, http.StatusNotFound, "not_found",
"no such endpoint; use GET /api/keys to list, DELETE /api/keys/{key} to remove, PUT /api/keys/{key} to update")
return
}
writeJSON(w, http.StatusOK, map[string]interface{}{"keys": g.core.ListKeys()})
case http.MethodPost:
var body struct {
Name string `json:"name"`
Role string `json:"role"`
Models []config.ModelScope `json:"models"`
Note string `json:"note"`
Name string `json:"name"`
Role string `json:"role"`
Models []config.ModelScope `json:"models"`
Note string `json:"note"`
}
if err := json.NewDecoder(r.Body).Decode(&body); err != nil {
writeError(w, http.StatusBadRequest, "invalid_request", "invalid json: "+err.Error())
return
}
if body.Role == "" {
body.Role = "user"
}
if body.Role != "admin" && body.Role != "user" {
writeError(w, http.StatusBadRequest, "invalid_request", "role must be admin or user")
return
}
body.Role = config.NormalizeRole(body.Role)
rec, err := g.core.CreateKey(body.Name, body.Role, body.Models, body.Note)
if err != nil {
writeError(w, http.StatusBadRequest, "key_error", err.Error())
@ -72,6 +70,9 @@ func (g *Gateway) handleKeysAPI(w http.ResponseWriter, r *http.Request) {
writeError(w, http.StatusBadRequest, "invalid_request", "invalid json: "+err.Error())
return
}
// Each scope entry carries its own token/request caps, so replacing the
// scope replaces the budgets with it — there is no separate key-wide
// quota that could drift out of sync with the models.
rec, err := g.core.UpdateKey(path, body.Name, body.Role, body.Models, body.Note)
if err != nil {
writeError(w, http.StatusBadRequest, "key_error", err.Error())
@ -143,8 +144,9 @@ func (g *Gateway) allowedModels(ctx context.Context) []config.ModelScope {
func (g *Gateway) handleAutoAPI(w http.ResponseWriter, r *http.Request) {
if r.Method == http.MethodGet {
writeJSON(w, http.StatusOK, map[string]interface{}{
"rules": g.core.AutoRules(),
"states": g.core.AutoSlotStates(),
"rules": g.core.AutoRules(),
"image_rules": g.core.AutoImageRules(),
"states": g.core.AutoSlotStates(),
})
return
}
@ -155,18 +157,31 @@ func (g *Gateway) handleAutoAPI(w http.ResponseWriter, r *http.Request) {
switch r.Method {
case http.MethodPut, http.MethodPost:
var body struct {
Rules []config.ModelScope `json:"rules"`
Rules []config.ModelScope `json:"rules"`
ImageRules []config.ModelScope `json:"image_rules"`
}
if err := json.NewDecoder(r.Body).Decode(&body); err != nil {
writeError(w, http.StatusBadRequest, "invalid_request", "invalid json: "+err.Error())
return
}
if err := g.core.SaveAutoRules(body.Rules); err != nil {
writeError(w, http.StatusBadRequest, "auto_error", err.Error())
return
if body.Rules != nil {
if err := g.core.SaveAutoRules(body.Rules); err != nil {
writeError(w, http.StatusBadRequest, "auto_error", err.Error())
return
}
}
writeJSON(w, http.StatusOK, map[string]interface{}{"ok": true, "rules": g.core.AutoRules()})
if body.ImageRules != nil {
if err := g.core.SaveAutoImageRules(body.ImageRules); err != nil {
writeError(w, http.StatusBadRequest, "auto_error", err.Error())
return
}
}
writeJSON(w, http.StatusOK, map[string]interface{}{
"ok": true,
"rules": g.core.AutoRules(),
"image_rules": g.core.AutoImageRules(),
})
default:
writeError(w, http.StatusMethodNotAllowed, "method_not_allowed", "")
}
}
}

View File

@ -0,0 +1,116 @@
package gateway
import (
"encoding/json"
"net/http"
"net/http/httptest"
"strings"
"testing"
)
func TestLoginBruteForceLockout(t *testing.T) {
g := newTestGateway(t)
body := `{"key":"wrong-key-XXXX"}`
// 5 failures allowed (each answered 401), 6th+ locked out
var gotLocked bool
for i := 1; i <= 8; i++ {
req := httptest.NewRequest("POST", "/api/login", strings.NewReader(body))
rr := httptest.NewRecorder()
g.handleLoginAPI(rr, req)
if rr.Code != http.StatusUnauthorized {
t.Fatalf("attempt %d: status=%d want 401", i, rr.Code)
}
if i >= 6 {
// after lockout the response must be identical (no info leak)
if !strings.Contains(rr.Body.String(), "invalid gateway api key") {
t.Fatalf("lockout reply changed body: %s", rr.Body.String())
}
gotLocked = true
}
}
if !gotLocked {
t.Fatal("expected lockout by attempt 6")
}
// correct key from the SAME ip must also be rejected while locked
okBody := `{"key":"sk-test"}`
req := httptest.NewRequest("POST", "/api/login", strings.NewReader(okBody))
rr := httptest.NewRecorder()
g.handleLoginAPI(rr, req)
if rr.Code != http.StatusUnauthorized {
t.Fatalf("locked-out valid key: status=%d want 401", rr.Code)
}
// a different IP is unaffected
req = httptest.NewRequest("POST", "/api/login", strings.NewReader(okBody))
req.RemoteAddr = "10.9.9.9:5555"
rr = httptest.NewRecorder()
g.handleLoginAPI(rr, req)
if rr.Code != http.StatusOK {
t.Fatalf("other IP valid key: status=%d want 200 body=%s", rr.Code, rr.Body.String())
}
var resp map[string]interface{}
_ = json.Unmarshal(rr.Body.Bytes(), &resp)
if resp["ok"] != true {
t.Fatalf("login ok flag missing: %s", rr.Body.String())
}
}
func TestLoginSuccessResetsFailures(t *testing.T) {
g := newTestGateway(t)
mkBad := func() *http.Request {
return httptest.NewRequest("POST", "/api/login", strings.NewReader(`{"key":"nope"}`))
}
for i := 0; i < loginMaxFails-1; i++ { // one below the lockout threshold
rr := httptest.NewRecorder()
g.handleLoginAPI(rr, mkBad())
if rr.Code != http.StatusUnauthorized {
t.Fatal("setup failure")
}
}
// success clears the counter
rr := httptest.NewRecorder()
g.handleLoginAPI(rr, httptest.NewRequest("POST", "/api/login", strings.NewReader(`{"key":"sk-test"}`)))
if rr.Code != http.StatusOK {
t.Fatalf("valid login failed: %d", rr.Code)
}
// now up to loginMaxFails-1 more failures are still accepted before lock
for i := 0; i < loginMaxFails-1; i++ {
last := httptest.NewRecorder()
g.handleLoginAPI(last, mkBad())
}
last := httptest.NewRecorder()
g.handleLoginAPI(last, mkBad())
// count is at loginMaxFails-1 failures since reset; this next one reaches
// the threshold but the attempt itself is still evaluated normally (401),
// and the FOLLOWING one must be locked.
if last.Code != http.StatusUnauthorized {
t.Fatalf("threshold attempt should still answer 401: %d", last.Code)
}
next := httptest.NewRecorder()
g.handleLoginAPI(next, mkBad())
if next.Code != http.StatusUnauthorized {
t.Fatal("want 401")
}
// verify lock actually engaged via the guard state
g.loginMu.Lock()
f := g.loginFails["192.0.2.1"]
engaged := f != nil && f.count >= loginMaxFails && f.until > 0
g.loginMu.Unlock()
if !engaged {
t.Fatal("expected engaged lockout state")
}
}
func TestClientIPExtraction(t *testing.T) {
r := httptest.NewRequest("POST", "/api/login", nil)
r.RemoteAddr = "203.0.113.7:12345"
if got := clientIP(r); got != "203.0.113.7" {
t.Errorf("clientIP=%q want 203.0.113.7", got)
}
r.Header.Set("X-Real-IP", "198.51.100.9")
if got := clientIP(r); got != "198.51.100.9" {
t.Errorf("clientIP with X-Real-IP=%q want proxy value", got)
}
}

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

@ -11,6 +11,7 @@ import (
"fmt"
"io/fs"
"log"
"net"
"net/http"
"net/url"
"strings"
@ -31,6 +32,87 @@ type Gateway struct {
stats *Stats
probeMu sync.Mutex
lastProbe time.Time
// loginGuard throttles /api/login brute-force attempts: per-IP failure
// counters with exponential backoff, plus a global cap so a distributed
// spray cannot outrun the per-IP window. Guarded by loginMu.
loginMu sync.Mutex
loginFails map[string]*loginFail
loginWin int64 // unix sec of current global window start
loginGlob int // failures in the global window
}
// loginFail tracks consecutive failed logins from one source IP.
type loginFail struct {
count int
until int64 // locked-out until this unix second (0 = not locked)
lastAt int64
}
const (
loginMaxFails = 5 // failures before lockout kicks in
loginLockBase = 30 // first lockout: 30s
loginLockCap = 30 * 60 // repeated lockouts cap at 30min
loginDecay = 15 * 60 // counters decay after 15min of quiet
loginGlobalMax = 100 // global failures per minute across all IPs
loginGlobalWindow = 60 // global window length (sec)
)
// loginAllow checks (and records) a login attempt for ip. It returns false
// when the attempt must be rejected: the IP is in exponential lockout or the
// global failure budget for the current window is exhausted.
func (g *Gateway) loginAllow(ip string, now int64) bool {
g.loginMu.Lock()
defer g.loginMu.Unlock()
if g.loginFails == nil {
g.loginFails = map[string]*loginFail{}
}
// global window rollover
if now-g.loginWin >= loginGlobalWindow {
g.loginWin = now
g.loginGlob = 0
}
f := g.loginFails[ip]
if f != nil && f.count > 0 && now-f.lastAt > loginDecay {
delete(g.loginFails, ip) // quiet long enough: forgive and forget
f = nil
}
if f != nil && f.until > now {
return false // still locked out
}
if g.loginGlob >= loginGlobalMax {
return false // global spray budget exhausted for this window
}
g.loginGlob++
return true
}
// loginRecord notes the outcome of one attempt: failures escalate toward an
// exponential lockout; a success clears the IP's counter entirely.
func (g *Gateway) loginRecord(ip string, ok bool, now int64) {
g.loginMu.Lock()
defer g.loginMu.Unlock()
if ok {
if g.loginFails != nil {
delete(g.loginFails, ip)
}
return
}
f := g.loginFails[ip]
if f == nil {
f = &loginFail{}
g.loginFails[ip] = f
}
f.count++
f.lastAt = now
if f.count >= loginMaxFails {
// lock duration doubles with each extra burst of failures, capped
epoch := int64(f.count/loginMaxFails) - 1
dur := int64(loginLockBase) << min(epoch, 6)
if dur > loginLockCap {
dur = loginLockCap
}
f.until = now + dur
}
}
func New(c *core.Core, gatewayKeys []string) (*Gateway, error) {
@ -38,7 +120,13 @@ func New(c *core.Core, gatewayKeys []string) (*Gateway, error) {
if err != nil {
return nil, err
}
st := NewStats(10000)
// The record ring is deliberately small: aggregates are built from the full
// audit history at startup (see Stats.LoadAudit), so the ring only has to
// cover the status page's 5-minute source windows and the dashboard's first
// screen. Everything older is paged from disk by /api/stats/records.
// A large ring here would put the ~25 MB of resident Req structs straight
// back, which is exactly what the on-demand paging removed.
st := NewStats(defaultRingSize)
if cfg := c.Config(); cfg != nil && cfg.RuntimeFile != "" {
st.LoadAudit(cfg.RuntimeFile + ".audit.jsonl")
}
@ -113,6 +201,17 @@ func (s *statusRecorder) Write(b []byte) (int, error) {
return s.ResponseWriter.Write(b)
}
// Flush forwards to the underlying writer so SSE handlers can stream
// incrementally through the access-log wrapper. Without this method the
// w.(http.Flusher) assertion inside the streaming chat handlers fails (the
// embedded ResponseWriter interface does not carry Flush into the method set),
// and every chunk stays buffered until the response ends.
func (s *statusRecorder) Flush() {
if f, ok := s.ResponseWriter.(http.Flusher); ok {
f.Flush()
}
}
func (g *Gateway) routes(w http.ResponseWriter, r *http.Request) {
switch {
case r.URL.Path == "/v1/chat/completions":
@ -123,14 +222,20 @@ func (g *Gateway) routes(w http.ResponseWriter, r *http.Request) {
g.handleModels(w, r)
case r.URL.Path == "/api/adapters" || strings.HasPrefix(r.URL.Path, "/api/adapters/"):
g.handleAdaptersAPI(w, r)
case r.URL.Path == "/api/v1" || strings.HasPrefix(r.URL.Path, "/api/v1/"):
g.apiV1Routes(w, r)
case r.URL.Path == "/api/sources" || strings.HasPrefix(r.URL.Path, "/api/sources/"):
g.handleSourcesAPI(w, r)
case r.URL.Path == "/api/source_templates" || strings.HasPrefix(r.URL.Path, "/api/source_templates/"):
g.handleSourceTemplatesAPI(w, r)
case r.URL.Path == "/api/chat":
g.handleChat(w, r)
case r.URL.Path == "/api/status":
g.handleStatusAPI(w, r)
case r.URL.Path == "/api/status/reset":
g.handleResetHealth(w, r)
case r.URL.Path == "/api/stats/records":
g.handleStatsRecordsAPI(w, r)
case r.URL.Path == "/api/stats" || strings.HasPrefix(r.URL.Path, "/api/stats/"):
g.handleStatsAPI(w, r)
case r.URL.Path == "/api/keys" || strings.HasPrefix(r.URL.Path, "/api/keys/"):
@ -325,34 +430,59 @@ document.getElementById('key').addEventListener('keydown',function(e){if(e.key==
apply();
</script></body></html>`
// handleLoginAPI validates the gateway key and issues a session cookie.
// handleLoginAPI validates the gateway key and issues a session cookie,
// guarded by per-IP + global brute-force throttling (see loginGuard).
func (g *Gateway) handleLoginAPI(w http.ResponseWriter, r *http.Request) {
if r.Method != http.MethodPost {
writeError(w, http.StatusMethodNotAllowed, "method_not_allowed", "use POST")
return
}
now := time.Now().Unix()
if !g.loginAllow(clientIP(r), now) {
// Don't reveal whether the key was right: answer the same 401 the
// failure path uses, but skip the expensive key lookup.
writeError(w, http.StatusUnauthorized, "invalid_api_key", "invalid gateway api key")
return
}
var body struct {
Key string `json:"key"`
}
if err := json.NewDecoder(r.Body).Decode(&body); err != nil {
g.loginRecord(clientIP(r), false, now)
writeError(w, http.StatusBadRequest, "invalid_request", "invalid json")
return
}
if _, ok := g.core.FindKey(body.Key); !ok {
g.loginRecord(clientIP(r), false, now)
writeError(w, http.StatusUnauthorized, "invalid_api_key", "invalid gateway api key")
return
}
g.loginRecord(clientIP(r), true, now)
http.SetCookie(w, &http.Cookie{
Name: "gw_key",
Value: body.Key,
Path: "/",
MaxAge: 86400 * 30,
Name: "gw_key",
Value: body.Key,
Path: "/",
MaxAge: 86400 * 30,
HttpOnly: true,
SameSite: http.SameSiteLaxMode,
})
writeJSON(w, http.StatusOK, map[string]interface{}{"ok": true})
}
// clientIP returns the caller's IP for rate limiting. Behind a trusted
// reverse proxy (nginx) the X-Real-IP header is preferred; otherwise the
// remote address is used with the port stripped.
func clientIP(r *http.Request) string {
if x := r.Header.Get("X-Real-IP"); x != "" {
return x
}
host, _, err := net.SplitHostPort(r.RemoteAddr)
if err != nil {
return r.RemoteAddr
}
return host
}
func (g *Gateway) handleModels(w http.ResponseWriter, r *http.Request) {
if r.Method != http.MethodGet {
writeError(w, http.StatusMethodNotAllowed, "method_not_allowed", "use GET")
@ -456,8 +586,25 @@ func (g *Gateway) handleStatusAPI(w http.ResponseWriter, r *http.Request) {
"gateway_keys": ks,
}
if reqRole(r.Context()) == "admin" {
resp["sources"] = g.core.Registry().Status()
sts := g.core.Registry().Status()
recent := g.stats.SourceRecent(300)
avgs := g.stats.SourceAverages(300)
for i := range sts {
if v, ok := recent[sts[i].Name]; ok {
sts[i].RecentOK = v[0]
sts[i].RecentErr = v[1]
}
if a, ok := avgs[sts[i].Name]; ok {
sts[i].AvgFirstByteMs = a.AvgFirstByteMs
sts[i].AvgTokPerS = a.AvgTokPerS
}
}
resp["sources"] = sts
resp["adapters"] = g.core.ListAdapters()
// Elastic Lua pool sizing: created/idle/in_use vs the ceiling, plus the
// current grow/shrink steps, so the algorithm is inspectable in the UI
// instead of being a black box.
resp["adapter_pools"] = g.core.VM().PoolStats()
}
writeJSON(w, http.StatusOK, resp)
}
@ -491,4 +638,4 @@ func writeJSON(w http.ResponseWriter, code int, v interface{}) {
if err := json.NewEncoder(w).Encode(v); err != nil {
log.Printf("[gateway] write json: %v", err)
}
}
}

File diff suppressed because it is too large Load Diff

View File

@ -92,21 +92,25 @@ func TestAuditRotationRecords(t *testing.T) {
func TestLoadAuditFullReplay(t *testing.T) {
dir := t.TempDir()
path := filepath.Join(dir, "audit.jsonl")
// timestamps relative to now: buckets are whole unix hours, so window
// assertions must not depend on the minute-of-hour of the test run.
// timestamps anchored to whole unix hours (see aggregateLocked): window
// assertions must hold no matter which minute-of-hour the suite runs at,
// so rows are placed relative to hour boundaries, not raw offsets from now.
now := time.Now()
thisHour := now.Truncate(time.Hour)
twoBucketsAgo := thisHour.Add(-90 * time.Minute) // bucket H-2
prevBucketMid := thisHour.Add(-30 * time.Minute) // bucket H-1
lines := []string{
`{"obj":"access","time":1699999999000,"key":"k","method":"GET","path":"/api/stats","status":200}`,
fmt.Sprintf(`{"time":%d,"key":"k","type":"chat","model":"m","source":"s","prompt_tokens":100,"completion_tokens":50,"latency_ms":10,"ok":true,"status":200}`, now.Add(-65*time.Minute).UnixMilli()),
fmt.Sprintf(`{"time":%d,"key":"k","type":"chat","model":"m","source":"s","prompt_tokens":100,"completion_tokens":50,"latency_ms":10,"ok":true,"status":200}`, twoBucketsAgo.UnixMilli()),
`{this is not valid json`,
fmt.Sprintf(`{"time":%d,"key":"k","type":"stream","model":"m","source":"s","prompt_tokens":200,"completion_tokens":20,"latency_ms":20,"ok":false,"status":503}`, now.Add(-30*time.Minute).UnixMilli()),
fmt.Sprintf(`{"time":%d,"key":"k","type":"stream","model":"m","source":"s","prompt_tokens":200,"completion_tokens":20,"latency_ms":20,"ok":false,"status":503}`, prevBucketMid.UnixMilli()),
"garbage-not-json\n",
fmt.Sprintf(`{"time":%d,"key":"k","type":"chat","model":"m2","source":"s","prompt_tokens":7,"completion_tokens":3,"latency_ms":5,"ok":true,"status":200}`, now.Add(-30*time.Minute).UnixMilli()),
fmt.Sprintf(`{"time":%d,"key":"k","type":"chat","model":"m2","source":"s","prompt_tokens":7,"completion_tokens":3,"latency_ms":5,"ok":true,"status":200}`, prevBucketMid.UnixMilli()),
}
loaded := strings.Join(lines, "\n") + "\n" + strings.Repeat("x", 1<<18) + "\n"
// oversized row at the END proves the scanner tolerates >64KB lines and
// still finishes the replay instead of truncating silently.
loaded += fmt.Sprintf(`{"time":%d,"key":"k","type":"chat","model":"m","source":"s","prompt_tokens":1,"completion_tokens":1,"latency_ms":1,"ok":true,"status":200}`, now.Add(-time.Minute).UnixMilli()) + "\n"
loaded += fmt.Sprintf(`{"time":%d,"key":"k","type":"chat","model":"m","source":"s","prompt_tokens":1,"completion_tokens":1,"latency_ms":1,"ok":true,"status":200}`, now.UnixMilli()) + "\n"
if err := os.WriteFile(path, []byte(loaded), 0644); err != nil {
t.Fatal(err)
}
@ -126,9 +130,10 @@ func TestLoadAuditFullReplay(t *testing.T) {
t.Fatalf("ring must hold only real requests, got %d rows: %#v", len(s.recs), s.recs)
}
// quota window rebuilt from full history. All-time and multi-hour windows
// must see everything; a 1h window must NOT return all records (the old
// ms/seconds unit bug made any sec>0 window return everything), and must
// always include the row written one minute ago (current hour bucket).
// must see everything. A 1h window must NOT return all records (the old
// ms/seconds unit bug made any sec>0 window return everything): buckets
// are whole unix hours, so only rows in the current hour bucket qualify
// — exactly the just-now row (2 tokens).
if w := s.WindowTokens("m", "s", 0); w != 372 {
t.Fatalf("all-time window want 372, got %d", w)
}
@ -138,8 +143,8 @@ func TestLoadAuditFullReplay(t *testing.T) {
if w := s.WindowTokens("m", "s", 24*hourSec); w != 372 {
t.Fatalf("24h window want 372, got %d", w)
}
if w := s.WindowTokens("m", "s", hourSec); w < 2 || w >= 372 {
t.Fatalf("1h window want [2,372), got %d", w)
if w := s.WindowTokens("m", "s", hourSec); w != 2 {
t.Fatalf("1h window want 2 (current-bucket row only), got %d", w)
}
if w := s.WindowTokens("m2", "s", 24*hourSec); w != 10 {
t.Fatalf("m2 24h window want 10, got %d", w)
@ -151,4 +156,407 @@ func TestLoadAuditFullReplay(t *testing.T) {
if s.byStatus[200].Err != 0 || s.byStatus[503] == nil || s.byStatus[503].Reqs != 1 {
t.Fatalf("by_status wrong: %#v", s.byStatus)
}
}
}
// ---- on-demand log loading (plan 阶段 3) ----
// writeAuditFile writes n synthetic request rows (plus some non-request event
// rows, which must be skipped) and returns the path.
func writeAuditFile(t *testing.T, path string, n int, key string, startMs int64) {
t.Helper()
f, err := os.Create(path)
if err != nil {
t.Fatalf("create audit: %v", err)
}
defer f.Close()
for i := 0; i < n; i++ {
fmt.Fprintf(f, `{"time":%d,"key":%q,"type":"chat","model":"m","source":"s","prompt_tokens":1,"completion_tokens":1,"latency_ms":5,"ok":true,"status":200}`+"\n",
startMs+int64(i)*1000, key)
if i%10 == 0 {
// an access-log/event row: must never appear as a record
fmt.Fprintf(f, `{"obj":"access","time":%d,"path":"/v1/models"}`+"\n", startMs+int64(i)*1000)
}
}
}
// TestLoadAuditAggregatesFullHistory is the correctness half of the log work:
// the dashboard's totals must cover EVERY audit record, however large the file,
// while the in-memory record ring stays bounded.
func TestLoadAuditAggregatesFullHistory(t *testing.T) {
dir := t.TempDir()
path := filepath.Join(dir, "audit.jsonl")
const n = 20000
writeAuditFile(t, path, n, "k1", 1_700_000_000_000)
s := NewStats(0)
s.LoadAudit(path)
if s.ReplayPartial() {
t.Fatal("a readable history must not be reported as partial")
}
snap := s.Snapshot(0, "")
total, _ := snap["total"].(Stat)
if total.Reqs != n {
t.Fatalf("total reqs = %d, want %d (aggregates must cover the whole file)", total.Reqs, n)
}
if total.Tokens != int64(n)*2 { // 1 prompt + 1 completion per row
t.Fatalf("total tokens = %d, want %d", total.Tokens, n*2)
}
// ...while the resident ring is capped and the snapshot ships one screen
recs, _ := snap["records"].([]Req)
if len(recs) == 0 {
t.Fatal("the newest records must stay resident")
}
if len(recs) > firstScreenRecords {
t.Fatalf("snapshot shipped %d records, must be capped at the first screen (%d)",
len(recs), firstScreenRecords)
}
s.mu.Lock()
ring := len(s.recs)
s.mu.Unlock()
if ring > defaultRingSize {
t.Fatalf("resident ring holds %d records, must be capped at %d", ring, defaultRingSize)
}
// the ring must hold the NEWEST records, not the first ones scanned
newest := recs[len(recs)-1]
wantNewest := int64(1_700_000_000_000) + int64(n-1)*1000
if newest.Time != wantNewest {
t.Fatalf("newest resident record time = %d, want %d", newest.Time, wantNewest)
}
}
// TestLoadAuditAggregatesAcrossRotatedFiles: history split over rotated .old
// files must still add up, and must be replayed oldest-first so the hourly
// quota buckets keep the intended trailing window.
func TestLoadAuditAggregatesAcrossRotatedFiles(t *testing.T) {
dir := t.TempDir()
path := filepath.Join(dir, "audit.jsonl")
writeAuditFile(t, path, 40, "k1", 1_700_000_500_000) // newest
writeAuditFile(t, path+".100.old", 40, "k1", 1_700_000_300_000) // older
writeAuditFile(t, path+".050.old", 40, "k1", 1_700_000_100_000) // oldest
s := NewStats(0)
s.LoadAudit(path)
snap := s.Snapshot(0, "")
total, _ := snap["total"].(Stat)
if total.Reqs != 120 {
t.Fatalf("total reqs = %d, want 120 across three files", total.Reqs)
}
recs, _ := snap["records"].([]Req)
if len(recs) == 0 {
t.Fatal("expected resident records")
}
// oldest-first replay + newest-last ring => the last resident record is the
// newest one in the live file
if got := recs[len(recs)-1].Time; got != 1_700_000_500_000+39*1000 {
t.Fatalf("newest resident record time = %d, want the live file's last row", got)
}
}
// TestReqRingKeepsNewest pins the ring semantics used by the full-history scan:
// bounded size, newest retained, oldest-first output.
func TestReqRingKeepsNewest(t *testing.T) {
r := newReqRing(3)
for i := 1; i <= 7; i++ {
r.push(Req{Time: int64(i)})
}
got := r.slice()
if len(got) != 3 {
t.Fatalf("ring len = %d, want 3", len(got))
}
for i, want := range []int64{5, 6, 7} {
if got[i].Time != want {
t.Fatalf("ring[%d] = %d, want %d (oldest-first, newest retained)", i, got[i].Time, want)
}
}
// a partially filled ring returns exactly what was pushed
r2 := newReqRing(5)
r2.push(Req{Time: 1})
r2.push(Req{Time: 2})
if got := r2.slice(); len(got) != 2 || got[0].Time != 1 || got[1].Time != 2 {
t.Fatalf("partial ring = %+v", got)
}
}
// TestLoadAuditFullFileWhenSmall: a small file is loaded completely and is not
// flagged partial.
func TestLoadAuditFullFileWhenSmall(t *testing.T) {
path := filepath.Join(t.TempDir(), "audit.jsonl")
writeAuditFile(t, path, 50, "k1", 1_700_000_000_000)
s := NewStats(0)
s.LoadAudit(path)
if s.ReplayPartial() {
t.Fatal("a small file must be replayed in full, not flagged partial")
}
snap := s.Snapshot(0, "")
total, _ := snap["total"].(Stat)
if total.Reqs != 50 {
t.Fatalf("total reqs = %d, want 50", total.Reqs)
}
}
// TestQuotaWindowsSurviveBoundedRing is the regression guard the plan calls out:
// capping the resident record ring must not break token-quota accounting, which
// reads the hour buckets rather than the ring.
func TestQuotaWindowsSurviveBoundedRing(t *testing.T) {
s := NewStats(10) // deliberately tiny ring
now := time.Now().UnixMilli()
for i := 0; i < 100; i++ {
s.Record(Req{Time: now, Key: "k", Model: "m", Source: "s", Prompt: 2, Compl: 3, OK: true, Status: 200})
}
if got := s.WindowTokens("m", "s", 24*hourSec); got != 500 {
t.Fatalf("window tokens = %d, want 500 (quota accounting must not depend on the ring)", got)
}
if got := s.ModelTokens("k")["m"]; got != 500 {
t.Fatalf("model tokens = %d, want 500", got)
}
if got := s.KeyTokens("k"); got != 500 {
t.Fatalf("key tokens = %d, want 500", got)
}
}
// TestAuditPagePagesBackwards walks the whole history one page at a time and
// asserts strict newest-first order with no gaps or repeats.
func TestAuditPagePagesBackwards(t *testing.T) {
path := filepath.Join(t.TempDir(), "audit.jsonl")
const n = 250
writeAuditFile(t, path, n, "k1", 1_700_000_000_000)
s := NewStats(0)
s.LoadAudit(path)
var seen []Req
cursor := ""
for i := 0; i < 20; i++ { // bounded loop guards against a cursor that never advances
page := s.AuditPage(cursor, 40, "")
if len(page.Records) == 0 {
break
}
seen = append(seen, page.Records...)
if !page.HasMore {
break
}
if page.Next == cursor {
t.Fatal("cursor did not advance")
}
cursor = page.Next
}
if len(seen) != n {
t.Fatalf("paged %d records, want %d", len(seen), n)
}
for i := 1; i < len(seen); i++ {
if seen[i-1].Time < seen[i].Time {
t.Fatalf("records not newest-first at %d: %d then %d", i, seen[i-1].Time, seen[i].Time)
}
}
// newest record first, oldest last
if seen[0].Time != 1_700_000_000_000+int64(n-1)*1000 {
t.Fatalf("first record time = %d, want the newest", seen[0].Time)
}
if seen[len(seen)-1].Time != 1_700_000_000_000 {
t.Fatalf("last record time = %d, want the oldest", seen[len(seen)-1].Time)
}
}
// TestAuditPageFirstScreenIsSmall: the first page must be exactly one screen,
// independent of how large the audit file is.
func TestAuditPageFirstScreenIsSmall(t *testing.T) {
path := filepath.Join(t.TempDir(), "audit.jsonl")
writeAuditFile(t, path, 5000, "k1", 1_700_000_000_000)
s := NewStats(0)
s.LoadAudit(path)
page := s.AuditPage("", 100, "")
if len(page.Records) != 100 {
t.Fatalf("first page = %d records, want 100", len(page.Records))
}
if !page.HasMore || page.Next == "" {
t.Fatal("a large file must report more pages plus a cursor")
}
}
// TestAuditPageFiltersByKey: a user key may only page its own records.
func TestAuditPageFiltersByKey(t *testing.T) {
dir := t.TempDir()
path := filepath.Join(dir, "audit.jsonl")
writeAuditFile(t, path, 30, "mine", 1_700_000_000_000)
f, err := os.OpenFile(path, os.O_APPEND|os.O_WRONLY, 0644)
if err != nil {
t.Fatal(err)
}
for i := 0; i < 30; i++ {
fmt.Fprintf(f, `{"time":%d,"key":"other","type":"chat","model":"m","source":"s","ok":true,"status":200}`+"\n",
1_700_000_100_000+int64(i)*1000)
}
f.Close()
s := NewStats(0)
s.LoadAudit(path)
page := s.AuditPage("", 1000, "mine")
if len(page.Records) != 30 {
t.Fatalf("filtered page = %d records, want 30", len(page.Records))
}
for _, r := range page.Records {
if r.Key != "mine" {
t.Fatalf("key filter leaked record for %q", r.Key)
}
}
}
// TestAuditPageSpansRotatedFiles: paging must continue into .old files instead
// of stopping at the live file's start.
func TestAuditPageSpansRotatedFiles(t *testing.T) {
dir := t.TempDir()
path := filepath.Join(dir, "audit.jsonl")
writeAuditFile(t, path, 40, "k1", 1_700_000_500_000) // newest
writeAuditFile(t, path+".100.old", 40, "k1", 1_700_000_300_000) // older
writeAuditFile(t, path+".050.old", 40, "k1", 1_700_000_100_000) // oldest
s := NewStats(0)
s.LoadAudit(path)
var seen []Req
cursor := ""
for i := 0; i < 30; i++ {
page := s.AuditPage(cursor, 25, "")
seen = append(seen, page.Records...)
if !page.HasMore {
break
}
cursor = page.Next
}
if len(seen) != 120 {
t.Fatalf("paged %d records across rotated files, want 120", len(seen))
}
for i := 1; i < len(seen); i++ {
if seen[i-1].Time < seen[i].Time {
t.Fatalf("rotation boundary broke the ordering at %d", i)
}
}
}
// TestAuditPageRotatedCursor: a cursor whose file rotated away must restart from
// the newest page and say so, rather than silently returning nothing.
func TestAuditPageRotatedCursor(t *testing.T) {
path := filepath.Join(t.TempDir(), "audit.jsonl")
writeAuditFile(t, path, 20, "k1", 1_700_000_000_000)
s := NewStats(0)
s.LoadAudit(path)
page := s.AuditPage("/gone/audit.jsonl:12345", 10, "")
if !page.Rotated {
t.Fatal("a stale cursor must be reported as rotated")
}
if len(page.Records) == 0 {
t.Fatal("a stale cursor must fall back to the newest page")
}
}
// TestAuditPageEmptyAndNoFile cover the degenerate inputs.
func TestAuditPageEmptyAndNoFile(t *testing.T) {
s := NewStats(0)
if page := s.AuditPage("", 10, ""); len(page.Records) != 0 || page.HasMore {
t.Fatalf("no audit path must yield an empty page, got %+v", page)
}
path := filepath.Join(t.TempDir(), "audit.jsonl")
if err := os.WriteFile(path, nil, 0644); err != nil {
t.Fatal(err)
}
s2 := NewStats(0)
s2.LoadAudit(path)
if page := s2.AuditPage("", 10, ""); len(page.Records) != 0 || page.HasMore {
t.Fatalf("empty audit file must yield an empty page, got %+v", page)
}
}
// TestAuditPageHandlesLongLines: a record larger than the reverse-read chunk
// must still be decoded (the reader carries partial lines across chunks).
func TestAuditPageHandlesLongLines(t *testing.T) {
path := filepath.Join(t.TempDir(), "audit.jsonl")
f, err := os.Create(path)
if err != nil {
t.Fatal(err)
}
big := strings.Repeat("e", revChunk+1024) // one record spanning >1 chunk
fmt.Fprintf(f, `{"time":1,"key":"k","type":"chat","model":"m","source":"s","ok":false,"status":500,"error":%q}`+"\n", big)
fmt.Fprintf(f, `{"time":2,"key":"k","type":"chat","model":"m","source":"s","ok":true,"status":200}`+"\n")
f.Close()
s := NewStats(0)
s.LoadAudit(path)
page := s.AuditPage("", 10, "")
if len(page.Records) != 2 {
t.Fatalf("got %d records, want 2 (a >chunk record must survive reverse reads)", len(page.Records))
}
if page.Records[0].Time != 2 || page.Records[1].Time != 1 {
t.Fatalf("ordering broken: %d then %d", page.Records[0].Time, page.Records[1].Time)
}
if len(page.Records[1].Err) != len(big) {
t.Fatalf("oversized field truncated: %d bytes, want %d", len(page.Records[1].Err), len(big))
}
}
// TestStreamAuditRecordsIsWindowed: export streaming must honour the time window
// and the key filter, and must hand records over one at a time.
func TestStreamAuditRecordsIsWindowed(t *testing.T) {
path := filepath.Join(t.TempDir(), "audit.jsonl")
writeAuditFile(t, path, 100, "k1", 1_700_000_000_000)
s := NewStats(0)
s.LoadAudit(path)
from := int64(1_700_000_010_000)
to := int64(1_700_000_019_000)
var got []Req
if err := s.StreamAuditRecords(from, to, "", func(r Req) error {
got = append(got, r)
return nil
}); err != nil {
t.Fatalf("stream: %v", err)
}
if len(got) != 10 {
t.Fatalf("streamed %d records, want 10 in [%d,%d]", len(got), from, to)
}
for _, r := range got {
if r.Time < from || r.Time > to {
t.Fatalf("record outside the window: %d", r.Time)
}
}
}
// TestStreamAuditRecordsAbortStops: a callback error (client disconnect) must
// stop the walk immediately instead of reading the rest of the history.
func TestStreamAuditRecordsAbortStops(t *testing.T) {
path := filepath.Join(t.TempDir(), "audit.jsonl")
writeAuditFile(t, path, 500, "k1", 1_700_000_000_000)
s := NewStats(0)
s.LoadAudit(path)
seen := 0
stop := fmt.Errorf("client gone")
err := s.StreamAuditRecords(0, 0, "", func(Req) error {
seen++
if seen == 5 {
return stop
}
return nil
})
if err != stop {
t.Fatalf("abort must surface the callback error, got %v", err)
}
if seen != 5 {
t.Fatalf("walk continued after abort: %d records", seen)
}
}
// TestAuditRotationThresholds pins the smaller-file/more-files trade-off that
// makes reverse paging cheap while keeping the same total budget.
func TestAuditRotationThresholds(t *testing.T) {
if auditRotateBytes != 16<<20 {
t.Fatalf("auditRotateBytes = %d, want 16 MiB", auditRotateBytes)
}
if auditKeepOld != 16 {
t.Fatalf("auditKeepOld = %d, want 16", auditKeepOld)
}
if total := auditRotateBytes * int64(auditKeepOld); total != 256<<20 {
t.Fatalf("total audit budget = %d, want 256 MiB", total)
}
}

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

File diff suppressed because it is too large Load Diff

View File

@ -0,0 +1,406 @@
package gateway
import (
"fmt"
"regexp"
"strings"
"testing"
)
// The WebUI talks to this package through a tiny `api(path, options)` helper
// that wraps fetch(). fetch() defaults to GET when no method is given, so an
// options object without `method` silently turns a DELETE/PUT into a GET —
// the request hits the wrong branch of the handler and the user sees a
// confusing error like "use GET /api/keys" while nothing is deleted.
//
// This has now bitten twice: ab20f1b fixed delSource/delTemplate, and the very
// same pattern was still live in delKey (deleting a gateway key was impossible)
// and delAdapter. Reviewing 4000+ lines of inline JS by eye clearly does not
// catch it, so it is pinned here instead.
// apiCallRe finds the start of every api(...) call in the UI source.
var apiCallRe = regexp.MustCompile(`\bapi\(`)
// uiSource returns the embedded WebUI document.
func uiSource(t *testing.T) string {
t.Helper()
b, err := uiFS.ReadFile("ui/index.html")
if err != nil {
t.Fatalf("read embedded ui: %v", err)
}
return string(b)
}
// extractCall returns the source of the balanced api(...) call starting at the
// opening parenthesis index, plus the line number it starts on.
func extractCall(src string, openIdx int) (string, bool) {
depth := 0
for i := openIdx; i < len(src); i++ {
switch src[i] {
case '(':
depth++
case ')':
depth--
if depth == 0 {
return src[openIdx : i+1], true
}
case '\n':
// A call spanning more than ~25 lines is not one of ours; bail out
// rather than scanning the rest of the document.
if i-openIdx > 4000 {
return "", false
}
}
}
return "", false
}
func lineOf(src string, idx int) int {
return strings.Count(src[:idx], "\n") + 1
}
// TestUIAPICallsDeclareMethod asserts that every api() call passing an options
// object also declares an HTTP method (or is a GET that only passes an
// AbortSignal). Without this, fetch defaults to GET and mutating endpoints are
// never reached.
func TestUIAPICallsDeclareMethod(t *testing.T) {
src := uiSource(t)
var offenders []string
for _, m := range apiCallRe.FindAllStringIndex(src, -1) {
open := m[1] - 1 // index of '('
call, ok := extractCall(src, open)
if !ok {
continue
}
// No options object at all => a plain GET, which is fine.
if !strings.Contains(call, "{") {
continue
}
if strings.Contains(call, "method:") {
continue
}
// A read that only carries an AbortSignal is a deliberate GET.
if strings.Contains(call, "signal:") && !strings.Contains(call, "body:") {
continue
}
offenders = append(offenders, fmt.Sprintf("line %d: %s", lineOf(src, open), oneLine(call)))
}
if len(offenders) > 0 {
t.Fatalf("api() called with an options object but no method (fetch will send GET):\n %s",
strings.Join(offenders, "\n "))
}
}
// TestUIDeleteHelpersUseDelete pins the specific helpers that remove things.
// Each must issue a DELETE; a GET here is the exact bug users reported as
// "cannot delete key: use GET /api/keys".
func TestUIDeleteHelpersUseDelete(t *testing.T) {
src := uiSource(t)
for _, fn := range []string{"delKey", "delAdapter", "delSource", "delTemplate"} {
body, ok := jsFunctionBody(src, fn)
if !ok {
t.Errorf("helper %s() not found in the WebUI", fn)
continue
}
if !strings.Contains(body, "api(") {
t.Errorf("%s() does not call api()", fn)
continue
}
if !strings.Contains(body, `method: "DELETE"`) {
t.Errorf("%s() must issue a DELETE, got:\n%s", fn, body)
}
}
}
// TestUIMutatingHelpersUseWriteMethods covers the write (non-delete) helpers:
// each must declare an explicit POST/PUT/PATCH.
func TestUIMutatingHelpersUseWriteMethods(t *testing.T) {
src := uiSource(t)
want := map[string][]string{
"createKey": {`method: "POST"`},
"putScope": {`method: "PUT"`},
"uploadAdapter": {`method: "POST"`},
"saveAuto": {`method: "PUT"`, `method: "POST"`},
}
for fn, methods := range want {
body, ok := jsFunctionBody(src, fn)
if !ok {
continue // helper renamed or absent: the method test above still guards it
}
found := false
for _, m := range methods {
if strings.Contains(body, m) {
found = true
break
}
}
if !found {
t.Errorf("%s() must declare one of %v", fn, methods)
}
}
}
// jsFunctionBody returns the source of `function name(` / `async function name(`
// up to its closing brace, using brace balancing.
func jsFunctionBody(src, name string) (string, bool) {
for _, prefix := range []string{"async function " + name + "(", "function " + name + "("} {
idx := strings.Index(src, prefix)
if idx < 0 {
continue
}
open := strings.Index(src[idx:], "{")
if open < 0 {
continue
}
open += idx
depth := 0
for i := open; i < len(src); i++ {
switch src[i] {
case '{':
depth++
case '}':
depth--
if depth == 0 {
return src[idx : i+1], true
}
}
}
}
return "", false
}
func oneLine(s string) string {
s = strings.Join(strings.Fields(s), " ")
if len(s) > 140 {
return s[:140] + "…"
}
return s
}
// TestUIRecordsPagingWiring guards the record table's paging plumbing. Two
// separate bugs made "scroll to the bottom" silently stop loading:
//
// 1. paintRecords rebuilt the whole table on every 5s poll whenever the row
// count did not exceed the first screen, wiping loaded pages and resetting
// scroll position.
// 2. IntersectionObserver only fires on TRANSITIONS. With a short list — or a
// page whose rows all duplicated the first screen — the sentinel stayed
// visible and never fired again.
//
// The fixes are a build-once flag, a scroll-position fallback, an initial
// viewport fill, and chaining when a page yields no new rows. None of that is
// reachable from Go, so the wiring is asserted structurally.
func TestUIRecordsPagingWiring(t *testing.T) {
src := uiSource(t)
paint, ok := jsFunctionBody(src, "paintRecords")
if !ok {
t.Fatal("paintRecords() not found")
}
if !strings.Contains(paint, "recsState.built") {
t.Error("paintRecords must build the table once (recsState.built) instead of rebuilding on every poll")
}
if !strings.Contains(paint, "afterbegin") {
t.Error("paintRecords must PREPEND newer rows on later polls, not rebuild the table")
}
if !strings.Contains(paint, "fillRecordsViewport") {
t.Error("paintRecords must fill the viewport so the sentinel can transition")
}
attach, ok := jsFunctionBody(src, "attachRecsObserver")
if !ok {
t.Fatal("attachRecsObserver() not found")
}
if !strings.Contains(attach, "IntersectionObserver") {
t.Error("attachRecsObserver must still use IntersectionObserver")
}
if !strings.Contains(attach, `addEventListener("scroll"`) {
t.Error("attachRecsObserver needs a scroll fallback: an already-visible sentinel never fires again")
}
load, ok := jsFunctionBody(src, "loadMoreRecords")
if !ok {
t.Fatal("loadMoreRecords() not found")
}
if !strings.Contains(load, "maxChain") {
t.Error("loadMoreRecords must chain when a page adds no new rows (the first page overlaps the first screen)")
}
if !strings.Contains(load, "signal: ctrl.signal") {
t.Error("loadMoreRecords must remain abortable")
}
fill, ok := jsFunctionBody(src, "fillRecordsViewport")
if !ok {
t.Fatal("fillRecordsViewport() not found")
}
for _, want := range []string{"scrollHeight", "clientHeight"} {
if !strings.Contains(fill, want) {
t.Errorf("fillRecordsViewport must compare %s to decide whether the list overflows", want)
}
}
rel, ok := jsFunctionBody(src, "releaseRecords")
if !ok {
t.Fatal("releaseRecords() not found")
}
for _, want := range []string{"observer.disconnect()", "removeEventListener", "abort()", "recsState.built = false"} {
if !strings.Contains(rel, want) {
t.Errorf("releaseRecords must clean up %s", want)
}
}
}
// TestUIThemeMatrix pins the theme system: every accent must define both a light
// and a dark variable block, so switching light/dark can never leave a theme
// half-defined (which shows up as unreadable text rather than as an error).
func TestUIThemeMatrix(t *testing.T) {
src := uiSource(t)
accents := []string{"sakura", "ocean", "violet", "mono"}
for _, a := range accents {
light := fmt.Sprintf(`html[data-accent=%q]{`, a)
dark := fmt.Sprintf(`html[data-theme="dark"][data-accent=%q]{`, a)
if !strings.Contains(src, light) {
t.Errorf("accent %q has no light-mode variable block (%s)", a, light)
}
// sakura is the default palette: its light values live in :root, and the
// generic dark block covers it.
if a != "sakura" && !strings.Contains(src, dark) {
t.Errorf("accent %q has no dark-mode variable block (%s)", a, dark)
}
if !strings.Contains(src, fmt.Sprintf(`data-accent="%s" title=`, a)) {
t.Errorf("accent %q has no picker button", a)
}
}
}
// TestUIMonoThemeIsFlat asserts the mono theme really is pure white / pure black:
// the surfaces must be opaque and the glass/blob decoration must be switched off,
// otherwise "#ffffff" renders as a grey translucent card over a gradient mesh.
func TestUIMonoThemeIsFlat(t *testing.T) {
src := uiSource(t)
lightBlock, ok := cssBlock(src, `html[data-accent="mono"]{`)
if !ok {
t.Fatal("mono light block not found")
}
for _, want := range []string{"--bg-s1:#ffffff", "--card:#ffffff", "--glass:0px"} {
if !strings.Contains(lightBlock, want) {
t.Errorf("mono light theme must set %s, got:\n%s", want, lightBlock)
}
}
if !strings.Contains(lightBlock, "--blob1:transparent") {
t.Error("mono light theme must disable the gradient blobs")
}
darkBlock, ok := cssBlock(src, `html[data-theme="dark"][data-accent="mono"]{`)
if !ok {
t.Fatal("mono dark block not found")
}
for _, want := range []string{"--bg-s1:#000000", "--glass:0px"} {
if !strings.Contains(darkBlock, want) {
t.Errorf("mono dark theme must set %s, got:\n%s", want, darkBlock)
}
}
if !strings.Contains(darkBlock, "--blob1:transparent") {
t.Error("mono dark theme must disable the gradient blobs")
}
// the flat-surface overrides must exist, or backdrop-filter turns #fff grey
for _, want := range []string{
`html[data-accent="mono"] #bgfx .blob{display:none}`,
`backdrop-filter:none`,
} {
if !strings.Contains(src, want) {
t.Errorf("mono theme needs the flat-surface override %q", want)
}
}
}
// TestUICanvasColorsResolveVars guards a class of silent theme bug: Canvas 2D
// does not understand CSS var(), so assigning "var(--x)" to strokeStyle is a
// no-op that leaves the previous colour (often black) in place — invisible on a
// dark background. Theme-driven canvas colours must go through cssVar().
func TestUICanvasColorsResolveVars(t *testing.T) {
src := uiSource(t)
if !strings.Contains(src, "function cssVar(") {
t.Fatal("cssVar() helper missing: canvas colours cannot resolve CSS variables without it")
}
for _, prop := range []string{"strokeStyle", "fillStyle"} {
bad := fmt.Sprintf(`ctx.%s = "var(--`, prop)
if strings.Contains(src, bad) {
t.Errorf("ctx.%s assigned a raw CSS var(): Canvas ignores it silently, use cssVar()", prop)
}
}
}
// cssBlock returns the text of the CSS rule starting at selector, up to its
// closing brace.
func cssBlock(src, selector string) (string, bool) {
i := strings.Index(src, selector)
if i < 0 {
return "", false
}
j := strings.Index(src[i:], "}")
if j < 0 {
return "", false
}
return src[i : i+j+1], true
}
// The WebUI ships two locale objects (`zh` and `en`) that every rendered string
// goes through. A key added to only one of them does not error: t() falls back
// to `t("key") || "literal"` in some call sites and to the bare key string in
// others, so the user sees either a hardcoded local string or a raw key name.
// Nothing in the Go test suite noticed — this was found by hand while adding
// thImgs. Pin the key sets so a one-sided edit fails here instead of shipping.
func TestUILocaleKeyParity(t *testing.T) {
src := uiSource(t)
iZh := strings.Index(src, "zh: {")
iEn := strings.Index(src, "en: {")
if iZh < 0 || iEn < 0 || iEn < iZh {
t.Fatalf("locale blocks not found (zh=%d en=%d)", iZh, iEn)
}
zh := src[iZh:iEn]
en := src[iEn:]
// End the EN block at its closing brace, so keys from later objects
// (config templates etc.) do not pollute the comparison.
if end := strings.Index(en, "\n },"); end > 0 {
en = en[:end]
}
keyRe := regexp.MustCompile(`(?m)^\s*([A-Za-z_][A-Za-z0-9_]*)\s*:`)
collect := func(seg string) map[string]bool {
m := map[string]bool{}
for _, g := range keyRe.FindAllStringSubmatch(seg, -1) {
m[g[1]] = true
}
return m
}
zk, ek := collect(zh), collect(en)
// The locale markers themselves are the two block headers.
delete(zk, "zh")
delete(ek, "en")
var missingEN, missingZH []string
for k := range zk {
if !ek[k] {
missingEN = append(missingEN, k)
}
}
for k := range ek {
if !zk[k] {
missingZH = append(missingZH, k)
}
}
if len(missingEN) > 0 {
t.Errorf("translation keys present in zh but missing in en: %v", missingEN)
}
if len(missingZH) > 0 {
t.Errorf("translation keys present in en but missing in zh: %v", missingZH)
}
if len(zk) < 100 {
t.Errorf("only %d zh keys parsed — the block boundary regexp has drifted "+
"and this test is no longer checking anything", len(zk))
}
}

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

@ -0,0 +1,212 @@
local adapter = {}
adapter.name = "agentrouter"
adapter.version = "1.0.0"
adapter.endpoint = "/chat/completions"
adapter.headers = {}
-- AgentRouter 的阿里云 WAF 按三重维度校验:出口 IP(必须海外)、
-- TLS 指纹(JS/Go HTTP 层与 curl 不同)、浏览器指纹 header 集。
-- 纯 UA 已不够:WAF 拦 12 种 agent CLI UA 中的 10 种(全 405 HTML),
-- 只有同时带 Origin: qwen.ai + Referer + X-Requested-With 才放行,
-- 且同 IP 短时间连续 3+ 次请求会被速率封禁,靠 adapter 侧低频调用 +
-- llmsproxy 的 prefFailStep 冷却自然限制,来源无法控制速率。
local default_ua = "QwenCode/0.2.0 (linux; x64)"
-- AgentRouter fronts Claude models (claude-opus-4-8), and Claude upstreams
-- enforce Anthropic's tool id rule ^[a-zA-Z0-9_-]{1,64}$ by rejecting the WHOLE
-- request. Plain OpenAI does not, so an id minted by a permissive model (e.g.
-- "bash:0" from moonshotai/kimi-k3) is replayed here by the client and kills
-- every turn. See anthropic.lua / openai.lua for the same helper: Lua adapters
-- have no shared prelude, so each carries its own copy.
local TOOL_ID_MAX = 64
local function safe_tool_id(id)
if type(id) ~= "string" or id == "" then return id end
local clean = string.gsub(id, "[^A-Za-z0-9_-]", "_")
if clean == id and #clean <= TOOL_ID_MAX then
return clean
end
local digest = string.sub(sha256_hex(id), 1, 8)
local keep = TOOL_ID_MAX - #digest - 1
if #clean > keep then clean = string.sub(clean, 1, keep) end
return clean .. "_" .. digest
end
function adapter.transform_request(raw_body)
local ok, req = pcall(json.decode, raw_body)
if not ok then return raw_body end
req.disable_thinking = nil
req.extra_body = nil
if req.messages then
for _, msg in ipairs(req.messages) do
msg.reasoning_content = nil
if msg.tool_call_id ~= nil then
msg.tool_call_id = safe_tool_id(msg.tool_call_id)
end
if type(msg.tool_calls) == "table" then
for _, tc in ipairs(msg.tool_calls) do
if type(tc) == "table" and tc.id ~= nil then
tc.id = safe_tool_id(tc.id)
end
end
end
end
end
return json.encode(req)
end
-- 注入 Authorization + 官方客户端 User-Agent
function adapter.build_headers(meta)
local ua = default_ua
local meta_ua = (meta.source or {}).meta and (meta.source).meta.user_agent
if meta_ua and meta_ua ~= "" then
ua = meta_ua
end
return {
["Content-Type"] = "application/json",
["Authorization"] = "Bearer " .. meta.api_key,
["User-Agent"] = ua,
-- QwenCode client fingerprint: WAF rejects bare UA without these.
["Accept"] = "application/json, text/plain, */*",
["Accept-Language"] = "zh-CN,zh;q=0.9",
["Origin"] = "https://qwen.ai",
["Referer"] = "https://qwen.ai/",
["X-Requested-With"] = "XMLHttpRequest",
["Connection"] = "keep-alive",
}
end
function adapter.transform_response(raw_body)
local ok, resp = pcall(json.decode, raw_body)
if not ok or resp == nil then return raw_body end
local unified = {
content = "",
finish_reason = "",
token_usage = { prompt = 0, completion = 0, total = 0 }
}
if type(resp.usage) == "table" then
unified.token_usage.prompt = resp.usage.prompt_tokens or 0
unified.token_usage.completion = resp.usage.completion_tokens or 0
unified.token_usage.total = resp.usage.total_tokens or 0
local hit = 0
if type(resp.usage.prompt_tokens_details) == "table" and resp.usage.prompt_tokens_details.cached_tokens ~= nil then
hit = resp.usage.prompt_tokens_details.cached_tokens
unified.token_usage.prompt_tokens_details = { cached_tokens = hit }
end
if (resp.usage.prompt_cache_hit_tokens or 0) > 0 then
unified.token_usage.prompt_cache_hit_tokens = resp.usage.prompt_cache_hit_tokens
unified.token_usage.prompt_cache_miss_tokens = resp.usage.prompt_cache_miss_tokens or 0
if hit == 0 then
unified.token_usage.prompt_tokens_details = { cached_tokens = resp.usage.prompt_cache_hit_tokens }
end
end
end
if type(resp.choices) == "table" and #resp.choices > 0 then
local ch = resp.choices[1]
if type(ch.message) == "table" then
unified.content = ch.message.content or ""
if ch.message.reasoning_content then
unified.reasoning_content = ch.message.reasoning_content
end
if type(ch.message.tool_calls) == "table" then
local tcs = {}
for _, tc in ipairs(ch.message.tool_calls) do
local args_ok, args = pcall(json.decode, tc["function"].arguments)
if not args_ok then args = {} end
table.insert(tcs, {
id = safe_tool_id(tc.id),
type = tc.type or "function",
name = tc["function"].name,
arguments = args
})
end
unified.tool_calls = tcs
end
end
unified.finish_reason = ch.finish_reason or ""
end
return json.encode(unified)
end
function adapter.transform_stream_chunk(raw_chunk)
local ok, chunk = pcall(json.decode, raw_chunk)
if not ok then return "" end
-- usage-only terminal chunk (empty choices + usage): forward it so the
-- gateway emits exact token counts (and cache fields when present).
local uses = nil
if type(chunk.usage) == "table" then
uses = {
prompt = chunk.usage.prompt_tokens or chunk.usage.prompt or 0,
completion = chunk.usage.completion_tokens or chunk.usage.completion or 0,
total = chunk.usage.total_tokens or chunk.usage.total or 0,
}
if type(chunk.usage.prompt_tokens_details) == "table" and chunk.usage.prompt_tokens_details.cached_tokens ~= nil then
uses.prompt_tokens_details = { cached_tokens = chunk.usage.prompt_tokens_details.cached_tokens }
elseif (chunk.usage.prompt_cache_hit_tokens or 0) > 0 then
uses.prompt_cache_hit_tokens = chunk.usage.prompt_cache_hit_tokens
uses.prompt_cache_miss_tokens = chunk.usage.prompt_cache_miss_tokens or 0
uses.prompt_tokens_details = { cached_tokens = chunk.usage.prompt_cache_hit_tokens }
end
end
if not chunk.choices or #chunk.choices == 0 then
if uses ~= nil then
return json.encode({ usage = uses, done = false })
end
return ""
end
local delta = chunk.choices[1].delta or {}
local fr = chunk.choices[1].finish_reason
local finish = (type(fr) == "string" and fr ~= "") and fr or nil
local unified = {
content = delta.content or "",
done = (finish ~= nil)
}
if finish then
unified.finish_reason = finish
end
if delta.reasoning_content then
unified.reasoning_content = delta.reasoning_content
end
if delta.tool_calls then
-- Sanitize outbound too: a dirty id must never enter a client session,
-- because the client replays it to every other source. Only the first
-- fragment of a streamed call carries an id; later argument fragments
-- have none and must stay id-less for index-based accumulation.
for _, tc in ipairs(delta.tool_calls) do
if type(tc) == "table" and tc.id ~= nil then
tc.id = safe_tool_id(tc.id)
end
end
unified.tool_calls = delta.tool_calls
end
if uses ~= nil then
unified.usage = uses
end
return json.encode(unified)
end
-- 错误收敛:前置 WAF 会返回整页 HTML(阿里云盾),JSON 时为 new-api 风格
function adapter.transform_error(status, body)
local low = string.lower(body or "")
if string.sub(low, 1, 9) == "<!doctype" or string.find(low, "<html", 1, true) then
return "blocked by AgentRouter WAF"
end
local ok, resp = pcall(json.decode, body)
if not ok or type(resp) ~= "table" then return nil end
local e = resp.error
if type(e) == "table" and type(e.message) == "string" then
return e.message
end
return nil
end
return adapter

View File

@ -1,144 +1,447 @@
local adapter = {}
adapter.name = "anthropic"
adapter.version = "2.0.0"
adapter.version = "3.0.0"
adapter.endpoint = "/v1/messages"
adapter.headers = {
["anthropic-version"] = "2023-06-01"
["anthropic-version"] = "2023-06-01",
}
function adapter.transform_request(raw_body)
local ok, req = pcall(json.decode, raw_body)
if not ok then return raw_body end
-- ============================================================
-- helpers
-- ============================================================
-- 将 OpenAI 风格 content(字符串或 [{type:*}] 数组)拆成文本/图片块
local function collect_blocks(content)
if type(content) == "string" then
return { { type = "text", text = content } }
end
local blocks = {}
for _, p in ipairs(content or {}) do
local ROLE_WHITELIST = {
system = true, user = true, assistant = true, tool = true,
}
-- Anthropic requires tool_use.id / tool_result.tool_use_id to match
-- ^[a-zA-Z0-9_-]{1,64}$ and rejects the WHOLE request otherwise with
-- REQUEST_BODY_INVALID / "Invalid tool use format". OpenAI imposes no such
-- rule, so an OpenAI-compatible upstream can hand a client an id like
-- "bash:0" (observed from moonshotai/kimi-k3); once that lands in the
-- client's history, every later replay kills every Claude slot at once and
-- an AUTO chain falls through all tiers.
--
-- safe_tool_id is pure and deterministic, so a tool_use block and its
-- matching tool_result are rewritten identically within one request. A
-- rewritten id keeps an 8-hex digest of the ORIGINAL id: without it two
-- distinct ids could collapse into one ("a:b" and "a_b"), which Anthropic
-- rejects as a duplicate/unpaired tool_use. Already-legal ids pass through
-- untouched so upstreams that round-trip their own ids are unaffected.
local TOOL_ID_MAX = 64
local function safe_tool_id(id)
if type(id) ~= "string" or id == "" then return "" end
local clean = string.gsub(id, "[^A-Za-z0-9_-]", "_")
if clean == id and #clean <= TOOL_ID_MAX then
return clean
end
local digest = string.sub(sha256_hex(id), 1, 8)
local keep = TOOL_ID_MAX - #digest - 1
if #clean > keep then clean = string.sub(clean, 1, keep) end
return clean .. "_" .. digest
end
-- map_usage converts one Anthropic usage object to the gateway's TokenUsage
-- shape. The two APIs disagree on what the prompt count MEANS:
--
-- Anthropic: input_tokens EXCLUDES cached blocks. cache_read_input_tokens and
-- cache_creation_input_tokens are separate, additive, billed input.
-- OpenAI: prompt_tokens INCLUDES its cached_tokens subset.
--
-- Mapping input_tokens straight onto prompt therefore did two wrong things at
-- once: it undercounted the billed prompt by the entire cache portion, and it
-- could report cached_tokens > prompt_tokens — a cache hit rate above 100%.
-- cache_creation_input_tokens was not read at all, so a cache-write turn lost
-- those billed tokens outright. Summing all three restores OpenAI semantics,
-- which is what clients (and the gateway's own audit trail) assume.
--
-- Returns nil when the object carries no counts, so callers can keep treating
-- "no usage in this chunk" as "say nothing" rather than reporting zeros.
local function map_usage(u)
if type(u) ~= "table" then return nil end
local fresh = u.input_tokens or 0
local read = u.cache_read_input_tokens or 0
local create = u.cache_creation_input_tokens or 0
local p = fresh + read + create
local c = u.output_tokens or 0
if p <= 0 and c <= 0 then return nil end
local uses = { prompt = p, completion = c, total = p + c }
-- Emit the split whenever Anthropic reports either cache field, even at 0:
-- a reported zero-hit ("cache missed") must stay distinguishable from "this
-- upstream never reports cache info at all".
if u.cache_read_input_tokens ~= nil or u.cache_creation_input_tokens ~= nil then
uses.prompt_tokens_details = { cached_tokens = read }
uses.prompt_cache_hit_tokens = read
uses.prompt_cache_miss_tokens = fresh + create
end
return uses
end
local function collect_blocks(content)
if type(content) == "string" then
if content == "" then return {} end
return { { type = "text", text = content } }
end
if type(content) ~= "table" then return {} end
local blocks = {}
for _, p in ipairs(content) do
if type(p) == "string" then
table.insert(blocks, { type = "text", text = p })
elseif type(p) == "table" then
if p.type == "text" then
table.insert(blocks, { type = "text", text = p.text })
table.insert(blocks, { type = "text", text = p.text or "" })
elseif p.type == "image_url" and type(p.image_url) == "table" and p.image_url.url then
local mt, b64 = string.match(p.image_url.url, "^data:([^,]+);base64,(.+)$")
local url = p.image_url.url
local mt, b64 = string.match(url, "^data:([^,]+);base64,(.+)$")
if b64 then
table.insert(blocks, { type = "image", source = { type = "base64", media_type = mt or "image/png", data = b64 } })
else
table.insert(blocks, { type = "image", source = { type = "url", url = p.image_url.url } })
table.insert(blocks, { type = "image", source = { type = "url", url = url } })
end
else
-- Unknown content type: preserve for forward compatibility
table.insert(blocks, p)
end
end
return blocks
end
local function text_of(content)
if type(content) == "string" then return content end
local t = ""
for _, p in ipairs(content or {}) do
if p.type == "text" and p.text then t = t .. p.text end
return blocks
end
local function text_of(content)
if type(content) == "string" then return content end
local t = ""
for _, p in ipairs(content or {}) do
if type(p) == "string" then
t = t .. p
elseif type(p) == "table" and p.type == "text" and p.text then
t = t .. p.text
end
return t
end
return t
end
--- Append blocks to the last user message (merge) or create a new one.
local function append_user(msgs, blocks)
if #msgs > 0 and msgs[#msgs].role == "user" then
for _, b in ipairs(blocks) do
table.insert(msgs[#msgs].content, b)
end
else
table.insert(msgs, { role = "user", content = blocks })
end
end
-- ============================================================
-- transform_request (OpenAI → Anthropic)
-- ============================================================
function adapter.transform_request(raw_body)
local ok, req = pcall(json.decode, raw_body)
if not ok or type(req) ~= "table" then return raw_body end
local msgs = {}
local system = ""
local pending_tool = {} -- accumulated tool_result blocks
for _, m in ipairs(req.messages or {}) do
if m.role == "system" then
local role = m.role or "user"
if role == "system" then
system = system .. text_of(m.content) .. "\n"
elseif role == "assistant" then
-- Build content array: text/image blocks from content + tool_use from tool_calls
local blocks = collect_blocks(m.content)
if type(m.tool_calls) == "table" then
for _, tc in ipairs(m.tool_calls) do
if type(tc) == "table" then
local fn = tc["function"] or {}
local args = fn.arguments
local input = {}
if type(args) == "string" and args ~= "" then
local ok2, parsed = pcall(json.decode, args)
if ok2 and type(parsed) == "table" then input = parsed end
elseif type(args) == "table" then
input = args
end
table.insert(blocks, {
type = "tool_use",
id = safe_tool_id(tc.id),
name = fn.name or "",
input = input,
})
end
end
end
table.insert(msgs, { role = "assistant", content = blocks })
elseif role == "tool" then
-- Accumulate consecutive tool results; will be flushed as one
-- user message with tool_result content blocks.
table.insert(pending_tool, {
type = "tool_result",
tool_use_id = safe_tool_id(m.tool_call_id),
content = text_of(m.content),
})
else
table.insert(msgs, { role = m.role, content = collect_blocks(m.content) })
-- user / any other role: flush pending tool results first
if #pending_tool > 0 then
append_user(msgs, pending_tool)
pending_tool = {}
end
append_user(msgs, collect_blocks(m.content))
end
end
local anthropic_req = {
model = req.model or "claude-sonnet-4-20250514",
max_tokens = req.max_tokens or 4096,
messages = msgs,
stream = req.stream or false,
}
if not req.disable_thinking then
anthropic_req.thinking = { type = "enabled", budget_tokens = 4096 }
-- Flush any trailing tool results
if #pending_tool > 0 then
append_user(msgs, pending_tool)
end
-- ── Build Anthropic request ──────────────────────────────
local anthropic_req = {
model = req.model or "claude-sonnet-4-20250514",
max_tokens = req.max_tokens or 4096,
messages = msgs,
stream = req.stream or false,
}
-- ── tools ────────────────────────────────────────────────
local has_tools = false
if req.tools and type(req.tools) == "table" then
local tools = {}
for _, t in ipairs(req.tools) do
if type(t) == "table" and t.type == "function"
and type(t["function"]) == "table" then
local fn = t["function"]
table.insert(tools, {
name = fn.name or "",
description = fn.description or "",
input_schema = fn.parameters or { type = "object", properties = {} },
})
end
end
if #tools > 0 then
anthropic_req.tools = tools
has_tools = true
end
end
-- ── tool_choice ──────────────────────────────────────────
-- OpenAI → Anthropic mapping:
-- "auto" → {type:"auto"}
-- "none" → remove tools entirely (Anthropic has no "none")
-- "required" → {type:"any"}
-- {type:"function", function:{name:"X"}} → {type:"tool", name:"X"}
if has_tools and req.tool_choice then
local tc = req.tool_choice
if type(tc) == "string" then
if tc == "auto" then
anthropic_req.tool_choice = { type = "auto" }
elseif tc == "none" then
anthropic_req.tools = nil
anthropic_req.tool_choice = nil
elseif tc == "required" then
anthropic_req.tool_choice = { type = "any" }
end
elseif type(tc) == "table" and tc.type == "function" then
local fn = tc["function"] or {}
anthropic_req.tool_choice = { type = "tool", name = fn.name or "" }
end
end
-- ── thinking (opt-in via extra_body) ──────────────────────
-- Default: OFF. Client sends extra_body.thinking to enable.
-- Example: {"extra_body": {"thinking": {"type": "enabled", "budget_tokens": 4096}}}
if type(req.extra_body) == "table" and type(req.extra_body.thinking) == "table" then
anthropic_req.thinking = req.extra_body.thinking
end
-- ── system ───────────────────────────────────────────────
if system ~= "" then
anthropic_req.system = system
-- Strip trailing newline from accumulation
anthropic_req.system = string.match(system, "^(.-)\n*$")
end
return json.encode(anthropic_req)
end
-- ============================================================
-- transform_response (Anthropic → OpenAI, non-streaming)
-- ============================================================
function adapter.transform_response(raw_body)
local ok, resp = pcall(json.decode, raw_body)
if not ok then return raw_body end
if not ok or resp == nil then return raw_body end
local unified = {
content = "",
finish_reason = "",
token_usage = { prompt = 0, completion = 0, total = 0 }
token_usage = { prompt = 0, completion = 0, total = 0 },
}
if resp.usage then
unified.token_usage.prompt = resp.usage.input_tokens or 0
unified.token_usage.completion = resp.usage.output_tokens or 0
unified.token_usage.total = (resp.usage.input_tokens or 0) + (resp.usage.output_tokens or 0)
local mapped = map_usage(resp.usage)
if mapped ~= nil then
unified.token_usage = mapped
end
if resp.content and #resp.content > 0 then
local tcs = {}
for _, block in ipairs(resp.content) do
if block.type == "text" then
unified.content = unified.content .. (block.text or "")
elseif block.type == "thinking" and block.thinking then
unified.reasoning_content = (unified.reasoning_content or "") .. block.thinking
elseif block.type == "tool_use" then
table.insert(tcs, {
id = safe_tool_id(block.id),
type = "function",
name = block.name or "",
arguments = block.input or {},
})
end
end
if #tcs > 0 then unified.tool_calls = tcs end
end
unified.finish_reason = resp.stop_reason or ""
if unified.finish_reason == "tool_use" then
unified.finish_reason = "tool_calls"
elseif unified.finish_reason == "max_tokens" then
unified.finish_reason = "length"
elseif unified.finish_reason == "end_turn"
or unified.finish_reason == "stop_sequence" then
unified.finish_reason = "stop"
end
return json.encode(unified)
end
-- ============================================================
-- transform_stream_chunk (Anthropic SSE → OpenAI SSE delta)
-- ============================================================
function adapter.transform_stream_chunk(raw_chunk)
local ok, chunk = pcall(json.decode, raw_chunk)
if not ok then return "" end
if chunk.type == "message_start" then return "" end
-- ── message_start: initial usage ─────────────────────────
if chunk.type == "message_start" then
if chunk.message then
-- Some Anthropic-compatible upstreams report a placeholder here and
-- only send the true prompt count in message_delta (justwoker:
-- 160 at message_start vs 6931 at message_delta). The gateway's
-- mergeUsage lets a later non-zero value win, so both are emitted.
local uses = map_usage(chunk.message.usage)
if uses ~= nil then
return json.encode({ usage = uses, done = false })
end
end
return ""
end
-- ── message_delta: stop_reason + final usage ─────────────
if chunk.type == "message_delta" then
return json.encode({ content = "", done = (chunk.delta and chunk.delta.stop_reason ~= nil) })
local finish = nil
if chunk.delta and chunk.delta.stop_reason ~= nil then
local sr = chunk.delta.stop_reason
if sr == "max_tokens" then
finish = "length"
elseif sr == "tool_use" then
finish = "tool_calls"
else
finish = "stop"
end
end
-- message_delta carries the FINAL usage, so it must map the cache
-- fields too; the old code dropped them here, which silently lost the
-- whole cache split on every streamed response.
local uses = map_usage(chunk.usage)
if uses ~= nil then
return json.encode({ content = "", done = (finish ~= nil), finish_reason = finish, usage = uses })
end
return json.encode({ content = "", done = (finish ~= nil), finish_reason = finish })
end
if chunk.type == "content_block_start" and chunk.content_block
and chunk.content_block.type == "tool_use" then
-- first fragment of a tool call: emit index + id + name, empty args
return json.encode({
content = "", done = false,
tool_calls = { {
index = chunk.index or 0,
id = chunk.content_block.id or "",
type = "function",
["function"] = { name = chunk.content_block.name or "", arguments = "" }
} }
})
-- ── content_block_start: begin text / thinking / tool_use ─
if chunk.type == "content_block_start" and chunk.content_block then
local cb = chunk.content_block
if cb.type == "tool_use" then
-- Pass Anthropic content_block index through; Go-side rewrites
-- to sequential OpenAI tool_call ordinal for multi-tool streams.
return json.encode({
content = "", done = false,
tool_calls = { {
index = chunk.index or 0,
id = safe_tool_id(cb.id),
type = "function",
["function"] = { name = cb.name or "", arguments = "" },
} },
})
end
return "" -- text / thinking block start: no OpenAI equivalent
end
-- ── content_block_delta: incremental content ──────────────
if chunk.type == "content_block_delta" and chunk.delta then
if chunk.delta.type == "input_json_delta" then
-- incremental JSON fragment; clients accumulate across chunks
local unified = { content = "", done = false, tool_calls = { {
index = chunk.index or 0,
id = "",
type = "function",
["function"] = { name = "", arguments = chunk.delta.partial_json or "" }
} } }
return json.encode(unified)
local d = chunk.delta
if d.type == "input_json_delta" then
-- Tool call argument fragment; Go-side rewrites index.
return json.encode({
content = "", done = false,
tool_calls = { {
index = chunk.index or 0,
id = "",
type = "function",
["function"] = { name = "", arguments = d.partial_json or "" },
} },
})
end
if chunk.delta.type == "thinking_delta" and chunk.delta.thinking then
return json.encode({ content = "", done = false, reasoning_content = chunk.delta.thinking })
if d.type == "thinking_delta" and d.thinking then
return json.encode({ content = "", done = false, reasoning_content = d.thinking })
end
if d.type == "text_delta" and d.text then
return json.encode({ content = d.text, done = false })
end
return json.encode({ content = chunk.delta.text or "", done = false })
end
-- ── content_block_stop / message_stop ─────────────────────
if chunk.type == "message_stop" then
return json.encode({ content = "", done = true })
end
if chunk.type == "content_block_stop" then
return json.encode({ content = "", done = false })
-- The message_delta event already emitted the true finish_reason.
-- Do NOT emit done=true here: an empty finish_reason would
-- overwrite the real one (tool_calls) in the Go gateway's
-- lastFinish tracker, causing the final SSE chunk to say
-- finish_reason=stop instead of tool_calls.
return ""
end
return ""
end
-- ============================================================
-- transform_error (Anthropic error → human-readable string)
-- ============================================================
function adapter.transform_error(status, body)
local ok, resp = pcall(json.decode, body)
if not ok or type(resp) ~= "table" then return nil end
-- Anthropic envelope: {type:"error", error:{type, message}}
if resp.type == "error" and type(resp.error) == "table"
and type(resp.error.message) == "string" then
return resp.error.message
end
-- Flat envelope: {error: {message: "..."}}
if type(resp.error) == "table" and type(resp.error.message) == "string" then
return resp.error.message
end
return nil
end
return adapter

View File

@ -49,6 +49,11 @@ function adapter.transform_response(raw_body)
unified.token_usage.prompt = resp.usage.prompt_tokens or 0
unified.token_usage.completion = resp.usage.completion_tokens or 0
unified.token_usage.total = resp.usage.total_tokens or 0
unified.token_usage.prompt_cache_hit_tokens = resp.usage.prompt_cache_hit_tokens or 0
unified.token_usage.prompt_cache_miss_tokens = resp.usage.prompt_cache_miss_tokens or 0
if resp.usage.prompt_tokens_details and resp.usage.prompt_tokens_details.cached_tokens ~= nil then
unified.token_usage.prompt_tokens_details = { cached_tokens = resp.usage.prompt_tokens_details.cached_tokens }
end
end
if type(resp.choices) == "table" and #resp.choices > 0 then
@ -82,14 +87,47 @@ end
function adapter.transform_stream_chunk(raw_chunk)
local ok, chunk = pcall(json.decode, raw_chunk)
if not ok then return "" end
if not chunk.choices or #chunk.choices == 0 then return "" end
-- OpenAI-style streams may attach usage to a chunk with empty choices
-- (the final usage chunk). 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_cache_hit_tokens = chunk.usage.prompt_cache_hit_tokens or 0,
prompt_cache_miss_tokens = chunk.usage.prompt_cache_miss_tokens or 0,
}
if chunk.usage.prompt_tokens_details and chunk.usage.prompt_tokens_details.cached_tokens ~= nil then
uses.prompt_tokens_details = { cached_tokens = chunk.usage.prompt_tokens_details.cached_tokens }
end
end
if not chunk.choices or #chunk.choices == 0 then
if uses ~= nil then
-- usage-only chunk is not a content/finish signal; the gateway
-- emits its own terminal stop chunk and merges this usage.
return json.encode({ usage = uses, done = false })
end
return ""
end
local delta = chunk.choices[1].delta or {}
local fr = chunk.choices[1].finish_reason
local finish = (type(fr) == "string" and fr ~= "") and fr or nil
local unified = {
content = delta.content or "",
done = (fr ~= nil)
done = (finish ~= nil)
}
if finish then
unified.finish_reason = finish
end
if uses ~= nil then
unified.usage = uses
end
if delta.reasoning_content then
unified.reasoning_content = delta.reasoning_content
end
@ -99,4 +137,15 @@ function adapter.transform_stream_chunk(raw_chunk)
return json.encode(unified)
end
-- 错误收敛:DeepSeek 走标准 OpenAI 信封 {error:{message,...}}
function adapter.transform_error(status, body)
local ok, resp = pcall(json.decode, body)
if not ok or type(resp) ~= "table" then return nil end
local e = resp.error
if type(e) == "table" and type(e.message) == "string" then
return e.message
end
return nil
end
return adapter

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
@ -68,20 +151,55 @@ function adapter.transform_response(raw_body)
unified.token_usage.prompt = resp.usageMetadata.promptTokenCount or 0
unified.token_usage.completion = resp.usageMetadata.candidatesTokenCount or 0
unified.token_usage.total = resp.usageMetadata.totalTokenCount or 0
-- Gemini reports context-cache reads as cachedContentTokenCount;
-- normalize into OpenAI-standard prompt_tokens_details.cached_tokens
-- so clients and the audit trail see the hit count. Emitted even when
-- 0 so a reported miss stays distinguishable from "not reported".
if resp.usageMetadata.cachedContentTokenCount ~= nil then
unified.token_usage.prompt_tokens_details = {
cached_tokens = resp.usageMetadata.cachedContentTokenCount
}
end
end
if resp.candidates and #resp.candidates > 0 then
local cand = resp.candidates[1]
local tools = {}
if cand.content and cand.content.parts then
for _, part in ipairs(cand.content.parts) do
if part.text then
unified.content = unified.content .. part.text
elseif part.functionCall then
-- Non-streaming tool calls used to be dropped here while
-- transform_stream_chunk handled them, so a non-streaming
-- agent turn looked like a plain text answer and the tool
-- loop died. Gemini's args are already an object.
local args = part.functionCall.args
if type(args) == "string" then
local aok, decoded = pcall(json.decode, args)
args = aok and decoded or {}
elseif type(args) ~= "table" then
args = {}
end
table.insert(tools, {
id = part.functionCall.id or ("call_" .. #tools),
type = "function",
name = part.functionCall.name or "",
arguments = args
})
end
end
end
if cand.finishReason then
unified.finish_reason = cand.finishReason
end
if #tools > 0 then
unified.tool_calls = tools
-- Gemini reports finishReason "STOP" even when it emitted a
-- functionCall; clients keyed on finish_reason would treat that as
-- a completed answer and never run the tool.
unified.finish_reason = "tool_calls"
end
end
return json.encode(unified)
@ -91,9 +209,45 @@ function adapter.transform_stream_chunk(raw_chunk)
local ok, chunk = pcall(json.decode, raw_chunk)
if not ok then return "" end
if not chunk.candidates or #chunk.candidates == 0 then return "" end
-- Gemini attaches usageMetadata to the final chunk (alongside or after
-- candidates). Keys map to Go's TokenUsage json tags (prompt/...).
local uses = nil
if type(chunk.usageMetadata) == "table" then
local p = chunk.usageMetadata.promptTokenCount or 0
local c = chunk.usageMetadata.candidatesTokenCount or 0
local t = chunk.usageMetadata.totalTokenCount or 0
if p > 0 or c > 0 or t > 0 then
uses = { prompt = p, completion = c, total = t }
-- Emit details whenever the field is present, even at 0, so a
-- reported cache miss stays distinguishable from "not reported".
if chunk.usageMetadata.cachedContentTokenCount ~= nil then
uses.prompt_tokens_details = {
cached_tokens = chunk.usageMetadata.cachedContentTokenCount
}
end
end
end
if not chunk.candidates or #chunk.candidates == 0 then
if uses ~= nil then
return json.encode({ usage = uses, done = false })
end
return ""
end
local cand = chunk.candidates[1]
local unified = { content = "", done = (cand.finishReason ~= nil) }
-- Gemini finishReason -> OpenAI finish_reason
local finish = nil
if cand.finishReason ~= nil then
if cand.finishReason == "MAX_TOKENS" then
finish = "length"
elseif cand.finishReason == "SAFETY" or cand.finishReason == "RECITATION"
or cand.finishReason == "BLOCKLIST" then
finish = "content_filter"
else
finish = "stop"
end
end
local unified = { content = "", done = (finish ~= nil), finish_reason = finish }
local reasoning = ""
local tools = {}
if cand.content and cand.content.parts then
@ -117,7 +271,26 @@ function adapter.transform_stream_chunk(raw_chunk)
end
if reasoning ~= "" then unified.reasoning_content = reasoning end
if #tools > 0 then unified.tool_calls = tools end
if uses ~= nil then
unified.usage = uses
end
return json.encode(unified)
end
-- 错误收敛:Gemini REST 信封 {error:{code, message, status}}
function adapter.transform_error(status, body)
local ok, resp = pcall(json.decode, body)
if not ok or type(resp) ~= "table" then return nil end
local e = resp.error
if type(e) == "table" then
if type(e.message) == "string" then
if type(e.status) == "string" then
return e.status .. ": " .. e.message
end
return e.message
end
end
return nil
end
return adapter

View File

@ -33,6 +33,18 @@ function adapter.transform_response(raw_body)
unified.token_usage.prompt = resp.usage.prompt_tokens or 0
unified.token_usage.completion = resp.usage.completion_tokens or 0
unified.token_usage.total = resp.usage.total_tokens or 0
local hit = 0
if type(resp.usage.prompt_tokens_details) == "table" and resp.usage.prompt_tokens_details.cached_tokens ~= nil then
hit = resp.usage.prompt_tokens_details.cached_tokens
unified.token_usage.prompt_tokens_details = { cached_tokens = hit }
end
if (resp.usage.prompt_cache_hit_tokens or 0) > 0 then
unified.token_usage.prompt_cache_hit_tokens = resp.usage.prompt_cache_hit_tokens
unified.token_usage.prompt_cache_miss_tokens = resp.usage.prompt_cache_miss_tokens or 0
if hit == 0 then
unified.token_usage.prompt_tokens_details = { cached_tokens = resp.usage.prompt_cache_hit_tokens }
end
end
end
if type(resp.choices) == "table" and #resp.choices > 0 then
@ -63,14 +75,49 @@ end
function adapter.transform_stream_chunk(raw_chunk)
local ok, chunk = pcall(json.decode, raw_chunk)
if not ok then return "" end
if not chunk.choices or #chunk.choices == 0 then return "" end
-- OpenAI-style streams may attach usage to a chunk with empty choices
-- (the final usage chunk). 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,
}
if type(chunk.usage.prompt_tokens_details) == "table" and chunk.usage.prompt_tokens_details.cached_tokens ~= nil then
uses.prompt_tokens_details = { cached_tokens = chunk.usage.prompt_tokens_details.cached_tokens }
elseif (chunk.usage.prompt_cache_hit_tokens or 0) > 0 then
uses.prompt_cache_hit_tokens = chunk.usage.prompt_cache_hit_tokens
uses.prompt_cache_miss_tokens = chunk.usage.prompt_cache_miss_tokens or 0
uses.prompt_tokens_details = { cached_tokens = chunk.usage.prompt_cache_hit_tokens }
end
end
if not chunk.choices or #chunk.choices == 0 then
if uses ~= nil then
-- usage-only chunk is not a content/finish signal; the gateway
-- emits its own terminal stop chunk and merges this usage.
return json.encode({ usage = uses, done = false })
end
return ""
end
local delta = chunk.choices[1].delta or {}
local fr = chunk.choices[1].finish_reason
local finish = (type(fr) == "string" and fr ~= "") and fr or nil
local unified = {
content = delta.content or "",
done = (fr ~= nil)
done = (finish ~= nil)
}
if finish then
unified.finish_reason = finish
end
if uses ~= nil then
unified.usage = uses
end
if delta.reasoning_content then
unified.reasoning_content = delta.reasoning_content
end
@ -80,4 +127,15 @@ function adapter.transform_stream_chunk(raw_chunk)
return json.encode(unified)
end
-- 错误收敛:GitHub Models 走标准 OpenAI 信封 {error:{message,...}}
function adapter.transform_error(status, body)
local ok, resp = pcall(json.decode, body)
if not ok or type(resp) ~= "table" then return nil end
local e = resp.error
if type(e) == "table" and type(e.message) == "string" then
return e.message
end
return nil
end
return adapter

View File

@ -32,6 +32,18 @@ function adapter.transform_response(raw_body)
unified.token_usage.prompt = resp.usage.prompt_tokens or 0
unified.token_usage.completion = resp.usage.completion_tokens or 0
unified.token_usage.total = resp.usage.total_tokens or 0
local hit = 0
if type(resp.usage.prompt_tokens_details) == "table" and resp.usage.prompt_tokens_details.cached_tokens ~= nil then
hit = resp.usage.prompt_tokens_details.cached_tokens
unified.token_usage.prompt_tokens_details = { cached_tokens = hit }
end
if (resp.usage.prompt_cache_hit_tokens or 0) > 0 then
unified.token_usage.prompt_cache_hit_tokens = resp.usage.prompt_cache_hit_tokens
unified.token_usage.prompt_cache_miss_tokens = resp.usage.prompt_cache_miss_tokens or 0
if hit == 0 then
unified.token_usage.prompt_tokens_details = { cached_tokens = resp.usage.prompt_cache_hit_tokens }
end
end
end
if type(resp.choices) == "table" and #resp.choices > 0 then
@ -62,14 +74,49 @@ end
function adapter.transform_stream_chunk(raw_chunk)
local ok, chunk = pcall(json.decode, raw_chunk)
if not ok then return "" end
if not chunk.choices or #chunk.choices == 0 then return "" end
-- OpenAI-style streams may attach usage to a chunk with empty choices
-- (the final usage chunk). 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,
}
if type(chunk.usage.prompt_tokens_details) == "table" and chunk.usage.prompt_tokens_details.cached_tokens ~= nil then
uses.prompt_tokens_details = { cached_tokens = chunk.usage.prompt_tokens_details.cached_tokens }
elseif (chunk.usage.prompt_cache_hit_tokens or 0) > 0 then
uses.prompt_cache_hit_tokens = chunk.usage.prompt_cache_hit_tokens
uses.prompt_cache_miss_tokens = chunk.usage.prompt_cache_miss_tokens or 0
uses.prompt_tokens_details = { cached_tokens = chunk.usage.prompt_cache_hit_tokens }
end
end
if not chunk.choices or #chunk.choices == 0 then
if uses ~= nil then
-- usage-only chunk is not a content/finish signal; the gateway
-- emits its own terminal stop chunk and merges this usage.
return json.encode({ usage = uses, done = false })
end
return ""
end
local delta = chunk.choices[1].delta or {}
local fr = chunk.choices[1].finish_reason
local finish = (type(fr) == "string" and fr ~= "") and fr or nil
local unified = {
content = delta.content or "",
done = (fr ~= nil)
done = (finish ~= nil)
}
if finish then
unified.finish_reason = finish
end
if uses ~= nil then
unified.usage = uses
end
if delta.reasoning_content then
unified.reasoning_content = delta.reasoning_content
end
@ -79,4 +126,15 @@ function adapter.transform_stream_chunk(raw_chunk)
return json.encode(unified)
end
-- 错误收敛:Groq 走标准 OpenAI 信封 {error:{message,...}}
function adapter.transform_error(status, body)
local ok, resp = pcall(json.decode, body)
if not ok or type(resp) ~= "table" then return nil end
local e = resp.error
if type(e) == "table" and type(e.message) == "string" then
return e.message
end
return nil
end
return adapter

View File

@ -64,6 +64,18 @@ function adapter.transform_response(raw_body)
unified.token_usage.prompt = resp.usage.prompt_tokens or 0
unified.token_usage.completion = resp.usage.completion_tokens or 0
unified.token_usage.total = resp.usage.total_tokens or 0
local hit = 0
if type(resp.usage.prompt_tokens_details) == "table" and resp.usage.prompt_tokens_details.cached_tokens ~= nil then
hit = resp.usage.prompt_tokens_details.cached_tokens
unified.token_usage.prompt_tokens_details = { cached_tokens = hit }
end
if (resp.usage.prompt_cache_hit_tokens or 0) > 0 then
unified.token_usage.prompt_cache_hit_tokens = resp.usage.prompt_cache_hit_tokens
unified.token_usage.prompt_cache_miss_tokens = resp.usage.prompt_cache_miss_tokens or 0
if hit == 0 then
unified.token_usage.prompt_tokens_details = { cached_tokens = resp.usage.prompt_cache_hit_tokens }
end
end
end
if type(resp.choices) == "table" and #resp.choices > 0 then
local ch = resp.choices[1]
@ -95,14 +107,49 @@ end
function adapter.transform_stream_chunk(raw_chunk)
local ok, chunk = pcall(json.decode, raw_chunk)
if not ok then return "" end
if not chunk.choices or #chunk.choices == 0 then return "" end
-- OpenAI-style streams may attach usage to a chunk with empty choices
-- (the final usage chunk). 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,
}
if type(chunk.usage.prompt_tokens_details) == "table" and chunk.usage.prompt_tokens_details.cached_tokens ~= nil then
uses.prompt_tokens_details = { cached_tokens = chunk.usage.prompt_tokens_details.cached_tokens }
elseif (chunk.usage.prompt_cache_hit_tokens or 0) > 0 then
uses.prompt_cache_hit_tokens = chunk.usage.prompt_cache_hit_tokens
uses.prompt_cache_miss_tokens = chunk.usage.prompt_cache_miss_tokens or 0
uses.prompt_tokens_details = { cached_tokens = chunk.usage.prompt_cache_hit_tokens }
end
end
if not chunk.choices or #chunk.choices == 0 then
if uses ~= nil then
-- usage-only chunk is not a content/finish signal; the gateway
-- emits its own terminal stop chunk and merges this usage.
return json.encode({ usage = uses, done = false })
end
return ""
end
local delta = chunk.choices[1].delta or {}
local fr = chunk.choices[1].finish_reason
local finish = (type(fr) == "string" and fr ~= "") and fr or nil
local unified = {
content = delta.content or "",
done = (fr ~= nil)
done = (finish ~= nil)
}
if finish then
unified.finish_reason = finish
end
if uses ~= nil then
unified.usage = uses
end
if delta.reasoning_content then
unified.reasoning_content = delta.reasoning_content
end
@ -112,4 +159,15 @@ function adapter.transform_stream_chunk(raw_chunk)
return json.encode(unified)
end
return adapter
-- 错误收敛:kimi 网关为 OpenAI 风格 {error:{message,...}}
function adapter.transform_error(status, body)
local ok, resp = pcall(json.decode, body)
if not ok or type(resp) ~= "table" then return nil end
local e = resp.error
if type(e) == "table" and type(e.message) == "string" then
return e.message
end
return nil
end
return adapter

View File

@ -32,6 +32,18 @@ function adapter.transform_response(raw_body)
unified.token_usage.prompt = resp.usage.prompt_tokens or 0
unified.token_usage.completion = resp.usage.completion_tokens or 0
unified.token_usage.total = resp.usage.total_tokens or 0
local hit = 0
if type(resp.usage.prompt_tokens_details) == "table" and resp.usage.prompt_tokens_details.cached_tokens ~= nil then
hit = resp.usage.prompt_tokens_details.cached_tokens
unified.token_usage.prompt_tokens_details = { cached_tokens = hit }
end
if (resp.usage.prompt_cache_hit_tokens or 0) > 0 then
unified.token_usage.prompt_cache_hit_tokens = resp.usage.prompt_cache_hit_tokens
unified.token_usage.prompt_cache_miss_tokens = resp.usage.prompt_cache_miss_tokens or 0
if hit == 0 then
unified.token_usage.prompt_tokens_details = { cached_tokens = resp.usage.prompt_cache_hit_tokens }
end
end
end
if type(resp.choices) == "table" and #resp.choices > 0 then
@ -62,14 +74,49 @@ end
function adapter.transform_stream_chunk(raw_chunk)
local ok, chunk = pcall(json.decode, raw_chunk)
if not ok then return "" end
if not chunk.choices or #chunk.choices == 0 then return "" end
-- OpenAI-style streams may attach usage to a chunk with empty choices
-- (the final usage chunk). 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,
}
if type(chunk.usage.prompt_tokens_details) == "table" and chunk.usage.prompt_tokens_details.cached_tokens ~= nil then
uses.prompt_tokens_details = { cached_tokens = chunk.usage.prompt_tokens_details.cached_tokens }
elseif (chunk.usage.prompt_cache_hit_tokens or 0) > 0 then
uses.prompt_cache_hit_tokens = chunk.usage.prompt_cache_hit_tokens
uses.prompt_cache_miss_tokens = chunk.usage.prompt_cache_miss_tokens or 0
uses.prompt_tokens_details = { cached_tokens = chunk.usage.prompt_cache_hit_tokens }
end
end
if not chunk.choices or #chunk.choices == 0 then
if uses ~= nil then
-- usage-only chunk is not a content/finish signal; the gateway
-- emits its own terminal stop chunk and merges this usage.
return json.encode({ usage = uses, done = false })
end
return ""
end
local delta = chunk.choices[1].delta or {}
local fr = chunk.choices[1].finish_reason
local finish = (type(fr) == "string" and fr ~= "") and fr or nil
local unified = {
content = delta.content or "",
done = (fr ~= nil)
done = (finish ~= nil)
}
if finish then
unified.finish_reason = finish
end
if uses ~= nil then
unified.usage = uses
end
if delta.reasoning_content then
unified.reasoning_content = delta.reasoning_content
end
@ -79,4 +126,14 @@ function adapter.transform_stream_chunk(raw_chunk)
return json.encode(unified)
end
-- 错误收敛:Mistral 用顶层 {message="...", type="..."}(非嵌套 error)
function adapter.transform_error(status, body)
local ok, resp = pcall(json.decode, body)
if not ok or type(resp) ~= "table" then return nil end
if type(resp.message) == "string" then
return resp.message
end
return nil
end
return adapter

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
@ -53,15 +97,44 @@ function adapter.transform_response(raw_body)
local ok, resp = pcall(json.decode, raw_body)
if not ok then return raw_body end
local p = resp.prompt_eval_count or 0
local c = resp.eval_count or 0
local unified = {
content = "",
finish_reason = resp.done_reason or "",
tool_calls = {},
usage = { prompt = 0, completion = 0, total = 0 }
-- key must be token_usage to match Go's UnifiedResponse json tag
token_usage = { prompt = p, completion = c, total = p + c }
}
if resp.message then
unified.content = resp.message.content or ""
-- Non-streaming tool calls were previously dropped: the field was
-- initialized to an empty table and never filled, while
-- transform_stream_chunk handled them. A non-streaming agent turn thus
-- looked like a plain answer and the tool loop stopped.
if type(resp.message.tool_calls) == "table" and #resp.message.tool_calls > 0 then
local tcs = {}
for _, tc in ipairs(resp.message.tool_calls) do
local fn = tc["function"] or {}
-- Ollama sends arguments as an object already
local args = fn.arguments
if type(args) == "string" then
local aok, decoded = pcall(json.decode, args)
args = aok and decoded or {}
elseif type(args) ~= "table" then
args = {}
end
table.insert(tcs, {
id = tc.id or ("call_" .. #tcs),
type = tc.type or "function",
name = fn.name or "",
arguments = args
})
end
unified.tool_calls = tcs
-- Ollama reports done_reason "stop" alongside tool calls
unified.finish_reason = "tool_calls"
end
end
return json.encode(unified)
@ -70,12 +143,43 @@ end
function adapter.transform_stream_chunk(raw_chunk)
local ok, chunk = pcall(json.decode, raw_chunk)
if not ok then return "" end
if not chunk.message then return "" end
-- Ollama's terminal chunk (done=true) carries token counts but may omit
-- message; pass them through so the gateway emits real usage.
local uses = nil
if chunk.done then
local p = chunk.prompt_eval_count or 0
local c = chunk.eval_count or 0
if p > 0 or c > 0 then
uses = { prompt = p, completion = c, total = p + c }
end
end
-- Ollama done_reason -> OpenAI finish_reason ("length" 透传,其余归一 stop)
local finish = nil
if chunk.done then
if chunk.done_reason == "length" then
finish = "length"
else
finish = "stop"
end
end
if not chunk.message then
if uses ~= nil then
return json.encode({ content = "", done = true, finish_reason = finish, usage = uses })
end
return json.encode({ content = "", done = true, finish_reason = finish })
end
local unified = {
content = chunk.message.content or "",
done = chunk.done or false
done = chunk.done or false,
finish_reason = finish
}
if uses ~= nil then
unified.usage = uses
end
if chunk.message.reasoning_content then
unified.reasoning_content = chunk.message.reasoning_content
end
@ -97,4 +201,15 @@ function adapter.transform_stream_chunk(raw_chunk)
return json.encode(unified)
end
-- 错误收敛:Ollama 常见 {error:"..."} 字符串(新版本也有对象形态)
function adapter.transform_error(status, body)
local ok, resp = pcall(json.decode, body)
if not ok or type(resp) ~= "table" then return nil end
if type(resp.error) == "string" then return resp.error end
if type(resp.error) == "table" and type(resp.error.message) == "string" then
return resp.error.message
end
return nil
end
return adapter

View File

@ -5,6 +5,36 @@ adapter.version = "2.0.0"
adapter.endpoint = "/chat/completions"
adapter.headers = {}
-- Claude-behind-OpenAI upstreams (tabitoken, 扇贝, …) convert /chat/completions
-- to the Anthropic Messages API internally, so they inherit Anthropic's
-- tool id rule: ^[a-zA-Z0-9_-]{1,64}$, enforced by rejecting the WHOLE request
-- with "Invalid tool use format" / "tool_use.id: String should match pattern".
-- Plain OpenAI has no such rule, so an OpenAI-compatible model can mint an id
-- like "bash:0" (observed from moonshotai/kimi-k3). In a fan-out router that id
-- is replayed to every other source, so one such id kills every Claude slot at
-- once and an AUTO request falls through all tiers.
--
-- safe_tool_id is pure and deterministic, so a tool_calls entry and its
-- matching tool_call_id are rewritten identically within one request. A
-- rewritten id keeps an 8-hex digest of the ORIGINAL id, without which two
-- distinct ids could collapse into one ("a:b" and "a_b") and become an
-- unpaired/duplicate tool call. Already-legal ids are returned untouched, so
-- well-behaved traffic is byte-identical to before.
-- (anthropic.lua carries the same helper; Lua adapters have no shared prelude.)
local TOOL_ID_MAX = 64
local function safe_tool_id(id)
if type(id) ~= "string" or id == "" then return id end
local clean = string.gsub(id, "[^A-Za-z0-9_-]", "_")
if clean == id and #clean <= TOOL_ID_MAX then
return clean
end
local digest = string.sub(sha256_hex(id), 1, 8)
local keep = TOOL_ID_MAX - #digest - 1
if #clean > keep then clean = string.sub(clean, 1, keep) end
return clean .. "_" .. digest
end
-- OpenAI /chat/completions format (pass-through, strip provider-specific fields)
function adapter.transform_request(raw_body)
local ok, req = pcall(json.decode, raw_body)
@ -14,6 +44,16 @@ function adapter.transform_request(raw_body)
if req.messages then
for _, msg in ipairs(req.messages) do
msg.reasoning_content = nil
if msg.tool_call_id ~= nil then
msg.tool_call_id = safe_tool_id(msg.tool_call_id)
end
if type(msg.tool_calls) == "table" then
for _, tc in ipairs(msg.tool_calls) do
if type(tc) == "table" and tc.id ~= nil then
tc.id = safe_tool_id(tc.id)
end
end
end
end
end
return json.encode(req)
@ -33,6 +73,22 @@ function adapter.transform_response(raw_body)
unified.token_usage.prompt = resp.usage.prompt_tokens or 0
unified.token_usage.completion = resp.usage.completion_tokens or 0
unified.token_usage.total = resp.usage.total_tokens or 0
-- Cache passthrough: OpenAI v2 prompt_tokens_details.cached_tokens
-- and DeepSeek-legacy prompt_cache_hit/miss_tokens. dsh reads
-- prompt_tokens_details.cached_tokens (falls back to the legacy
-- standalone field), so both shapes reach clients.
local hit = 0
if type(resp.usage.prompt_tokens_details) == "table" and resp.usage.prompt_tokens_details.cached_tokens ~= nil then
hit = resp.usage.prompt_tokens_details.cached_tokens
unified.token_usage.prompt_tokens_details = { cached_tokens = hit }
end
if (resp.usage.prompt_cache_hit_tokens or 0) > 0 then
unified.token_usage.prompt_cache_hit_tokens = resp.usage.prompt_cache_hit_tokens
unified.token_usage.prompt_cache_miss_tokens = resp.usage.prompt_cache_miss_tokens or 0
if hit == 0 then
unified.token_usage.prompt_tokens_details = { cached_tokens = resp.usage.prompt_cache_hit_tokens }
end
end
end
if type(resp.choices) == "table" and #resp.choices > 0 then
@ -48,7 +104,7 @@ function adapter.transform_response(raw_body)
local args_ok, args = pcall(json.decode, tc["function"].arguments)
if not args_ok then args = {} end
table.insert(tcs, {
id = tc.id,
id = safe_tool_id(tc.id),
type = tc.type or "function",
name = tc["function"].name,
arguments = args
@ -67,22 +123,106 @@ function adapter.transform_stream_chunk(raw_chunk)
local ok, chunk = pcall(json.decode, raw_chunk)
if not ok then return "" end
if not chunk.choices or #chunk.choices == 0 then return "" end
-- OpenAI-style streams may attach usage to a chunk with empty choices
-- (the final usage chunk). 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,
}
if type(chunk.usage.prompt_tokens_details) == "table" and chunk.usage.prompt_tokens_details.cached_tokens ~= nil then
uses.prompt_tokens_details = { cached_tokens = chunk.usage.prompt_tokens_details.cached_tokens }
elseif (chunk.usage.prompt_cache_hit_tokens or 0) > 0 then
uses.prompt_cache_hit_tokens = chunk.usage.prompt_cache_hit_tokens
uses.prompt_cache_miss_tokens = chunk.usage.prompt_cache_miss_tokens or 0
uses.prompt_tokens_details = { cached_tokens = chunk.usage.prompt_cache_hit_tokens }
end
end
if not chunk.choices or #chunk.choices == 0 then
if uses ~= nil then
-- usage-only chunk is not a content/finish signal; the gateway
-- emits its own terminal stop chunk and merges this usage.
return json.encode({ usage = uses, done = false })
end
return ""
end
local delta = chunk.choices[1].delta or {}
local fr = chunk.choices[1].finish_reason
local finish = (type(fr) == "string" and fr ~= "") and fr or nil
local unified = {
content = delta.content or "",
done = (fr ~= nil)
done = (finish ~= nil)
}
if finish then
unified.finish_reason = finish
end
if uses ~= nil then
unified.usage = uses
end
if delta.reasoning_content then
unified.reasoning_content = delta.reasoning_content
end
if delta.tool_calls then
-- pass raw streaming fragments through; OpenAI clients accumulate index+id+name+arguments
-- Sanitize on the way OUT too: an id this upstream happily minted (it
-- does not validate them) becomes a landmine once the client replays
-- it to a Claude upstream. Only the first fragment of a streamed call
-- carries an id; later argument fragments have none and are untouched.
for _, tc in ipairs(delta.tool_calls) do
if type(tc) == "table" and tc.id ~= nil then
tc.id = safe_tool_id(tc.id)
end
end
unified.tool_calls = delta.tool_calls
end
return json.encode(unified)
end
-- 错误收敛:标准 OpenAI 信封 {error:{message,...}}
function adapter.transform_error(status, body)
-- Non-JSON body (nginx HTML error pages, plain text): extract a short
-- human-readable reason instead of letting the raw body reach the log.
local ok, resp = pcall(json.decode, body)
if not ok or type(resp) ~= "table" then
-- HTML error page: pull the <title> text (e.g. "413 Request Entity Too Large")
local title = string.match(body or "", "<title>(.-)</title>")
if title and title ~= "" then return title end
-- Bare text: first non-empty line, capped
local line = string.match(body or "", "^%s*([^\r\n]+)")
if line and line ~= "" and not string.match(line, "^<") then
return string.sub(line, 1, 200)
end
return nil
end
-- Standard OpenAI envelope: {error:{message,...}} or {error:"..."}
local e = resp.error
if type(e) == "table" and type(e.message) == "string" then
return e.message
end
if type(e) == "string" then return e end
-- Flat envelope used by many OpenAI-compatible gateways:
-- {"code":20012,"message":"Model does not exist..."}
-- {"code":"INVALID_API_KEY","message":"Invalid API key"}
if type(resp.message) == "string" and resp.message ~= "" then
if resp.code ~= nil then
return tostring(resp.code) .. ": " .. resp.message
end
return resp.message
end
-- Some gateways use {detail:"..."} (FastAPI style)
if type(resp.detail) == "string" and resp.detail ~= "" then
return resp.detail
end
return nil
end
return adapter

View File

@ -4,18 +4,40 @@ adapter.name = "opencode"
adapter.version = "1.0.0"
adapter.endpoint = "/chat/completions"
-- opencode.ai zen 网关按 User-Agent 指纹识别官方客户端并把请求分到免费额度池;
-- 非官方 UA(curl/Go 默认等)会被分到匿名池并触发 FreeUsageLimitError。
-- 因此固定发送 opencode 客户端的 UA;配合源配置 api_key: "public"(官方无 key
-- 客户端实际发送 Bearer public)即可走免费池。
-- opencode.ai zen 网关按客户端指纹(UA + x-opencode-* 头)路由请求池:
-- 缺少 x-opencode-client/session/request/project 身份头的请求会被判为匿名
-- 客户端,x-preview-f-free 等模型在带 tools 时上游直接失败——流式返回单
-- chunk "finish_reason":"network_error" 空 content,非流式返回 503
-- "Endpoint is unavailable",表现为"空回复"。因此除 UA 外必须附带
-- x-opencode-* 身份头;session/request 每请求派生唯一值(见 build_headers)。
adapter.headers = {
["User-Agent"] = "opencode/0.1.0",
["User-Agent"] = "opencode/1.18.21 ai-sdk/provider-utils/4.0.23 runtime/bun/1.3.14",
}
-- 每请求生成身份头。沙箱无 os/math,用 meta.timestamp + 请求体哈希派生:
-- 同秒内重复请求 id 相同可接受(zen 只校验存在性,不校验格式)。
local function rand_id(prefix, seed)
return prefix .. string.sub(sha256_hex(seed), 1, 24)
end
function adapter.build_headers(meta)
local ts = tostring(meta.timestamp or "")
return {
["User-Agent"] = adapter.headers["User-Agent"],
["x-opencode-client"] = "cli",
-- project 固定:同网关实例共享一个工作区身份
["x-opencode-project"] = string.sub(sha256_hex("llmsproxy|" .. (meta.source and meta.source.name or "")), 1, 32),
["x-opencode-session"] = rand_id("ses_", "session|" .. ts),
["x-opencode-request"] = rand_id("msg_", "request|" .. ts .. "|" .. tostring(meta.body or "")),
}
end
-- OpenAI /chat/completions format (pass-through, strip provider-specific fields)
-- zen 上游 schema 只接受 text content part(无视觉/音频能力):多模态 part
-- (image_url / input_audio / file 等)一律剥离;因此失去全部 content 的
-- 消息整条丢弃,避免上游 "unknown variant `image_url`, expected `text`"。
-- (image_url / input_audio / file 等)一律剥离。剥离后 content 变空的消息
-- 若不再携带 tool_calls / tool_call_id 才整条丢弃(避免上游
-- "unknown variant `image_url`, expected `text`");带工具调用的必须保留,
-- 否则会把紧随其后的 tool 结果变成孤儿,模型会反复重发同一个调用。
-- zen 上游角色白名单只有 system / user / assistant / tool / latest_reminder:
-- OpenAI 的 developer(及 function 等)不在其中,直接透传会触发上游
-- "unknown variant `developer`, expected one of ..." 错误;统一归一化为 system。
@ -32,6 +54,14 @@ function adapter.transform_request(raw_body)
if not ok then return raw_body end
req.disable_thinking = nil
req.extra_body = nil
-- stream_options is only valid alongside stream:true. Sending it on a
-- non-streaming request is rejected by stricter upstreams (OpenCode Go:
-- "stream_options should be set along with stream"), which failed every
-- non-stream call. Only attach it when the request actually streams.
if req.stream then
if type(req.stream_options) ~= "table" then req.stream_options = {} end
req.stream_options.include_usage = true
end
if req.messages then
local kept = {}
for _, msg in ipairs(req.messages) do
@ -50,7 +80,24 @@ function adapter.transform_request(raw_body)
end
end
if #parts == 0 then
drop = true
-- Content collapsed to nothing after stripping unsupported
-- parts. A message that still carries a tool call must
-- NEVER be dropped: the very next message is its tool
-- result, and dropping the call orphans that result. The
-- model then sees a result for a call it never made and
-- re-issues the same tool call on every turn (observed as
-- an infinite "repeated tool call" loop).
--
-- This is the common shape for agent clients: an assistant
-- turn whose content is only [thinking, toolCall] serialises
-- to content:[] with tool_calls — exactly the case that used
-- to vanish here. Emit "" instead, which zen accepts.
if (type(msg.tool_calls) == "table" and #msg.tool_calls > 0)
or msg.tool_call_id ~= nil then
msg.content = ""
else
drop = true
end
else
msg.content = parts
end
@ -78,14 +125,27 @@ function adapter.transform_response(raw_body)
unified.token_usage.prompt = resp.usage.prompt_tokens or 0
unified.token_usage.completion = resp.usage.completion_tokens or 0
unified.token_usage.total = resp.usage.total_tokens or 0
local hit = 0
if type(resp.usage.prompt_tokens_details) == "table" and resp.usage.prompt_tokens_details.cached_tokens ~= nil then
hit = resp.usage.prompt_tokens_details.cached_tokens
unified.token_usage.prompt_tokens_details = { cached_tokens = hit }
end
if (resp.usage.prompt_cache_hit_tokens or 0) > 0 then
unified.token_usage.prompt_cache_hit_tokens = resp.usage.prompt_cache_hit_tokens
unified.token_usage.prompt_cache_miss_tokens = resp.usage.prompt_cache_miss_tokens or 0
if hit == 0 then
unified.token_usage.prompt_tokens_details = { cached_tokens = resp.usage.prompt_cache_hit_tokens }
end
end
end
if type(resp.choices) == "table" and #resp.choices > 0 then
local ch = resp.choices[1]
if type(ch.message) == "table" then
unified.content = ch.message.content or ""
if ch.message.reasoning_content then
unified.reasoning_content = ch.message.reasoning_content
local reasoning = ch.message.reasoning_content or ch.message.reasoning
if reasoning then
unified.reasoning_content = reasoning
end
if type(ch.message.tool_calls) == "table" then
local tcs = {}
@ -112,16 +172,53 @@ function adapter.transform_stream_chunk(raw_chunk)
local ok, chunk = pcall(json.decode, raw_chunk)
if not ok then return "" end
if not chunk.choices or #chunk.choices == 0 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,
}
if type(chunk.usage.prompt_tokens_details) == "table" and chunk.usage.prompt_tokens_details.cached_tokens ~= nil then
uses.prompt_tokens_details = { cached_tokens = chunk.usage.prompt_tokens_details.cached_tokens }
elseif (chunk.usage.prompt_cache_hit_tokens or 0) > 0 then
uses.prompt_cache_hit_tokens = chunk.usage.prompt_cache_hit_tokens
uses.prompt_cache_miss_tokens = chunk.usage.prompt_cache_miss_tokens or 0
uses.prompt_tokens_details = { cached_tokens = chunk.usage.prompt_cache_hit_tokens }
end
end
if not chunk.choices or #chunk.choices == 0 then
if uses ~= nil then
-- usage-only chunk is not a content/finish signal; the gateway
-- emits its own terminal stop chunk and merges this usage.
return json.encode({ usage = uses, done = false })
end
return ""
end
local delta = chunk.choices[1].delta or {}
local fr = chunk.choices[1].finish_reason
local finish = (type(fr) == "string" and fr ~= "") and fr or nil
local unified = {
content = delta.content or "",
done = (fr ~= nil)
done = (finish ~= nil)
}
if delta.reasoning_content then
unified.reasoning_content = delta.reasoning_content
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
@ -130,4 +227,17 @@ function adapter.transform_stream_chunk(raw_chunk)
return json.encode(unified)
end
-- 错误收敛(可选钩子):zen 错误信封固定为 {error={type,message}};
-- 免费池限流 FreeUsageLimitError 单独标注。返回 nil 走通用兜底。
function adapter.transform_error(status, body)
local ok, resp = pcall(json.decode, body)
if not ok or type(resp) ~= "table" then return nil end
local e = resp.error
if type(e) ~= "table" then return nil end
if e.type == "FreeUsageLimitError" then
return "zen free pool quota exhausted"
end
return e.message
end
return adapter

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,174 @@
-- sensenova API format adapter
-- streaming: delta only has reasoning_content, no content field
-- non-streaming: has both content and reasoning_content
local adapter = {}
adapter.name = "sensenova"
adapter.version = "1.0.0"
adapter.endpoint = "/chat/completions"
adapter.headers = {}
-- Same as openai - strip provider-specific fields
function adapter.transform_request(raw_body)
local ok, req = pcall(json.decode, raw_body)
if not ok then return raw_body end
req.disable_thinking = nil
req.extra_body = nil
if req.messages then
for _, msg in ipairs(req.messages) do
msg.reasoning_content = nil
end
end
return json.encode(req)
end
-- Same as openai - extract content from response
function adapter.transform_response(raw_body)
local ok, resp = pcall(json.decode, raw_body)
if not ok or resp == nil then return raw_body end
local unified = {
content = "",
finish_reason = "",
token_usage = { prompt = 0, completion = 0, total = 0 }
}
if type(resp.usage) == "table" then
unified.token_usage.prompt = resp.usage.prompt_tokens or 0
unified.token_usage.completion = resp.usage.completion_tokens or 0
unified.token_usage.total = resp.usage.total_tokens or 0
local hit = 0
if type(resp.usage.prompt_tokens_details) == "table" and resp.usage.prompt_tokens_details.cached_tokens ~= nil then
hit = resp.usage.prompt_tokens_details.cached_tokens
unified.token_usage.prompt_tokens_details = { cached_tokens = hit }
end
if (resp.usage.prompt_cache_hit_tokens or 0) > 0 then
unified.token_usage.prompt_cache_hit_tokens = resp.usage.prompt_cache_hit_tokens
unified.token_usage.prompt_cache_miss_tokens = resp.usage.prompt_cache_miss_tokens or 0
if hit == 0 then
unified.token_usage.prompt_tokens_details = { cached_tokens = resp.usage.prompt_cache_hit_tokens }
end
end
end
if type(resp.choices) == "table" and #resp.choices > 0 then
local ch = resp.choices[1]
if type(ch.message) == "table" then
unified.content = ch.message.content or ""
if ch.message.reasoning_content then
unified.reasoning_content = ch.message.reasoning_content
elseif ch.message.reasoning then
-- sensenova-6.8-flash-lite 用 reasoning 字段而不是 reasoning_content
unified.reasoning_content = ch.message.reasoning
end
-- Tool calls MUST be forwarded. Dropping them while keeping
-- finish_reason="tool_calls" makes the client replay an assistant
-- message whose function name/arguments are empty, and sensenova
-- then rejects the next turn with
-- 400 invalid tool_call function, function/name/arguments cannot be empty
-- i.e. a tool-using conversation dies on its second request.
if type(ch.message.tool_calls) == "table" and #ch.message.tool_calls > 0 then
local tcs = {}
for _, tc in ipairs(ch.message.tool_calls) do
local fn = tc["function"] or {}
-- arguments arrives as a JSON *string* on the wire; the
-- unified shape expects a decoded object.
local args = fn.arguments
if type(args) == "string" then
local args_ok, decoded = pcall(json.decode, args)
args = args_ok and decoded or {}
elseif type(args) ~= "table" then
args = {}
end
table.insert(tcs, {
id = tc.id,
type = tc.type or "function",
name = fn.name,
arguments = args
})
end
unified.tool_calls = tcs
end
end
unified.finish_reason = ch.finish_reason or ""
end
return json.encode(unified)
end
-- Sensenova-specific stream handling
-- Upstream puts content in reasoning_content only (no content field)
-- Also sends finish_reason="" (empty string) on every chunk
function adapter.transform_stream_chunk(raw_chunk)
local ok, chunk = pcall(json.decode, raw_chunk)
if not ok then return "" end
local uses = nil
if type(chunk.usage) == "table" then
uses = {
prompt = chunk.usage.prompt_tokens or chunk.usage.prompt or 0,
completion = chunk.usage.completion_tokens or chunk.usage.completion or 0,
total = chunk.usage.total_tokens or chunk.usage.total or 0,
}
if type(chunk.usage.prompt_tokens_details) == "table" and chunk.usage.prompt_tokens_details.cached_tokens ~= nil then
uses.prompt_tokens_details = { cached_tokens = chunk.usage.prompt_tokens_details.cached_tokens }
elseif (chunk.usage.prompt_cache_hit_tokens or 0) > 0 then
uses.prompt_cache_hit_tokens = chunk.usage.prompt_cache_hit_tokens
uses.prompt_cache_miss_tokens = chunk.usage.prompt_cache_miss_tokens or 0
uses.prompt_tokens_details = { cached_tokens = chunk.usage.prompt_cache_hit_tokens }
end
end
if not chunk.choices or #chunk.choices == 0 then
if uses ~= nil then
return json.encode({ usage = uses, done = false })
end
return ""
end
local delta = chunk.choices[1].delta or {}
-- Sensenova: delta has reasoning_content (deepseek-v4-flash) or reasoning
-- (sensenova-6.8-flash-lite) but no content field
local content = delta.content or ""
if content == "" and (delta.reasoning_content or delta.reasoning) then
content = delta.reasoning_content or delta.reasoning
end
-- Sensenova: finish_reason is "" on every chunk, "stop" on last
local fr = chunk.choices[1].finish_reason
local done = (fr == "stop" or fr == "length")
local unified = {
content = content,
done = done,
}
if chunk.choices[1].finish_reason and chunk.choices[1].finish_reason ~= "" then
unified.finish_reason = chunk.choices[1].finish_reason
end
if delta.tool_calls then
unified.tool_calls = delta.tool_calls
end
if uses ~= nil then
unified.usage = uses
end
return json.encode(unified)
end
-- 错误收敛:sensenova 为 OpenAI 风格 {error:{message,...}};
-- 配额类错误单独点出便于客户端识别重置周期。
function adapter.transform_error(status, body)
local ok, resp = pcall(json.decode, body)
if not ok or type(resp) ~= "table" then return nil end
local e = resp.error
if type(e) ~= "table" then return nil end
if status == 429 and type(e.code) == "string"
and e.code == "insufficient_quota" then
return "workspace quota exhausted (resets periodically)"
end
if type(e.message) == "string" then return e.message end
return nil
end
return adapter

View File

@ -0,0 +1,286 @@
-- trae 源适配器(trae-local-api 的 OpenAI 兼容代理)
-- trae-local-api 运行在 http://127.0.0.1:19900,接 Trae CN 账号的云资源。
-- 协议:OpenAI /v1/chat/completions
-- 特性:
-- - 排队:上游 busy 时排队,可能长时间无响应
-- - 非流式请求偶发返回 SSE 数据(upstream bug)
-- - 支持 thinking,由上游模型自动决定是否开启
local adapter = {}
adapter.name = "trae"
adapter.version = "1.0.0"
adapter.endpoint = "/v1/chat/completions"
adapter.headers = {}
function adapter.transform_request(raw_body)
local ok, req = pcall(json.decode, raw_body)
if not ok then return raw_body end
req.disable_thinking = nil
req.extra_body = nil
if req.messages then
for _, msg in ipairs(req.messages) do
msg.reasoning_content = nil
end
end
return json.encode(req)
end
-- parse_text_tool_calls extracts tool calls that an upstream emitted as PLAIN
-- TEXT instead of using the OpenAI tool_calls field.
--
-- trae-local-api's OpenAI endpoint (/v1/chat/completions) does not read the
-- request's `tools` array at all, so the relayed model is never told the tool
-- schema; it falls back to printing
-- <tool_call>
-- {"name": "get_weather", "arguments": {"location": "北京"}}
-- </tool_call>
-- into message.content, leaves message.tool_calls null, and reports
-- finish_reason="stop". A client following the OpenAI contract therefore never
-- sees a tool call: the agent loop terminates unexpectedly mid-conversation
-- (and a hand-written replay produces an empty function name next turn).
--
-- Both tag spellings are accepted: the same codebase's Anthropic endpoint
-- instructs models to emit <toolcall>, and models mix the two. Key names vary
-- too (arguments / params / input), so all are tried.
--
-- Returns (tool_calls_array_or_nil, content_with_blocks_removed).
local function parse_text_tool_calls(content)
if type(content) ~= "string" or content == "" then return nil, content end
local tcs = {}
local idx = 0
local function collect(pattern)
for payload in content:gmatch(pattern) do
local ok, obj = pcall(json.decode, payload)
if ok and type(obj) == "table" then
-- some models wrap it as {"function":{"name":..,"arguments":..}}
local fn = obj["function"]
local name = obj.name or (type(fn) == "table" and fn.name) or nil
if name then
local args = obj.arguments or obj.params or obj.input
if args == nil and type(fn) == "table" then
args = fn.arguments or fn.params
end
if type(args) == "string" then
local aok, decoded = pcall(json.decode, args)
args = aok and decoded or {}
elseif type(args) ~= "table" then
args = {}
end
idx = idx + 1
table.insert(tcs, {
id = obj.id or ("call_text_" .. idx),
type = "function",
name = name,
arguments = args
})
end
end
end
end
-- `<tag ...>` allows attributes; %s* handles the "</tool_call >" spacing
-- that trae-local-api's own prompt example uses.
collect("<tool_call[^>]*>%s*(.-)%s*</tool_call%s*>")
collect("<toolcall[^>]*>%s*(.-)%s*</toolcall%s*>")
-- Legacy / mimicked form: models that saw past assistant turns folded as
-- [Called tool: name({...})]
-- tend to emit the same shape instead of a real block. Recover it so
-- the agent loop keeps working; otherwise the client receives plain text
-- where a structured tool call should be. Capture the name and the JSON
-- object literal independently (the JSON is the only {...} run here).
if #tcs == 0 and content:find("%[Called tool") then
for name, argsraw in content:gmatch("%[Called tool%s*:%s*([%w_%-%.]+)%s*%((%b{})%s*%)%s*%]") do
local aok, decoded = pcall(json.decode, argsraw)
idx = idx + 1
table.insert(tcs, {
id = "call_legacy_" .. idx,
type = "function",
name = name,
arguments = aok and decoded or {}
})
end
end
if #tcs == 0 then return nil, content end
-- drop the blocks from user-visible content; keep any surrounding prose
local stripped = content:gsub("<tool_call[^>]*>%s*.-%s*</tool_call%s*>", "")
stripped = stripped:gsub("<toolcall[^>]*>%s*.-%s*</toolcall%s*>", "")
stripped = stripped:gsub("%[Called tool%s*:%s*[%w_%-%.]+%s*%(%b{}%s*%)%s*%]", "")
stripped = stripped:gsub("^%s+", ""):gsub("%s+$", "")
return tcs, stripped
end
function adapter.transform_response(raw_body)
-- trae-local-api 偶发在非流式请求中返回 SSE 格式数据,
-- 表现为多个 data: {...} 行或混合了 reasoning_chunk 等。
-- 剥掉 data: 前缀,取最后一个完整的 JSON 块(通常是
-- 最终的 stop chunk 或 usage chunk)。
local use_json = raw_body
if raw_body:match("^data:") or raw_body:match("\ndata:") then
-- 取最后一个 data: 行的 JSON
local last = nil
for line in raw_body:gmatch("data: ([^\n\r]+)") do
local trimmed = line:match("^%s*(.-)%s*$")
if trimmed and trimmed ~= "[DONE]" then
local ok2, obj = pcall(json.decode, trimmed)
if ok2 and type(obj) == "table" then
last = obj
end
end
end
if last then
use_json = json.encode(last)
end
-- 如果解析失败,回退到原始 body
end
local ok, resp = pcall(json.decode, use_json)
if not ok or resp == nil then return raw_body end
local unified = {
content = "",
finish_reason = "",
token_usage = { prompt = 0, completion = 0, total = 0 }
}
if type(resp.usage) == "table" then
unified.token_usage.prompt = resp.usage.prompt_tokens or 0
unified.token_usage.completion = resp.usage.completion_tokens or 0
unified.token_usage.total = resp.usage.total_tokens or 0
end
if type(resp.choices) == "table" and #resp.choices > 0 then
local ch = resp.choices[1]
if type(ch.message) == "table" then
unified.content = ch.message.content or ""
if ch.message.reasoning_content then
unified.reasoning_content = ch.message.reasoning_content
end
if type(ch.message.tool_calls) == "table" and #ch.message.tool_calls > 0 then
local tcs = {}
for _, tc in ipairs(ch.message.tool_calls) do
local fn = tc["function"] or {}
local args = fn.arguments
if type(args) == "string" then
local args_ok, decoded = pcall(json.decode, args)
args = args_ok and decoded or {}
elseif type(args) ~= "table" then
args = {}
end
table.insert(tcs, {
id = tc.id,
type = tc.type or "function",
name = fn.name,
arguments = args
})
end
unified.tool_calls = tcs
end
end
unified.finish_reason = ch.finish_reason or ""
-- Fallback: trae-local-api's OpenAI endpoint drops the request's
-- `tools` array entirely, so the relayed model is never told the tool
-- schema and instead PRINTS a <tool_call>{...}</tool_call> block into
-- content, leaving message.tool_calls null and finish_reason="stop".
-- A client following the OpenAI contract then sees a normal completion
-- and its agent loop terminates mid-conversation. Recover the
-- structured call so the loop can continue.
if unified.tool_calls == nil then
local recovered, cleaned = parse_text_tool_calls(unified.content)
if recovered then
unified.tool_calls = recovered
unified.content = cleaned
unified.finish_reason = "tool_calls"
end
end
end
return json.encode(unified)
end
function adapter.transform_stream_chunk(raw_chunk)
local ok, chunk = pcall(json.decode, raw_chunk)
if not ok then return "" end
local uses = nil
if type(chunk.usage) == "table" then
uses = {
prompt = chunk.usage.prompt_tokens or chunk.usage.prompt or 0,
completion = chunk.usage.completion_tokens or chunk.usage.completion or 0,
total = chunk.usage.total_tokens or chunk.usage.total or 0,
}
if type(chunk.usage.prompt_tokens_details) == "table" and chunk.usage.prompt_tokens_details.cached_tokens ~= nil then
uses.prompt_tokens_details = { cached_tokens = chunk.usage.prompt_tokens_details.cached_tokens }
elseif (chunk.usage.prompt_cache_hit_tokens or 0) > 0 then
uses.prompt_cache_hit_tokens = chunk.usage.prompt_cache_hit_tokens
uses.prompt_cache_miss_tokens = chunk.usage.prompt_cache_miss_tokens or 0
uses.prompt_tokens_details = { cached_tokens = chunk.usage.prompt_cache_hit_tokens }
end
end
if not chunk.choices or #chunk.choices == 0 then
if uses ~= nil then
return json.encode({ usage = uses, done = false })
end
return ""
end
local delta = chunk.choices[1].delta or {}
local fr = chunk.choices[1].finish_reason
local finish = (type(fr) == "string" and fr ~= "") and fr or nil
local unified = {
content = delta.content or "",
done = (finish ~= nil)
}
if finish then
unified.finish_reason = finish
end
if uses ~= nil then
unified.usage = uses
end
if delta.reasoning_content then
unified.reasoning_content = delta.reasoning_content
end
if delta.tool_calls then
unified.tool_calls = delta.tool_calls
end
return json.encode(unified)
end
function adapter.transform_error(status, body)
local ok, resp = pcall(json.decode, body)
if not ok or type(resp) ~= "table" then
-- Non-JSON body: extract title from HTML or first line
local title = string.match(body or "", "<title>(.-)</title>")
if title and title ~= "" then return title end
local line = string.match(body or "", "^%s*([^\r\n]+)")
if line and line ~= "" and not string.match(line, "^<") then
return string.sub(line, 1, 200)
end
return nil
end
local e = resp.error
if type(e) == "table" and type(e.message) == "string" then
-- trae quota / rate limit annotations
if e.code then
return tostring(e.code) .. ": " .. e.message
end
return e.message
end
if type(e) == "string" then return e end
if type(resp.message) == "string" and resp.message ~= "" then
if resp.code ~= nil then
return tostring(resp.code) .. ": " .. resp.message
end
return resp.message
end
return nil
end
return adapter

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

@ -4,9 +4,15 @@
//
// The runtime uses LuaJIT (github.com/aarzilli/golua) so every adapter runs on
// an independent interpreter state. Because a lua.State is NOT goroutine-safe,
// each adapter owns a pool of workers (states); the pool size equals the sum
// of max_concurrent over all sources that use the adapter, and a worker is
// checked out for the duration of a single hook call.
// each adapter owns a pool of workers (states) and a worker is checked out for
// the duration of a single hook call.
//
// Pools are ELASTIC: the sum of max_concurrent over the sources using an adapter
// is a ceiling, not a preallocation. States are booted on demand in batches
// sized by that ceiling, and reclaimed by a single background janitor at a rate
// scaled inversely to the adapter's live connection count. An idle gateway
// therefore holds close to zero Lua states regardless of how much aggregate
// concurrency is configured. See adapterPool for the sizing rules.
package lua
import (
@ -23,6 +29,7 @@ import (
"strconv"
"strings"
"sync"
"time"
// golua exposes the LuaJIT C API. Build with -tags luajit so the cgo
// LDFLAGS resolve to -lluajit-5.1 (see golua's lua.go).
@ -64,7 +71,20 @@ type staticInfo struct {
}
// adapterPool owns the worker states for one adapter. Idle workers are held in
// a slice guarded by a cond; up to target workers are created lazily.
// a slice guarded by a cond; workers are created lazily on demand and reclaimed
// by the VM janitor once demand drops.
//
// Sizing is elastic rather than fixed (plan 阶段 2). Three inputs drive it:
//
// - maxW — the adapter's maximum concurrency (Σ max_concurrent of every source
// using it). It is the hard ceiling AND sets the growth step, so a hot,
// high-concurrency adapter ramps in batches while a low-concurrency one
// creeps up one state at a time.
// - created / len(idle) — how many states already exist, i.e. how much room is
// left to grow and how much is available to reclaim.
// - inUse — how many states are checked out right now, i.e. the live
// connection count for this adapter. It sets the shrink step: the fewer
// connections, the more aggressively idle states are released.
type adapterPool struct {
name string
script string
@ -73,24 +93,164 @@ type adapterPool struct {
mu sync.Mutex
cond *sync.Cond
idle []*worker
created int
target int
created int // live states: len(idle) + inUse (+ states being booted)
inUse int // states checked out right now == live connections
waiting int // goroutines blocked waiting for a free state
maxW int // hard ceiling = Σ max_concurrent of the sources using this adapter
used bool
closed bool
lastGrow time.Time
idleRounds int // consecutive janitor rounds that saw reclaimable slack
peakInUse int // high-water mark of inUse, for observability
}
const (
// growDivisor turns the adapter's max concurrency into the growth step CAP:
// step <= clamp(ceil(maxW / growDivisor), 1, growStepCap). The actual batch
// is additionally bounded by how many callers are queued (see growPlanLocked),
// so a wide adapter may ramp in big jumps but never warms states nobody is
// waiting for.
growDivisor = 8
growStepCap = 8
// growCooldown keeps a burst of misses from batching repeatedly while the
// previous batch is still booting.
growCooldown = time.Second
// residentWorkers is how many states an adapter keeps warm once it has
// served at least one request. Booting is milliseconds, but keeping one warm
// removes that from the critical path of the next request. Adapters that
// were never used keep nothing.
residentWorkers = 1
// idleHeadroom is the spare idle state kept above the live connection count,
// so the next concurrent request does not have to boot.
idleHeadroom = 1
// shrinkInterval is how often the VM janitor reclaims idle states.
shrinkInterval = 30 * time.Second
// shrinkGraceRounds is how many consecutive janitor rounds must see slack
// before anything is released. It guards against tearing a pool down between
// two bursts — NOT against ordinary traffic: a pool serving one request at a
// time still has reclaimable slack, and requiring "no traffic at all" would
// pin a burst's leftover states forever on any busy gateway. The counter is
// therefore reset by GROWTH (real contention), not by a mere checkout.
shrinkGraceRounds = 2
)
// PoolStats is a snapshot of one adapter pool's elastic sizing, surfaced by the
// status API so the algorithm is observable instead of a black box.
type PoolStats struct {
Name string `json:"name"`
Created int `json:"created"`
Idle int `json:"idle"`
InUse int `json:"in_use"`
Waiting int `json:"waiting"`
Max int `json:"max"`
Resident int `json:"resident"`
GrowStep int `json:"grow_step"`
ShrinkStep int `json:"shrink_step"`
PeakInUse int `json:"peak_in_use"`
}
func newAdapterPool(name, script string, static staticInfo) *adapterPool {
p := &adapterPool{name: name, script: script, static: static, target: 1}
p := &adapterPool{name: name, script: script, static: static, maxW: 1}
p.cond = sync.NewCond(&p.mu)
return p
}
func (p *adapterPool) setTarget(n int) {
// setMax updates the pool's hard ceiling (Σ max_concurrent of its sources).
// Lowering it does not kill live states: the janitor reclaims the excess.
func (p *adapterPool) setMax(n int) {
if n < 1 {
n = 1
}
p.mu.Lock()
p.target = n
p.maxW = n
p.cond.Broadcast()
p.mu.Unlock()
}
func ceilDiv(a, b int) int {
if b <= 0 {
return a
}
return (a + b - 1) / b
}
// growStepLocked is how many states to warm on a miss, derived from the
// adapter's maximum concurrency. Caller holds p.mu.
func (p *adapterPool) growStepLocked() int {
step := ceilDiv(p.maxW, growDivisor)
if step < 1 {
return 1
}
if step > growStepCap {
return growStepCap
}
return step
}
// residentLocked is the floor the janitor will not shrink below. Caller holds p.mu.
func (p *adapterPool) residentLocked() int {
if !p.used {
return 0
}
if residentWorkers > p.maxW {
return p.maxW
}
return residentWorkers
}
// shrinkStepLocked is how many idle states to release this round. The step is
// inversely proportional to the live connection count: with no connections the
// slack collapses in one round, while a busy adapter gives up one state at a
// time so the hot path keeps its warm states. Caller holds p.mu.
func (p *adapterPool) shrinkStepLocked() int {
excess := p.created - p.keepLocked()
if excess <= 0 {
return 0
}
step := ceilDiv(excess, 1+p.inUse)
if step > excess {
step = excess
}
if step > len(p.idle) {
step = len(p.idle)
}
if step < 0 {
return 0
}
return step
}
// keepLocked is the number of live states this pool should retain right now.
// Caller holds p.mu.
func (p *adapterPool) keepLocked() int {
keep := p.residentLocked()
if h := p.inUse + idleHeadroom; h > keep {
keep = h
}
if keep > p.maxW {
keep = p.maxW
}
return keep
}
func (p *adapterPool) stats() PoolStats {
p.mu.Lock()
defer p.mu.Unlock()
return PoolStats{
Name: p.name,
Created: p.created,
Idle: len(p.idle),
InUse: p.inUse,
Waiting: p.waiting,
Max: p.maxW,
Resident: p.residentLocked(),
GrowStep: p.growStepLocked(),
ShrinkStep: p.shrinkStepLocked(),
PeakInUse: p.peakInUse,
}
}
// boot creates a fresh LuaJIT state, registers the shared globals and executes
// the adapter script, storing the returned adapter table in the state.
func (p *adapterPool) boot() (*worker, error) {
@ -110,6 +270,96 @@ func (p *adapterPool) boot() (*worker, error) {
return &worker{L: L}, nil
}
// growPlanLocked reserves capacity for a batch prewarm and returns how many
// EXTRA states to boot in the background (the caller already reserved one for
// itself).
//
// Two independent limits apply, and BOTH matter:
//
// - the adapter's max concurrency caps the step (growStepLocked), so a
// high-concurrency adapter is allowed to ramp in bigger jumps than a
// narrow one;
// - the number of goroutines actually BLOCKED waiting for a state bounds it
// to real demand.
//
// Sizing on the ceiling alone over-provisions badly: an adapter with
// max_concurrent=76 has a step of 8, so two concurrent requests would warm 8
// states, and the janitor would throw 7 of them away a minute later — boot,
// discard, repeat, with RSS oscillating for no benefit. Prewarming only for
// goroutines that are genuinely queued keeps the ramp cheap without the churn.
//
// Caller holds p.mu.
func (p *adapterPool) growPlanLocked(contended bool) int {
if !contended || p.waiting <= 0 {
return 0
}
if time.Since(p.lastGrow) < growCooldown {
return 0
}
extra := p.growStepLocked() - 1
if extra > p.waiting {
extra = p.waiting // never warm more than the queue needs
}
if extra <= 0 {
return 0
}
if room := p.maxW - p.created; extra > room {
extra = room
}
if extra <= 0 {
return 0
}
// Reserve immediately so concurrent acquires respect the ceiling while the
// batch is still booting.
p.created += extra
p.lastGrow = time.Now()
return extra
}
// prewarm boots n reserved states and publishes them as idle. A boot failure
// releases its reservation; it is not fatal, since the synchronous path
// surfaces adapter errors already.
func (p *adapterPool) prewarm(n int) {
for i := 0; i < n; i++ {
p.mu.Lock()
closed := p.closed
p.mu.Unlock()
if closed {
p.releaseReservation(n - i)
return
}
w, err := p.boot()
if err != nil {
p.releaseReservation(n - i)
return
}
p.mu.Lock()
if p.closed {
p.created--
p.mu.Unlock()
w.close()
p.releaseReservation(n - i - 1)
return
}
p.idle = append(p.idle, w)
p.cond.Signal()
p.mu.Unlock()
}
}
func (p *adapterPool) releaseReservation(n int) {
if n <= 0 {
return
}
p.mu.Lock()
p.created -= n
if p.created < 0 {
p.created = 0
}
p.cond.Broadcast()
p.mu.Unlock()
}
func (p *adapterPool) acquire() (*worker, error) {
p.mu.Lock()
for {
@ -120,32 +370,68 @@ func (p *adapterPool) acquire() (*worker, error) {
if n := len(p.idle); n > 0 {
w := p.idle[n-1]
p.idle = p.idle[:n-1]
p.checkoutLocked()
p.mu.Unlock()
return w, nil
}
if p.created < p.target {
if p.created < p.maxW {
// Every existing state busy on a miss = real concurrency, not a
// sequential caller reusing one warm state.
contended := p.created > 0 && p.inUse >= p.created
p.created++
break
// Growing is the signal that capacity is genuinely short, so any
// pending shrink decision is stale.
p.idleRounds = 0
extra := p.growPlanLocked(contended)
p.mu.Unlock()
if extra > 0 {
go p.prewarm(extra)
}
w, err := p.boot()
if err != nil {
p.mu.Lock()
p.created--
p.cond.Signal()
p.mu.Unlock()
return nil, err
}
p.mu.Lock()
p.checkoutLocked()
p.mu.Unlock()
return w, nil
}
p.waiting++
p.cond.Wait()
p.waiting--
}
p.mu.Unlock()
}
w, err := p.boot()
if err != nil {
p.mu.Lock()
p.created--
p.cond.Signal()
p.mu.Unlock()
return nil, err
// checkoutLocked accounts one state as handed out. Caller holds p.mu.
func (p *adapterPool) checkoutLocked() {
p.inUse++
p.used = true
if p.inUse > p.peakInUse {
p.peakInUse = p.inUse
}
return w, nil
}
func (p *adapterPool) release(w *worker) {
w.L.SetTop(0)
p.mu.Lock()
if p.inUse > 0 {
p.inUse--
}
if p.closed {
p.created--
p.mu.Unlock()
w.close()
return
}
// A lowered ceiling is honoured on the spot rather than waiting for the
// janitor: returning a state the pool may no longer keep closes it.
if p.created > p.maxW {
p.created--
p.cond.Signal()
p.mu.Unlock()
w.close()
return
@ -155,11 +441,42 @@ func (p *adapterPool) release(w *worker) {
p.mu.Unlock()
}
// reclaim releases idle states down to the current keep target, using a step
// scaled by the live connection count. It returns how many states were closed.
func (p *adapterPool) reclaim() int {
p.mu.Lock()
if p.closed {
p.mu.Unlock()
return 0
}
step := p.shrinkStepLocked()
if step <= 0 {
p.idleRounds = 0
p.mu.Unlock()
return 0
}
p.idleRounds++
if p.idleRounds < shrinkGraceRounds {
p.mu.Unlock()
return 0
}
victims := append([]*worker(nil), p.idle[len(p.idle)-step:]...)
p.idle = p.idle[:len(p.idle)-step]
p.created -= step
p.idleRounds = 0
p.mu.Unlock()
for _, w := range victims {
w.close()
}
return step
}
func (p *adapterPool) shutdown() {
p.mu.Lock()
p.closed = true
idle := p.idle
p.idle = nil
p.created -= len(idle)
p.cond.Broadcast()
p.mu.Unlock()
for _, w := range idle {
@ -173,13 +490,84 @@ type VM struct {
mu sync.RWMutex
dir string
pools map[string]*adapterPool
// janitor drives elastic shrinking. One goroutine serves every pool (not
// one per adapter), so an idle gateway costs a single sleeping goroutine.
janitorOnce sync.Once
janitorStop chan struct{}
janitorDone chan struct{}
}
func NewVM(dir string) *VM {
return &VM{dir: dir, pools: map[string]*adapterPool{}}
return &VM{
dir: dir,
pools: map[string]*adapterPool{},
janitorStop: make(chan struct{}),
janitorDone: make(chan struct{}),
}
}
// startJanitor launches the single shrink loop. Idempotent.
func (v *VM) startJanitor() {
v.janitorOnce.Do(func() {
go v.janitorLoop(shrinkInterval)
})
}
func (v *VM) janitorLoop(every time.Duration) {
defer close(v.janitorDone)
t := time.NewTicker(every)
defer t.Stop()
for {
select {
case <-v.janitorStop:
return
case <-t.C:
v.reclaimIdle()
}
}
}
// reclaimIdle runs one shrink round over every pool. Exported behaviour is
// tested through ReclaimIdleNow.
func (v *VM) reclaimIdle() int {
v.mu.RLock()
pools := make([]*adapterPool, 0, len(v.pools))
for _, p := range v.pools {
pools = append(pools, p)
}
v.mu.RUnlock()
n := 0
for _, p := range pools {
n += p.reclaim()
}
return n
}
// ReclaimIdleNow forces one shrink round and reports how many worker states
// were closed. Intended for tests and for an explicit "release idle memory"
// action; normal operation relies on the janitor.
func (v *VM) ReclaimIdleNow() int { return v.reclaimIdle() }
// PoolStats returns the elastic sizing snapshot of every loaded adapter,
// sorted by name.
func (v *VM) PoolStats() []PoolStats {
v.mu.RLock()
pools := make([]*adapterPool, 0, len(v.pools))
for _, p := range v.pools {
pools = append(pools, p)
}
v.mu.RUnlock()
out := make([]PoolStats, 0, len(pools))
for _, p := range pools {
out = append(out, p.stats())
}
sort.Slice(out, func(i, j int) bool { return out[i].Name < out[j].Name })
return out
}
func (v *VM) Start() error {
v.startJanitor()
if v.dir == "" {
return nil
}
@ -214,6 +602,12 @@ func (v *VM) Start() error {
}
func (v *VM) Stop() {
select {
case <-v.janitorStop:
// already stopped
default:
close(v.janitorStop)
}
v.mu.Lock()
pools := v.pools
v.pools = map[string]*adapterPool{}
@ -223,17 +617,17 @@ func (v *VM) Stop() {
}
}
// ConfigureConcurrency sets each adapter's worker target to the sum of
// max_concurrent of every source using it. Values below one are clamped to 1.
// ConfigureConcurrency sets each adapter's worker CEILING to the sum of
// max_concurrent of every source using it. It is a ceiling, not a preallocation:
// states are booted on demand and reclaimed when demand drops, so a config with
// a large aggregate concurrency no longer implies a large resident pool.
// Values below one are clamped to 1.
func (v *VM) ConfigureConcurrency(adapterConcurrency map[string]int) {
v.mu.RLock()
defer v.mu.RUnlock()
for name, n := range adapterConcurrency {
if p, ok := v.pools[name]; ok {
if n < 1 {
n = 1
}
p.setTarget(n)
p.setMax(n)
}
}
}
@ -260,16 +654,19 @@ func (v *VM) LoadAdapter(path string) error {
}
// LoadAdapterSource registers an adapter from source code. Replacing an
// existing adapter shuts the old pool down and re-boots lazily.
// existing adapter shuts the old pool down and re-boots lazily, carrying the
// configured ceiling over so a hot-reload does not reset elastic sizing.
func (v *VM) LoadAdapterSource(name, code string) error {
static, err := inspectScript(name, code)
if err != nil {
return err
}
target := 1
maxW := 1
v.mu.RLock()
if p, ok := v.pools[static.name]; ok {
target = p.target
p.mu.Lock()
maxW = p.maxW
p.mu.Unlock()
}
v.mu.RUnlock()
@ -278,9 +675,10 @@ func (v *VM) LoadAdapterSource(name, code string) error {
p.shutdown()
}
p := newAdapterPool(static.name, code, *static)
p.setTarget(target)
p.setMax(maxW)
v.pools[static.name] = p
v.mu.Unlock()
v.startJanitor()
return nil
}
@ -407,6 +805,46 @@ func (v *VM) BuildHeaders(name string, meta map[string]interface{}) (map[string]
return headers, nil
}
// TransformError executes the optional adapter.transform_error(status, body)
// hook: per-source condensing of an upstream error response into a short
// client-facing reason. ok=false means the adapter defines no hook (or it
// failed) — the caller falls back to the generic Go-side condenser.
func (v *VM) TransformError(name string, status int, body string) (reason string, ok bool, err error) {
p := v.pool(name)
if p == nil {
return "", false, fmt.Errorf("adapter %s not loaded", name)
}
w, err := p.acquire()
if err != nil {
return "", false, err
}
defer p.release(w)
L := w.L
L.SetTop(0)
defer L.SetTop(0)
L.GetGlobal(adapterGlobal)
if L.IsNil(-1) {
return "", false, nil
}
L.GetField(-1, "transform_error")
if !L.IsFunction(-1) {
return "", false, nil
}
L.SetTop(0)
L.GetGlobal(adapterGlobal)
L.GetField(-1, "transform_error")
L.PushInteger(int64(status))
L.PushString(body)
if callErr := L.Call(2, 1); callErr != nil {
return "", false, fmt.Errorf("transform_error: %w", callErr)
}
if L.Type(-1) != golua.LUA_TSTRING {
return "", false, nil
}
return L.ToString(-1), true, nil
}
func (v *VM) Endpoint(name string) string {
p := v.pool(name)
if p == nil {

File diff suppressed because it is too large Load Diff

File diff suppressed because it is too large Load Diff

View File

@ -7,7 +7,9 @@ import (
"fmt"
"net/http"
"net/http/httptest"
"os"
"path/filepath"
"strings"
"sync/atomic"
"testing"
"time"
@ -196,8 +198,8 @@ func TestProviderAuthFailureSelfHeals(t *testing.T) {
if st.FailCount() != backoffCapN {
t.Fatalf("auth failure must jump to capped count, got %d", st.FailCount())
}
if until := st.CooldownUntil(); until <= time.Now().Add(25*time.Minute).Unix() {
t.Fatalf("auth failure must cool near the cap (until=%d)", until)
if until := st.CooldownUntil(); until <= time.Now().Add(authCooldown-time.Minute).Unix() {
t.Fatalf("auth failure must cool for authCooldown (until=%d)", until)
}
if st.Pref() != -2*int64(prefFailStep) {
t.Fatalf("auth failure pref penalty must be doubled, got %d", st.Pref())
@ -240,8 +242,81 @@ func TestModelStateCooldownAndRecovery(t *testing.T) {
if st.FailCount() != 0 {
t.Fatalf("fail count after success = %d", st.FailCount())
}
if st.Pref() != 1-int64(prefFailStep) {
t.Fatalf("pref after one failure (-5) then success (+1) = %d, want %d", st.Pref(), 1-int64(prefFailStep))
// failure penalties prefFailStep, success rewards prefSuccessStep:
// -5 then +2 = -3
if st.Pref() != int64(prefSuccessStep)-int64(prefFailStep) {
t.Fatalf("pref after one failure then success = %d, want %d", st.Pref(), int64(prefSuccessStep)-int64(prefFailStep))
}
}
func TestModelStateRateLimitShortCooldown(t *testing.T) {
// 429 must NOT use the exponential backoff schedule: a rate-limited but
// quota-rich source has to re-enter rotation after the short fixed window.
p := newTestProvider(t, src("mock", "http://127.0.0.1:1", "openai", "m"))
st := p.state("m")
// repeated 429s must stay at the fixed short cooldown, never escalate
for i := 0; i < 15; i++ {
p.ReportStatus("m", 429)
}
until := st.CooldownUntil()
want := time.Now().Add(rateLimitCooldown).Unix()
if until < want-2 || until > want+2 {
t.Fatalf("429 cooldown = %d, want ~%d (fixed %v, not exponential)", until, want, rateLimitCooldown)
}
if st.FailCount() != 15 {
t.Fatalf("fail count = %d, want 15 (counted but not escalating)", st.FailCount())
}
if st.Pref() != -15*int64(prefQuotaStep) {
t.Fatalf("pref after 15 x 429 = %d, want %d (a 429 is a quota-class event, not a fault)", st.Pref(), -15*int64(prefQuotaStep))
}
if p.ModelAvailable("m") {
t.Fatal("model must be cooling right after a 429")
}
}
func TestThrottleSpacingAndCancel(t *testing.T) {
p := newTestProvider(t, src("mock", "http://127.0.0.1:1", "openai", "m"))
p.cfg.RPM = 120 // gap = 500ms
p.rpmGap = time.Minute / time.Duration(p.cfg.RPM)
start := time.Now()
for i := 0; i < 3; i++ {
if err := p.Throttle(context.Background()); err != nil {
t.Fatalf("throttle %d: %v", i, err)
}
}
elapsed := time.Since(start)
// first call passes immediately, the next two wait one gap each
want := 2 * time.Minute / time.Duration(p.cfg.RPM)
if elapsed < want-time.Duration(100*time.Millisecond) || elapsed > want+time.Second {
t.Fatalf("3 throttled calls took %v, want ~%v", elapsed, want)
}
// unlimited source: Throttle is a no-op
p2 := newTestProvider(t, src("free", "http://127.0.0.1:1", "openai", "m"))
if err := p2.Throttle(context.Background()); err != nil {
t.Fatalf("unlimited throttle: %v", err)
}
// cancelled context aborts a pending window reservation
p3 := newTestProvider(t, src("slow", "http://127.0.0.1:1", "openai", "m"))
p3.cfg.RPM = 6 // 10s gap
p3.rpmGap = time.Minute / time.Duration(p3.cfg.RPM)
// first call passes immediately (fresh provider) and reserves the next window
if err := p3.Throttle(context.Background()); err != nil {
t.Fatalf("first throttle: %v", err)
}
ctx, cancel := context.WithTimeout(context.Background(), 50*time.Millisecond)
defer cancel()
// second call must wait ~10s for the reserved window; the short deadline aborts it
start = time.Now()
err := p3.Throttle(ctx)
if !errors.Is(err, context.DeadlineExceeded) {
t.Fatalf("cancelled throttle = %v, want DeadlineExceeded", err)
}
if time.Since(start) > time.Second {
t.Fatalf("cancel took %v, want fast abort", time.Since(start))
}
}
@ -345,7 +420,8 @@ func TestChatClientCancelNotRecorded(t *testing.T) {
}
// TestChatStreamEmptyBodyNotSuccess: a 200 that yields zero chunks and no
// [DONE] is a failure before the first chunk — the slot must back off.
// [DONE] must fail the candidate (error return, no delivered chunks) so the
// scheduler falls through to the next source, and record the failure.
func TestChatStreamEmptyBodyNotSuccess(t *testing.T) {
up := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
w.Header().Set("Content-Type", "text/event-stream")
@ -357,15 +433,12 @@ func TestChatStreamEmptyBodyNotSuccess(t *testing.T) {
Model: "m",
Messages: []types.ChatMessage{{Role: "user", Content: types.StringContent("x")}},
})
if err != nil {
t.Fatalf("stream: %v", err)
}
got := 0
for range ch {
got++
}
if got != 0 {
t.Fatalf("want empty stream, got %d chunks", got)
if err == nil {
if ch != nil {
for range ch {
}
}
t.Fatal("empty stream must fail the candidate for scheduler fallback")
}
eventually(t, 2*time.Second, func() bool {
_, fc, _ := p.ModelHealthInfo("m")
@ -373,6 +446,62 @@ func TestChatStreamEmptyBodyNotSuccess(t *testing.T) {
}, "empty stream must record a failure")
}
// TestChatStreamErrorFinishFailsCandidate: zen free pool answers HTTP 200
// with a single-chunk stream whose only payload is finish_reason:"network_error"
// and an empty delta. That must fail the candidate (not deliver a laundered
// empty reply), so the scheduler can fall through to the next source.
func TestChatStreamErrorFinishFailsCandidate(t *testing.T) {
up := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
w.Header().Set("Content-Type", "text/event-stream")
fmt.Fprint(w, `data: {"choices":[{"index":0,"finish_reason":"network_error","delta":{"role":"assistant","content":""}}]}`+"\n\n")
fmt.Fprint(w, "data: [DONE]\n\n")
}))
defer up.Close()
p := newTestProvider(t, src("mock", up.URL, "opencode", "m"))
_, err := p.ChatStream(context.Background(), &types.ChatRequest{
Model: "m",
Messages: []types.ChatMessage{{Role: "user", Content: types.StringContent("x")}},
})
if err == nil {
t.Fatal("network_error-only stream must fail the candidate")
}
eventually(t, 2*time.Second, func() bool {
_, fc, _ := p.ModelHealthInfo("m")
return fc >= 1
}, "error-only stream must record a failure")
}
// TestChatStreamInstantEmptyStopStillDelivered: a legitimate completion that
// ends immediately with the standard finish_reason:"stop" and zero content is
// NOT an upstream error and must still reach the client.
func TestChatStreamInstantEmptyStopStillDelivered(t *testing.T) {
up := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
w.Header().Set("Content-Type", "text/event-stream")
fmt.Fprint(w, `data: {"choices":[{"index":0,"finish_reason":"stop","delta":{}}]}`+"\n\n")
fmt.Fprint(w, "data: [DONE]\n\n")
}))
defer up.Close()
p := newTestProvider(t, src("mock", up.URL, "openai", "m"))
ch, err := p.ChatStream(context.Background(), &types.ChatRequest{
Model: "m",
Messages: []types.ChatMessage{{Role: "user", Content: types.StringContent("x")}},
})
if err != nil {
t.Fatalf("standard stop stream must be delivered: %v", err)
}
n := 0
for range ch {
n++
}
if n == 0 {
t.Fatal("expected at least the terminating chunk")
}
eventually(t, 2*time.Second, func() bool {
_, fc, _ := p.ModelHealthInfo("m")
return fc == 0
}, "standard stop must count as success")
}
// TestImageAutoUsesImageModel: AUTO image generation on a mixed source must
// send the image-kind model id, never the best chat model.
func TestImageAutoUsesImageModel(t *testing.T) {
@ -401,3 +530,656 @@ func TestImageAutoUsesImageModel(t *testing.T) {
t.Fatalf("AUTO image must use the image-kind model, got %q", gotModel)
}
}
func TestStandardSSEChunkFinishReason(t *testing.T) {
out := standardSSEChunk(`{"choices":[{"index":0,"finish_reason":"tool_calls","delta":{}}]}`)
if !strings.Contains(out, `"done":true`) || !strings.Contains(out, `"finish_reason":"tool_calls"`) {
t.Fatalf("real finish reason must pass through: %s", out)
}
out = standardSSEChunk(`{"choices":[{"index":0,"finish_reason":"","delta":{"content":"hi"}}]}`)
if strings.Contains(out, `"Done":true`) {
t.Fatalf("empty-string finish reason is not a finish signal: %s", out)
}
}
// TestAdapterHookCondensesError: the adapter transform_error hook owns the
// per-source error format; its reason must reach the client verbatim.
func TestAdapterHookCondensesError(t *testing.T) {
up := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
w.Header().Set("Content-Type", "application/json")
w.WriteHeader(503)
fmt.Fprint(w, `{"error":{"type":"FreeUsageLimitError","message":"Rate limit exceeded. Please try again later."}}`)
}))
defer up.Close()
p := newTestProvider(t, src("mock", up.URL, "opencode", "m"))
_, err := p.Chat(context.Background(), &types.ChatRequest{
Model: "m",
Messages: []types.ChatMessage{{Role: "user", Content: types.StringContent("x")}},
})
if err == nil || !strings.Contains(err.Error(), "zen free pool quota exhausted") {
t.Fatalf("adapter hook reason must surface, got %v", err)
}
if strings.Contains(err.Error(), "Rate limit exceeded") {
t.Fatalf("raw upstream body must not leak past the hook: %v", err)
}
}
// TestUnknownErrorFallbackWithoutHook: an adapter without transform_error
// gets the uniform core fallback; the raw body must not leak to clients.
func TestUnknownErrorFallbackWithoutHook(t *testing.T) {
dir := filepath.Join(t.TempDir(), "adapters")
if err := os.MkdirAll(dir, 0o755); err != nil {
t.Fatal(err)
}
minimal := `return { name="custom", endpoint="/chat/completions",
transform_request=function(raw) return raw end,
transform_response=function(raw) return raw end }
`
if err := os.WriteFile(filepath.Join(dir, "custom.lua"), []byte(minimal), 0o644); err != nil {
t.Fatal(err)
}
vm := lua.NewVM(dir)
if err := vm.Start(); err != nil {
t.Fatalf("vm: %v", err)
}
defer vm.Stop()
up := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
w.WriteHeader(500)
fmt.Fprint(w, `{"weird":{"shape":["no","message"]}}`)
}))
defer up.Close()
p := New(src("mock", up.URL, "custom", "m"), vm)
_, err := p.Chat(context.Background(), &types.ChatRequest{
Model: "m",
Messages: []types.ChatMessage{{Role: "user", Content: types.StringContent("x")}},
})
if err == nil || !strings.Contains(err.Error(), "unknown error") {
t.Fatalf("hook-less adapter must fall back to unknown error, got %v", err)
}
if strings.Contains(err.Error(), "shape") {
t.Fatalf("raw body must not leak: %v", err)
}
}
// TestToolIndexRemapper verifies sparse upstream tool-call indices are
// compacted to dense 0-based OpenAI ordinals (Issue 2). Anthropic emits the
// content_block ordinal, so a thinking block at index 0 pushes the first
// tool_use to index 1 — OpenAI clients accumulating by index would leave a
// hole at 0 and mis-assemble multi-tool responses.
func TestToolIndexRemapper(t *testing.T) {
t.Run("single tool after thinking block", func(t *testing.T) {
r := newToolIndexRemapper()
// content_block index 1 (thinking was 0) -> dense 0
out := r.remap(json.RawMessage(`[{"index":1,"id":"toolu_a","type":"function","function":{"name":"calc","arguments":""}}]`))
var got []map[string]interface{}
if err := json.Unmarshal(out, &got); err != nil {
t.Fatalf("unmarshal: %v (%s)", err, out)
}
if got[0]["index"] != float64(0) {
t.Fatalf("index = %v, want 0: %s", got[0]["index"], out)
}
// id/name/arguments must survive verbatim
if got[0]["id"] != "toolu_a" {
t.Fatalf("id lost: %s", out)
}
fn := got[0]["function"].(map[string]interface{})
if fn["name"] != "calc" {
t.Fatalf("function.name lost: %s", out)
}
// subsequent argument fragments on the SAME upstream index reuse ordinal 0
out2 := r.remap(json.RawMessage(`[{"index":1,"id":"","type":"function","function":{"name":"","arguments":"{\"a\":1}"}}]`))
var got2 []map[string]interface{}
if err := json.Unmarshal(out2, &got2); err != nil {
t.Fatalf("unmarshal2: %v", err)
}
if got2[0]["index"] != float64(0) {
t.Fatalf("fragment index = %v, want stable 0: %s", got2[0]["index"], out2)
}
})
t.Run("multiple tools get distinct dense ordinals", func(t *testing.T) {
r := newToolIndexRemapper()
// thinking=0, tool_use=1, tool_use=2 -> 0, 1
a := r.remap(json.RawMessage(`[{"index":1,"id":"t1","type":"function","function":{"name":"f1","arguments":""}}]`))
b := r.remap(json.RawMessage(`[{"index":2,"id":"t2","type":"function","function":{"name":"f2","arguments":""}}]`))
var ga, gb []map[string]interface{}
json.Unmarshal(a, &ga)
json.Unmarshal(b, &gb)
if ga[0]["index"] != float64(0) {
t.Fatalf("first tool index = %v, want 0", ga[0]["index"])
}
if gb[0]["index"] != float64(1) {
t.Fatalf("second tool index = %v, want 1", gb[0]["index"])
}
// interleaved fragments keep their own ordinals
a2 := r.remap(json.RawMessage(`[{"index":1,"id":"","type":"function","function":{"name":"","arguments":"x"}}]`))
var ga2 []map[string]interface{}
json.Unmarshal(a2, &ga2)
if ga2[0]["index"] != float64(0) {
t.Fatalf("t1 fragment index = %v, want 0", ga2[0]["index"])
}
})
t.Run("already dense indices pass through untouched", func(t *testing.T) {
r := newToolIndexRemapper()
in := json.RawMessage(`[{"index":0,"id":"t","type":"function","function":{"name":"f","arguments":""}}]`)
out := r.remap(in)
if string(out) != string(in) {
t.Fatalf("dense input was rewritten:\n in=%s\nout=%s", in, out)
}
})
t.Run("empty and malformed input is safe", func(t *testing.T) {
r := newToolIndexRemapper()
if got := r.remap(nil); got != nil {
t.Fatalf("nil -> %s", got)
}
bad := json.RawMessage(`not json`)
if got := r.remap(bad); string(got) != string(bad) {
t.Fatalf("malformed input must pass through: %s", got)
}
})
}
// ---- half-cooldown probe gate (plan 阶段 1) ----
// TestProbeGateHalfCooldown pins the core rule: a cooling model is fully silent
// in the first half of its cooldown window and probeable exactly once in the
// second half.
func TestProbeGateHalfCooldown(t *testing.T) {
var st ModelState
// open a 100s window centred so that "now" sits in the first half
now := time.Now().Unix()
st.cooldownFrom.Store(now - 10)
st.cooldownUntil.Store(now + 90)
if st.TryProbe() {
t.Fatal("first half of the cooldown window must stay silent")
}
// move into the second half
st.cooldownFrom.Store(now - 90)
st.cooldownUntil.Store(now + 10)
if !st.TryProbe() {
t.Fatal("second half of the cooldown window must allow one probe")
}
if st.TryProbe() {
t.Fatal("probe permit must be single-token")
}
if !st.Probing() {
t.Fatal("Probing() must report the claimed permit")
}
st.ProbeDone()
if st.Probing() {
t.Fatal("ProbeDone must release the permit")
}
if !st.TryProbe() {
t.Fatal("permit must be reusable after ProbeDone")
}
}
// TestProbeGateNoProbeWhenHealthy: a healthy model needs no probe permit — the
// caller schedules it normally.
func TestProbeGateNoProbeWhenHealthy(t *testing.T) {
var idle ModelState
if idle.TryProbe() {
t.Fatal("a healthy model must not be probed (schedule it normally)")
}
}
// TestProbeGateRescuesPrefFloor documents a fix that comes with the probe gate:
// a slot whose preference sank to prefMin used to be refused by ModelAvailable
// forever, even long after every cooldown expired — a permanent blacklist. It is
// now probeable, so one success brings it back.
func TestProbeGateRescuesPrefFloor(t *testing.T) {
var dead ModelState
dead.pref.Store(prefMin)
if !dead.TryProbe() {
t.Fatal("a slot stuck at prefMin must be probeable, not blacklisted")
}
dead.RecordSuccess()
dead.ProbeDone()
if dead.Pref() <= prefMin {
t.Fatalf("one probe success must lift the slot off the floor, pref=%d", dead.Pref())
}
}
// TestProbeSuccessRestoresImmediately is the user-visible payoff: a recovered
// upstream must not sit out the rest of its cooldown.
func TestProbeSuccessRestoresImmediately(t *testing.T) {
var st ModelState
for i := 0; i < backoffCapN; i++ {
st.RecordFailure(false)
}
if cd := st.CooldownUntil() - time.Now().Unix(); cd < int64(backoffCap.Seconds())-2 {
t.Fatalf("10 failures must cool for backoffCap, got %ds", cd)
}
// jump into the probe window: the midpoint of [from, until) must be in the
// past, so push the window's start back past its full length
now := time.Now().Unix()
st.cooldownFrom.Store(now - int64(backoffCap.Seconds()) - 10)
if !st.TryProbe() {
t.Fatal("probe must be allowed past the midpoint")
}
if !st.Probing() {
t.Fatal("probe permit must be held while the probe runs")
}
st.RecordSuccess()
st.ProbeDone()
if !st.Available() {
t.Fatal("probe success must clear the cooldown immediately")
}
if st.FailCount() != 0 {
t.Fatalf("probe success must reset failCount, got %d", st.FailCount())
}
if st.CooldownUntil() != 0 || st.CooldownFrom() != 0 {
t.Fatal("probe success must clear both cooldown bounds")
}
}
// TestProbeFailureDefersNextProbe: a failed probe must not immediately retry —
// the new window's midpoint pushes the next probe out.
func TestProbeFailureDefersNextProbe(t *testing.T) {
var st ModelState
st.RecordFailure(false) // 5s window
now := time.Now().Unix()
st.cooldownFrom.Store(now - 100)
st.cooldownUntil.Store(now + 10)
if !st.TryProbe() {
t.Fatal("probe must be allowed past the midpoint")
}
st.RecordFailure(false) // probe failed: a fresh, longer window opens
st.ProbeDone()
if st.TryProbe() {
t.Fatal("a failed probe must defer the next probe to the new midpoint")
}
if from := st.CooldownFrom(); from < now {
t.Fatalf("failed probe must reopen the window (from=%d, now=%d)", from, now)
}
}
// TestAuthCooldownIsSeparateFromCap documents that credential failures cool on
// their own (longer) schedule instead of reusing the exponential cap.
func TestAuthCooldownIsSeparateFromCap(t *testing.T) {
var st ModelState
st.RecordFailure(true)
cd := st.CooldownUntil() - time.Now().Unix()
if cd < int64(authCooldown.Seconds())-2 || cd > int64(authCooldown.Seconds())+2 {
t.Fatalf("auth failure must cool for authCooldown (%v), got %ds", authCooldown, cd)
}
if st.FailCount() != backoffCapN {
t.Fatalf("auth failure must persist the capped count, got %d", st.FailCount())
}
}
// TestQuotaExhaustedDoesNotEscalate is the second half of the user's complaint:
// running out of allowance must not be treated as a failure streak.
func TestQuotaExhaustedDoesNotEscalate(t *testing.T) {
var st ModelState
st.RecordQuotaExhausted(2 * time.Minute)
if st.FailCount() != 0 {
t.Fatalf("quota exhaustion must not bump failCount, got %d", st.FailCount())
}
cd := st.CooldownUntil() - time.Now().Unix()
if cd < 118 || cd > 122 {
t.Fatalf("quota cooldown must honour the reset window, got %ds", cd)
}
// an unbounded window is clamped so a topped-up allowance is still noticed
var big ModelState
big.RecordQuotaExhausted(72 * time.Hour)
if cd := big.CooldownUntil() - time.Now().Unix(); cd > int64(maxQuotaCooldown.Seconds())+2 {
t.Fatalf("quota cooldown must be capped at maxQuotaCooldown, got %ds", cd)
}
// no hint at all falls back to the short window
var none ModelState
none.RecordQuotaExhausted(0)
if cd := none.CooldownUntil() - time.Now().Unix(); cd > int64(rateLimitCooldown.Seconds())+2 {
t.Fatalf("unknown reset window must fall back to rateLimitCooldown, got %ds", cd)
}
}
// TestQuotaReasonRouting checks that a 429 carrying a quota phrase is cooled as
// quota exhaustion while a plain rate limit keeps the short fixed window.
func TestQuotaReasonRouting(t *testing.T) {
p := newTestProvider(t, src("q", "http://127.0.0.1:1", "openai", "m"))
p.ReportStatusReason("m", 429, "429 insufficient_quota: you exceeded your current quota")
st := p.state("m")
if st.FailCount() != 0 {
t.Fatalf("quota rejection must not bump failCount, got %d", st.FailCount())
}
if cd := st.CooldownUntil() - time.Now().Unix(); cd <= int64(rateLimitCooldown.Seconds()) {
t.Fatalf("quota rejection must cool longer than a plain 429, got %ds", cd)
}
p2 := newTestProvider(t, src("r", "http://127.0.0.1:1", "openai", "m"))
p2.ReportStatusReason("m", 429, "rate limit reached, slow down")
st2 := p2.state("m")
if cd := st2.CooldownUntil() - time.Now().Unix(); cd > int64(rateLimitCooldown.Seconds())+2 {
t.Fatalf("plain 429 must keep the short window, got %ds", cd)
}
}
// TestModelSchedulableProbeHandoff covers the provider-level gate the scheduler
// consumes.
func TestModelSchedulableProbeHandoff(t *testing.T) {
p := newTestProvider(t, src("s", "http://127.0.0.1:1", "openai", "m"))
if ok, probe := p.ModelSchedulable("m"); !ok || probe {
t.Fatalf("healthy model must be normally schedulable, got ok=%v probe=%v", ok, probe)
}
st := p.state("m")
st.RecordFailure(false)
now := time.Now().Unix()
st.cooldownFrom.Store(now - 100)
st.cooldownUntil.Store(now + 10)
ok, probe := p.ModelSchedulable("m")
if !ok || !probe {
t.Fatalf("cooling model past midpoint must be probeable, got ok=%v probe=%v", ok, probe)
}
if ok2, _ := p.ModelSchedulable("m"); ok2 {
t.Fatal("only one probe permit may be outstanding")
}
p.ProbeDone("m")
if ok3, probe3 := p.ModelSchedulable("m"); !ok3 || !probe3 {
t.Fatal("permit must be reclaimable after ProbeDone")
}
}
// TestProbeSlotsFollowsMaxConcurrent pins the user's example: max_concurrent=10
// yields exactly one probe slot.
func TestProbeSlotsFollowsMaxConcurrent(t *testing.T) {
for _, tc := range []struct{ mc, want int }{
{0, 1}, {1, 1}, {4, 1}, {10, 1}, {20, 2}, {100, 2},
} {
s := src("p", "http://127.0.0.1:1", "openai", "m")
s.MaxConcurrent = tc.mc
p := newTestProvider(t, s)
if got := p.ProbeSlots(); got != tc.want {
t.Fatalf("max_concurrent=%d: probe slots = %d, want %d", tc.mc, got, tc.want)
}
}
}
// TestModelProbeInfo checks the UI-facing view of the probe window.
func TestModelProbeInfo(t *testing.T) {
p := newTestProvider(t, src("i", "http://127.0.0.1:1", "openai", "m"))
if from, after, probing := p.ModelProbeInfo("m"); from != 0 || after != 0 || probing {
t.Fatalf("idle model must report no probe window, got %d %d %v", from, after, probing)
}
st := p.state("m")
now := time.Now().Unix()
st.cooldownFrom.Store(now)
st.cooldownUntil.Store(now + 100)
from, after, _ := p.ModelProbeInfo("m")
if from != now || after != now+50 {
t.Fatalf("probe window = (%d,%d), want (%d,%d)", from, after, now, now+50)
}
}
// TestUpstreamRecoveryWithoutWaitingOutCooldown is the end-to-end shape of the
// user's complaint: an upstream that 500s ten times lands in the capped
// cooldown, and used to be unusable for a full 30 minutes even after it
// recovered. With the probe gate the slot is silent for the first half of the
// window and then serves the very next request itself, clearing the cooldown.
//
// Real time is not waited out: the window bounds are moved so the test observes
// the same state the scheduler would see at cap/2.
func TestUpstreamRecoveryWithoutWaitingOutCooldown(t *testing.T) {
var healthy atomic.Bool
var hits atomic.Int64
up := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
hits.Add(1)
if !healthy.Load() {
w.WriteHeader(500)
return
}
fmt.Fprint(w, `{"choices":[{"message":{"content":"back"},"finish_reason":"stop"}]}`)
}))
defer up.Close()
p := newTestProvider(t, src("recov", up.URL, "openai", "m"))
req := func() *types.ChatRequest {
return &types.ChatRequest{
Model: "m",
Messages: []types.ChatMessage{{Role: "user", Content: types.StringContent("x")}},
}
}
// 1. drive the slot into the capped cooldown
for i := 0; i < backoffCapN; i++ {
if _, err := p.Chat(context.Background(), req()); err == nil {
t.Fatalf("attempt %d: expected upstream failure", i)
}
}
st := p.state("m")
if p.ModelAvailable("m") {
t.Fatal("10 consecutive 5xx must take the slot out of normal rotation")
}
cd := st.CooldownUntil() - time.Now().Unix()
if cd < int64(backoffCap.Seconds())-2 {
t.Fatalf("cooldown = %ds, want ~%v", cd, backoffCap)
}
// the old 30-minute lockout is gone
if cd > int64((6 * time.Minute).Seconds()) {
t.Fatalf("cooldown %ds is too long: the cap must be minutes, not half an hour", cd)
}
// 2. upstream recovers, but we are still in the FIRST half of the window:
// the slot must stay completely silent (no probe traffic at all).
healthy.Store(true)
quiet := hits.Load()
if ok, probe := p.ModelSchedulable("m"); ok || probe {
t.Fatalf("first half of the window must stay silent, got ok=%v probe=%v", ok, probe)
}
if hits.Load() != quiet {
t.Fatal("no request may reach the upstream during the silent half")
}
// 3. advance past the midpoint: exactly one probe is admitted.
st.cooldownFrom.Store(time.Now().Unix() - int64(backoffCap.Seconds()) - 1)
ok, isProbe := p.ModelSchedulable("m")
if !ok || !isProbe {
t.Fatalf("past the midpoint one probe must be admitted, got ok=%v probe=%v", ok, isProbe)
}
if ok2, _ := p.ModelSchedulable("m"); ok2 {
t.Fatal("only ONE probe may be in flight while cooling")
}
// 4. the probe is a real request: its success restores the slot fully.
resp, err := p.Chat(context.Background(), req())
p.ProbeDone("m")
if err != nil {
t.Fatalf("probe request failed: %v", err)
}
if resp.Content != "back" {
t.Fatalf("probe response = %q, want \"back\"", resp.Content)
}
if !p.ModelAvailable("m") {
t.Fatal("a successful probe must return the slot to normal rotation immediately")
}
if st.CooldownUntil() != 0 || st.FailCount() != 0 {
t.Fatalf("probe success must clear cooldown/failCount, got until=%d fails=%d",
st.CooldownUntil(), st.FailCount())
}
if ok, probe := p.ModelSchedulable("m"); !ok || probe {
t.Fatalf("recovered slot must schedule normally, got ok=%v probe=%v", ok, probe)
}
}
// TestRecoveryProbeBudgetIsBounded quantifies the cost of the shorter cap
// against a still-dead upstream: one probe per window, not a retry storm.
func TestRecoveryProbeBudgetIsBounded(t *testing.T) {
var hits atomic.Int64
up := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
hits.Add(1)
w.WriteHeader(500)
}))
defer up.Close()
p := newTestProvider(t, src("dead", up.URL, "openai", "m"))
req := &types.ChatRequest{
Model: "m",
Messages: []types.ChatMessage{{Role: "user", Content: types.StringContent("x")}},
}
for i := 0; i < backoffCapN; i++ {
_, _ = p.Chat(context.Background(), req)
}
baseline := hits.Load()
st := p.state("m")
// simulate three consecutive probe windows against a dead upstream
for i := 0; i < 3; i++ {
st.cooldownFrom.Store(time.Now().Unix() - int64(backoffCap.Seconds()) - 1)
ok, isProbe := p.ModelSchedulable("m")
if !ok || !isProbe {
t.Fatalf("window %d: expected a probe permit", i)
}
// a second caller in the same window gets nothing
if ok2, _ := p.ModelSchedulable("m"); ok2 {
t.Fatalf("window %d: more than one probe admitted", i)
}
if _, err := p.Chat(context.Background(), req); err == nil {
t.Fatalf("window %d: expected the probe to fail", i)
}
p.ProbeDone("m")
}
if extra := hits.Load() - baseline; extra != 3 {
t.Fatalf("3 probe windows must cost exactly 3 upstream requests, got %d", extra)
}
}
// TestQuotaPenaltyIsLighterThanFailure is the regression for the report
// "今天一次失败没看到,分数目前已经 -10 了" on an allowance-metered source
// (sensenova). RecordQuotaExhausted deliberately does not touch failCount — the
// WebUI therefore showed fails=0 and cooling=false — yet it used to charge the
// SAME score penalty as a real failure, so two ordinary quota resets dragged the
// slot to -10 with no visible fault. An exhausted allowance means "no budget
// right now", not "this model is broken": the cooldown already keeps it out of
// rotation, so the score penalty must be mild.
func TestQuotaPenaltyIsLighterThanFailure(t *testing.T) {
quota := &ModelState{}
quota.RecordQuotaExhausted(time.Minute)
fail := &ModelState{}
fail.RecordFailure(false)
qp, fp := quota.Pref(), fail.Pref()
if qp <= fp {
t.Fatalf("a quota reset must cost less score than a real failure: quota=%d failure=%d", qp, fp)
}
if qp != -prefQuotaStep {
t.Errorf("quota penalty = %d, want -%d", qp, prefQuotaStep)
}
if fp != -prefFailStep {
t.Errorf("failure penalty = %d, want -%d", fp, prefFailStep)
}
// The reported symptom: several quota resets in a day must NOT approach the
// prefMin floor, since each one is a normal event for a metered source.
s := &ModelState{}
for i := 0; i < 5; i++ {
s.RecordQuotaExhausted(time.Minute)
}
if got := s.Pref(); got <= prefMin/2 {
t.Errorf("5 quota resets sank the score to %d (floor %d); a metered source would be written off", got, prefMin)
}
if s.FailCount() != 0 {
t.Errorf("quota exhaustion must not inflate failCount, got %d", s.FailCount())
}
}
// TestRateLimitPenaltyIsMild covers the sibling case: a 429 means "too fast",
// not "broken", so a popular-but-healthy slot must not sink for being throttled.
func TestRateLimitPenaltyIsMild(t *testing.T) {
s := &ModelState{}
s.RecordRateLimit()
if got := s.Pref(); got != -prefQuotaStep {
t.Errorf("429 penalty = %d, want -%d (same class as a quota event)", got, prefQuotaStep)
}
}
// TestSuccessRecoversFasterThanBefore pins the penalty/reward asymmetry. Reward
// used to be +1 against a -5 penalty, so a slot at -10 needed TEN consecutive
// successes just to reach neutral — in practice it never got there, because a
// low score makes the slot unattractive to the scheduler in the first place.
func TestSuccessRecoversFasterThanBefore(t *testing.T) {
s := &ModelState{}
s.RecordFailure(false)
s.RecordFailure(false) // -10, the score from the bug report
n := 0
for s.Pref() < 0 {
s.RecordSuccess()
n++
if n > 50 {
t.Fatal("score never recovered")
}
}
// with +2 per success, -10 needs 5; the old +1 needed 10
if n > 5 {
t.Errorf("recovery from -10 took %d successes, want <= 5", n)
}
// a real failure must still outweigh a single success, or failures stop
// meaning anything
if prefSuccessStep >= prefFailStep {
t.Errorf("success reward %d must stay below the failure penalty %d", prefSuccessStep, prefFailStep)
}
}
// TestPrefDecayRehabilitatesIdleSlot closes the starvation loop: a penalised
// slot is unattractive, so the scheduler stops picking it, so it never earns the
// successes that would rehabilitate it. Without decay the score stays negative
// forever even though the upstream may be perfectly healthy.
func TestPrefDecayRehabilitatesIdleSlot(t *testing.T) {
s := &ModelState{}
s.RecordFailure(false)
s.RecordFailure(false)
if s.Pref() != -2*prefFailStep {
t.Fatalf("setup: pref=%d", s.Pref())
}
// backdate the last touch by three decay intervals
step := int64(prefDecayAfter / time.Second)
s.prefTouched.Store(time.Now().Unix() - 3*step)
got := s.Pref()
if got != -2*prefFailStep+3 {
t.Errorf("after 3 idle intervals pref=%d, want %d", got, -2*prefFailStep+3)
}
// decay must never overshoot into positive territory
s.prefTouched.Store(time.Now().Unix() - 1000*step)
if got := s.Pref(); got != 0 {
t.Errorf("decay overshot to %d, want to stop at 0", got)
}
// a positive score must not be dragged down for being idle
good := &ModelState{}
good.RecordSuccess()
good.RecordSuccess()
before := good.Pref()
good.prefTouched.Store(time.Now().Unix() - 1000*step)
if after := good.Pref(); after != before {
t.Errorf("idle decay must not touch a positive score: %d -> %d", before, after)
}
}
// TestDecayLiftsSlotOffTheFloor verifies decay actually restores schedulability:
// a slot at prefMin is refused by ModelAvailable, and before decay existed only
// a probe could rescue it.
func TestDecayLiftsSlotOffTheFloor(t *testing.T) {
s := &ModelState{}
for i := 0; i < 10; i++ {
s.RecordFailure(false)
}
if s.Pref() != prefMin {
t.Fatalf("setup: want floor %d, got %d", prefMin, s.Pref())
}
// clear the cooldown so only the score gates the slot
s.cooldownUntil.Store(0)
s.cooldownFrom.Store(0)
step := int64(prefDecayAfter / time.Second)
s.prefTouched.Store(time.Now().Unix() - 2*step)
if got := s.Pref(); got <= prefMin {
t.Errorf("decay failed to lift the slot off the floor: %d", got)
}
}

View File

@ -211,6 +211,18 @@ type SourceStatus struct {
FailCount int `json:"fail_count,omitempty"`
BackoffUntil int64 `json:"backoff_until,omitempty"`
Permanent bool `json:"permanent,omitempty"`
// RecentOK / RecentErr count real gateway requests served by this source
// within the last 300s (drives the status column so a source that is
// actually serving traffic can never be shown as down just because a
// probe was rate-limited by the upstream).
RecentOK int64 `json:"recent_ok,omitempty"`
RecentErr int64 `json:"recent_err,omitempty"`
// AvgFirstByteMs is the mean time-to-first-byte (ms) for recently
// successful requests served by this source (0 when unmeasured).
AvgFirstByteMs int64 `json:"avg_first_byte_ms,omitempty"`
// AvgTokPerS is the aggregate completion throughput (tokens/s) for
// recently successful requests served by this source (0 when no data).
AvgTokPerS int64 `json:"avg_tok_per_s,omitempty"`
}
// ProbeAll runs a live reachability check for every provider (in parallel).
@ -260,4 +272,4 @@ func (r *Registry) Status() []SourceStatus {
out = append(out, s)
}
return out
}
}

View File

@ -42,6 +42,12 @@ type Provider interface {
Name() string
ModelFor(reqModel string) string
ModelAvailable(model string) bool
// ModelSchedulable is the probe-aware availability gate: ok reports
// whether the model may take a request now, isProbe marks that it is only
// allowed as a cooldown probe (the caller must release the permit with
// ProbeDone once the attempt finished).
ModelSchedulable(model string) (ok bool, isProbe bool)
ProbeDone(model string)
Pref(model string) int64
Chat(ctx context.Context, req *types.ChatRequest) (*types.UnifiedResponse, error)
ChatStream(ctx context.Context, req *types.ChatRequest) (<-chan types.UnifiedChunk, error)
@ -176,18 +182,84 @@ type tierResult struct {
hard []TierError // hard failures seen in this pass (nil = none)
}
// candidate is one schedulable slot in a tier pass. probe marks a slot that is
// still cooling but past its half-cooldown mark and holding the probe permit:
// it is tried only after every normal slot, so probe traffic is what the tier
// falls back to instead of what it prefers.
type candidate struct {
slot *Slot
probe bool
}
// collectCands partitions a tier's slots into normal and probe candidates,
// normal first. Quota-exhausted slots are dropped outright. Every probe
// candidate returned holds a probe permit, so the caller MUST call
// releaseProbes on the result exactly once.
func collectCands(slots []*Slot, exhausted func(*Slot) bool) []candidate {
var normal, probes []candidate
for _, sl := range slots {
if exhausted != nil && exhausted(sl) {
continue
}
ok, isProbe := sl.Prov.ModelSchedulable(sl.Model)
if !ok {
continue
}
if isProbe {
probes = append(probes, candidate{slot: sl, probe: true})
continue
}
normal = append(normal, candidate{slot: sl})
}
return append(normal, probes...)
}
// releaseProbes hands every claimed probe permit back, whether or not the probe
// slot was actually used.
func releaseProbes(cands []candidate) {
for _, c := range cands {
if c.probe {
c.slot.Prov.ProbeDone(c.slot.Model)
}
}
}
// normalCount is how many leading candidates are normal (non-probe). The
// round-robin cursor rotates only over those: probe slots are a strictly
// ordered tail, never a rotation target.
func normalCount(cands []candidate) int {
for i, c := range cands {
if c.probe {
return i
}
}
return len(cands)
}
// runTier executes one tier pass starting at the round-robin base index.
// Cooldown is the only hard skip (re-verified per slot); a busy slot is
// skipped without any penalty; a hard failure is recorded and the pass moves
// on to the next slot (plan 2.3: "单请求内不重试已失败槽" — the failed slot is
// not retried, the others still are). hard == nil and no success means every
// candidate was merely busy/cooling, so the caller may wait a bounded time.
func runTier(ctx context.Context, tn *TierNode, cands []*Slot, base int64, req *types.ChatRequest, stream bool) tierResult {
//
// Normal candidates rotate by base; probe candidates form a fixed tail tried
// only after every normal slot failed or was busy.
func runTier(ctx context.Context, tn *TierNode, cands []candidate, base int64, req *types.ChatRequest, stream bool) tierResult {
n := len(cands)
norm := normalCount(cands)
var hard []TierError
for i := 0; i < n; i++ {
sl := cands[(int(base)+i)%n]
if !sl.Prov.ModelAvailable(sl.Model) {
var c candidate
if i < norm {
c = cands[(int(base)+i)%norm] // rotate within the normal head
} else {
c = cands[i] // probe tail keeps its order
}
sl := c.slot
// A probe candidate is intentionally NOT re-checked here: it is cooling
// by definition, and its permit was already claimed.
if !c.probe && !sl.Prov.ModelAvailable(sl.Model) {
continue
}
r := *req
@ -233,76 +305,40 @@ func (s *Scheduler) chainDrive(ctx context.Context, chain *Chain, req *types.Cha
}
var ce ChainErr
for _, tn := range chain.Tiers {
// initial filter: quota-exhausted and cooling slots are dropped
var cands []*Slot
for _, sl := range tn.Slots {
if exhausted != nil && exhausted(sl) {
continue
}
if !sl.Prov.ModelAvailable(sl.Model) {
continue
}
cands = append(cands, sl)
}
// initial filter: quota-exhausted slots are dropped, cooling slots are
// dropped unless they qualify as half-cooldown probes (appended last).
cands := collectCands(tn.Slots, exhausted)
if len(cands) == 0 {
ce.Skipped = append(ce.Skipped, fmt.Sprintf("tier %d: no schedulable slot (cooling or quota exhausted)", tn.Tier))
continue
}
// preference orders a same-tier run; stable so equal prefs keep order
sort.SliceStable(cands, func(i, j int) bool {
return cands[i].Prov.Pref(cands[i].Model) > cands[j].Prov.Pref(cands[j].Model)
})
// No Pref sort: load balancing is done by round-robin cursor.
// Persistently failing slots are excluded by ModelSchedulable
// (which checks Pref > prefMin).
base := tn.NextStart()
res := runTier(ctx, tn, cands, base, req, stream)
if res.resp != nil || res.chunks != nil {
releaseProbes(cands)
return res.resp, res.chunks, res.src, res.model, nil
}
if ctx.Err() != nil {
releaseProbes(cands)
return nil, nil, "", "", ctx.Err()
}
if len(res.hard) > 0 {
ce.Tiers = append(ce.Tiers, res.hard...)
releaseProbes(cands)
continue // hard failures: fall through to the next tier, no waiting
}
// every candidate was busy or cooling: bounded poll before downgrading
deadline := time.Now().Add(busyWait)
timer := time.NewTimer(busyPoll)
defer timer.Stop()
for {
select {
case <-ctx.Done():
return nil, nil, "", "", ctx.Err()
case <-timer.C:
if err := s.pollBusyTier(ctx, tn, cands, base, req, stream, &ce); err != nil {
releaseProbes(cands)
if r, ok := err.(*tierSuccess); ok {
return r.res.resp, r.res.chunks, r.res.src, r.res.model, nil
}
done := time.Now().After(deadline)
if done {
ce.Skipped = append(ce.Skipped, fmt.Sprintf("tier %d: no free slot within %v", tn.Tier, busyWait))
break
}
// refresh candidates: cooldowns may have expired meanwhile
var again []*Slot
for _, sl := range cands {
if sl.Prov.ModelAvailable(sl.Model) {
again = append(again, sl)
}
}
if len(again) == 0 {
ce.Skipped = append(ce.Skipped, fmt.Sprintf("tier %d: no free slot within %v", tn.Tier, busyWait))
break
}
res = runTier(ctx, tn, again, base, req, stream)
if res.resp != nil || res.chunks != nil {
return res.resp, res.chunks, res.src, res.model, nil
}
if ctx.Err() != nil {
return nil, nil, "", "", ctx.Err()
}
if len(res.hard) > 0 {
ce.Tiers = append(ce.Tiers, res.hard...)
break // hard failure while waiting: stop waiting, fall through
}
timer.Reset(busyPoll)
return nil, nil, "", "", err
}
releaseProbes(cands)
}
if len(ce.Tiers) == 0 && len(ce.Skipped) == 0 {
return nil, nil, "", "", fmt.Errorf("no auto slot configured")
@ -310,6 +346,57 @@ func (s *Scheduler) chainDrive(ctx context.Context, chain *Chain, req *types.Cha
return nil, nil, "", "", &ce
}
// tierSuccess carries a successful result out of pollBusyTier through the error
// return. It is never surfaced to callers of chainDrive.
type tierSuccess struct{ res tierResult }
func (t *tierSuccess) Error() string { return "tier success" }
// pollBusyTier waits a bounded time for a fully-busy tier to free a slot,
// retrying the pass while cooldowns expire. It returns nil when the tier should
// be abandoned (caller falls through to the next tier), a *tierSuccess when a
// retry succeeded, or a context error.
func (s *Scheduler) pollBusyTier(ctx context.Context, tn *TierNode, cands []candidate, base int64, req *types.ChatRequest, stream bool, ce *ChainErr) error {
deadline := time.Now().Add(busyWait)
timer := time.NewTimer(busyPoll)
defer timer.Stop()
for {
select {
case <-ctx.Done():
return ctx.Err()
case <-timer.C:
}
if time.Now().After(deadline) {
ce.Skipped = append(ce.Skipped, fmt.Sprintf("tier %d: no free slot within %v", tn.Tier, busyWait))
return nil
}
// refresh candidates: cooldowns may have expired meanwhile. Probe
// candidates keep their already-claimed permit and stay eligible.
var again []candidate
for _, c := range cands {
if c.probe || c.slot.Prov.ModelAvailable(c.slot.Model) {
again = append(again, c)
}
}
if len(again) == 0 {
ce.Skipped = append(ce.Skipped, fmt.Sprintf("tier %d: no free slot within %v", tn.Tier, busyWait))
return nil
}
res := runTier(ctx, tn, again, base, req, stream)
if res.resp != nil || res.chunks != nil {
return &tierSuccess{res: res}
}
if ctx.Err() != nil {
return ctx.Err()
}
if len(res.hard) > 0 {
ce.Tiers = append(ce.Tiers, res.hard...)
return nil // hard failure while waiting: stop waiting, fall through
}
timer.Reset(busyPoll)
}
}
// ChainChat runs a non-streaming AUTO request down the chain. exhausted, when
// non-nil, decides slot token-quota exhaustion. Returns the response, the
// serving source and the exact model id used; on total failure a *ChainErr
@ -327,6 +414,65 @@ func (s *Scheduler) ChainChatStream(ctx context.Context, chain *Chain, req *type
return chunks, src, model, err
}
// ChainImage runs an image-generation AUTO request down the chain: tiers
// ascending (tier 1 highest priority), per-tier round-robin, same-tier order
// by preference. Each slot's model is pinned to its own image id (ModelFor),
// so a fallback switches per source. Cooling-down slots are skipped. Returns
// the response, serving source and the exact model id used; on total failure
// a *ChainErr summarizing every tier.
func (s *Scheduler) ChainImage(ctx context.Context, chain *Chain, req *types.ImageGenRequest) (*types.UnifiedResponse, string, string, error) {
if chain == nil || len(chain.Tiers) == 0 {
return nil, "", "", fmt.Errorf("no image auto slot configured")
}
var ce ChainErr
for _, tn := range chain.Tiers {
cands := collectCands(tn.Slots, nil)
if len(cands) == 0 {
ce.Skipped = append(ce.Skipped, fmt.Sprintf("tier %d: no schedulable slot (cooling)", tn.Tier))
continue
}
// No Pref sort: load balancing is done by round-robin cursor.
base := tn.NextStart()
norm := normalCount(cands)
var hard []TierError
for i := 0; i < len(cands); i++ {
var c candidate
if i < norm {
c = cands[(int(base)+i)%norm]
} else {
c = cands[i]
}
sl := c.slot
if !c.probe && !sl.Prov.ModelAvailable(sl.Model) {
continue
}
r := *req
r.Model = sl.Model
resp, err := sl.Prov.Image(ctx, &r)
if ctx.Err() != nil {
releaseProbes(cands)
return nil, "", "", ctx.Err()
}
if err == nil {
releaseProbes(cands)
return resp, sl.Source, sl.Model, nil
}
if errors.Is(err, types.ErrBusy) {
continue
}
hard = append(hard, TierError{Tier: tn.Tier, Source: sl.Source, Model: sl.Model, Err: err})
}
releaseProbes(cands)
if len(hard) > 0 {
ce.Tiers = append(ce.Tiers, hard...)
}
}
if len(ce.Tiers) == 0 && len(ce.Skipped) == 0 {
return nil, "", "", fmt.Errorf("no image auto slot configured")
}
return nil, "", "", &ce
}
// ---- direct scheduling ----
// Chat runs a chat request across cands, falling back on failure. Each
@ -345,11 +491,15 @@ func (s *Scheduler) Chat(ctx context.Context, cands []Provider, req *types.ChatR
p := cands[i]
r := *req
r.Model = p.ModelFor(req.Model)
if !p.ModelAvailable(r.Model) {
ok, isProbe := p.ModelSchedulable(r.Model)
if !ok {
lastErr = fmt.Errorf("provider %s: model %q cooling down", p.Name(), r.Model)
continue
}
resp, err := p.Chat(ctx, &r)
if isProbe {
p.ProbeDone(r.Model)
}
if ctx.Err() != nil {
return nil, "", "", ctx.Err()
}
@ -377,11 +527,15 @@ func (s *Scheduler) ChatStream(ctx context.Context, cands []Provider, req *types
p := cands[i]
r := *req
r.Model = p.ModelFor(req.Model)
if !p.ModelAvailable(r.Model) {
ok, isProbe := p.ModelSchedulable(r.Model)
if !ok {
lastErr = fmt.Errorf("provider %s: model %q cooling down", p.Name(), r.Model)
continue
}
resp, err := p.ChatStream(ctx, &r)
if isProbe {
p.ProbeDone(r.Model)
}
if err == nil {
return resp, p.Name(), r.Model, nil
}
@ -400,11 +554,16 @@ func (s *Scheduler) Image(ctx context.Context, cands []Provider, req *types.Imag
var lastErr error
for i := 0; i < attempts && i < len(cands); i++ {
p := cands[i]
if !p.ModelAvailable(p.ModelFor(req.Model)) {
lastErr = fmt.Errorf("provider %s: model %q cooling down", p.Name(), p.ModelFor(req.Model))
im := p.ModelFor(req.Model)
ok, isProbe := p.ModelSchedulable(im)
if !ok {
lastErr = fmt.Errorf("provider %s: model %q cooling down", p.Name(), im)
continue
}
resp, err := p.Image(ctx, req)
if isProbe {
p.ProbeDone(im)
}
if err == nil {
return resp, p.Name(), nil
}

View File

@ -22,6 +22,14 @@ type fakeProvider struct {
busy atomic.Bool
fail atomic.Bool
chatHits atomic.Int64
// probeable makes an unavailable provider eligible as a half-cooldown
// probe candidate. probePermit is the single-token permit; probeClaims
// counts how many times it was handed out and probeDones how many times it
// was returned, so tests can assert the permit is always released.
probeable atomic.Bool
probePermit atomic.Bool
probeClaims atomic.Int64
probeDones atomic.Int64
}
func fakeProv(name, model string) *fakeProvider {
@ -36,6 +44,25 @@ func (f *fakeProvider) ModelFor(reqModel string) string { return f.model }
func (f *fakeProvider) ModelAvailable(model string) bool { return f.available.Load() }
func (f *fakeProvider) ModelSchedulable(model string) (bool, bool) {
if f.available.Load() {
return true, false
}
if !f.probeable.Load() {
return false, false
}
if f.probePermit.CompareAndSwap(false, true) {
f.probeClaims.Add(1)
return true, true
}
return false, false
}
func (f *fakeProvider) ProbeDone(model string) {
f.probeDones.Add(1)
f.probePermit.Store(false)
}
func (f *fakeProvider) Pref(model string) int64 { return f.pref.Load() }
func (f *fakeProvider) Chat(ctx context.Context, req *types.ChatRequest) (*types.UnifiedResponse, error) {
@ -177,14 +204,21 @@ func TestChainPreferenceSinksButStaysReachable(t *testing.T) {
if err != nil {
t.Fatalf("chain: %v", err)
}
// higher pref (good) is tried first and hard-fails; the pass moves on to
// the negative-pref slot, which sinks but stays reachable: with the real
// provider its success would RecordSuccess (+1 pref, self-heal)
// Round-robin starts at index 0 (neg). neg is available (pref=-5 > -20)
// and succeeds. good is never tried because neg already won.
if src != "neg" || model != "n" || resp.Content != "neg" {
t.Fatalf("served src=%q model=%q content=%q", src, model, resp.Content)
}
if good.chatHits.Load() == 0 {
t.Fatal("higher-pref slot must be tried first")
if good.chatHits.Load() != 0 {
t.Fatal("neg was tried first and succeeded; good must not be attempted")
}
// Second request: cursor advances. neg wins again (good hard-fails).
resp2, src2, _, err2 := s.ChainChat(context.Background(), ch, chatReq(), nil)
if err2 != nil {
t.Fatalf("second chain: %v", err2)
}
if src2 != "neg" || resp2.Content != "neg" {
t.Fatalf("second req src=%q content=%q", src2, resp2.Content)
}
}
@ -315,3 +349,149 @@ func TestChainResetCooldownAfterSwap(t *testing.T) {
t.Fatalf("first start = %d, want 0", base)
}
}
// ---- half-cooldown probe candidates (plan 阶段 1.3) ----
// TestChainProbeIsLastResort pins the ordering rule: a probe candidate must not
// steal traffic from a healthy slot in the same tier.
func TestChainProbeIsLastResort(t *testing.T) {
healthy := fakeProv("healthy", "h")
cooling := fakeProv("cooling", "c")
cooling.available.Store(false)
cooling.probeable.Store(true)
ch := BuildChain([]Rule{
{Tier: 0, Model: "c", Source: "cooling"}, // configured FIRST on purpose
{Tier: 0, Model: "h", Source: "healthy"},
}, bySource(healthy, cooling))
s := New(0)
for i := 0; i < 3; i++ {
_, src, _, err := s.ChainChat(context.Background(), ch, chatReq(), nil)
if err != nil {
t.Fatalf("iter %d: %v", i, err)
}
if src != "healthy" {
t.Fatalf("iter %d served by %q; a probe slot must never outrank a healthy one", i, src)
}
}
if cooling.chatHits.Load() != 0 {
t.Fatalf("probe slot was hit %d times while a healthy slot existed", cooling.chatHits.Load())
}
if got := cooling.probeDones.Load(); got != cooling.probeClaims.Load() {
t.Fatalf("probe permits leaked: claims=%d dones=%d", cooling.probeClaims.Load(), got)
}
}
// TestChainProbeServesWhenNothingElseCan is the recovery path: with every normal
// slot gone, the cooling slot's probe carries the request instead of the tier
// failing outright.
func TestChainProbeServesWhenNothingElseCan(t *testing.T) {
cooling := fakeProv("cooling", "c")
cooling.available.Store(false)
cooling.probeable.Store(true)
ch := BuildChain([]Rule{{Tier: 0, Model: "c", Source: "cooling"}}, bySource(cooling))
s := New(0)
resp, src, model, err := s.ChainChat(context.Background(), ch, chatReq(), nil)
if err != nil {
t.Fatalf("probe must serve the request: %v", err)
}
if src != "cooling" || model != "c" {
t.Fatalf("served by src=%q model=%q, want cooling/c", src, model)
}
if resp == nil || resp.Content != "cooling" {
t.Fatalf("bad response: %#v", resp)
}
if cooling.chatHits.Load() != 1 {
t.Fatalf("probe must issue exactly one request, got %d", cooling.chatHits.Load())
}
if cooling.probeClaims.Load() != 1 || cooling.probeDones.Load() != 1 {
t.Fatalf("permit accounting: claims=%d dones=%d, want 1/1", cooling.probeClaims.Load(), cooling.probeDones.Load())
}
}
// TestChainProbePermitReleasedOnFailure: a failed probe must still hand its
// permit back, otherwise the slot could never be probed again.
func TestChainProbePermitReleasedOnFailure(t *testing.T) {
cooling := fakeProv("cooling", "c")
cooling.available.Store(false)
cooling.probeable.Store(true)
cooling.fail.Store(true)
ch := BuildChain([]Rule{{Tier: 0, Model: "c", Source: "cooling"}}, bySource(cooling))
s := New(0)
if _, _, _, err := s.ChainChat(context.Background(), ch, chatReq(), nil); err == nil {
t.Fatal("expected the failing probe to surface an error")
}
if cooling.probeClaims.Load() != 1 || cooling.probeDones.Load() != 1 {
t.Fatalf("permit accounting: claims=%d dones=%d, want 1/1", cooling.probeClaims.Load(), cooling.probeDones.Load())
}
// permit is free again for the next attempt
if _, _, _, err := s.ChainChat(context.Background(), ch, chatReq(), nil); err == nil {
t.Fatal("expected the second probe to fail too")
}
if cooling.probeClaims.Load() != 2 {
t.Fatalf("permit must be reclaimable, claims=%d", cooling.probeClaims.Load())
}
}
// TestChainProbeDoesNotBlockTierFallthrough: a cooling tier-1 slot that is not
// probeable must still let the request fall through to tier 2.
func TestChainProbeDoesNotBlockTierFallthrough(t *testing.T) {
cold := fakeProv("cold", "c")
cold.available.Store(false) // cooling and NOT probeable
backup := fakeProv("backup", "b")
ch := BuildChain([]Rule{
{Tier: 1, Model: "c", Source: "cold"},
{Tier: 2, Model: "b", Source: "backup"},
}, bySource(cold, backup))
s := New(0)
_, src, _, err := s.ChainChat(context.Background(), ch, chatReq(), nil)
if err != nil || src != "backup" {
t.Fatalf("want fallthrough to backup, got src=%q err=%v", src, err)
}
if cold.chatHits.Load() != 0 {
t.Fatal("non-probeable cooling slot must not be hit")
}
}
// TestDirectProbeReleasesPermit covers the non-AUTO paths (Chat/ChatStream/Image).
func TestDirectProbeReleasesPermit(t *testing.T) {
cooling := fakeProv("cooling", "m")
cooling.available.Store(false)
cooling.probeable.Store(true)
s := New(0)
_, src, _, err := s.Chat(context.Background(), []Provider{cooling}, chatReq())
if err != nil || src != "cooling" {
t.Fatalf("direct probe must serve: src=%q err=%v", src, err)
}
if cooling.probeClaims.Load() != 1 || cooling.probeDones.Load() != 1 {
t.Fatalf("permit accounting: claims=%d dones=%d", cooling.probeClaims.Load(), cooling.probeDones.Load())
}
}
// TestChainProbeRoundRobinUnaffected: adding a probe tail must not disturb the
// round-robin rotation over the healthy head.
func TestChainProbeRoundRobinUnaffected(t *testing.T) {
a, b := fakeProv("s1", "a"), fakeProv("s2", "b")
cooling := fakeProv("s3", "c")
cooling.available.Store(false)
cooling.probeable.Store(true)
ch := BuildChain([]Rule{
{Tier: 0, Model: "a", Source: "s1"},
{Tier: 0, Model: "b", Source: "s2"},
{Tier: 0, Model: "c", Source: "s3"},
}, bySource(a, b, cooling))
s := New(0)
var got []string
for i := 0; i < 4; i++ {
_, src, _, err := s.ChainChat(context.Background(), ch, chatReq(), nil)
if err != nil {
t.Fatalf("iter %d: %v", i, err)
}
got = append(got, src)
}
want := []string{"s1", "s2", "s1", "s2"}
for i := range want {
if got[i] != want[i] {
t.Fatalf("rr order = %v, want %v (probe tail must not join the rotation)", got, want)
}
}
}

View File

@ -3,8 +3,10 @@
package types
import (
"bytes"
"encoding/json"
"errors"
"strings"
"time"
)
@ -18,15 +20,19 @@ var ErrBusy = errors.New("provider busy")
// ---- OpenAI wire request (gateway input) ----
type ChatRequest struct {
Model string `json:"model"`
Messages []ChatMessage `json:"messages"`
Temperature *float64 `json:"temperature,omitempty"`
MaxTokens int `json:"max_tokens,omitempty"`
Stream bool `json:"stream,omitempty"`
Tools []interface{} `json:"tools,omitempty"`
ToolChoice interface{} `json:"tool_choice,omitempty"`
DisableThinking bool `json:"disable_thinking"`
ExtraBody map[string]interface{} `json:"-"`
Model string `json:"model"`
Messages []ChatMessage `json:"messages"`
Temperature *float64 `json:"temperature,omitempty"`
MaxTokens int `json:"max_tokens,omitempty"`
Stream bool `json:"stream,omitempty"`
Tools []interface{} `json:"tools,omitempty"`
ToolChoice interface{} `json:"tool_choice,omitempty"`
DisableThinking bool `json:"disable_thinking"`
ExtraBody map[string]interface{} `json:"-"`
// ClientSession 是客户端自带的会话标识(请求头 x-session-affinity 等)。
// 不参与序列化:它只用于让适配器生成稳定的上游会话,不应泄漏给上游 body。
ClientSession string `json:"-"`
}
func (r *ChatRequest) MarshalJSON() ([]byte, error) {
@ -60,6 +66,39 @@ type ChatMessage struct {
func StringContent(s string) json.RawMessage { b, _ := json.Marshal(s); return b }
// isEmptyJSONArray reports whether raw is the literal empty JSON array [].
func isEmptyJSONArray(raw json.RawMessage) bool {
t := bytes.TrimSpace(raw)
return len(t) == 2 && t[0] == '[' && t[1] == ']'
}
// UnmarshalJSON decodes a chat message, normalising an empty-array content
// ("content":[]) to an empty string and dropping an empty tool_calls array.
//
// Why this exists: Lua adapters json.decode the request and re-encode it, and
// an empty Lua table is indistinguishable from an empty JSON array — the
// encoder emits {} for both. Agent clients serialise an assistant turn that
// carries tool_calls and no text as content:[], so every pass-through adapter
// turned it into content:{} — a shape that is not valid OpenAI (content is
// string | array of parts | null) and that real upstreams reject with
// "400 invalid arguments". Normalising at the decode boundary fixes every
// adapter at once, including ones added later.
func (m *ChatMessage) UnmarshalJSON(b []byte) error {
type alias ChatMessage
var a alias
if err := json.Unmarshal(b, &a); err != nil {
return err
}
*m = ChatMessage(a)
if isEmptyJSONArray(m.Content) {
m.Content = StringContent("")
}
if isEmptyJSONArray(m.ToolCalls) {
m.ToolCalls = nil
}
return nil
}
type ToolCall struct {
ID string `json:"id"`
Type string `json:"type"`
@ -70,6 +109,7 @@ type ToolCall struct {
// ---- Unified internal representation (what adapters produce) ----
type UnifiedResponse struct {
Model string `json:"model,omitempty"`
Content string `json:"content"`
ReasoningContent string `json:"reasoning_content,omitempty"`
FinishReason string `json:"finish_reason,omitempty"`
@ -77,12 +117,83 @@ type UnifiedResponse struct {
ToolCalls []ToolCall `json:"tool_calls,omitempty"`
// ImageData used by image-generation adapters.
ImageData []ImageData `json:"image_data,omitempty"`
// Cost is the upstream-reported charge for this request, passed through
// verbatim (OpenCode Zen/Go put a decimal **string** at the response top
// level). Deliberately not parsed or summed by the gateway: it is an
// upstream fact, and only some upstreams report it at all.
Cost string `json:"cost,omitempty"`
}
type TokenUsage struct {
Prompt int `json:"prompt"`
Completion int `json:"completion"`
Total int `json:"total"`
// PromptTokensDetails mirrors the OpenAI v2 usage.prompt_tokens_details
// object so cache-hit counts reported by OpenAI-compatible upstreams
// (and by adapters that normalize their own cache fields into it) pass
// through to clients that read it — dsh reads cached_tokens from here.
PromptTokensDetails *PromptTokensDetails `json:"prompt_tokens_details,omitempty"`
// PromptCacheHit / PromptCacheMiss carry the DeepSeek-legacy standalone
// fields; dsh falls back to prompt_cache_hit_tokens when
// prompt_tokens_details.cached_tokens is absent.
PromptCacheHit int `json:"prompt_cache_hit_tokens,omitempty"`
PromptCacheMiss int `json:"prompt_cache_miss_tokens,omitempty"`
// CompletionTokensDetails mirrors the OpenAI v2 usage.completion_tokens_details
// object. OpenCode Zen/Go report how much of the completion was thinking
// tokens here; without it clients cannot tell how much of the billed output
// was reasoning rather than visible answer.
CompletionTokensDetails *CompletionTokensDetails `json:"completion_tokens_details,omitempty"`
}
// CompletionTokensDetails is the OpenAI v2 completion_tokens_details object.
type CompletionTokensDetails struct {
ReasoningTokens int `json:"reasoning_tokens"`
}
// PromptTokensDetails is the OpenAI v2 prompt_tokens_details object. Only
// CachedTokens is emitted (omitempty drops the whole object when zero).
type PromptTokensDetails struct {
// CachedTokens is always emitted (even 0) so clients can distinguish
// "upstream reports cache, this request missed" from "no cache data".
CachedTokens int `json:"cached_tokens"`
}
// MarshalJSON emits both the legacy short keys (prompt/completion/total, used
// by the internal unified representation and older clients) and the OpenAI
// standard keys (prompt_tokens/completion_tokens/total_tokens). Standard
// clients such as DSH and DevEco Code read the *_tokens fields.
func (t TokenUsage) MarshalJSON() ([]byte, error) {
// Auto-generate prompt_tokens_details from legacy DeepSeek fields when
// the upstream adapter only set the standalone hit count (openai.lua
// does this in Lua, but other adapters or the standardSSEChunk fallback
// may not). dsh reads prompt_tokens_details.cached_tokens first.
pdetails := t.PromptTokensDetails
if pdetails == nil && t.PromptCacheHit > 0 {
pdetails = &PromptTokensDetails{CachedTokens: t.PromptCacheHit}
}
return json.Marshal(struct {
PromptTokens int `json:"prompt_tokens"`
CompletionTokens int `json:"completion_tokens"`
TotalTokens int `json:"total_tokens"`
Prompt int `json:"prompt"`
Completion int `json:"completion"`
Total int `json:"total"`
PromptTokensDetails *PromptTokensDetails `json:"prompt_tokens_details,omitempty"`
PromptCacheHitTokens int `json:"prompt_cache_hit_tokens,omitempty"`
PromptCacheMissTokens int `json:"prompt_cache_miss_tokens,omitempty"`
CompletionDetails *CompletionTokensDetails `json:"completion_tokens_details,omitempty"`
}{
PromptTokens: t.Prompt,
CompletionTokens: t.Completion,
TotalTokens: t.Total,
Prompt: t.Prompt,
Completion: t.Completion,
Total: t.Total,
PromptTokensDetails: pdetails,
PromptCacheHitTokens: t.PromptCacheHit,
PromptCacheMissTokens: t.PromptCacheMiss,
CompletionDetails: t.CompletionTokensDetails,
})
}
type ImageData struct {
@ -112,10 +223,22 @@ type ImageGenResponse struct {
// streaming tool_calls array (incremental fragments with an index field), which
// OpenAI-compatible clients accumulate themselves.
type UnifiedChunk struct {
Content string `json:"content"`
Done bool `json:"done"`
Content string `json:"content"`
Done bool `json:"done"`
// FinishReason carries the upstream finish/stop reason ("tool_calls",
// "length", ...) when the adapter provides it; the gateway emits it on
// the terminating chunk instead of the default "stop".
FinishReason string `json:"finish_reason,omitempty"`
ToolCalls json.RawMessage `json:"tool_calls,omitempty"`
ReasoningContent string `json:"reasoning_content,omitempty"`
// Usage carries the upstream token usage when the stream chunk provides
// it (OpenAI-style streams attach usage to some chunks; the final one
// 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
@ -128,4 +251,15 @@ type BuildMeta struct {
Source map[string]interface{} `json:"source"`
}
func Now() int64 { return time.Now().Unix() }
func Now() int64 { return time.Now().Unix() }
// OneLine flattens an error string to a single line capped at n chars.
// Shared by the provider layer (upstream error reasons) and the gateway
// layer (client-facing AUTO-chain summaries).
func OneLine(s string, n int) string {
s = strings.Join(strings.Fields(s), " ")
if len(s) > n {
s = s[:n] + "..."
}
return s
}

View File

@ -0,0 +1,154 @@
package types
import (
"encoding/json"
"strings"
"testing"
)
func TestTokenUsageMarshalCacheFields(t *testing.T) {
// DeepSeek-style: prompt_cache_hit_tokens + prompt_tokens_details auto-generated
u := TokenUsage{
Prompt: 100,
Completion: 50,
Total: 150,
PromptCacheHit: 80,
PromptCacheMiss: 20,
}
b, _ := json.Marshal(u)
var out map[string]interface{}
json.Unmarshal(b, &out)
check := func(name string, want interface{}) {
got := out[name]
if got != want {
t.Errorf("%s = %v (want %v)", name, got, want)
}
}
check("prompt_tokens", 100.0)
check("completion_tokens", 50.0)
check("prompt_cache_hit_tokens", 80.0)
check("prompt_cache_miss_tokens", 20.0)
details := out["prompt_tokens_details"]
if details == nil {
t.Error("prompt_tokens_details is missing (should be set from prompt_cache_hit_tokens?)")
} else {
d := details.(map[string]interface{})
if d["cached_tokens"] != 80.0 {
t.Errorf("cached_tokens = %v (want 80)", d["cached_tokens"])
}
}
// OpenAI-style: prompt_tokens_details.cached_tokens
u2 := TokenUsage{
Prompt: 200,
Completion: 100,
Total: 300,
PromptTokensDetails: &PromptTokensDetails{
CachedTokens: 150,
},
}
b2, _ := json.Marshal(u2)
json.Unmarshal(b2, &out)
check2 := func(name string, want interface{}) {
got := out[name]
if got != want {
t.Errorf("OpenAI %s = %v (want %v)", name, got, want)
}
}
check2("prompt_tokens", 200.0)
details2 := out["prompt_tokens_details"]
if details2 == nil {
t.Error("OpenAI prompt_tokens_details is missing")
} else {
d := details2.(map[string]interface{})
if d["cached_tokens"] != 150.0 {
t.Errorf("OpenAI cached_tokens = %v (want 150)", d["cached_tokens"])
}
}
t.Logf("DeepSeek JSON: %s", string(b))
t.Logf("OpenAI JSON: %s", string(b2))
}
// TestChatMessageNormalizesEmptyArrayContent pins the fix for the malformed
// content:{} that pass-through adapters produced for an assistant turn carrying
// tool_calls with no text. Lua adapters cannot tell an empty JSON array from an
// empty object, so the decode boundary normalises it for every adapter at once.
func TestChatMessageNormalizesEmptyArrayContent(t *testing.T) {
var msg ChatMessage
raw := `{"role":"assistant","content":[],"tool_calls":[{"id":"call_1","type":"function","function":{"name":"read","arguments":"{}"}}]}`
if err := json.Unmarshal([]byte(raw), &msg); err != nil {
t.Fatal(err)
}
if string(msg.Content) != `""` {
t.Fatalf("empty-array content must normalise to \"\", got %s", msg.Content)
}
if len(msg.ToolCalls) == 0 {
t.Fatal("tool_calls must survive normalisation")
}
// A non-empty content array must be left byte-identical (multimodal path).
var mm ChatMessage
multi := `{"role":"user","content":[{"type":"text","text":"hi"}]}`
if err := json.Unmarshal([]byte(multi), &mm); err != nil {
t.Fatal(err)
}
if string(mm.Content) != `[{"type":"text","text":"hi"}]` {
t.Fatalf("multimodal content must pass through untouched, got %s", mm.Content)
}
// A plain string content is untouched.
var sm ChatMessage
if err := json.Unmarshal([]byte(`{"role":"user","content":"plain"}`), &sm); err != nil {
t.Fatal(err)
}
if string(sm.Content) != `"plain"` {
t.Fatalf("string content must pass through untouched, got %s", sm.Content)
}
// An omitted content stays omitted (omitempty semantics preserved).
var om ChatMessage
if err := json.Unmarshal([]byte(`{"role":"assistant","tool_calls":[]}`), &om); err != nil {
t.Fatal(err)
}
if om.Content != nil {
t.Fatalf("absent content must stay absent, got %s", om.Content)
}
if om.ToolCalls != nil {
t.Fatalf("empty tool_calls array must be dropped, got %s", om.ToolCalls)
}
}
// 上游报告输出里有多少是思考 token;缺少这个字段,客户端无法判断
// completion_tokens 里多少是可见回答、多少是思考(两者都按输出计费)。
func TestTokenUsageMarshalsCompletionTokensDetails(t *testing.T) {
u := TokenUsage{
Prompt: 37, Completion: 40, Total: 77,
CompletionTokensDetails: &CompletionTokensDetails{ReasoningTokens: 40},
}
b, err := json.Marshal(u)
if err != nil {
t.Fatal(err)
}
var got map[string]interface{}
if err := json.Unmarshal(b, &got); err != nil {
t.Fatal(err)
}
ctd, ok := got["completion_tokens_details"].(map[string]interface{})
if !ok {
t.Fatalf("completion_tokens_details 未输出: %s", b)
}
if n, _ := ctd["reasoning_tokens"].(float64); int(n) != 40 {
t.Errorf("reasoning_tokens = %v, want 40 (%s)", ctd["reasoning_tokens"], b)
}
// 标准 OpenAI 字段同时存在,老客户端不受影响。
if _, ok := got["completion_tokens"]; !ok {
t.Errorf("标准 completion_tokens 丢失: %s", b)
}
// 未设置时不得凭空造字段。
b2, _ := json.Marshal(TokenUsage{Prompt: 1, Completion: 2, Total: 3})
if strings.Contains(string(b2), "completion_tokens_details") {
t.Errorf("无数据时不应输出 completion_tokens_details: %s", b2)
}
}

90
packaging/core-dist.sh Executable file
View File

@ -0,0 +1,90 @@
#!/usr/bin/env bash
# Build the llmsproxy CORE distribution packages (deb + rpm + tar.gz).
#
# Unlike the GUI packages (Electron), the core is a single static-ish binary
# that server admins install with dpkg/rpm. The systemd unit bakes in the
# memory tuning (GOGC=50, MALLOC_ARENA_MAX=2) measured on the production box.
#
# Requires: go, nfpm (PATH), dpkg-deb/rpmbuild for verification
# Outputs: cmd/build/dist/llmsproxy_<ver>_amd64.deb
# cmd/build/dist/llmsproxy-<ver>.x86_64.rpm
# cmd/build/dist/llmsproxy-<ver>-linux-amd64.tar.gz
set -euo pipefail
ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)"
cd "$ROOT"
VERSION="${1:-$(git describe --tags --dirty 2>/dev/null | sed 's/^v//' || echo 0.1.0)}"
DIST="$ROOT/cmd/build/dist"
BUILD_BIN="$ROOT/cmd/build/llmsproxy"
log() { printf '\033[1;34m[core-dist]\033[0m %s\n' "$*"; }
# nfpm is installed via `go install` -> lives in GOPATH/bin, which may not be
# on the caller's PATH (non-interactive shells, cron). Resolve it explicitly.
NFPM_BIN="$(command -v nfpm || true)"
[ -z "$NFPM_BIN" ] && [ -n "$(go env GOPATH)" ] && {
[ -x "$(go env GOPATH)/bin/nfpm" ] && NFPM_BIN="$(go env GOPATH)/bin/nfpm"
}
[ -n "$NFPM_BIN" ] || {
echo "[core-dist] nfpm missing. Install: go install github.com/goreleaser/nfpm/v2/cmd/nfpm@v2.41.2"
exit 1
}
mkdir -p "$DIST"
# 1) build the core binary (luajit tag, -s -w like deploy.sh)
log "building core binary (luajit, stripped)..."
CGO_ENABLED=1 go build -tags luajit -trimpath -ldflags="-s -w" -o "$BUILD_BIN" ./cmd/llmsproxy
# 2) sanity gates: binary must contain the expected symbols and not be degenerate.
# NOTE under `set -o pipefail`, `strings | grep -q` is BROKEN: grep -q exits
# on first hit and closes the pipe, strings gets SIGPIPE -> pipeline returns
# 141 -> the check spuriously fails. Stage strings output in a temp file.
SYMS_TMP="$(mktemp)"
strings "$BUILD_BIN" > "$SYMS_TMP"
for sym in ProbeSlots reqRing shrinkStepLocked mergeUsage; do
grep -qF "$sym" "$SYMS_TMP" || {
echo "[core-dist] FATAL: binary missing expected symbol '$sym'"; rm -f "$SYMS_TMP"; exit 1
}
done
rm -f "$SYMS_TMP"
[ "$(stat -c%s "$BUILD_BIN")" -gt 5000000 ] || {
echo "[core-dist] FATAL: core binary implausibly small"; exit 1
}
log " ok: $(stat -c%s "$BUILD_BIN") bytes, symbols present"
# 3) deb + rpm via nfpm. nfpm v2 here does not render {{ .Env.VERSION }}, so
# stamp the version into a throwaway copy of the config; -t places the
# generated package in the dist dir.
log "packaging deb + rpm (v$VERSION)..."
NFPM_CFG="$(mktemp)" && sed "s/__VERSION__/$VERSION/g" "$ROOT/packaging/nfpm.yaml" > "$NFPM_CFG"
"$NFPM_BIN" package --config "$NFPM_CFG" --packager deb -t "$DIST"
"$NFPM_BIN" package --config "$NFPM_CFG" --packager rpm -t "$DIST"
rm -f "$NFPM_CFG"
# 4) tar.gz
log "packaging tar.gz..."
TAR_DIR="$DIST/llmsproxy-$VERSION-linux-amd64"
rm -rf "$TAR_DIR"; mkdir -p "$TAR_DIR"
cp "$BUILD_BIN" "$TAR_DIR/llmsproxy"
cp "$ROOT/config.yaml" "$TAR_DIR/config.example.yaml"
cp "$ROOT/packaging/llmsproxy.service" "$TAR_DIR/llmsproxy.service"
mkdir -p "$TAR_DIR/adapters"
cp "$ROOT"/internal/lua/adapters/*.lua "$TAR_DIR/adapters/"
tar -C "$DIST" -czf "$DIST/llmsproxy-$VERSION-linux-amd64.tar.gz" \
"$(basename "$TAR_DIR")"
rm -rf "$TAR_DIR"
# 5) verify artifacts are healthy (the same gate that would have caught the
# 264KB exe fiasco)
log "verifying packages..."
for f in "$DIST"/llmsproxy_*.deb "$DIST"/llmsproxy-*.rpm "$DIST"/*.tar.gz; do
[ -f "$f" ] || continue
sz=$(stat -c%s "$f")
[ "$sz" -gt 100000 ] || { echo "[core-dist] FATAL: $f suspiciously small ($sz B)"; exit 1; }
log " ✓ $f ($((sz/1024)) KB)"
done
dpkg-deb -I "$DIST"/llmsproxy_*.deb 2>/dev/null | grep -E 'Package|Version' | head -2
log "done: $DIST"

View File

@ -0,0 +1,24 @@
[Unit]
Description=LLMSProxy - unified OpenAI-compatible multi-source LLM gateway
After=network.target
[Service]
Type=simple
# Memory tuning (measured, see README "内存占用"):
# MALLOC_ARENA_MAX=2 caps glibc per-thread malloc arenas. LuaJIT allocates
# through cgo -> glibc malloc, and glibc defaults to 8*nproc arenas, so every
# OS thread that touches malloc reserved its own ~1 MB arena that is never
# returned. Measured: 8-12 arenas -> 0.
# GOGC=50 halves the Go heap growth target. On its own it does NOT help (the
# saved heap is immediately eaten by more glibc arenas); combined with
# MALLOC_ARENA_MAX it cut settled RSS by ~19%. This gateway is I/O bound, so
# the extra GC cycles are free.
Environment=GOGC=50
Environment=MALLOC_ARENA_MAX=2
ExecStart=/usr/bin/llmsproxy -config /etc/llmsproxy/config.yaml
WorkingDirectory=/etc/llmsproxy
Restart=always
RestartSec=5
[Install]
WantedBy=multi-user.target

59
packaging/nfpm.yaml Normal file
View File

@ -0,0 +1,59 @@
# nfpm config for the llmsproxy core gateway binary.
# The version placeholder below is substituted by packaging/core-dist.sh (nfpm v2
# here does not honour {{ .Env.VERSION }} template rendering), so do not edit it
# by hand.
name: llmsproxy
arch: amd64
platform: linux
version: __VERSION__
epoch: "1"
section: net
priority: optional
maintainer: JianFeeeee <dev@modelrouter.local>
description: |
Unified OpenAI-compatible multi-source LLM gateway.
Routes requests across multiple upstream LLM providers with Lua adapters,
model scheduling, quota windows, an elastic Lua worker pool, and a Web UI.
vendor: ModelRouter
homepage: https://gitcode.com/JianFeeeee/ModelRouter
license: MIT
contents:
# binary
- src: ./cmd/build/llmsproxy
dst: /usr/bin/llmsproxy
file_info:
mode: 0755
# systemd unit (memory tuning baked in)
- src: ./packaging/llmsproxy.service
dst: /etc/systemd/system/llmsproxy.service
file_info:
mode: 0644
# example config (seeded to /etc/llmsproxy on first install by postinst)
- src: ./config.yaml
dst: /usr/share/llmsproxy/config.example.yaml
file_info:
mode: 0644
# built-in Lua adapters (deployed to /etc/llmsproxy/adapters)
- src: ./internal/lua/adapters/*.lua
dst: /usr/share/llmsproxy/adapters/
file_info:
mode: 0644
# README pointer
- src: ./README.md
dst: /usr/share/doc/llmsproxy/README.md
file_info:
mode: 0644
scripts:
postinstall: ./packaging/postinst.sh
preremove: ./packaging/prerm.sh
rpm:
group: Applications/Internet
summary: Unified OpenAI-compatible multi-source LLM gateway
compression: gzip
deb:
compression: gzip
fields:
Recommends: ca-certificates

22
packaging/postinst.sh Executable file
View File

@ -0,0 +1,22 @@
#!/bin/sh
# llmsproxy core package postinst
set -e
# Runtime state belongs under /etc/llmsproxy just like the production layout:
# runtime.json holds keys/created sources and is NOT touched on reinstall.
mkdir -p /etc/llmsproxy/adapters
if [ -f /etc/llmsproxy/config.yaml ]; then
echo "llmsproxy: keeping existing /etc/llmsproxy/config.yaml"
else
# 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
fi
# systemd
if command -v systemctl >/dev/null 2>&1; then
systemctl daemon-reload >/dev/null 2>&1 || true
echo "llmsproxy: use 'systemctl enable --now llmsproxy' to start"
fi
exit 0

14
packaging/prerm.sh Executable file
View File

@ -0,0 +1,14 @@
#!/bin/sh
# llmsproxy core package prerm
set -e
if command -v systemctl >/dev/null 2>&1; then
# stop the service if this package's unit is the one installed
if systemctl list-unit-files llmsproxy.service >/dev/null 2>&1 && \
systemctl is-active llmsproxy.service >/dev/null 2>&1; then
systemctl stop llmsproxy.service >/dev/null 2>&1 || true
systemctl disable llmsproxy.service >/dev/null 2>&1 || true
fi
systemctl daemon-reload >/dev/null 2>&1 || true
fi
exit 0

49
packaging/verify-dist.sh Executable file
View File

@ -0,0 +1,49 @@
#!/usr/bin/env bash
# Verify packaged GUI artifacts are healthy before a release is declared done.
#
# Background: on 2026-08-28 (1.3.0) the host wine broke (missing syswow64, so
# wine -- the NSIS self-extract step -- failed), and electron-builder silently
# wrote a 264 KB installer shell that did NOT embed the 86 MB payload. No gate
# caught it, so a broken Windows installer shipped. This script is that gate.
#
# It enforces a per-artifact minimum size, so a degenerate installer (or a
# build that only emitted a shell without its embedded payload) fails loudly.
# A size check is the right signal here: a healthy NSIS exe is ~86 MB and the
# broken 1.3.0 one was 264 KB (326x) — electron-builder compresses the embedded
# app.7z, so a plaintext-payload string probe would be unreliable.
set -uo pipefail
GUI_DIST="${1:-cmd/build/gui-dist}"
VERSION="${2:-}"
log() { printf '[verify-dist] %s\n' "$*"; }
fail() { printf '[verify-dist] FATAL: %s\n' "$*" >&2; exit 1; }
[ -d "$GUI_DIST" ] || fail "no dist dir: $GUI_DIST"
# (glob, min-bytes) — nsis.7z is NOT a required artifact: with `--win nsis` the
# payload is embedded inside the Setup exe and only the exe is the formal product
# (the .nsis.7z appears in intermediate/--dir builds only).
checks=(
"ModelRouter*.AppImage:10000000"
"modelrouter-*.deb:10000000"
"modelrouter-*.rpm:10000000"
"ModelRouter Setup*.exe:5000000"
)
for entry in "${checks[@]}"; do
glob="${entry%%:*}"; min="${entry#*:}"
# gui-dist accumulates installers from previous versions; check the MOST
# RECENT artifact (last in sort) so a stale file never masks a broken build.
# `find -name` (not a bare glob) so the space in "ModelRouter Setup *.exe"
# is preserved as one file pattern.
match="$(find "$GUI_DIST" -maxdepth 1 -name "$glob" -type f | sort -V | tail -1)"
[ -z "$match" ] && fail "no artifact matching '$glob' in $GUI_DIST"
sz=$(stat -c%s "$match")
if [ "$sz" -lt "$min" ]; then
fail "$(basename "$match") is only $sz bytes (min $min) — likely a degenerate shell like the 1.3.0 264KB exe"
fi
log " ✓ $(basename "$match") ($((sz/1024)) KB)"
done
log "all artifacts healthy"