mirror of
https://gitcode.com/JianFeeeee/ModelRouter.git
synced 2026-11-07 17:57:23 +00:00
Compare commits
33 Commits
| Author | SHA1 | Date | |
|---|---|---|---|
| 624fd74b45 | |||
| 94cbcb6771 | |||
| a5384d9fb6 | |||
| ab20f1be48 | |||
| 8334cffbc9 | |||
| a42ff62d06 | |||
| 66585549f1 | |||
| 747dff5b76 | |||
| 5d045f97a8 | |||
| e48baa1bb3 | |||
| 519b18518a | |||
| 77b00bfe3f | |||
| 21ec8f59d8 | |||
| ec89daad62 | |||
| 24609289e8 | |||
| 18c2385c61 | |||
| 2a5c7bbbc7 | |||
| 045ecf47bc | |||
| 334b984c25 | |||
| f3d5ba6cea | |||
| 98847adfd6 | |||
| 5bb94db08c | |||
| ef396f9b47 | |||
| c8a300c998 | |||
| 210243f724 | |||
| b2183df1e8 | |||
| 6eac80bc6c | |||
| e3a36bc7ca | |||
| 31b1ed2c96 | |||
| e19cff248f | |||
| bb3af3bdb3 | |||
| 9c99ded8c0 | |||
| 690f55c6f1 |
1
.gitignore
vendored
1
.gitignore
vendored
@ -20,3 +20,4 @@ cmd/gui/*.log
|
||||
|
||||
# local ops scripts (machine-specific, not for sharing)
|
||||
scripts/
|
||||
cmd/gui/dist/
|
||||
|
||||
10
.omo/run-continuation/ses_fcda7f5e1ffe9g78A7lVoLMxn7.json
Normal file
10
.omo/run-continuation/ses_fcda7f5e1ffe9g78A7lVoLMxn7.json
Normal file
@ -0,0 +1,10 @@
|
||||
{
|
||||
"sessionID": "ses_fcda7f5e1ffe9g78A7lVoLMxn7",
|
||||
"updatedAt": "2026-08-24T11:19:43.949Z",
|
||||
"sources": {
|
||||
"background-task": {
|
||||
"state": "idle",
|
||||
"updatedAt": "2026-08-24T11:19:43.949Z"
|
||||
}
|
||||
}
|
||||
}
|
||||
58
README.md
58
README.md
@ -184,16 +184,24 @@ 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 并附上每档失败摘要。
|
||||
|
||||
### 源(Source)与适配器(Adapter)的关系
|
||||
|
||||
@ -224,6 +232,23 @@ 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
|
||||
的具体源。
|
||||
|
||||
### disable_thinking
|
||||
|
||||
请求体带 `"disable_thinking": true` 时,网关透传给各适配器;DeepSeek 适配器将其
|
||||
@ -296,14 +321,23 @@ return {
|
||||
一键启停、系统托盘常驻、开机自启、静默启动,内容区一比一内嵌完整 WebUI
|
||||
(状态/对话/密钥/优先级/源/适配器 六个页面,免登录)。服务器用户继续使用纯 Go 二进制。
|
||||
|
||||
### 两种安装包:Headless 与 Desktop
|
||||
|
||||
发布到 Release 的安装包分为两类,面向不同用户:
|
||||
|
||||
- **Headless(服务器版)**:纯核心二进制(无 Electron),配置驱动,适合服务器、
|
||||
容器、systemd 等无人值守场景。Linux 为 `.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**(单二进制) | 极轻量(~10MB、~15MB RSS),零依赖单进程,直接跑在 systemd/任意容器里,远程管理 |
|
||||
| 个人桌面日常使用 / 多设备内网共享 | **Desktop**(GUI) | 免登录内嵌 WebUI、系统托盘一键启停、开机自启、静默后台,适合不懂命令行的使用者 |
|
||||
| Windows 桌面 | **Desktop** | 纯后端在 Windows 上需自行注册服务,GUI 提供原生托盘/自启体验 |
|
||||
| CI 一键出三平台安装包 | **Desktop 打包脚本** | `make gui-*` 系列出 deb/AppImage/NSIS,可进发行流水线 |
|
||||
|
||||
两者完全同源:GUI 内嵌的就是纯后端的同一份 `llmsproxy` 核心(LuaJIT 版),
|
||||
配置/适配器格式完全一致,可随时互换。
|
||||
@ -321,8 +355,10 @@ make gui-win-docker # 构建 Windows nsis 安装包(全 docker 自足,无
|
||||
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)。rpm 需系统 `rpmbuild`。
|
||||
发布时会把产物重命名为 `ModelRouter-Desktop-*`(GUI)与 `ModelRouter-Headless-*`(纯后端)
|
||||
两类上传到 GitCode Release。
|
||||
|
||||
#### Windows 打包(docker 固化方案,推荐)
|
||||
|
||||
|
||||
72
README_EN.md
72
README_EN.md
@ -176,12 +176,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 +197,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 +236,28 @@ 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.
|
||||
|
||||
### disable_thinking
|
||||
|
||||
With `"disable_thinking": true` in the request body, the gateway passes it to
|
||||
@ -308,14 +339,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. Linux ships as
|
||||
`.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) | ~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 | **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.
|
||||
@ -332,9 +374,11 @@ make gui-win # Windows NSIS (host mingw + wine)
|
||||
make gui-win-docker # Windows NSIS, fully dockerized (no host mingw needed)
|
||||
```
|
||||
|
||||
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). rpm
|
||||
needs system `rpmbuild`. 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)
|
||||
|
||||
|
||||
@ -263,7 +263,6 @@ function waitForCoreReady(timeoutMs = 10000) {
|
||||
let coreLastExit = 0;
|
||||
|
||||
function stopCore() {
|
||||
app.isQuitting = true;
|
||||
coreReady = false;
|
||||
if (coreProc && coreProc.exitCode === null) {
|
||||
try {
|
||||
|
||||
6
cmd/gui/package-lock.json
generated
6
cmd/gui/package-lock.json
generated
@ -1,12 +1,12 @@
|
||||
{
|
||||
"name": "modelrouter-gui",
|
||||
"version": "1.0.0",
|
||||
"version": "1.2.0",
|
||||
"lockfileVersion": 3,
|
||||
"requires": true,
|
||||
"packages": {
|
||||
"": {
|
||||
"name": "modelrouter-gui",
|
||||
"version": "1.0.0",
|
||||
"version": "1.2.0",
|
||||
"license": "Proprietary",
|
||||
"devDependencies": {
|
||||
"electron": "^33.0.0",
|
||||
@ -4045,4 +4045,4 @@
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
@ -1,6 +1,6 @@
|
||||
{
|
||||
"name": "modelrouter-gui",
|
||||
"version": "1.0.0",
|
||||
"version": "1.2.0",
|
||||
"description": "ModelRouter Desktop — embedded ModelRouter core with tray",
|
||||
"author": "ModelRouter",
|
||||
"main": "main.js",
|
||||
|
||||
@ -26,16 +26,6 @@ const state = {
|
||||
|
||||
const $ = (s) => document.querySelector(s);
|
||||
|
||||
function esc(s) {
|
||||
return String(s == null ? "" : s).replace(
|
||||
/[&<>"']/g,
|
||||
(c) =>
|
||||
({ "&": "&", "<": "<", ">": ">", '"': """, "'": "'" })[
|
||||
c
|
||||
],
|
||||
);
|
||||
}
|
||||
|
||||
function toast(msg, isErr, ms) {
|
||||
const t = $("#toast");
|
||||
t.textContent = msg;
|
||||
|
||||
271
deploy.sh
Executable file
271
deploy.sh
Executable file
@ -0,0 +1,271 @@
|
||||
#!/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. 健康检查确认新进程正常
|
||||
#
|
||||
# 注意事项:
|
||||
# - 本服务依赖自身(推理依赖),restart 前后会短暂中断,
|
||||
# 但重试机制(systemd RestartSec=5 + 客户端 retry)可自愈
|
||||
# - 适配器文件替换后需重启才能生效(VM 启动时批量加载)
|
||||
# ============================================================
|
||||
set -euo pipefail
|
||||
|
||||
# ---------- 配置 ----------
|
||||
readonly TARGET_BIN="/usr/local/bin/llmsproxy"
|
||||
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 BACKUP_BIN="${TARGET_BIN}.bak"
|
||||
|
||||
# ---------- 日志 ----------
|
||||
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"
|
||||
}
|
||||
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 权限"
|
||||
|
||||
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_binary() {
|
||||
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
|
||||
}
|
||||
|
||||
# ---------- 构建 ----------
|
||||
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
|
||||
warn "新二进制与旧版本完全一致(sha256 相同),跳过部署"
|
||||
exit 0
|
||||
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})"
|
||||
}
|
||||
|
||||
# ---------- 同步适配器 ----------
|
||||
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 等),也能让旧版本继续服务,避免推理能力丢失。
|
||||
rollback() {
|
||||
err "健康检查失败,启动回滚流程"
|
||||
if [[ ! -f "$BACKUP_BIN" ]]; then
|
||||
err "无回滚备份 ($BACKUP_BIN 不存在),无法自动恢复"
|
||||
err "请手动检查: systemctl status llmsproxy.service && journalctl -u llmsproxy -n 50"
|
||||
return 1
|
||||
fi
|
||||
|
||||
log "恢复旧二进制"
|
||||
cp -f "$BACKUP_BIN" "$TARGET_BIN"
|
||||
chmod 0755 "$TARGET_BIN"
|
||||
|
||||
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_ADAPTERS"
|
||||
log "═════════════════════════════════════════════"
|
||||
}
|
||||
|
||||
# ---------- 主流程 ----------
|
||||
main() {
|
||||
log "═══ llmsproxy 原子部署开始 ═══"
|
||||
precheck
|
||||
build
|
||||
sync_adapters
|
||||
backup_old_binary
|
||||
atomic_replace_binary
|
||||
restart_service
|
||||
if ! healthcheck; then
|
||||
rollback || true
|
||||
exit 1
|
||||
fi
|
||||
print_summary
|
||||
}
|
||||
|
||||
main "$@"
|
||||
@ -55,6 +55,7 @@ state may be shared across workers (`log` to stdout is the only side effect).
|
||||
| `version` | string | no | Version, shown in the WebUI |
|
||||
| `endpoint` | string | no | Upstream path, default `/chat/completions`; overridable by `source.endpoint` / `source.image_endpoint` |
|
||||
| `headers` | table | no | Static default request headers; used as fallback when no `build_headers` hook is defined |
|
||||
| `transform_error(status, body)` | function | no | Error-response condensing: return a one-line reason; on nil / absence clients uniformly receive `unknown error` (raw body goes to server logs only) |
|
||||
|
||||
These fields are extracted statically at load time (compile-once); reading them never
|
||||
occupies a pooled worker.
|
||||
@ -147,6 +148,27 @@ end
|
||||
|
||||
## 4. Optional Hooks
|
||||
|
||||
### `transform_error(status, body) -> string | nil`
|
||||
|
||||
Condense this source's error response into a one-line reason. Every upstream
|
||||
formats errors differently — that is adapter territory: all built-in adapters
|
||||
implement their own envelope parsing (zen's `{error={type,message}}`,
|
||||
Anthropic's `{type="error",error={...}}`, Gemini's
|
||||
`{error={code,message,status}}`, Ollama's string `{error="..."}`, etc.).
|
||||
When the hook is absent or returns nil the core does not guess: clients get
|
||||
`api error <status>: unknown error` and the raw body is logged server-side only.
|
||||
|
||||
```lua
|
||||
function adapter.transform_error(status, body)
|
||||
local ok, resp = pcall(json.decode, body)
|
||||
if not ok or type(resp) ~= "table" then return nil end
|
||||
if resp.error and resp.error.type == "FreeUsageLimitError" then
|
||||
return "zen free pool quota exhausted"
|
||||
end
|
||||
return resp.error and resp.error.message or nil
|
||||
end
|
||||
```
|
||||
|
||||
### `build_headers(meta) -> table<string,string>`
|
||||
|
||||
Dynamically generate / sign request headers (e.g. KimiCode's HMAC signature). If the
|
||||
|
||||
@ -50,6 +50,7 @@ return adapter
|
||||
| `version` | string | 否 | 版本号,用于 WebUI 展示 |
|
||||
| `endpoint` | string | 否 | 上游请求路径,默认 `/chat/completions`;可被 `source.endpoint` / `source.image_endpoint` 覆盖 |
|
||||
| `headers` | table | 否 | 静态默认请求头;若未定义 `build_headers` 钩子则作为请求头回退 |
|
||||
| `transform_error(status, body)` | function | 否 | 错误响应收敛:返回一行短原因;返回 nil / 未定义时客户端统一收到 `unknown error`(原始响应体只进服务端日志) |
|
||||
|
||||
这些字段在加载时静态提取(compile-once),之后读它们不会占用池内 worker。
|
||||
|
||||
@ -138,6 +139,26 @@ end
|
||||
|
||||
## 4. 可选钩子
|
||||
|
||||
### `transform_error(status, body) -> string | nil`
|
||||
|
||||
把该源特有的错误响应收敛成一行短原因。每个上游的错误格式不同——这是适配器层
|
||||
的职责:内置各适配器均实现了自己的信封解析(如 zen 的
|
||||
`{error={type,message}}`、Anthropic 的 `{type="error",error={...}}`、Gemini 的
|
||||
`{error={code,message,status}}`、Ollama 的字符串 `{error="..."}` 等)。
|
||||
未定义本钩子或返回 nil 时,核心不猜测格式,客户端统一收到
|
||||
`api error <status>: unknown error`,原始响应体仅记录在服务端日志。
|
||||
|
||||
```lua
|
||||
function adapter.transform_error(status, body)
|
||||
local ok, resp = pcall(json.decode, body)
|
||||
if not ok or type(resp) ~= "table" then return nil end
|
||||
if resp.error and resp.error.type == "FreeUsageLimitError" then
|
||||
return "zen free pool quota exhausted"
|
||||
end
|
||||
return resp.error and resp.error.message or nil
|
||||
end
|
||||
```
|
||||
|
||||
### `build_headers(meta) -> table<string,string>`
|
||||
|
||||
动态生成/签名请求头(如 KimiCode 的 HMAC 签名)。若脚本未定义此函数,Go 层回退使用静态
|
||||
|
||||
@ -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)
|
||||
}
|
||||
}
|
||||
|
||||
@ -20,12 +20,12 @@ 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 优先级页编辑)
|
||||
Auto []ModelScope `yaml:"auto,omitempty"` // AUTO 调度链规则(WebUI 优先级页编辑,chat)
|
||||
AutoImage []ModelScope `yaml:"auto_image,omitempty"` // AUTO 生图调度链规则(WebUI 优先级页·生图)
|
||||
Keys []GWKey `yaml:"keys,omitempty"` // 网关密钥(WebUI 密钥页管理)
|
||||
}
|
||||
|
||||
@ -62,6 +62,7 @@ type Source struct {
|
||||
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
|
||||
@ -313,11 +314,32 @@ 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"`
|
||||
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"
|
||||
|
||||
@ -318,3 +318,55 @@ sources:
|
||||
t.Fatalf("after no-op: %v %v", len(cfg.Sources), err)
|
||||
}
|
||||
}
|
||||
|
||||
// TestSeedTemplatesOnce verifies preset templates are seeded exactly once:
|
||||
// a user's edit survives a restart, and a deliberately deleted preset does
|
||||
// not silently come back.
|
||||
func TestSeedTemplatesOnce(t *testing.T) {
|
||||
dir := t.TempDir()
|
||||
path := filepath.Join(dir, "runtime.json")
|
||||
presets := []SourceTemplate{
|
||||
{Name: "DeepSeek", BaseURL: "https://api.deepseek.com", Adapter: "deepseek",
|
||||
Models: []Model{{ID: "deepseek-chat", Kind: "chat"}}},
|
||||
{Name: "OpenAI", BaseURL: "https://api.openai.com/v1", Adapter: "openai",
|
||||
Models: []Model{{ID: "gpt-4o", Kind: "chat"}}},
|
||||
}
|
||||
|
||||
// first run: both presets land in the store
|
||||
s1 := NewStore(path)
|
||||
if err := s1.Load(); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := s1.SeedTemplates(presets); err != nil {
|
||||
t.Fatalf("seed: %v", err)
|
||||
}
|
||||
if got := len(s1.ListTemplates()); got != 2 {
|
||||
t.Fatalf("after seed len = %d, want 2", got)
|
||||
}
|
||||
|
||||
// user edits one preset, deletes the other
|
||||
edited := presets[0]
|
||||
edited.BaseURL = "https://my-proxy.internal/v1"
|
||||
if err := s1.UpsertTemplate(edited); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if _, err := s1.RemoveTemplate("OpenAI"); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
// restart: re-seeding must be a no-op for both
|
||||
s2 := NewStore(path)
|
||||
if err := s2.Load(); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := s2.SeedTemplates(presets); err != nil {
|
||||
t.Fatalf("reseed: %v", err)
|
||||
}
|
||||
tpls := s2.ListTemplates()
|
||||
if len(tpls) != 1 {
|
||||
t.Fatalf("after restart len = %d, want 1 (deleted preset must not return): %+v", len(tpls), tpls)
|
||||
}
|
||||
if tpls[0].Name != "DeepSeek" || tpls[0].BaseURL != "https://my-proxy.internal/v1" {
|
||||
t.Fatalf("user edit lost on restart: %+v", tpls[0])
|
||||
}
|
||||
}
|
||||
|
||||
212
internal/config/presets.go
Normal file
212
internal/config/presets.go
Normal file
@ -0,0 +1,212 @@
|
||||
package config
|
||||
|
||||
// Built-in preset source templates. Seeded into the runtime store on first
|
||||
// startup so the WebUI can present "从模板创建" for popular providers — the
|
||||
// user only needs to supply a name and API key.
|
||||
//
|
||||
// Model IDs verified against provider docs 2026-08-28. Providers ship new IDs
|
||||
// frequently; treat these as a starting point and edit in the WebUI as needed.
|
||||
var PresetTemplates = []SourceTemplate{
|
||||
// ── 国内热门 ────────────────────────────────────────────────────────────
|
||||
{
|
||||
Name: "DeepSeek",
|
||||
BaseURL: "https://api.deepseek.com",
|
||||
Adapter: "deepseek",
|
||||
MaxConcurrent: 8,
|
||||
Models: []Model{
|
||||
{ID: "deepseek-v4-flash", Priority: 100, Kind: "chat"},
|
||||
{ID: "deepseek-v4-pro", Priority: 90, Kind: "chat"},
|
||||
{ID: "deepseek-v4-flash-vision-exp", Priority: 60, Kind: "chat"},
|
||||
},
|
||||
},
|
||||
{
|
||||
Name: "智谱 GLM",
|
||||
BaseURL: "https://open.bigmodel.cn/api/paas/v4",
|
||||
Adapter: "openai",
|
||||
MaxConcurrent: 8,
|
||||
Models: []Model{
|
||||
{ID: "glm-5.3", Priority: 100, Kind: "chat"},
|
||||
{ID: "glm-5.3-flash", Priority: 90, Kind: "chat"},
|
||||
{ID: "glm-5.2", Priority: 80, Kind: "chat"},
|
||||
{ID: "glm-5.1", Priority: 60, Kind: "chat"},
|
||||
},
|
||||
},
|
||||
{
|
||||
Name: "Kimi / Moonshot",
|
||||
BaseURL: "https://api.moonshot.cn/v1",
|
||||
Adapter: "openai",
|
||||
MaxConcurrent: 8,
|
||||
Models: []Model{
|
||||
{ID: "kimi-k3", Priority: 100, Kind: "chat"},
|
||||
{ID: "kimi-k2.7-code", Priority: 90, Kind: "chat"},
|
||||
{ID: "kimi-k2.7-code-highspeed", Priority: 85, Kind: "chat"},
|
||||
{ID: "kimi-k2.6", Priority: 70, Kind: "chat"},
|
||||
},
|
||||
},
|
||||
{
|
||||
Name: "SiliconFlow",
|
||||
BaseURL: "https://api.siliconflow.cn/v1",
|
||||
Adapter: "openai",
|
||||
MaxConcurrent: 8,
|
||||
Models: []Model{
|
||||
{ID: "deepseek-ai/DeepSeek-V4-Pro", Priority: 100, Kind: "chat"},
|
||||
{ID: "deepseek-ai/DeepSeek-V4-Flash", Priority: 90, Kind: "chat"},
|
||||
{ID: "moonshotai/Kimi-K2.6", Priority: 80, Kind: "chat"},
|
||||
{ID: "zai-org/GLM-5.1", Priority: 70, Kind: "chat"},
|
||||
{ID: "Kwai-Kolors/Kolors", Priority: 50, Kind: "image"},
|
||||
},
|
||||
},
|
||||
{
|
||||
Name: "讯飞星火 (Spark)",
|
||||
BaseURL: "https://spark-api-open.xf-yun.com/v1",
|
||||
Adapter: "openai",
|
||||
MaxConcurrent: 4,
|
||||
Models: []Model{
|
||||
{ID: "4.0Ultra", Priority: 100, Kind: "chat"},
|
||||
{ID: "generalv3.5", Priority: 80, Kind: "chat"},
|
||||
{ID: "max-32k", Priority: 70, Kind: "chat"},
|
||||
},
|
||||
},
|
||||
{
|
||||
Name: "阿里云 通义千问",
|
||||
BaseURL: "https://dashscope.aliyuncs.com/compatible-mode/v1",
|
||||
Adapter: "openai",
|
||||
MaxConcurrent: 8,
|
||||
Models: []Model{
|
||||
{ID: "qwen3.8-max", Priority: 100, Kind: "chat"},
|
||||
{ID: "qwen3.8-flash", Priority: 90, Kind: "chat"},
|
||||
{ID: "qwen3.7-plus", Priority: 85, Kind: "chat"},
|
||||
{ID: "deepseek-v4-pro", Priority: 70, Kind: "chat"},
|
||||
{ID: "kimi-k3", Priority: 70, Kind: "chat"},
|
||||
},
|
||||
},
|
||||
{
|
||||
Name: "百度 千帆",
|
||||
BaseURL: "https://qianfan.baidubce.com/v2",
|
||||
Adapter: "openai",
|
||||
MaxConcurrent: 8,
|
||||
Models: []Model{
|
||||
{ID: "ernie-5.0-turbo", Priority: 100, Kind: "chat"},
|
||||
{ID: "ernie-5.0", Priority: 90, Kind: "chat"},
|
||||
{ID: "deepseek-v4-pro", Priority: 70, Kind: "chat"},
|
||||
},
|
||||
},
|
||||
{
|
||||
Name: "豆包 (火山方舟)",
|
||||
BaseURL: "https://ark.cn-beijing.volces.com/api/v3",
|
||||
Adapter: "openai",
|
||||
MaxConcurrent: 8,
|
||||
Models: []Model{
|
||||
{ID: "doubao-seed-2.0-pro", Priority: 100, Kind: "chat"},
|
||||
{ID: "doubao-seed-2.0", Priority: 90, Kind: "chat"},
|
||||
{ID: "doubao-seed-2.0-flash", Priority: 80, Kind: "chat"},
|
||||
},
|
||||
},
|
||||
{
|
||||
Name: "MiniMax",
|
||||
BaseURL: "https://api.minimax.chat/v1",
|
||||
Adapter: "openai",
|
||||
MaxConcurrent: 8,
|
||||
Models: []Model{
|
||||
{ID: "MiniMax-M3", Priority: 100, Kind: "chat"},
|
||||
{ID: "MiniMax-M2.5", Priority: 80, Kind: "chat"},
|
||||
},
|
||||
},
|
||||
// ── 国际主流 ────────────────────────────────────────────────────────────
|
||||
{
|
||||
Name: "OpenAI",
|
||||
BaseURL: "https://api.openai.com/v1",
|
||||
Adapter: "openai",
|
||||
MaxConcurrent: 8,
|
||||
Models: []Model{
|
||||
{ID: "gpt-5.6-sol", Priority: 100, Kind: "chat"},
|
||||
{ID: "gpt-5.6-terra", Priority: 90, Kind: "chat"},
|
||||
{ID: "gpt-5.6-luna", Priority: 80, Kind: "chat"},
|
||||
},
|
||||
},
|
||||
{
|
||||
Name: "Anthropic Claude",
|
||||
BaseURL: "https://api.anthropic.com",
|
||||
Adapter: "anthropic",
|
||||
MaxConcurrent: 8,
|
||||
Models: []Model{
|
||||
{ID: "claude-opus-5", Priority: 100, Kind: "chat"},
|
||||
{ID: "claude-fable-5", Priority: 95, Kind: "chat"},
|
||||
{ID: "claude-sonnet-5", Priority: 85, Kind: "chat"},
|
||||
{ID: "claude-haiku-4-5", Priority: 70, Kind: "chat"},
|
||||
},
|
||||
},
|
||||
{
|
||||
Name: "Google Gemini",
|
||||
BaseURL: "https://generativelanguage.googleapis.com/v1beta",
|
||||
Adapter: "gemini",
|
||||
MaxConcurrent: 8,
|
||||
Models: []Model{
|
||||
{ID: "gemini-3.6-flash", Priority: 100, Kind: "chat"},
|
||||
{ID: "gemini-3.5-flash", Priority: 90, Kind: "chat"},
|
||||
{ID: "gemini-3.5-flash-lite", Priority: 80, Kind: "chat"},
|
||||
{ID: "gemini-3.1-pro", Priority: 85, Kind: "chat"},
|
||||
},
|
||||
},
|
||||
{
|
||||
Name: "Groq",
|
||||
BaseURL: "https://api.groq.com/openai/v1",
|
||||
Adapter: "groq",
|
||||
MaxConcurrent: 8,
|
||||
Models: []Model{
|
||||
{ID: "openai/gpt-oss-120b", Priority: 100, Kind: "chat"},
|
||||
{ID: "openai/gpt-oss-20b", Priority: 80, Kind: "chat"},
|
||||
{ID: "groq/compound", Priority: 70, Kind: "chat"},
|
||||
},
|
||||
},
|
||||
{
|
||||
Name: "Mistral",
|
||||
BaseURL: "https://api.mistral.ai/v1",
|
||||
Adapter: "mistral",
|
||||
MaxConcurrent: 8,
|
||||
Models: []Model{
|
||||
{ID: "mistral-medium-latest", Priority: 100, Kind: "chat"},
|
||||
{ID: "mistral-large-latest", Priority: 90, Kind: "chat"},
|
||||
{ID: "codestral-latest", Priority: 85, Kind: "chat"},
|
||||
},
|
||||
},
|
||||
{
|
||||
Name: "OpenRouter",
|
||||
BaseURL: "https://openrouter.ai/api/v1",
|
||||
Adapter: "openai",
|
||||
MaxConcurrent: 8,
|
||||
Headers: map[string]string{
|
||||
"HTTP-Referer": "http://localhost:8080",
|
||||
"X-Title": "ModelRouter",
|
||||
},
|
||||
Models: []Model{
|
||||
{ID: "openrouter/auto", Priority: 100, Kind: "chat"},
|
||||
{ID: "deepseek/deepseek-v4-flash", Priority: 90, Kind: "chat"},
|
||||
{ID: "anthropic/claude-opus-5", Priority: 85, Kind: "chat"},
|
||||
{ID: "openai/gpt-5.6-sol", Priority: 85, Kind: "chat"},
|
||||
},
|
||||
},
|
||||
{
|
||||
Name: "GitHub Models",
|
||||
BaseURL: "https://models.inference.ai.azure.com",
|
||||
Adapter: "github",
|
||||
MaxConcurrent: 8,
|
||||
Models: []Model{
|
||||
{ID: "gpt-5.6-sol", Priority: 100, Kind: "chat"},
|
||||
{ID: "gpt-5.6-luna", Priority: 80, Kind: "chat"},
|
||||
},
|
||||
},
|
||||
// ── 本地模型 ────────────────────────────────────────────────────────────
|
||||
{
|
||||
Name: "Ollama (localhost)",
|
||||
BaseURL: "http://127.0.0.1:11434",
|
||||
Adapter: "ollama",
|
||||
Endpoint: "/api/chat",
|
||||
MaxConcurrent: 4,
|
||||
Models: []Model{
|
||||
{ID: "qwen3", Priority: 80, Kind: "chat"},
|
||||
{ID: "deepseek-r1", Priority: 80, Kind: "chat"},
|
||||
{ID: "llama3.3", Priority: 70, Kind: "chat"},
|
||||
},
|
||||
},
|
||||
}
|
||||
@ -128,15 +128,3 @@ func (b *SecretBox) MustDecrypt(v string) string {
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
// resolvedAPIKey resolves a source key from api_key_env (env var) or api_key.
|
||||
func resolvedAPIKey(src Source, box *SecretBox) (string, error) {
|
||||
if src.APIKeyEnv != "" {
|
||||
v := os.Getenv(src.APIKeyEnv)
|
||||
if v == "" {
|
||||
return "", fmt.Errorf("source %s: env %s is empty or unset", src.Name, src.APIKeyEnv)
|
||||
}
|
||||
return v, nil
|
||||
}
|
||||
return box.MustDecrypt(src.APIKey), nil
|
||||
}
|
||||
@ -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 {
|
||||
|
||||
@ -32,7 +32,8 @@ type Core struct {
|
||||
store *config.Store
|
||||
scheduler *scheduler.Scheduler
|
||||
registry *provider.Registry
|
||||
autoChain atomic.Pointer[scheduler.Chain]
|
||||
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.
|
||||
@ -68,6 +69,11 @@ 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)
|
||||
}
|
||||
return c, nil
|
||||
}
|
||||
|
||||
@ -161,6 +167,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 +199,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 }
|
||||
@ -294,6 +315,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 +355,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.
|
||||
@ -420,6 +469,7 @@ func (c *Core) rebuildRegistry() error {
|
||||
c.registry.Replace(providers)
|
||||
}
|
||||
c.buildAutoChain()
|
||||
c.buildAutoImageChain()
|
||||
return nil
|
||||
}
|
||||
|
||||
@ -466,6 +516,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"`
|
||||
@ -601,6 +692,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")
|
||||
|
||||
@ -72,6 +72,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 +106,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 +128,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) {
|
||||
@ -139,26 +197,20 @@ func (g *Gateway) handleStatsAPI(w http.ResponseWriter, r *http.Request) {
|
||||
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"})
|
||||
_ = 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"})
|
||||
for _, rec := range g.stats.AuditRecords(from, to, key) {
|
||||
_ = cw.Write([]string{
|
||||
time.UnixMilli(rec.Time).Format(time.RFC3339),
|
||||
@ -172,6 +224,9 @@ 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,
|
||||
})
|
||||
}
|
||||
@ -179,31 +234,24 @@ func (g *Gateway) handleStatsAPI(w http.ResponseWriter, r *http.Request) {
|
||||
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 +261,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 +271,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,
|
||||
})
|
||||
}
|
||||
|
||||
@ -5,6 +5,7 @@ import (
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"log"
|
||||
"net/http"
|
||||
"strings"
|
||||
"sync/atomic"
|
||||
@ -387,13 +388,6 @@ func (g *Gateway) handleChat(w http.ResponseWriter, r *http.Request) {
|
||||
g.singleChat(w, ctx, cands, inner, effective, rec)
|
||||
}
|
||||
|
||||
func normalizeModel(m string) string {
|
||||
if isAuto(m) {
|
||||
return ""
|
||||
}
|
||||
return m
|
||||
}
|
||||
|
||||
func firstModel(p *provider.Provider) string {
|
||||
ms := p.Models()
|
||||
if len(ms) > 0 {
|
||||
@ -523,21 +517,51 @@ 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
|
||||
// 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.
|
||||
func clientUpstreamErr(err error) string {
|
||||
log.Printf("[gateway] upstream failure surfaced to client: %v", err)
|
||||
var ce *scheduler.ChainErr
|
||||
if errors.As(err, &ce) {
|
||||
parts := make([]string, 0, len(ce.Tiers)+len(ce.Skipped))
|
||||
for _, t := range ce.Tiers {
|
||||
parts = append(parts, types.OneLine(fmt.Sprintf("%s/%s: %v", t.Source, t.Model, t.Err), 80))
|
||||
}
|
||||
for _, sk := range ce.Skipped {
|
||||
parts = append(parts, types.OneLine(sk, 80))
|
||||
}
|
||||
msg := strings.Join(parts, "; ")
|
||||
if len(msg) > 300 {
|
||||
msg = msg[:300] + "..."
|
||||
}
|
||||
return fmt.Sprintf("all %d auto providers failed: %s", len(parts), msg)
|
||||
}
|
||||
rec.OK = true
|
||||
rec.Status = http.StatusOK
|
||||
return types.OneLine(err.Error(), 160)
|
||||
}
|
||||
|
||||
// failChat maps a scheduling failure onto the audit record and answers the
|
||||
// client. A failed AUTO chain additionally pins its first failed tier onto
|
||||
// the record; direct-path errors never match *scheduler.ChainErr, so the
|
||||
// extraction is safely shared by all four entry points. Callers own the
|
||||
// writeRec call (inline for non-stream paths, deferred for stream paths).
|
||||
func (g *Gateway) failChat(w http.ResponseWriter, rec *Req, err error) {
|
||||
rec.OK = false
|
||||
rec.Status = upstreamErrStatus(err)
|
||||
rec.Err = err.Error()
|
||||
var ce *scheduler.ChainErr
|
||||
if errors.As(err, &ce) && len(ce.Tiers) > 0 {
|
||||
rec.Source = ce.Tiers[0].Source
|
||||
rec.Model = ce.Tiers[0].Model
|
||||
}
|
||||
writeError(w, rec.Status, "upstream_error", clientUpstreamErr(err))
|
||||
}
|
||||
|
||||
// recordChatUsage fills token accounting for a finished non-streaming
|
||||
// request: exact upstream numbers win, byte estimates fill the gaps.
|
||||
func recordChatUsage(rec *Req, req *types.ChatRequest, resp *types.UnifiedResponse) {
|
||||
rec.Prompt = int64(resp.TokenUsage.Prompt)
|
||||
if rec.Prompt == 0 {
|
||||
rec.Prompt = estimatePromptTokens(req)
|
||||
@ -546,9 +570,29 @@ func (g *Gateway) singleChat(w http.ResponseWriter, ctx context.Context, cands [
|
||||
if rec.Compl == 0 {
|
||||
rec.Compl = estimateTextTokens(resp.Content, resp.ReasoningContent, resp.ToolCalls)
|
||||
}
|
||||
rec.Source = usedSrc
|
||||
rec.Model = usedModel
|
||||
g.writeRec(rec)
|
||||
// Cache accounting: whichever source format the adapter normalized into
|
||||
// (prompt_tokens_details.cached_tokens or legacy prompt_cache_hit_tokens),
|
||||
// read it back so the request record carries the hit/miss split.
|
||||
if d := resp.TokenUsage.PromptTokensDetails; d != nil && d.CachedTokens > 0 {
|
||||
rec.CacheHit = int64(d.CachedTokens)
|
||||
}
|
||||
if resp.TokenUsage.PromptTokensDetails != nil {
|
||||
// Upstream reported cache details (even a 0 hit) — tag the row so
|
||||
// the UI can show 0% rather than “—”.
|
||||
rec.CacheReported = true
|
||||
} else if resp.TokenUsage.PromptCacheHit > 0 {
|
||||
rec.CacheHit = int64(resp.TokenUsage.PromptCacheHit)
|
||||
rec.CacheReported = true
|
||||
}
|
||||
if resp.TokenUsage.PromptCacheMiss > 0 {
|
||||
rec.CacheMiss = int64(resp.TokenUsage.PromptCacheMiss)
|
||||
}
|
||||
}
|
||||
|
||||
// writeChatCompletion renders a unified response as an OpenAI
|
||||
// chat.completion object. modelName is the id clients see as the serving
|
||||
// model: the requested id for direct routes, the exact slot model for AUTO.
|
||||
func writeChatCompletion(w http.ResponseWriter, resp *types.UnifiedResponse, modelName string) {
|
||||
msg := RespMessage{Role: "assistant", Content: resp.Content}
|
||||
if resp.ReasoningContent != "" {
|
||||
msg.ReasoningContent = resp.ReasoningContent
|
||||
@ -560,7 +604,7 @@ 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}},
|
||||
}
|
||||
if resp.TokenUsage.Total > 0 || resp.TokenUsage.Prompt > 0 || resp.TokenUsage.Completion > 0 {
|
||||
@ -569,6 +613,30 @@ func (g *Gateway) singleChat(w http.ResponseWriter, ctx context.Context, cands [
|
||||
writeJSON(w, http.StatusOK, out)
|
||||
}
|
||||
|
||||
// singleChat runs a direct (model-pinned) non-streaming request across the
|
||||
// candidate list, falling back on failure.
|
||||
func (g *Gateway) singleChat(w http.ResponseWriter, ctx context.Context, cands []*provider.Provider, req *types.ChatRequest, effective string, rec *Req) {
|
||||
rec.LatMs = 0
|
||||
t0 := time.Now()
|
||||
resp, usedSrc, usedModel, err := g.core.Scheduler().Chat(ctx, toScheduler(cands), req)
|
||||
rec.LatMs = time.Since(t0).Milliseconds()
|
||||
if err != nil {
|
||||
g.failChat(w, rec, err)
|
||||
g.writeRec(rec)
|
||||
return
|
||||
}
|
||||
rec.OK = true
|
||||
rec.Status = http.StatusOK
|
||||
recordChatUsage(rec, req, resp)
|
||||
rec.Source = usedSrc
|
||||
rec.Model = usedModel
|
||||
// Non-streaming: the whole response arrives at once, so TTFB equals
|
||||
// the total latency.
|
||||
rec.FirstByteMs = rec.LatMs
|
||||
g.writeRec(rec)
|
||||
writeChatCompletion(w, resp, effective)
|
||||
}
|
||||
|
||||
// writeRec records a finished request (audit + aggregates).
|
||||
func (g *Gateway) writeRec(rec *Req) {
|
||||
if rec == nil {
|
||||
@ -603,30 +671,28 @@ func mergeUsage(prev, cur *types.TokenUsage) *types.TokenUsage {
|
||||
out.Completion = cur.Completion
|
||||
}
|
||||
out.Total = out.Prompt + out.Completion
|
||||
if cur.PromptTokensDetails != nil {
|
||||
// Keep the details object even when CachedTokens is 0: a reported
|
||||
// zero-hit is meaningful ("cache missed") and must stay
|
||||
// distinguishable from "upstream never reported cache info".
|
||||
// Dropping it here made streaming rows show “—” instead of 0%.
|
||||
out.PromptTokensDetails = cur.PromptTokensDetails
|
||||
}
|
||||
if cur.PromptCacheHit > 0 {
|
||||
out.PromptCacheHit = cur.PromptCacheHit
|
||||
}
|
||||
if cur.PromptCacheMiss > 0 {
|
||||
out.PromptCacheMiss = cur.PromptCacheMiss
|
||||
}
|
||||
return &out
|
||||
}
|
||||
|
||||
func (g *Gateway) streamChat(w http.ResponseWriter, ctx context.Context, cands []*provider.Provider, req *types.ChatRequest, effective string, rec *Req) {
|
||||
rec.LatMs = 0
|
||||
t0 := time.Now()
|
||||
rec.OK = true
|
||||
rec.Status = http.StatusOK
|
||||
defer func() {
|
||||
rec.LatMs = time.Since(t0).Milliseconds()
|
||||
g.writeRec(rec)
|
||||
}()
|
||||
chunks, _, usedModel, err := g.core.Scheduler().ChatStream(ctx, toScheduler(cands), req)
|
||||
if err != nil {
|
||||
rec.OK = false
|
||||
rec.Status = upstreamErrStatus(err)
|
||||
rec.Err = err.Error()
|
||||
writeError(w, rec.Status, "upstream_error", err.Error())
|
||||
return
|
||||
}
|
||||
if usedModel != "" {
|
||||
rec.Model = usedModel
|
||||
}
|
||||
rec.Prompt = estimatePromptTokens(req)
|
||||
// pumpStream writes the full SSE sequence for a started stream: role
|
||||
// preamble, one chunk per unified delta, the terminating finish_reason, the
|
||||
// OpenAI-standard final usage chunk (empty choices) and [DONE]. modelName
|
||||
// follows writeChatCompletion's rule (requested id for direct routes, exact
|
||||
// slot model for AUTO).
|
||||
func (g *Gateway) pumpStream(w http.ResponseWriter, rec *Req, chunks <-chan types.UnifiedChunk, modelName string, t0 time.Time) {
|
||||
w.Header().Set("Content-Type", "text/event-stream")
|
||||
w.Header().Set("Cache-Control", "no-cache")
|
||||
w.Header().Set("Connection", "keep-alive")
|
||||
@ -650,18 +716,22 @@ 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
|
||||
lastFinish := ""
|
||||
for ck := range chunks {
|
||||
if ck.Usage != nil {
|
||||
lastUsage = mergeUsage(lastUsage, ck.Usage)
|
||||
}
|
||||
chunk := ChatChunk{
|
||||
ID: id, Object: "chat.completion.chunk", Created: created, Model: effective,
|
||||
ID: id, Object: "chat.completion.chunk", Created: created, Model: modelName,
|
||||
}
|
||||
delta := RespMessage{Role: "assistant", Content: ck.Content}
|
||||
if ck.ReasoningContent != "" {
|
||||
@ -672,8 +742,12 @@ func (g *Gateway) streamChat(w http.ResponseWriter, ctx context.Context, cands [
|
||||
}
|
||||
choice := ChunkChoice{Index: 0, Delta: delta}
|
||||
if ck.Done {
|
||||
stop := "stop"
|
||||
choice.FinishReason = &stop
|
||||
fin := ck.FinishReason
|
||||
if fin == "" {
|
||||
fin = "stop"
|
||||
}
|
||||
lastFinish = fin
|
||||
choice.FinishReason = &fin
|
||||
}
|
||||
chunk.Choices = []ChunkChoice{choice}
|
||||
rec.Compl += int64(len(ck.Content)+len(ck.ReasoningContent)+len(ck.ToolCalls)) / 3
|
||||
@ -681,10 +755,13 @@ 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
|
||||
@ -692,6 +769,20 @@ func (g *Gateway) streamChat(w http.ResponseWriter, ctx context.Context, cands [
|
||||
var tut *types.TokenUsage
|
||||
if lastUsage != nil {
|
||||
tut = lastUsage
|
||||
// Write the upstream's cache accounting back onto the request record
|
||||
// so the audit trail carries the hit/miss split for streaming too.
|
||||
if d := tut.PromptTokensDetails; d != nil {
|
||||
rec.CacheReported = true
|
||||
if d.CachedTokens > 0 {
|
||||
rec.CacheHit = int64(d.CachedTokens)
|
||||
}
|
||||
} else if tut.PromptCacheHit > 0 {
|
||||
rec.CacheHit = int64(tut.PromptCacheHit)
|
||||
rec.CacheReported = true
|
||||
}
|
||||
if tut.PromptCacheMiss > 0 {
|
||||
rec.CacheMiss = int64(tut.PromptCacheMiss)
|
||||
}
|
||||
} else if rec.Prompt+rec.Compl > 0 {
|
||||
u := types.TokenUsage{
|
||||
Prompt: int(rec.Prompt),
|
||||
@ -702,7 +793,7 @@ func (g *Gateway) streamChat(w http.ResponseWriter, ctx context.Context, cands [
|
||||
}
|
||||
if tut != nil {
|
||||
send(ChatChunk{
|
||||
ID: id, Object: "chat.completion.chunk", Created: created, Model: effective,
|
||||
ID: id, Object: "chat.completion.chunk", Created: created, Model: modelName,
|
||||
Choices: []ChunkChoice{},
|
||||
Usage: tut,
|
||||
})
|
||||
@ -713,6 +804,33 @@ func (g *Gateway) streamChat(w http.ResponseWriter, ctx context.Context, cands [
|
||||
}
|
||||
}
|
||||
|
||||
// streamChat runs a direct (model-pinned) streaming request across the
|
||||
// candidate list, falling back early on connect errors.
|
||||
func (g *Gateway) streamChat(w http.ResponseWriter, ctx context.Context, cands []*provider.Provider, req *types.ChatRequest, effective string, rec *Req) {
|
||||
rec.LatMs = 0
|
||||
t0 := time.Now()
|
||||
rec.OK = true
|
||||
rec.Status = http.StatusOK
|
||||
defer func() {
|
||||
rec.LatMs = time.Since(t0).Milliseconds()
|
||||
g.writeRec(rec)
|
||||
}()
|
||||
chunks, usedSrc, usedModel, err := g.core.Scheduler().ChatStream(ctx, toScheduler(cands), req)
|
||||
if err != nil {
|
||||
g.failChat(w, rec, err)
|
||||
return
|
||||
}
|
||||
if usedModel != "" {
|
||||
rec.Model = usedModel
|
||||
}
|
||||
// Audit accuracy: pin the source that actually served the stream (after a
|
||||
// failover it differs from the first candidate). Direct streams previously
|
||||
// discarded it.
|
||||
rec.Source = usedSrc
|
||||
rec.Prompt = estimatePromptTokens(req)
|
||||
g.pumpStream(w, rec, chunks, effective, t0)
|
||||
}
|
||||
|
||||
// singleChatAuto runs a non-streaming AUTO request down the chain (see
|
||||
// scheduler.ChainChat): tiers ascending (tier 1 = highest priority first),
|
||||
// per-tier round-robin ordered by preference, cooldown as the only hard skip,
|
||||
@ -725,48 +843,18 @@ func (g *Gateway) singleChatAuto(w http.ResponseWriter, ctx context.Context, cha
|
||||
resp, usedSrc, usedModel, err := g.core.Scheduler().ChainChat(ctx, chain, req, quotaExhausted)
|
||||
rec.LatMs = time.Since(t0).Milliseconds()
|
||||
if err != nil {
|
||||
rec.OK = false
|
||||
rec.Status = upstreamErrStatus(err)
|
||||
rec.Err = err.Error()
|
||||
if ce, ok := err.(*scheduler.ChainErr); ok && len(ce.Tiers) > 0 {
|
||||
rec.Source = ce.Tiers[0].Source
|
||||
rec.Model = ce.Tiers[0].Model
|
||||
}
|
||||
g.failChat(w, rec, err)
|
||||
g.writeRec(rec)
|
||||
writeError(w, rec.Status, "upstream_error", err.Error())
|
||||
return
|
||||
}
|
||||
rec.OK = true
|
||||
rec.Status = http.StatusOK
|
||||
rec.Prompt = int64(resp.TokenUsage.Prompt)
|
||||
if rec.Prompt == 0 {
|
||||
rec.Prompt = estimatePromptTokens(req)
|
||||
}
|
||||
rec.Compl = int64(resp.TokenUsage.Completion)
|
||||
if rec.Compl == 0 {
|
||||
rec.Compl = estimateTextTokens(resp.Content, resp.ReasoningContent, resp.ToolCalls)
|
||||
}
|
||||
recordChatUsage(rec, req, resp)
|
||||
rec.Source = usedSrc
|
||||
rec.Model = usedModel
|
||||
rec.FirstByteMs = rec.LatMs
|
||||
g.writeRec(rec)
|
||||
msg := RespMessage{Role: "assistant", Content: resp.Content}
|
||||
if resp.ReasoningContent != "" {
|
||||
msg.ReasoningContent = resp.ReasoningContent
|
||||
}
|
||||
if len(resp.ToolCalls) > 0 {
|
||||
msg.ToolCalls = toolCallsWire(resp.ToolCalls)
|
||||
}
|
||||
out := ChatCompletion{
|
||||
ID: newID(),
|
||||
Object: "chat.completion",
|
||||
Created: time.Now().Unix(),
|
||||
Model: usedModel,
|
||||
Choices: []ChatChoice{{Index: 0, Message: msg, FinishReason: resp.FinishReason}},
|
||||
}
|
||||
if resp.TokenUsage.Total > 0 || resp.TokenUsage.Prompt > 0 || resp.TokenUsage.Completion > 0 {
|
||||
out.Usage = &resp.TokenUsage
|
||||
}
|
||||
writeJSON(w, http.StatusOK, out)
|
||||
writeChatCompletion(w, resp, usedModel)
|
||||
}
|
||||
|
||||
// streamChatAuto streams an AUTO request down the chain. A slot is abandoned
|
||||
@ -784,14 +872,7 @@ func (g *Gateway) streamChatAuto(w http.ResponseWriter, ctx context.Context, cha
|
||||
}()
|
||||
chunks, usedSrc, usedModel, err := g.core.Scheduler().ChainChatStream(ctx, chain, req, quotaExhausted)
|
||||
if err != nil {
|
||||
rec.OK = false
|
||||
rec.Status = upstreamErrStatus(err)
|
||||
rec.Err = err.Error()
|
||||
if ce, ok := err.(*scheduler.ChainErr); ok && len(ce.Tiers) > 0 {
|
||||
rec.Source = ce.Tiers[0].Source
|
||||
rec.Model = ce.Tiers[0].Model
|
||||
}
|
||||
writeError(w, rec.Status, "upstream_error", err.Error())
|
||||
g.failChat(w, rec, err)
|
||||
return
|
||||
}
|
||||
if usedModel != "" {
|
||||
@ -799,88 +880,7 @@ func (g *Gateway) streamChatAuto(w http.ResponseWriter, ctx context.Context, cha
|
||||
}
|
||||
rec.Source = usedSrc
|
||||
rec.Prompt = estimatePromptTokens(req)
|
||||
w.Header().Set("Content-Type", "text/event-stream")
|
||||
w.Header().Set("Cache-Control", "no-cache")
|
||||
w.Header().Set("Connection", "keep-alive")
|
||||
w.WriteHeader(http.StatusOK)
|
||||
flusher, _ := w.(http.Flusher)
|
||||
id := newID()
|
||||
created := time.Now().Unix()
|
||||
send := func(obj interface{}) bool {
|
||||
b, err := json.Marshal(obj)
|
||||
if err != nil {
|
||||
return false
|
||||
}
|
||||
if _, err := fmt.Fprintf(w, "data: %s\n\n", b); err != nil {
|
||||
return false
|
||||
}
|
||||
if flusher != nil {
|
||||
flusher.Flush()
|
||||
}
|
||||
return true
|
||||
}
|
||||
if !send(ChatChunk{
|
||||
ID: id, Object: "chat.completion.chunk", Created: created, Model: rec.Model,
|
||||
Choices: []ChunkChoice{{Index: 0, Delta: RespMessage{Role: "assistant"}}},
|
||||
}) {
|
||||
return
|
||||
}
|
||||
var lastUsage *types.TokenUsage
|
||||
for ck := range chunks {
|
||||
if ck.Usage != nil {
|
||||
lastUsage = mergeUsage(lastUsage, ck.Usage)
|
||||
}
|
||||
chunk := ChatChunk{
|
||||
ID: id, Object: "chat.completion.chunk", Created: created, Model: rec.Model,
|
||||
}
|
||||
delta := RespMessage{Role: "assistant", Content: ck.Content}
|
||||
if ck.ReasoningContent != "" {
|
||||
delta.ReasoningContent = ck.ReasoningContent
|
||||
}
|
||||
if len(ck.ToolCalls) > 0 {
|
||||
delta.ToolCalls = ck.ToolCalls
|
||||
}
|
||||
choice := ChunkChoice{Index: 0, Delta: delta}
|
||||
if ck.Done {
|
||||
stop := "stop"
|
||||
choice.FinishReason = &stop
|
||||
}
|
||||
chunk.Choices = []ChunkChoice{choice}
|
||||
rec.Compl += int64(len(ck.Content)+len(ck.ReasoningContent)+len(ck.ToolCalls)) / 3
|
||||
if !send(chunk) {
|
||||
return
|
||||
}
|
||||
}
|
||||
stop := "stop"
|
||||
send(ChatChunk{
|
||||
ID: id, Object: "chat.completion.chunk", Created: created, Model: rec.Model,
|
||||
Choices: []ChunkChoice{{Index: 0, Delta: RespMessage{}, FinishReason: &stop}},
|
||||
})
|
||||
// Final usage chunk (OpenAI standard: empty choices + usage before [DONE]).
|
||||
// Prefer the upstream's exact usage if the stream carried it; fall back to
|
||||
// the gateway's estimate otherwise.
|
||||
var tut *types.TokenUsage
|
||||
if lastUsage != nil {
|
||||
tut = lastUsage
|
||||
} else if rec.Prompt+rec.Compl > 0 {
|
||||
u := types.TokenUsage{
|
||||
Prompt: int(rec.Prompt),
|
||||
Completion: int(rec.Compl),
|
||||
Total: int(rec.Prompt + rec.Compl),
|
||||
}
|
||||
tut = &u
|
||||
}
|
||||
if tut != nil {
|
||||
send(ChatChunk{
|
||||
ID: id, Object: "chat.completion.chunk", Created: created, Model: rec.Model,
|
||||
Choices: []ChunkChoice{},
|
||||
Usage: tut,
|
||||
})
|
||||
}
|
||||
fmt.Fprintf(w, "data: [DONE]\n\n")
|
||||
if flusher != nil {
|
||||
flusher.Flush()
|
||||
}
|
||||
g.pumpStream(w, rec, chunks, usedModel, t0)
|
||||
}
|
||||
|
||||
func (g *Gateway) handleImage(w http.ResponseWriter, r *http.Request) {
|
||||
@ -901,11 +901,45 @@ 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 msg := g.checkModelScope(r.Context(), "AUTO"); msg != "" {
|
||||
writeError(w, http.StatusForbidden, "model_not_allowed", msg)
|
||||
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
|
||||
rec.Compl = int64(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)
|
||||
@ -930,12 +964,15 @@ 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))
|
||||
|
||||
62
internal/gateway/chat_cache_test.go
Normal file
62
internal/gateway/chat_cache_test.go
Normal file
@ -0,0 +1,62 @@
|
||||
package gateway
|
||||
|
||||
import (
|
||||
"testing"
|
||||
|
||||
"llmsproxy/internal/types"
|
||||
)
|
||||
|
||||
// TestMergeUsageKeepsZeroCacheDetails locks the streaming counterpart of the
|
||||
// "distinguish missed from not reported" contract: a usage chunk that reports
|
||||
// prompt_tokens_details with 0 cached tokens must survive the merge, otherwise
|
||||
// streaming rows lose cache_reported and the UI shows "—" instead of 0%.
|
||||
func TestMergeUsageKeepsZeroCacheDetails(t *testing.T) {
|
||||
t.Run("zero-hit details survive merge", func(t *testing.T) {
|
||||
prev := &types.TokenUsage{Prompt: 10, Completion: 1, Total: 11}
|
||||
cur := &types.TokenUsage{
|
||||
Prompt: 10, Completion: 5, Total: 15,
|
||||
PromptTokensDetails: &types.PromptTokensDetails{CachedTokens: 0},
|
||||
}
|
||||
got := mergeUsage(prev, cur)
|
||||
if got.PromptTokensDetails == nil {
|
||||
t.Fatal("zero-hit prompt_tokens_details was dropped by mergeUsage")
|
||||
}
|
||||
if got.PromptTokensDetails.CachedTokens != 0 {
|
||||
t.Fatalf("cached_tokens = %d, want 0", got.PromptTokensDetails.CachedTokens)
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("non-zero hit still wins", func(t *testing.T) {
|
||||
prev := &types.TokenUsage{Prompt: 10, Completion: 1, Total: 11}
|
||||
cur := &types.TokenUsage{
|
||||
Prompt: 10, Completion: 5, Total: 15,
|
||||
PromptTokensDetails: &types.PromptTokensDetails{CachedTokens: 64},
|
||||
}
|
||||
got := mergeUsage(prev, cur)
|
||||
if got.PromptTokensDetails == nil || got.PromptTokensDetails.CachedTokens != 64 {
|
||||
t.Fatalf("cached_tokens lost: %+v", got.PromptTokensDetails)
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("absent details do not overwrite an earlier report", func(t *testing.T) {
|
||||
prev := &types.TokenUsage{
|
||||
Prompt: 10, Completion: 1, Total: 11,
|
||||
PromptTokensDetails: &types.PromptTokensDetails{CachedTokens: 32},
|
||||
}
|
||||
cur := &types.TokenUsage{Prompt: 10, Completion: 5, Total: 15}
|
||||
got := mergeUsage(prev, cur)
|
||||
if got.PromptTokensDetails == nil || got.PromptTokensDetails.CachedTokens != 32 {
|
||||
t.Fatalf("earlier cache report clobbered by a later chunk without details: %+v",
|
||||
got.PromptTokensDetails)
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("no cache data anywhere stays nil", func(t *testing.T) {
|
||||
prev := &types.TokenUsage{Prompt: 10, Completion: 1, Total: 11}
|
||||
cur := &types.TokenUsage{Prompt: 10, Completion: 5, Total: 15}
|
||||
got := mergeUsage(prev, cur)
|
||||
if got.PromptTokensDetails != nil {
|
||||
t.Fatalf("fabricated cache details: %+v", got.PromptTokensDetails)
|
||||
}
|
||||
})
|
||||
}
|
||||
@ -139,7 +139,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())
|
||||
@ -671,3 +671,71 @@ func TestHasScopeModelWithSourcePrefix(t *testing.T) {
|
||||
t.Error("AUTO scope should allow any prefixed model")
|
||||
}
|
||||
}
|
||||
|
||||
// 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)
|
||||
}
|
||||
}
|
||||
|
||||
@ -143,8 +143,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,17 +156,30 @@ 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", "")
|
||||
}
|
||||
|
||||
114
internal/gateway/login_guard_test.go
Normal file
114
internal/gateway/login_guard_test.go
Normal file
@ -0,0 +1,114 @@
|
||||
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)
|
||||
}
|
||||
}
|
||||
@ -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) {
|
||||
@ -113,6 +195,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":
|
||||
@ -125,6 +218,8 @@ func (g *Gateway) routes(w http.ResponseWriter, r *http.Request) {
|
||||
g.handleAdaptersAPI(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":
|
||||
@ -325,23 +420,34 @@ document.getElementById('key').addEventListener('keydown',function(e){if(e.key==
|
||||
apply();
|
||||
</script></body></html>`
|
||||
|
||||
// handleLoginAPI validates the gateway key and issues a session cookie.
|
||||
// handleLoginAPI validates the gateway key and issues a session cookie,
|
||||
// guarded by per-IP + global brute-force throttling (see loginGuard).
|
||||
func (g *Gateway) handleLoginAPI(w http.ResponseWriter, r *http.Request) {
|
||||
if r.Method != http.MethodPost {
|
||||
writeError(w, http.StatusMethodNotAllowed, "method_not_allowed", "use POST")
|
||||
return
|
||||
}
|
||||
now := time.Now().Unix()
|
||||
if !g.loginAllow(clientIP(r), now) {
|
||||
// Don't reveal whether the key was right: answer the same 401 the
|
||||
// failure path uses, but skip the expensive key lookup.
|
||||
writeError(w, http.StatusUnauthorized, "invalid_api_key", "invalid gateway api key")
|
||||
return
|
||||
}
|
||||
var body struct {
|
||||
Key string `json:"key"`
|
||||
}
|
||||
if err := json.NewDecoder(r.Body).Decode(&body); err != nil {
|
||||
g.loginRecord(clientIP(r), false, now)
|
||||
writeError(w, http.StatusBadRequest, "invalid_request", "invalid json")
|
||||
return
|
||||
}
|
||||
if _, ok := g.core.FindKey(body.Key); !ok {
|
||||
g.loginRecord(clientIP(r), false, now)
|
||||
writeError(w, http.StatusUnauthorized, "invalid_api_key", "invalid gateway api key")
|
||||
return
|
||||
}
|
||||
g.loginRecord(clientIP(r), true, now)
|
||||
http.SetCookie(w, &http.Cookie{
|
||||
Name: "gw_key",
|
||||
Value: body.Key,
|
||||
@ -353,6 +459,20 @@ func (g *Gateway) handleLoginAPI(w http.ResponseWriter, r *http.Request) {
|
||||
writeJSON(w, http.StatusOK, map[string]interface{}{"ok": true})
|
||||
}
|
||||
|
||||
// clientIP returns the caller's IP for rate limiting. Behind a trusted
|
||||
// reverse proxy (nginx) the X-Real-IP header is preferred; otherwise the
|
||||
// remote address is used with the port stripped.
|
||||
func clientIP(r *http.Request) string {
|
||||
if x := r.Header.Get("X-Real-IP"); x != "" {
|
||||
return x
|
||||
}
|
||||
host, _, err := net.SplitHostPort(r.RemoteAddr)
|
||||
if err != nil {
|
||||
return r.RemoteAddr
|
||||
}
|
||||
return host
|
||||
}
|
||||
|
||||
func (g *Gateway) handleModels(w http.ResponseWriter, r *http.Request) {
|
||||
if r.Method != http.MethodGet {
|
||||
writeError(w, http.StatusMethodNotAllowed, "method_not_allowed", "use GET")
|
||||
@ -458,11 +578,16 @@ func (g *Gateway) handleStatusAPI(w http.ResponseWriter, r *http.Request) {
|
||||
if reqRole(r.Context()) == "admin" {
|
||||
sts := g.core.Registry().Status()
|
||||
recent := g.stats.SourceRecent(300)
|
||||
avgs := g.stats.SourceAverages(300)
|
||||
for i := range sts {
|
||||
if v, ok := recent[sts[i].Name]; ok {
|
||||
sts[i].RecentOK = v[0]
|
||||
sts[i].RecentErr = v[1]
|
||||
}
|
||||
if a, ok := avgs[sts[i].Name]; ok {
|
||||
sts[i].AvgFirstByteMs = a.AvgFirstByteMs
|
||||
sts[i].AvgTokPerS = a.AvgTokPerS
|
||||
}
|
||||
}
|
||||
resp["sources"] = sts
|
||||
resp["adapters"] = g.core.ListAdapters()
|
||||
|
||||
@ -26,6 +26,19 @@ type Req struct {
|
||||
Compl int64 `json:"completion_tokens"`
|
||||
// LatMs total handling time ms
|
||||
LatMs int64 `json:"latency_ms"`
|
||||
// FirstByteMs time-to-first-byte for streaming (ms from request start
|
||||
// to the first SSE chunk sent to the client); for non-streaming it
|
||||
// equals LatMs. 0 when unmeasured (legacy records).
|
||||
FirstByteMs int64 `json:"first_byte_ms,omitempty"`
|
||||
// CacheHit / CacheMiss carry the upstream prompt-cache accounting
|
||||
// (DeepSeek-style hit/miss tokens) when the upstream reports it.
|
||||
// Both 0 = upstream gave no cache data.
|
||||
CacheHit int64 `json:"cache_hit_tokens,omitempty"`
|
||||
CacheMiss int64 `json:"cache_miss_tokens,omitempty"`
|
||||
// CacheReported marks that the upstream usage reported cache
|
||||
// accounting at all (even when the hit count is 0). The WebUI shows
|
||||
// "0%" instead of "—" for such rows.
|
||||
CacheReported bool `json:"cache_reported,omitempty"`
|
||||
OK bool `json:"ok"`
|
||||
// Status http status code
|
||||
Status int `json:"status"`
|
||||
@ -35,14 +48,15 @@ type Req struct {
|
||||
|
||||
// Stat aggregates counters for one dimension row.
|
||||
type Stat struct {
|
||||
Reqs int64 `json:"reqs"`
|
||||
OK int64 `json:"ok"`
|
||||
Err int64 `json:"err"`
|
||||
Tokens int64 `json:"tokens"`
|
||||
Prompt int64 `json:"prompt_tokens"`
|
||||
Compl int64 `json:"completion_tokens"`
|
||||
LatSum int64 `json:"latency_sum_ms"`
|
||||
LatMax int64 `json:"latency_max_ms"`
|
||||
Reqs int64 `json:"reqs"`
|
||||
OK int64 `json:"ok"`
|
||||
Err int64 `json:"err"`
|
||||
Tokens int64 `json:"tokens"`
|
||||
Prompt int64 `json:"prompt_tokens"`
|
||||
Compl int64 `json:"completion_tokens"`
|
||||
LatSum int64 `json:"latency_sum_ms"`
|
||||
LatMax int64 `json:"latency_max_ms"`
|
||||
FirstByteSum int64 `json:"first_byte_sum_ms,omitempty"`
|
||||
}
|
||||
|
||||
type agrRow struct {
|
||||
@ -130,6 +144,9 @@ func incStatus(a *Stat, name string, r Req) {
|
||||
if r.LatMs > a.LatMax {
|
||||
a.LatMax = r.LatMs
|
||||
}
|
||||
if r.FirstByteMs > 0 {
|
||||
a.FirstByteSum += r.FirstByteMs
|
||||
}
|
||||
}
|
||||
|
||||
func (s *Stats) LoadAudit(path string) {
|
||||
@ -380,26 +397,6 @@ type StatsRow struct {
|
||||
Stat
|
||||
}
|
||||
|
||||
// Records returns request records filtered by unix-millisecond time range and key.
|
||||
func (s *Stats) Records(from, to int64, key string) []Req {
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
out := make([]Req, 0, len(s.recs))
|
||||
for _, r := range s.recs {
|
||||
if key != "" && r.Key != key {
|
||||
continue
|
||||
}
|
||||
if from > 0 && r.Time < from {
|
||||
continue
|
||||
}
|
||||
if to > 0 && r.Time > to {
|
||||
continue
|
||||
}
|
||||
out = append(out, r)
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
// AuditRecords returns every request row from the audit file plus its rotated
|
||||
// .old files within the [from,to] unix-millisecond window, optionally for one
|
||||
// masked key id. Unlike Records it is not bounded by the in-memory ring, so it
|
||||
@ -472,6 +469,65 @@ func (s *Stats) SourceRecent(windowSec int64) map[string][2]int64 {
|
||||
return out
|
||||
}
|
||||
|
||||
// SourceAvg carries per-source performance averages for the status page.
|
||||
type SourceAvg struct {
|
||||
// AvgFirstByteMs is the mean time-to-first-byte over successful
|
||||
// requests in the window (0 when no measured samples).
|
||||
AvgFirstByteMs int64 `json:"avg_first_byte_ms"`
|
||||
// AvgTokPerS is the aggregate completion throughput:
|
||||
// sum(completion_tokens) / sum(latency_seconds) (0 when no samples).
|
||||
AvgTokPerS int64 `json:"avg_tok_per_s"`
|
||||
// Samples is the number of successful requests the averages cover.
|
||||
Samples int64 `json:"samples"`
|
||||
}
|
||||
|
||||
// SourceAverages computes TTFB and tokens/s averages per source from the
|
||||
// in-memory ring within the window (unix seconds). Only successful chat/
|
||||
// stream rows count; image and failed rows are skipped.
|
||||
func (s *Stats) SourceAverages(windowSec int64) map[string]SourceAvg {
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
cut := time.Now().Unix() - windowSec
|
||||
type acc struct {
|
||||
fbSum, latSum, complSum, n int64
|
||||
}
|
||||
accs := map[string]*acc{}
|
||||
for _, r := range s.recs {
|
||||
if !r.OK || r.Source == "" || r.Time/1000 < cut {
|
||||
continue
|
||||
}
|
||||
if r.Type != "chat" && r.Type != "stream" {
|
||||
continue
|
||||
}
|
||||
if r.LatMs <= 0 {
|
||||
continue
|
||||
}
|
||||
a := accs[r.Source]
|
||||
if a == nil {
|
||||
a = &acc{}
|
||||
accs[r.Source] = a
|
||||
}
|
||||
a.latSum += r.LatMs
|
||||
a.complSum += r.Compl
|
||||
a.fbSum += r.FirstByteMs
|
||||
if r.FirstByteMs > 0 {
|
||||
a.n++
|
||||
}
|
||||
}
|
||||
out := make(map[string]SourceAvg, len(accs))
|
||||
for src, a := range accs {
|
||||
avg := SourceAvg{Samples: a.n}
|
||||
if a.n > 0 {
|
||||
avg.AvgFirstByteMs = a.fbSum / a.n
|
||||
}
|
||||
if a.latSum > 0 && a.complSum > 0 {
|
||||
avg.AvgTokPerS = a.complSum * 1000 / a.latSum
|
||||
}
|
||||
out[src] = avg
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
// Snapshot returns the whole dashboard payload; when key != "" the records
|
||||
// and aggregate views are restricted to that gateway key.
|
||||
func (s *Stats) Snapshot(limit int, key string) map[string]interface{} {
|
||||
|
||||
@ -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)
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
164
internal/lua/adapters/agentrouter.lua
Normal file
164
internal/lua/adapters/agentrouter.lua
Normal file
@ -0,0 +1,164 @@
|
||||
local adapter = {}
|
||||
|
||||
adapter.name = "agentrouter"
|
||||
adapter.version = "1.0.0"
|
||||
adapter.endpoint = "/chat/completions"
|
||||
adapter.headers = {}
|
||||
|
||||
-- AgentRouter 的 WAF 会按客户端指纹白名单校验,只放行官方客户端。
|
||||
-- 这里通过 build_headers 注入官方客户端 User-Agent 以通过校验。
|
||||
-- 默认 QwenCode 指纹(实测可通);可在 source.meta.user_agent 覆盖成 Claude Code 指纹。
|
||||
|
||||
local default_ua = "QwenCode/0.2.0 (linux; x64)"
|
||||
|
||||
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
|
||||
|
||||
-- 注入 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,
|
||||
}
|
||||
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 = 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
|
||||
unified.tool_calls = delta.tool_calls
|
||||
end
|
||||
if uses ~= nil then
|
||||
unified.usage = uses
|
||||
end
|
||||
return json.encode(unified)
|
||||
end
|
||||
|
||||
-- 错误收敛:前置 WAF 会返回整页 HTML(阿里云盾),JSON 时为 new-api 风格
|
||||
function adapter.transform_error(status, body)
|
||||
local low = string.lower(body or "")
|
||||
if string.sub(low, 1, 9) == "<!doctype" or string.find(low, "<html", 1, true) then
|
||||
return "blocked by AgentRouter WAF"
|
||||
end
|
||||
local ok, resp = pcall(json.decode, body)
|
||||
if not ok or type(resp) ~= "table" then return nil end
|
||||
local e = resp.error
|
||||
if type(e) == "table" and type(e.message) == "string" then
|
||||
return e.message
|
||||
end
|
||||
return nil
|
||||
end
|
||||
|
||||
return adapter
|
||||
@ -1,120 +1,310 @@
|
||||
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,
|
||||
}
|
||||
|
||||
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 = tc.id or "",
|
||||
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 = m.tool_call_id or "",
|
||||
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.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)
|
||||
unified.token_usage.total = (resp.usage.input_tokens or 0) + (resp.usage.output_tokens or 0)
|
||||
-- Emit details whenever Anthropic reports the field, even at 0, so a
|
||||
-- reported cache miss stays distinguishable from "not reported".
|
||||
if resp.usage.cache_read_input_tokens ~= nil then
|
||||
unified.token_usage.prompt_tokens_details = {
|
||||
cached_tokens = resp.usage.cache_read_input_tokens
|
||||
}
|
||||
end
|
||||
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 = block.id or "",
|
||||
type = "function",
|
||||
name = block.name or "",
|
||||
arguments = block.input or {},
|
||||
})
|
||||
end
|
||||
end
|
||||
if #tcs > 0 then unified.tool_calls = tcs end
|
||||
end
|
||||
|
||||
unified.finish_reason = resp.stop_reason or ""
|
||||
if unified.finish_reason == "tool_use" then
|
||||
unified.finish_reason = "tool_calls"
|
||||
elseif unified.finish_reason == "max_tokens" then
|
||||
unified.finish_reason = "length"
|
||||
elseif unified.finish_reason == "end_turn"
|
||||
or unified.finish_reason == "stop_sequence" then
|
||||
unified.finish_reason = "stop"
|
||||
end
|
||||
|
||||
return json.encode(unified)
|
||||
end
|
||||
|
||||
-- ============================================================
|
||||
-- transform_stream_chunk (Anthropic SSE → OpenAI SSE delta)
|
||||
-- ============================================================
|
||||
|
||||
function adapter.transform_stream_chunk(raw_chunk)
|
||||
local ok, chunk = pcall(json.decode, raw_chunk)
|
||||
if not ok then return "" end
|
||||
|
||||
-- ── message_start: initial usage ─────────────────────────
|
||||
if chunk.type == "message_start" then
|
||||
local uses = nil
|
||||
if chunk.message and type(chunk.message.usage) == "table" then
|
||||
local u = chunk.message.usage
|
||||
local p = u.input_tokens or 0
|
||||
local c = u.output_tokens or 0
|
||||
if p > 0 or c > 0 then
|
||||
uses = { prompt = p, completion = c, total = p + c }
|
||||
local uses = { prompt = p, completion = c, total = p + c }
|
||||
if u.cache_read_input_tokens ~= nil then
|
||||
uses.prompt_tokens_details = {
|
||||
cached_tokens = u.cache_read_input_tokens
|
||||
}
|
||||
end
|
||||
return json.encode({ usage = uses, done = false })
|
||||
end
|
||||
end
|
||||
if uses ~= nil then
|
||||
return json.encode({ usage = uses, done = false })
|
||||
end
|
||||
return ""
|
||||
end
|
||||
|
||||
-- ── message_delta: stop_reason + final usage ─────────────
|
||||
if chunk.type == "message_delta" then
|
||||
local finish = nil
|
||||
if chunk.delta and chunk.delta.stop_reason ~= nil then
|
||||
local sr = chunk.delta.stop_reason
|
||||
if sr == "max_tokens" then
|
||||
finish = "length"
|
||||
elseif sr == "tool_use" then
|
||||
finish = "tool_calls"
|
||||
else
|
||||
finish = "stop"
|
||||
end
|
||||
end
|
||||
local uses = nil
|
||||
if type(chunk.usage) == "table" then
|
||||
local u = chunk.usage
|
||||
@ -124,49 +314,87 @@ function adapter.transform_stream_chunk(raw_chunk)
|
||||
uses = { prompt = p, completion = c, total = p + c }
|
||||
end
|
||||
end
|
||||
local done = (chunk.delta and chunk.delta.stop_reason ~= nil)
|
||||
if uses ~= nil then
|
||||
-- completion is final here; prompt is merged from message_start
|
||||
return json.encode({ content = "", done = done, usage = uses })
|
||||
return json.encode({ content = "", done = (finish ~= nil), finish_reason = finish, usage = uses })
|
||||
end
|
||||
return json.encode({ content = "", done = done })
|
||||
return json.encode({ content = "", done = (finish ~= nil), finish_reason = finish })
|
||||
end
|
||||
if chunk.type == "content_block_start" and chunk.content_block
|
||||
and chunk.content_block.type == "tool_use" then
|
||||
-- first fragment of a tool call: emit index + id + name, empty args
|
||||
return json.encode({
|
||||
content = "", done = false,
|
||||
tool_calls = { {
|
||||
index = chunk.index or 0,
|
||||
id = chunk.content_block.id or "",
|
||||
type = "function",
|
||||
["function"] = { name = chunk.content_block.name or "", arguments = "" }
|
||||
} }
|
||||
})
|
||||
|
||||
-- ── content_block_start: begin text / thinking / tool_use ─
|
||||
if chunk.type == "content_block_start" and chunk.content_block then
|
||||
local cb = chunk.content_block
|
||||
if cb.type == "tool_use" then
|
||||
-- Pass Anthropic content_block index through; Go-side rewrites
|
||||
-- to sequential OpenAI tool_call ordinal for multi-tool streams.
|
||||
return json.encode({
|
||||
content = "", done = false,
|
||||
tool_calls = { {
|
||||
index = chunk.index or 0,
|
||||
id = cb.id or "",
|
||||
type = "function",
|
||||
["function"] = { name = cb.name or "", arguments = "" },
|
||||
} },
|
||||
})
|
||||
end
|
||||
return "" -- text / thinking block start: no OpenAI equivalent
|
||||
end
|
||||
|
||||
-- ── content_block_delta: incremental content ──────────────
|
||||
if chunk.type == "content_block_delta" and chunk.delta then
|
||||
if chunk.delta.type == "input_json_delta" then
|
||||
-- incremental JSON fragment; clients accumulate across chunks
|
||||
local unified = { content = "", done = false, tool_calls = { {
|
||||
index = chunk.index or 0,
|
||||
id = "",
|
||||
type = "function",
|
||||
["function"] = { name = "", arguments = chunk.delta.partial_json or "" }
|
||||
} } }
|
||||
return json.encode(unified)
|
||||
local d = chunk.delta
|
||||
if d.type == "input_json_delta" then
|
||||
-- Tool call argument fragment; Go-side rewrites index.
|
||||
return json.encode({
|
||||
content = "", done = false,
|
||||
tool_calls = { {
|
||||
index = chunk.index or 0,
|
||||
id = "",
|
||||
type = "function",
|
||||
["function"] = { name = "", arguments = d.partial_json or "" },
|
||||
} },
|
||||
})
|
||||
end
|
||||
if chunk.delta.type == "thinking_delta" and chunk.delta.thinking then
|
||||
return json.encode({ content = "", done = false, reasoning_content = chunk.delta.thinking })
|
||||
if d.type == "thinking_delta" and d.thinking then
|
||||
return json.encode({ content = "", done = false, reasoning_content = d.thinking })
|
||||
end
|
||||
if d.type == "text_delta" and d.text then
|
||||
return json.encode({ content = d.text, done = false })
|
||||
end
|
||||
return json.encode({ content = chunk.delta.text or "", done = false })
|
||||
end
|
||||
|
||||
-- ── content_block_stop / message_stop ─────────────────────
|
||||
if chunk.type == "message_stop" then
|
||||
return json.encode({ content = "", done = true })
|
||||
end
|
||||
if chunk.type == "content_block_stop" then
|
||||
return json.encode({ content = "", done = false })
|
||||
-- The message_delta event already emitted the true finish_reason.
|
||||
-- Do NOT emit done=true here: an empty finish_reason would
|
||||
-- overwrite the real one (tool_calls) in the Go gateway's
|
||||
-- lastFinish tracker, causing the final SSE chunk to say
|
||||
-- finish_reason=stop instead of tool_calls.
|
||||
return ""
|
||||
end
|
||||
|
||||
return ""
|
||||
end
|
||||
|
||||
-- ============================================================
|
||||
-- transform_error (Anthropic error → human-readable string)
|
||||
-- ============================================================
|
||||
|
||||
function adapter.transform_error(status, body)
|
||||
local ok, resp = pcall(json.decode, body)
|
||||
if not ok or type(resp) ~= "table" then return nil end
|
||||
|
||||
-- Anthropic envelope: {type:"error", error:{type, message}}
|
||||
if resp.type == "error" and type(resp.error) == "table"
|
||||
and type(resp.error.message) == "string" then
|
||||
return resp.error.message
|
||||
end
|
||||
|
||||
-- Flat envelope: {error: {message: "..."}}
|
||||
if type(resp.error) == "table" and type(resp.error.message) == "string" then
|
||||
return resp.error.message
|
||||
end
|
||||
|
||||
return nil
|
||||
end
|
||||
|
||||
return adapter
|
||||
|
||||
@ -49,6 +49,11 @@ function adapter.transform_response(raw_body)
|
||||
unified.token_usage.prompt = resp.usage.prompt_tokens or 0
|
||||
unified.token_usage.completion = resp.usage.completion_tokens or 0
|
||||
unified.token_usage.total = resp.usage.total_tokens or 0
|
||||
unified.token_usage.prompt_cache_hit_tokens = resp.usage.prompt_cache_hit_tokens or 0
|
||||
unified.token_usage.prompt_cache_miss_tokens = resp.usage.prompt_cache_miss_tokens or 0
|
||||
if resp.usage.prompt_tokens_details and resp.usage.prompt_tokens_details.cached_tokens ~= nil then
|
||||
unified.token_usage.prompt_tokens_details = { cached_tokens = resp.usage.prompt_tokens_details.cached_tokens }
|
||||
end
|
||||
end
|
||||
|
||||
if type(resp.choices) == "table" and #resp.choices > 0 then
|
||||
@ -92,7 +97,12 @@ function adapter.transform_stream_chunk(raw_chunk)
|
||||
prompt = chunk.usage.prompt_tokens or chunk.usage.prompt or 0,
|
||||
completion = chunk.usage.completion_tokens or chunk.usage.completion or 0,
|
||||
total = chunk.usage.total_tokens or chunk.usage.total or 0,
|
||||
prompt_cache_hit_tokens = chunk.usage.prompt_cache_hit_tokens or 0,
|
||||
prompt_cache_miss_tokens = chunk.usage.prompt_cache_miss_tokens or 0,
|
||||
}
|
||||
if chunk.usage.prompt_tokens_details and chunk.usage.prompt_tokens_details.cached_tokens ~= nil then
|
||||
uses.prompt_tokens_details = { cached_tokens = chunk.usage.prompt_tokens_details.cached_tokens }
|
||||
end
|
||||
end
|
||||
|
||||
if not chunk.choices or #chunk.choices == 0 then
|
||||
@ -106,10 +116,15 @@ function adapter.transform_stream_chunk(raw_chunk)
|
||||
local delta = chunk.choices[1].delta or {}
|
||||
local fr = chunk.choices[1].finish_reason
|
||||
|
||||
local finish = (type(fr) == "string" and fr ~= "") and fr or nil
|
||||
|
||||
local unified = {
|
||||
content = delta.content or "",
|
||||
done = (fr ~= nil)
|
||||
done = (finish ~= nil)
|
||||
}
|
||||
if finish then
|
||||
unified.finish_reason = finish
|
||||
end
|
||||
if uses ~= nil then
|
||||
unified.usage = uses
|
||||
end
|
||||
@ -122,4 +137,15 @@ function adapter.transform_stream_chunk(raw_chunk)
|
||||
return json.encode(unified)
|
||||
end
|
||||
|
||||
-- 错误收敛:DeepSeek 走标准 OpenAI 信封 {error:{message,...}}
|
||||
function adapter.transform_error(status, body)
|
||||
local ok, resp = pcall(json.decode, body)
|
||||
if not ok or type(resp) ~= "table" then return nil end
|
||||
local e = resp.error
|
||||
if type(e) == "table" and type(e.message) == "string" then
|
||||
return e.message
|
||||
end
|
||||
return nil
|
||||
end
|
||||
|
||||
return adapter
|
||||
|
||||
@ -68,6 +68,15 @@ 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
|
||||
@ -100,6 +109,13 @@ function adapter.transform_stream_chunk(raw_chunk)
|
||||
local t = chunk.usageMetadata.totalTokenCount or 0
|
||||
if p > 0 or c > 0 or t > 0 then
|
||||
uses = { prompt = p, completion = c, total = t }
|
||||
-- Emit details whenever the field is present, even at 0, so a
|
||||
-- reported cache miss stays distinguishable from "not reported".
|
||||
if chunk.usageMetadata.cachedContentTokenCount ~= nil then
|
||||
uses.prompt_tokens_details = {
|
||||
cached_tokens = chunk.usageMetadata.cachedContentTokenCount
|
||||
}
|
||||
end
|
||||
end
|
||||
end
|
||||
|
||||
@ -110,7 +126,19 @@ function adapter.transform_stream_chunk(raw_chunk)
|
||||
return ""
|
||||
end
|
||||
local cand = chunk.candidates[1]
|
||||
local unified = { content = "", done = (cand.finishReason ~= nil) }
|
||||
-- Gemini finishReason -> OpenAI finish_reason
|
||||
local finish = nil
|
||||
if cand.finishReason ~= nil then
|
||||
if cand.finishReason == "MAX_TOKENS" then
|
||||
finish = "length"
|
||||
elseif cand.finishReason == "SAFETY" or cand.finishReason == "RECITATION"
|
||||
or cand.finishReason == "BLOCKLIST" then
|
||||
finish = "content_filter"
|
||||
else
|
||||
finish = "stop"
|
||||
end
|
||||
end
|
||||
local unified = { content = "", done = (finish ~= nil), finish_reason = finish }
|
||||
local reasoning = ""
|
||||
local tools = {}
|
||||
if cand.content and cand.content.parts then
|
||||
@ -140,4 +168,20 @@ function adapter.transform_stream_chunk(raw_chunk)
|
||||
return json.encode(unified)
|
||||
end
|
||||
|
||||
-- 错误收敛:Gemini REST 信封 {error:{code, message, status}}
|
||||
function adapter.transform_error(status, body)
|
||||
local ok, resp = pcall(json.decode, body)
|
||||
if not ok or type(resp) ~= "table" then return nil end
|
||||
local e = resp.error
|
||||
if type(e) == "table" then
|
||||
if type(e.message) == "string" then
|
||||
if type(e.status) == "string" then
|
||||
return e.status .. ": " .. e.message
|
||||
end
|
||||
return e.message
|
||||
end
|
||||
end
|
||||
return nil
|
||||
end
|
||||
|
||||
return adapter
|
||||
|
||||
@ -33,6 +33,18 @@ function adapter.transform_response(raw_body)
|
||||
unified.token_usage.prompt = resp.usage.prompt_tokens or 0
|
||||
unified.token_usage.completion = resp.usage.completion_tokens or 0
|
||||
unified.token_usage.total = resp.usage.total_tokens or 0
|
||||
local hit = 0
|
||||
if type(resp.usage.prompt_tokens_details) == "table" and resp.usage.prompt_tokens_details.cached_tokens ~= nil then
|
||||
hit = resp.usage.prompt_tokens_details.cached_tokens
|
||||
unified.token_usage.prompt_tokens_details = { cached_tokens = hit }
|
||||
end
|
||||
if (resp.usage.prompt_cache_hit_tokens or 0) > 0 then
|
||||
unified.token_usage.prompt_cache_hit_tokens = resp.usage.prompt_cache_hit_tokens
|
||||
unified.token_usage.prompt_cache_miss_tokens = resp.usage.prompt_cache_miss_tokens or 0
|
||||
if hit == 0 then
|
||||
unified.token_usage.prompt_tokens_details = { cached_tokens = resp.usage.prompt_cache_hit_tokens }
|
||||
end
|
||||
end
|
||||
end
|
||||
|
||||
if type(resp.choices) == "table" and #resp.choices > 0 then
|
||||
@ -74,6 +86,13 @@ function adapter.transform_stream_chunk(raw_chunk)
|
||||
completion = chunk.usage.completion_tokens or chunk.usage.completion or 0,
|
||||
total = chunk.usage.total_tokens or chunk.usage.total or 0,
|
||||
}
|
||||
if type(chunk.usage.prompt_tokens_details) == "table" and chunk.usage.prompt_tokens_details.cached_tokens ~= nil then
|
||||
uses.prompt_tokens_details = { cached_tokens = chunk.usage.prompt_tokens_details.cached_tokens }
|
||||
elseif (chunk.usage.prompt_cache_hit_tokens or 0) > 0 then
|
||||
uses.prompt_cache_hit_tokens = chunk.usage.prompt_cache_hit_tokens
|
||||
uses.prompt_cache_miss_tokens = chunk.usage.prompt_cache_miss_tokens or 0
|
||||
uses.prompt_tokens_details = { cached_tokens = chunk.usage.prompt_cache_hit_tokens }
|
||||
end
|
||||
end
|
||||
|
||||
if not chunk.choices or #chunk.choices == 0 then
|
||||
@ -87,10 +106,15 @@ function adapter.transform_stream_chunk(raw_chunk)
|
||||
local delta = chunk.choices[1].delta or {}
|
||||
local fr = chunk.choices[1].finish_reason
|
||||
|
||||
local finish = (type(fr) == "string" and fr ~= "") and fr or nil
|
||||
|
||||
local unified = {
|
||||
content = delta.content or "",
|
||||
done = (fr ~= nil)
|
||||
done = (finish ~= nil)
|
||||
}
|
||||
if finish then
|
||||
unified.finish_reason = finish
|
||||
end
|
||||
if uses ~= nil then
|
||||
unified.usage = uses
|
||||
end
|
||||
@ -103,4 +127,15 @@ function adapter.transform_stream_chunk(raw_chunk)
|
||||
return json.encode(unified)
|
||||
end
|
||||
|
||||
-- 错误收敛:GitHub Models 走标准 OpenAI 信封 {error:{message,...}}
|
||||
function adapter.transform_error(status, body)
|
||||
local ok, resp = pcall(json.decode, body)
|
||||
if not ok or type(resp) ~= "table" then return nil end
|
||||
local e = resp.error
|
||||
if type(e) == "table" and type(e.message) == "string" then
|
||||
return e.message
|
||||
end
|
||||
return nil
|
||||
end
|
||||
|
||||
return adapter
|
||||
|
||||
@ -32,6 +32,18 @@ function adapter.transform_response(raw_body)
|
||||
unified.token_usage.prompt = resp.usage.prompt_tokens or 0
|
||||
unified.token_usage.completion = resp.usage.completion_tokens or 0
|
||||
unified.token_usage.total = resp.usage.total_tokens or 0
|
||||
local hit = 0
|
||||
if type(resp.usage.prompt_tokens_details) == "table" and resp.usage.prompt_tokens_details.cached_tokens ~= nil then
|
||||
hit = resp.usage.prompt_tokens_details.cached_tokens
|
||||
unified.token_usage.prompt_tokens_details = { cached_tokens = hit }
|
||||
end
|
||||
if (resp.usage.prompt_cache_hit_tokens or 0) > 0 then
|
||||
unified.token_usage.prompt_cache_hit_tokens = resp.usage.prompt_cache_hit_tokens
|
||||
unified.token_usage.prompt_cache_miss_tokens = resp.usage.prompt_cache_miss_tokens or 0
|
||||
if hit == 0 then
|
||||
unified.token_usage.prompt_tokens_details = { cached_tokens = resp.usage.prompt_cache_hit_tokens }
|
||||
end
|
||||
end
|
||||
end
|
||||
|
||||
if type(resp.choices) == "table" and #resp.choices > 0 then
|
||||
@ -73,6 +85,13 @@ function adapter.transform_stream_chunk(raw_chunk)
|
||||
completion = chunk.usage.completion_tokens or chunk.usage.completion or 0,
|
||||
total = chunk.usage.total_tokens or chunk.usage.total or 0,
|
||||
}
|
||||
if type(chunk.usage.prompt_tokens_details) == "table" and chunk.usage.prompt_tokens_details.cached_tokens ~= nil then
|
||||
uses.prompt_tokens_details = { cached_tokens = chunk.usage.prompt_tokens_details.cached_tokens }
|
||||
elseif (chunk.usage.prompt_cache_hit_tokens or 0) > 0 then
|
||||
uses.prompt_cache_hit_tokens = chunk.usage.prompt_cache_hit_tokens
|
||||
uses.prompt_cache_miss_tokens = chunk.usage.prompt_cache_miss_tokens or 0
|
||||
uses.prompt_tokens_details = { cached_tokens = chunk.usage.prompt_cache_hit_tokens }
|
||||
end
|
||||
end
|
||||
|
||||
if not chunk.choices or #chunk.choices == 0 then
|
||||
@ -86,10 +105,15 @@ function adapter.transform_stream_chunk(raw_chunk)
|
||||
local delta = chunk.choices[1].delta or {}
|
||||
local fr = chunk.choices[1].finish_reason
|
||||
|
||||
local finish = (type(fr) == "string" and fr ~= "") and fr or nil
|
||||
|
||||
local unified = {
|
||||
content = delta.content or "",
|
||||
done = (fr ~= nil)
|
||||
done = (finish ~= nil)
|
||||
}
|
||||
if finish then
|
||||
unified.finish_reason = finish
|
||||
end
|
||||
if uses ~= nil then
|
||||
unified.usage = uses
|
||||
end
|
||||
@ -102,4 +126,15 @@ function adapter.transform_stream_chunk(raw_chunk)
|
||||
return json.encode(unified)
|
||||
end
|
||||
|
||||
-- 错误收敛:Groq 走标准 OpenAI 信封 {error:{message,...}}
|
||||
function adapter.transform_error(status, body)
|
||||
local ok, resp = pcall(json.decode, body)
|
||||
if not ok or type(resp) ~= "table" then return nil end
|
||||
local e = resp.error
|
||||
if type(e) == "table" and type(e.message) == "string" then
|
||||
return e.message
|
||||
end
|
||||
return nil
|
||||
end
|
||||
|
||||
return adapter
|
||||
|
||||
@ -64,6 +64,18 @@ function adapter.transform_response(raw_body)
|
||||
unified.token_usage.prompt = resp.usage.prompt_tokens or 0
|
||||
unified.token_usage.completion = resp.usage.completion_tokens or 0
|
||||
unified.token_usage.total = resp.usage.total_tokens or 0
|
||||
local hit = 0
|
||||
if type(resp.usage.prompt_tokens_details) == "table" and resp.usage.prompt_tokens_details.cached_tokens ~= nil then
|
||||
hit = resp.usage.prompt_tokens_details.cached_tokens
|
||||
unified.token_usage.prompt_tokens_details = { cached_tokens = hit }
|
||||
end
|
||||
if (resp.usage.prompt_cache_hit_tokens or 0) > 0 then
|
||||
unified.token_usage.prompt_cache_hit_tokens = resp.usage.prompt_cache_hit_tokens
|
||||
unified.token_usage.prompt_cache_miss_tokens = resp.usage.prompt_cache_miss_tokens or 0
|
||||
if hit == 0 then
|
||||
unified.token_usage.prompt_tokens_details = { cached_tokens = resp.usage.prompt_cache_hit_tokens }
|
||||
end
|
||||
end
|
||||
end
|
||||
if type(resp.choices) == "table" and #resp.choices > 0 then
|
||||
local ch = resp.choices[1]
|
||||
@ -106,6 +118,13 @@ function adapter.transform_stream_chunk(raw_chunk)
|
||||
completion = chunk.usage.completion_tokens or chunk.usage.completion or 0,
|
||||
total = chunk.usage.total_tokens or chunk.usage.total or 0,
|
||||
}
|
||||
if type(chunk.usage.prompt_tokens_details) == "table" and chunk.usage.prompt_tokens_details.cached_tokens ~= nil then
|
||||
uses.prompt_tokens_details = { cached_tokens = chunk.usage.prompt_tokens_details.cached_tokens }
|
||||
elseif (chunk.usage.prompt_cache_hit_tokens or 0) > 0 then
|
||||
uses.prompt_cache_hit_tokens = chunk.usage.prompt_cache_hit_tokens
|
||||
uses.prompt_cache_miss_tokens = chunk.usage.prompt_cache_miss_tokens or 0
|
||||
uses.prompt_tokens_details = { cached_tokens = chunk.usage.prompt_cache_hit_tokens }
|
||||
end
|
||||
end
|
||||
|
||||
if not chunk.choices or #chunk.choices == 0 then
|
||||
@ -119,10 +138,15 @@ function adapter.transform_stream_chunk(raw_chunk)
|
||||
local delta = chunk.choices[1].delta or {}
|
||||
local fr = chunk.choices[1].finish_reason
|
||||
|
||||
local finish = (type(fr) == "string" and fr ~= "") and fr or nil
|
||||
|
||||
local unified = {
|
||||
content = delta.content or "",
|
||||
done = (fr ~= nil)
|
||||
done = (finish ~= nil)
|
||||
}
|
||||
if finish then
|
||||
unified.finish_reason = finish
|
||||
end
|
||||
if uses ~= nil then
|
||||
unified.usage = uses
|
||||
end
|
||||
@ -135,4 +159,15 @@ function adapter.transform_stream_chunk(raw_chunk)
|
||||
return json.encode(unified)
|
||||
end
|
||||
|
||||
return adapter
|
||||
-- 错误收敛:kimi 网关为 OpenAI 风格 {error:{message,...}}
|
||||
function adapter.transform_error(status, body)
|
||||
local ok, resp = pcall(json.decode, body)
|
||||
if not ok or type(resp) ~= "table" then return nil end
|
||||
local e = resp.error
|
||||
if type(e) == "table" and type(e.message) == "string" then
|
||||
return e.message
|
||||
end
|
||||
return nil
|
||||
end
|
||||
|
||||
return adapter
|
||||
|
||||
@ -32,6 +32,18 @@ function adapter.transform_response(raw_body)
|
||||
unified.token_usage.prompt = resp.usage.prompt_tokens or 0
|
||||
unified.token_usage.completion = resp.usage.completion_tokens or 0
|
||||
unified.token_usage.total = resp.usage.total_tokens or 0
|
||||
local hit = 0
|
||||
if type(resp.usage.prompt_tokens_details) == "table" and resp.usage.prompt_tokens_details.cached_tokens ~= nil then
|
||||
hit = resp.usage.prompt_tokens_details.cached_tokens
|
||||
unified.token_usage.prompt_tokens_details = { cached_tokens = hit }
|
||||
end
|
||||
if (resp.usage.prompt_cache_hit_tokens or 0) > 0 then
|
||||
unified.token_usage.prompt_cache_hit_tokens = resp.usage.prompt_cache_hit_tokens
|
||||
unified.token_usage.prompt_cache_miss_tokens = resp.usage.prompt_cache_miss_tokens or 0
|
||||
if hit == 0 then
|
||||
unified.token_usage.prompt_tokens_details = { cached_tokens = resp.usage.prompt_cache_hit_tokens }
|
||||
end
|
||||
end
|
||||
end
|
||||
|
||||
if type(resp.choices) == "table" and #resp.choices > 0 then
|
||||
@ -73,6 +85,13 @@ function adapter.transform_stream_chunk(raw_chunk)
|
||||
completion = chunk.usage.completion_tokens or chunk.usage.completion or 0,
|
||||
total = chunk.usage.total_tokens or chunk.usage.total or 0,
|
||||
}
|
||||
if type(chunk.usage.prompt_tokens_details) == "table" and chunk.usage.prompt_tokens_details.cached_tokens ~= nil then
|
||||
uses.prompt_tokens_details = { cached_tokens = chunk.usage.prompt_tokens_details.cached_tokens }
|
||||
elseif (chunk.usage.prompt_cache_hit_tokens or 0) > 0 then
|
||||
uses.prompt_cache_hit_tokens = chunk.usage.prompt_cache_hit_tokens
|
||||
uses.prompt_cache_miss_tokens = chunk.usage.prompt_cache_miss_tokens or 0
|
||||
uses.prompt_tokens_details = { cached_tokens = chunk.usage.prompt_cache_hit_tokens }
|
||||
end
|
||||
end
|
||||
|
||||
if not chunk.choices or #chunk.choices == 0 then
|
||||
@ -86,10 +105,15 @@ function adapter.transform_stream_chunk(raw_chunk)
|
||||
local delta = chunk.choices[1].delta or {}
|
||||
local fr = chunk.choices[1].finish_reason
|
||||
|
||||
local finish = (type(fr) == "string" and fr ~= "") and fr or nil
|
||||
|
||||
local unified = {
|
||||
content = delta.content or "",
|
||||
done = (fr ~= nil)
|
||||
done = (finish ~= nil)
|
||||
}
|
||||
if finish then
|
||||
unified.finish_reason = finish
|
||||
end
|
||||
if uses ~= nil then
|
||||
unified.usage = uses
|
||||
end
|
||||
@ -102,4 +126,14 @@ function adapter.transform_stream_chunk(raw_chunk)
|
||||
return json.encode(unified)
|
||||
end
|
||||
|
||||
-- 错误收敛:Mistral 用顶层 {message="...", type="..."}(非嵌套 error)
|
||||
function adapter.transform_error(status, body)
|
||||
local ok, resp = pcall(json.decode, body)
|
||||
if not ok or type(resp) ~= "table" then return nil end
|
||||
if type(resp.message) == "string" then
|
||||
return resp.message
|
||||
end
|
||||
return nil
|
||||
end
|
||||
|
||||
return adapter
|
||||
|
||||
@ -85,16 +85,27 @@ function adapter.transform_stream_chunk(raw_chunk)
|
||||
end
|
||||
end
|
||||
|
||||
-- Ollama done_reason -> OpenAI finish_reason ("length" 透传,其余归一 stop)
|
||||
local finish = nil
|
||||
if chunk.done then
|
||||
if chunk.done_reason == "length" then
|
||||
finish = "length"
|
||||
else
|
||||
finish = "stop"
|
||||
end
|
||||
end
|
||||
|
||||
if not chunk.message then
|
||||
if uses ~= nil then
|
||||
return json.encode({ content = "", done = true, usage = uses })
|
||||
return json.encode({ content = "", done = true, finish_reason = finish, usage = uses })
|
||||
end
|
||||
return ""
|
||||
return json.encode({ content = "", done = true, finish_reason = finish })
|
||||
end
|
||||
|
||||
local unified = {
|
||||
content = chunk.message.content or "",
|
||||
done = chunk.done or false
|
||||
done = chunk.done or false,
|
||||
finish_reason = finish
|
||||
}
|
||||
if uses ~= nil then
|
||||
unified.usage = uses
|
||||
@ -120,4 +131,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
|
||||
|
||||
@ -33,6 +33,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
|
||||
@ -77,6 +93,13 @@ function adapter.transform_stream_chunk(raw_chunk)
|
||||
completion = chunk.usage.completion_tokens or chunk.usage.completion or 0,
|
||||
total = chunk.usage.total_tokens or chunk.usage.total or 0,
|
||||
}
|
||||
if type(chunk.usage.prompt_tokens_details) == "table" and chunk.usage.prompt_tokens_details.cached_tokens ~= nil then
|
||||
uses.prompt_tokens_details = { cached_tokens = chunk.usage.prompt_tokens_details.cached_tokens }
|
||||
elseif (chunk.usage.prompt_cache_hit_tokens or 0) > 0 then
|
||||
uses.prompt_cache_hit_tokens = chunk.usage.prompt_cache_hit_tokens
|
||||
uses.prompt_cache_miss_tokens = chunk.usage.prompt_cache_miss_tokens or 0
|
||||
uses.prompt_tokens_details = { cached_tokens = chunk.usage.prompt_cache_hit_tokens }
|
||||
end
|
||||
end
|
||||
|
||||
if not chunk.choices or #chunk.choices == 0 then
|
||||
@ -90,10 +113,15 @@ function adapter.transform_stream_chunk(raw_chunk)
|
||||
local delta = chunk.choices[1].delta or {}
|
||||
local fr = chunk.choices[1].finish_reason
|
||||
|
||||
local finish = (type(fr) == "string" and fr ~= "") and fr or nil
|
||||
|
||||
local unified = {
|
||||
content = delta.content or "",
|
||||
done = (fr ~= nil)
|
||||
done = (finish ~= nil)
|
||||
}
|
||||
if finish then
|
||||
unified.finish_reason = finish
|
||||
end
|
||||
if uses ~= nil then
|
||||
unified.usage = uses
|
||||
end
|
||||
@ -106,4 +134,46 @@ function adapter.transform_stream_chunk(raw_chunk)
|
||||
return json.encode(unified)
|
||||
end
|
||||
|
||||
-- 错误收敛:标准 OpenAI 信封 {error:{message,...}}
|
||||
function adapter.transform_error(status, body)
|
||||
-- Non-JSON body (nginx HTML error pages, plain text): extract a short
|
||||
-- human-readable reason instead of letting the raw body reach the log.
|
||||
local ok, resp = pcall(json.decode, body)
|
||||
if not ok or type(resp) ~= "table" then
|
||||
-- HTML error page: pull the <title> text (e.g. "413 Request Entity Too Large")
|
||||
local title = string.match(body or "", "<title>(.-)</title>")
|
||||
if title and title ~= "" then return title end
|
||||
-- Bare text: first non-empty line, capped
|
||||
local line = string.match(body or "", "^%s*([^\r\n]+)")
|
||||
if line and line ~= "" and not string.match(line, "^<") then
|
||||
return string.sub(line, 1, 200)
|
||||
end
|
||||
return nil
|
||||
end
|
||||
|
||||
-- Standard OpenAI envelope: {error:{message,...}} or {error:"..."}
|
||||
local e = resp.error
|
||||
if type(e) == "table" and type(e.message) == "string" then
|
||||
return e.message
|
||||
end
|
||||
if type(e) == "string" then return e end
|
||||
|
||||
-- Flat envelope used by many OpenAI-compatible gateways:
|
||||
-- {"code":20012,"message":"Model does not exist..."}
|
||||
-- {"code":"INVALID_API_KEY","message":"Invalid API key"}
|
||||
if type(resp.message) == "string" and resp.message ~= "" then
|
||||
if resp.code ~= nil then
|
||||
return tostring(resp.code) .. ": " .. resp.message
|
||||
end
|
||||
return resp.message
|
||||
end
|
||||
|
||||
-- Some gateways use {detail:"..."} (FastAPI style)
|
||||
if type(resp.detail) == "string" and resp.detail ~= "" then
|
||||
return resp.detail
|
||||
end
|
||||
|
||||
return nil
|
||||
end
|
||||
|
||||
return adapter
|
||||
|
||||
@ -4,14 +4,34 @@ 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 的
|
||||
@ -32,6 +52,8 @@ function adapter.transform_request(raw_body)
|
||||
if not ok then return raw_body end
|
||||
req.disable_thinking = nil
|
||||
req.extra_body = nil
|
||||
if type(req.stream_options) ~= "table" then req.stream_options = {} end
|
||||
req.stream_options.include_usage = true
|
||||
if req.messages then
|
||||
local kept = {}
|
||||
for _, msg in ipairs(req.messages) do
|
||||
@ -78,14 +100,27 @@ function adapter.transform_response(raw_body)
|
||||
unified.token_usage.prompt = resp.usage.prompt_tokens or 0
|
||||
unified.token_usage.completion = resp.usage.completion_tokens or 0
|
||||
unified.token_usage.total = resp.usage.total_tokens or 0
|
||||
local hit = 0
|
||||
if type(resp.usage.prompt_tokens_details) == "table" and resp.usage.prompt_tokens_details.cached_tokens ~= nil then
|
||||
hit = resp.usage.prompt_tokens_details.cached_tokens
|
||||
unified.token_usage.prompt_tokens_details = { cached_tokens = hit }
|
||||
end
|
||||
if (resp.usage.prompt_cache_hit_tokens or 0) > 0 then
|
||||
unified.token_usage.prompt_cache_hit_tokens = resp.usage.prompt_cache_hit_tokens
|
||||
unified.token_usage.prompt_cache_miss_tokens = resp.usage.prompt_cache_miss_tokens or 0
|
||||
if hit == 0 then
|
||||
unified.token_usage.prompt_tokens_details = { cached_tokens = resp.usage.prompt_cache_hit_tokens }
|
||||
end
|
||||
end
|
||||
end
|
||||
|
||||
if type(resp.choices) == "table" and #resp.choices > 0 then
|
||||
local ch = resp.choices[1]
|
||||
if type(ch.message) == "table" then
|
||||
unified.content = ch.message.content or ""
|
||||
if ch.message.reasoning_content then
|
||||
unified.reasoning_content = ch.message.reasoning_content
|
||||
local reasoning = ch.message.reasoning_content or ch.message.reasoning
|
||||
if reasoning then
|
||||
unified.reasoning_content = reasoning
|
||||
end
|
||||
if type(ch.message.tool_calls) == "table" then
|
||||
local tcs = {}
|
||||
@ -123,6 +158,13 @@ function adapter.transform_stream_chunk(raw_chunk)
|
||||
completion = chunk.usage.completion_tokens or chunk.usage.completion or 0,
|
||||
total = chunk.usage.total_tokens or chunk.usage.total or 0,
|
||||
}
|
||||
if type(chunk.usage.prompt_tokens_details) == "table" and chunk.usage.prompt_tokens_details.cached_tokens ~= nil then
|
||||
uses.prompt_tokens_details = { cached_tokens = chunk.usage.prompt_tokens_details.cached_tokens }
|
||||
elseif (chunk.usage.prompt_cache_hit_tokens or 0) > 0 then
|
||||
uses.prompt_cache_hit_tokens = chunk.usage.prompt_cache_hit_tokens
|
||||
uses.prompt_cache_miss_tokens = chunk.usage.prompt_cache_miss_tokens or 0
|
||||
uses.prompt_tokens_details = { cached_tokens = chunk.usage.prompt_cache_hit_tokens }
|
||||
end
|
||||
end
|
||||
|
||||
if not chunk.choices or #chunk.choices == 0 then
|
||||
@ -136,15 +178,22 @@ function adapter.transform_stream_chunk(raw_chunk)
|
||||
local delta = chunk.choices[1].delta or {}
|
||||
local fr = chunk.choices[1].finish_reason
|
||||
|
||||
local finish = (type(fr) == "string" and fr ~= "") and fr or nil
|
||||
|
||||
local unified = {
|
||||
content = delta.content or "",
|
||||
done = (fr ~= nil)
|
||||
done = (finish ~= nil)
|
||||
}
|
||||
if finish then
|
||||
unified.finish_reason = finish
|
||||
end
|
||||
if uses ~= nil then
|
||||
unified.usage = uses
|
||||
end
|
||||
if delta.reasoning_content then
|
||||
unified.reasoning_content = delta.reasoning_content
|
||||
-- zen 用 reasoning 字段承载推理文本(OpenAI 惯例是 reasoning_content)
|
||||
local reasoning = delta.reasoning_content or delta.reasoning
|
||||
if reasoning then
|
||||
unified.reasoning_content = reasoning
|
||||
end
|
||||
if delta.tool_calls then
|
||||
-- pass raw streaming fragments through; OpenAI clients accumulate index+id+name+arguments
|
||||
@ -153,4 +202,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
|
||||
|
||||
142
internal/lua/adapters/sensenova.lua
Normal file
142
internal/lua/adapters/sensenova.lua
Normal file
@ -0,0 +1,142 @@
|
||||
-- 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
|
||||
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 but no content
|
||||
local content = delta.content or ""
|
||||
if content == "" and delta.reasoning_content then
|
||||
content = delta.reasoning_content
|
||||
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
|
||||
@ -407,6 +407,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 {
|
||||
|
||||
@ -115,12 +115,31 @@ func TestOpenCodeAdapterFingerprint(t *testing.T) {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer vm.Stop()
|
||||
hdrs, err := vm.BuildHeaders("opencode", map[string]interface{}{"api_key": "public"})
|
||||
hdrs, err := vm.BuildHeaders("opencode", map[string]interface{}{
|
||||
"api_key": "public",
|
||||
"timestamp": 1700000000,
|
||||
"body": `{"model":"x"}`,
|
||||
"url": "https://opencode.ai/zen/v1/chat/completions",
|
||||
"method": "POST",
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("build headers: %v", err)
|
||||
}
|
||||
if ua := hdrs["User-Agent"]; ua != "opencode/0.1.0" {
|
||||
t.Fatalf("opencode adapter must send the opencode client UA, got %q", ua)
|
||||
// zen 按客户端指纹路由请求池:UA 必须是 opencode 真实格式,
|
||||
// 且必须携带 x-opencode-* 身份头,否则带 tools 的请求会被上游
|
||||
// 判为匿名客户端并返回 network_error 空流(表现为空回复)。
|
||||
if ua := hdrs["User-Agent"]; !strings.HasPrefix(ua, "opencode/") ||
|
||||
!strings.Contains(ua, "ai-sdk/provider-utils/") {
|
||||
t.Fatalf("opencode adapter must send the real opencode client UA, got %q", ua)
|
||||
}
|
||||
for _, h := range []string{"x-opencode-client", "x-opencode-project", "x-opencode-session", "x-opencode-request"} {
|
||||
if hdrs[h] == "" {
|
||||
t.Fatalf("opencode adapter must send %q (zen client fingerprint), got %v", h, hdrs)
|
||||
}
|
||||
}
|
||||
hdrs2, _ := vm.BuildHeaders("opencode", map[string]interface{}{"api_key": "public", "timestamp": 1700000001, "body": `{"model":"y"}`, "method": "POST"})
|
||||
if hdrs["x-opencode-request"] == hdrs2["x-opencode-request"] {
|
||||
t.Fatalf("x-opencode-request must vary per request, got %q twice", hdrs["x-opencode-request"])
|
||||
}
|
||||
if hdrs["Authorization"] != "" {
|
||||
t.Fatalf("opencode adapter must not hardcode Authorization (config api_key supplies it), got %q", hdrs["Authorization"])
|
||||
@ -254,7 +273,8 @@ func TestDisableThinkingPassthrough(t *testing.T) {
|
||||
t.Fatalf("thinking.type = %v", thinking["type"])
|
||||
}
|
||||
|
||||
// anthropic: disable_thinking removes the thinking block
|
||||
// anthropic: thinking is opt-in (Issue 3) — never emitted by default, and
|
||||
// disable_thinking is not a trigger either.
|
||||
out2, err := vm.Transform("anthropic", "transform_request", body)
|
||||
if err != nil {
|
||||
t.Fatalf("anthropic transform: %v", err)
|
||||
@ -266,8 +286,17 @@ func TestDisableThinkingPassthrough(t *testing.T) {
|
||||
if err != nil {
|
||||
t.Fatalf("anthropic transform: %v", err)
|
||||
}
|
||||
if !strings.Contains(out3, "enabled") {
|
||||
t.Fatalf("anthropic should enable thinking by default: %s", out3)
|
||||
if strings.Contains(out3, "thinking") {
|
||||
t.Fatalf("anthropic must not enable thinking by default: %s", out3)
|
||||
}
|
||||
// opt-in path: extra_body.thinking is forwarded verbatim
|
||||
out4, err := vm.Transform("anthropic", "transform_request",
|
||||
`{"model":"x","messages":[{"role":"user","content":"hi"}],"extra_body":{"thinking":{"type":"enabled","budget_tokens":2048}}}`)
|
||||
if err != nil {
|
||||
t.Fatalf("anthropic transform: %v", err)
|
||||
}
|
||||
if !strings.Contains(out4, `"type":"enabled"`) || !strings.Contains(out4, "2048") {
|
||||
t.Fatalf("anthropic should forward extra_body.thinking: %s", out4)
|
||||
}
|
||||
}
|
||||
|
||||
@ -403,3 +432,317 @@ func TestOpenCodeAdapterKeepsWhitelistedRolesAndDropsMultimodal(t *testing.T) {
|
||||
t.Fatalf("unexpected roles: %s", out)
|
||||
}
|
||||
}
|
||||
|
||||
func TestAdaptersPassFinishReason(t *testing.T) {
|
||||
vm := NewVM(freshAdapterDir(t))
|
||||
if err := vm.Start(); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer vm.Stop()
|
||||
// OpenAI-style chunk with a real finish reason must surface finish_reason
|
||||
// and done=true; an empty-string finish_reason (sensenova sends "" on
|
||||
// every chunk) must NOT terminate the stream.
|
||||
chunk := `{"choices":[{"index":0,"finish_reason":"tool_calls","delta":{"content":""}}]}`
|
||||
empty := `{"choices":[{"index":0,"finish_reason":"","delta":{"content":"x"}}]}`
|
||||
for _, name := range []string{"openai", "deepseek", "github", "groq", "kimicode", "mistral", "opencode"} {
|
||||
out, err := vm.Transform(name, "transform_stream_chunk", chunk)
|
||||
if err != nil {
|
||||
t.Fatalf("%s: %v", name, err)
|
||||
}
|
||||
if !strings.Contains(out, `"finish_reason":"tool_calls"`) {
|
||||
t.Fatalf("%s: finish_reason lost: %s", name, out)
|
||||
}
|
||||
if !strings.Contains(out, `"done":true`) {
|
||||
t.Fatalf("%s: done not set on real finish: %s", name, out)
|
||||
}
|
||||
out, err = vm.Transform(name, "transform_stream_chunk", empty)
|
||||
if err != nil {
|
||||
t.Fatalf("%s: %v", name, err)
|
||||
}
|
||||
if strings.Contains(out, `"done":true`) {
|
||||
t.Fatalf("%s: empty finish_reason must not end stream: %s", name, out)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// TestAnthropicToolRoundTrip verifies the OpenAI->Anthropic request mapping for
|
||||
// a full agent tool-call round (Issue 1): assistant tool_calls become
|
||||
// tool_use content blocks and role:"tool" results become user tool_result
|
||||
// blocks merged into ONE user message. This is the regression that made agent
|
||||
// clients (dsh/Claude Code/Cursor) repeatedly re-invoke the same tool.
|
||||
func TestAnthropicToolRoundTrip(t *testing.T) {
|
||||
vm := NewVM(freshAdapterDir(t))
|
||||
if err := vm.Start(); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer vm.Stop()
|
||||
|
||||
round1 := `{"model":"x","tools":[{"type":"function","function":{"name":"calc","description":"multiply","parameters":{"type":"object","properties":{"a":{"type":"integer"},"b":{"type":"integer"}},"required":["a","b"]}}}],"messages":[{"role":"user","content":"what is 17*23?"}]}`
|
||||
out1, err := vm.Transform("anthropic", "transform_request", round1)
|
||||
if err != nil {
|
||||
t.Fatalf("round1 transform: %v", err)
|
||||
}
|
||||
var r1 struct {
|
||||
Tools []struct {
|
||||
Name string `json:"name"`
|
||||
InputSchema map[string]interface{} `json:"input_schema"`
|
||||
} `json:"tools"`
|
||||
Messages []map[string]interface{} `json:"messages"`
|
||||
}
|
||||
if err := json.Unmarshal([]byte(out1), &r1); err != nil {
|
||||
t.Fatalf("unmarshal r1: %v (%s)", err, out1)
|
||||
}
|
||||
if len(r1.Tools) != 1 || r1.Tools[0].Name != "calc" {
|
||||
t.Fatalf("tools not mapped: %s", out1)
|
||||
}
|
||||
|
||||
// Round 2: assistant tool_calls + tool result
|
||||
round2 := `{"model":"x","messages":[
|
||||
{"role":"user","content":"what is 17*23?"},
|
||||
{"role":"assistant","content":"","tool_calls":[
|
||||
{"id":"call_1","type":"function","function":{"name":"calc","arguments":"{\"a\":17,\"b\":23}"}}
|
||||
]},
|
||||
{"role":"tool","tool_call_id":"call_1","content":"391"}
|
||||
]}`
|
||||
out2, err := vm.Transform("anthropic", "transform_request", round2)
|
||||
if err != nil {
|
||||
t.Fatalf("round2 transform: %v", err)
|
||||
}
|
||||
var r2 struct {
|
||||
Messages []struct {
|
||||
Role string `json:"role"`
|
||||
Content []struct {
|
||||
Type string `json:"type"`
|
||||
ID string `json:"id"`
|
||||
Name string `json:"name"`
|
||||
Input map[string]interface{} `json:"input"`
|
||||
ToolUseID string `json:"tool_use_id"`
|
||||
ContentText string `json:"content"`
|
||||
} `json:"content"`
|
||||
} `json:"messages"`
|
||||
}
|
||||
if err := json.Unmarshal([]byte(out2), &r2); err != nil {
|
||||
t.Fatalf("unmarshal r2: %v (%s)", err, out2)
|
||||
}
|
||||
// Message 2 (index 1) must be assistant with one tool_use block
|
||||
am := r2.Messages[1]
|
||||
if am.Role != "assistant" {
|
||||
t.Fatalf("msg[1].role = %q, want assistant", am.Role)
|
||||
}
|
||||
var toolUse struct {
|
||||
Type string `json:"type"`
|
||||
ID string `json:"id"`
|
||||
Name string `json:"name"`
|
||||
Input map[string]interface{} `json:"input"`
|
||||
}
|
||||
for _, b := range am.Content {
|
||||
if b.Type == "tool_use" {
|
||||
toolUse = struct {
|
||||
Type string `json:"type"`
|
||||
ID string `json:"id"`
|
||||
Name string `json:"name"`
|
||||
Input map[string]interface{} `json:"input"`
|
||||
}{b.Type, b.ID, b.Name, b.Input}
|
||||
}
|
||||
}
|
||||
if toolUse.ID != "call_1" || toolUse.Name != "calc" {
|
||||
t.Fatalf("tool_use not mapped: %+v", am.Content)
|
||||
}
|
||||
if toolUse.Input["a"] != float64(17) || toolUse.Input["b"] != float64(23) {
|
||||
t.Fatalf("tool_use input args not decoded from JSON string: %+v", toolUse.Input)
|
||||
}
|
||||
// Message 3 (index 2) must be user with one tool_result block
|
||||
um := r2.Messages[2]
|
||||
if um.Role != "user" {
|
||||
t.Fatalf("msg[2].role = %q, want user (tool_result)", um.Role)
|
||||
}
|
||||
if len(um.Content) != 1 || um.Content[0].Type != "tool_result" || um.Content[0].ToolUseID != "call_1" {
|
||||
t.Fatalf("tool_result not mapped: %+v", um.Content)
|
||||
}
|
||||
|
||||
// Round 3: consecutive tool results must merge into ONE user message
|
||||
round3 := `{"model":"x","messages":[
|
||||
{"role":"user","content":"do both"},
|
||||
{"role":"assistant","content":"","tool_calls":[
|
||||
{"id":"c1","type":"function","function":{"name":"calc","arguments":"{\"a\":1,\"b\":2}"}},
|
||||
{"id":"c2","type":"function","function":{"name":"calc","arguments":"{\"a\":3,\"b\":4}"}}
|
||||
]},
|
||||
{"role":"tool","tool_call_id":"c1","content":"2"},
|
||||
{"role":"tool","tool_call_id":"c2","content":"12"}
|
||||
]}`
|
||||
out3, err := vm.Transform("anthropic", "transform_request", round3)
|
||||
if err != nil {
|
||||
t.Fatalf("round3 transform: %v", err)
|
||||
}
|
||||
var r3 struct {
|
||||
Messages []map[string]interface{} `json:"messages"`
|
||||
}
|
||||
if err := json.Unmarshal([]byte(out3), &r3); err != nil {
|
||||
t.Fatalf("unmarshal r3: %v (%s)", err, out3)
|
||||
}
|
||||
// messages: user, assistant(tool_use x2), user(tool_result x2 merged)
|
||||
if len(r3.Messages) != 3 {
|
||||
t.Fatalf("round3 len(messages) = %d, want 3 (merged tool results): %s", len(r3.Messages), out3)
|
||||
}
|
||||
last := r3.Messages[3-1]
|
||||
blocks := last["content"].([]interface{})
|
||||
if len(blocks) != 2 {
|
||||
t.Fatalf("last user message content blocks = %d, want 2 merged tool_results: %s", len(blocks), out3)
|
||||
}
|
||||
}
|
||||
|
||||
// TestOpenAITransformErrorEnvelopes covers Issue 6: the openai adapter must
|
||||
// condense every shape of upstream error body that real OpenAI-compatible
|
||||
// gateways emit, not just the {error:{message}} envelope. Unhandled shapes
|
||||
// used to fall through to Go's generic "unknown error" and dump raw HTML /
|
||||
// JSON into the server log.
|
||||
func TestOpenAITransformErrorEnvelopes(t *testing.T) {
|
||||
vm := NewVM(freshAdapterDir(t))
|
||||
if err := vm.Start(); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer vm.Stop()
|
||||
|
||||
cases := []struct {
|
||||
name string
|
||||
status int
|
||||
body string
|
||||
want string
|
||||
}{
|
||||
{"standard openai envelope", 429,
|
||||
`{"error":{"message":"Rate limit reached","type":"rate_limit"}}`,
|
||||
"Rate limit reached"},
|
||||
{"error as bare string", 400,
|
||||
`{"error":"bad request"}`,
|
||||
"bad request"},
|
||||
{"flat numeric code (qijiar/siliconflow style)", 400,
|
||||
`{"code":20012,"message":"Model does not exist. Please check it carefully.","data":null}`,
|
||||
"20012: Model does not exist. Please check it carefully."},
|
||||
{"flat string code (remotezen style)", 401,
|
||||
`{"code":"INVALID_API_KEY","message":"Invalid API key"}`,
|
||||
"INVALID_API_KEY: Invalid API key"},
|
||||
{"fastapi detail", 422,
|
||||
`{"detail":"validation failed"}`,
|
||||
"validation failed"},
|
||||
{"nginx html error page", 413,
|
||||
`<html> <head><title>413 Request Entity Too Large</title></head> <body> <center><h1>413 Request Entity Too Large</h1></center> <hr><center>nginx/1.18.0 (Ubuntu)</center> </body> </html>`,
|
||||
"413 Request Entity Too Large"},
|
||||
{"plain text body", 502,
|
||||
"upstream connect error",
|
||||
"upstream connect error"},
|
||||
}
|
||||
for _, tc := range cases {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
got, ok, err := vm.TransformError("openai", tc.status, tc.body)
|
||||
if err != nil {
|
||||
t.Fatalf("TransformError: %v", err)
|
||||
}
|
||||
if !ok {
|
||||
t.Fatalf("hook returned no reason for %s body: %s", tc.name, tc.body)
|
||||
}
|
||||
if got != tc.want {
|
||||
t.Fatalf("reason = %q, want %q", got, tc.want)
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
// A body carrying no usable message must fall through (ok=false) so the
|
||||
// Go-side generic condenser stays in charge instead of inventing text.
|
||||
t.Run("no usable message falls through", func(t *testing.T) {
|
||||
if _, ok, _ := vm.TransformError("openai", 500, `{"foo":"bar"}`); ok {
|
||||
t.Fatal("expected fallthrough for a body with no message field")
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
// TestAdaptersReportZeroCacheHit covers the "distinguish missed from not
|
||||
// reported" contract on BOTH adapter paths: whenever an upstream reports a
|
||||
// cache field, the adapter must emit prompt_tokens_details even when the hit
|
||||
// count is 0, so the gateway can record cache_reported=true and the UI shows
|
||||
// 0% instead of "—". Adapters that dropped the 0 case made a reported miss
|
||||
// indistinguishable from an upstream that never reported cache info.
|
||||
func TestAdaptersReportZeroCacheHit(t *testing.T) {
|
||||
vm := NewVM(freshAdapterDir(t))
|
||||
if err := vm.Start(); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer vm.Stop()
|
||||
|
||||
// OpenAI-shaped upstreams: usage.prompt_tokens_details.cached_tokens = 0
|
||||
openaiLike := []string{"openai", "deepseek", "sensenova", "opencode",
|
||||
"agentrouter", "github", "groq", "kimicode", "mistral"}
|
||||
respBody := `{"choices":[{"message":{"content":"hi"},"finish_reason":"stop"}],
|
||||
"usage":{"prompt_tokens":10,"completion_tokens":2,"total_tokens":12,
|
||||
"prompt_tokens_details":{"cached_tokens":0}}}`
|
||||
streamBody := `{"choices":[],"usage":{"prompt_tokens":10,"completion_tokens":2,
|
||||
"total_tokens":12,"prompt_tokens_details":{"cached_tokens":0}}}`
|
||||
for _, name := range openaiLike {
|
||||
out, err := vm.Transform(name, "transform_response", respBody)
|
||||
if err != nil {
|
||||
t.Fatalf("%s transform_response: %v", name, err)
|
||||
}
|
||||
if !strings.Contains(out, "prompt_tokens_details") {
|
||||
t.Fatalf("%s: zero cached_tokens dropped in transform_response: %s", name, out)
|
||||
}
|
||||
out, err = vm.Transform(name, "transform_stream_chunk", streamBody)
|
||||
if err != nil {
|
||||
t.Fatalf("%s transform_stream_chunk: %v", name, err)
|
||||
}
|
||||
if !strings.Contains(out, "prompt_tokens_details") {
|
||||
t.Fatalf("%s: zero cached_tokens dropped in transform_stream_chunk: %s", name, out)
|
||||
}
|
||||
}
|
||||
|
||||
// gemini: usageMetadata.cachedContentTokenCount = 0
|
||||
gResp := `{"candidates":[{"content":{"parts":[{"text":"hi"}]},"finishReason":"STOP"}],
|
||||
"usageMetadata":{"promptTokenCount":10,"candidatesTokenCount":2,
|
||||
"totalTokenCount":12,"cachedContentTokenCount":0}}`
|
||||
out, err := vm.Transform("gemini", "transform_response", gResp)
|
||||
if err != nil {
|
||||
t.Fatalf("gemini transform_response: %v", err)
|
||||
}
|
||||
if !strings.Contains(out, "prompt_tokens_details") {
|
||||
t.Fatalf("gemini: cachedContentTokenCount not normalized in transform_response: %s", out)
|
||||
}
|
||||
gStream := `{"usageMetadata":{"promptTokenCount":10,"candidatesTokenCount":2,
|
||||
"totalTokenCount":12,"cachedContentTokenCount":0}}`
|
||||
out, err = vm.Transform("gemini", "transform_stream_chunk", gStream)
|
||||
if err != nil {
|
||||
t.Fatalf("gemini transform_stream_chunk: %v", err)
|
||||
}
|
||||
if !strings.Contains(out, "prompt_tokens_details") {
|
||||
t.Fatalf("gemini: zero cachedContentTokenCount dropped in stream: %s", out)
|
||||
}
|
||||
|
||||
// anthropic: usage.cache_read_input_tokens = 0
|
||||
aResp := `{"content":[{"type":"text","text":"hi"}],"stop_reason":"end_turn",
|
||||
"usage":{"input_tokens":10,"output_tokens":2,"cache_read_input_tokens":0}}`
|
||||
out, err = vm.Transform("anthropic", "transform_response", aResp)
|
||||
if err != nil {
|
||||
t.Fatalf("anthropic transform_response: %v", err)
|
||||
}
|
||||
if !strings.Contains(out, "prompt_tokens_details") {
|
||||
t.Fatalf("anthropic: zero cache_read_input_tokens dropped in response: %s", out)
|
||||
}
|
||||
aStream := `{"type":"message_start","message":{"usage":{"input_tokens":10,
|
||||
"output_tokens":2,"cache_read_input_tokens":0}}}`
|
||||
out, err = vm.Transform("anthropic", "transform_stream_chunk", aStream)
|
||||
if err != nil {
|
||||
t.Fatalf("anthropic transform_stream_chunk: %v", err)
|
||||
}
|
||||
if !strings.Contains(out, "prompt_tokens_details") {
|
||||
t.Fatalf("anthropic: zero cache_read_input_tokens dropped in stream: %s", out)
|
||||
}
|
||||
|
||||
// Negative control: an upstream that reports NO cache field at all must
|
||||
// not fabricate details (that would flip "not reported" into a fake 0%).
|
||||
noCache := `{"choices":[{"message":{"content":"hi"},"finish_reason":"stop"}],
|
||||
"usage":{"prompt_tokens":10,"completion_tokens":2,"total_tokens":12}}`
|
||||
out, err = vm.Transform("openai", "transform_response", noCache)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if strings.Contains(out, "prompt_tokens_details") {
|
||||
t.Fatalf("openai fabricated cache details when upstream reported none: %s", out)
|
||||
}
|
||||
}
|
||||
|
||||
@ -10,6 +10,7 @@ import (
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"io"
|
||||
"log"
|
||||
"net"
|
||||
"net/http"
|
||||
"strings"
|
||||
@ -47,6 +48,13 @@ const (
|
||||
backoffCapN = 10
|
||||
)
|
||||
|
||||
// rateLimitCooldown is the fixed cooldown applied on HTTP 429. A rate-limit
|
||||
// rejection means "too fast", not "broken": unlike transport/5xx failures it
|
||||
// must NOT escalate exponentially to backoffCap, or a quota-rich source gets
|
||||
// locked out for up to 30 minutes while its allowance is untouched. A short
|
||||
// fixed window lets the slot re-enter the rotation quickly.
|
||||
const rateLimitCooldown = 30 * time.Second
|
||||
|
||||
func clampPref(a *atomic.Int64, lo, hi int64) {
|
||||
for {
|
||||
cur := a.Load()
|
||||
@ -100,6 +108,17 @@ func (s *ModelState) RecordFailure(auth bool) {
|
||||
clampPref(&s.pref, prefMin, prefMax)
|
||||
}
|
||||
|
||||
// RecordRateLimit records an HTTP 429: a short fixed cooldown instead of the
|
||||
// exponential schedule, so a merely-throttled source recovers in seconds and
|
||||
// its remaining quota stays usable. The preference penalty still applies so
|
||||
// other slots are preferred while cooling.
|
||||
func (s *ModelState) RecordRateLimit() {
|
||||
s.failCount.Add(1)
|
||||
s.cooldownUntil.Store(time.Now().Add(rateLimitCooldown).Unix())
|
||||
s.pref.Add(-prefFailStep)
|
||||
clampPref(&s.pref, prefMin, prefMax)
|
||||
}
|
||||
|
||||
// RecordSuccess resets the failure counter and cooldown and bumps the
|
||||
// preference score by one.
|
||||
func (s *ModelState) RecordSuccess() {
|
||||
@ -139,6 +158,12 @@ type Provider struct {
|
||||
mu sync.Mutex
|
||||
sem chan struct{}
|
||||
states map[string]*ModelState // key = model id
|
||||
// proactive rate limiting: nextOK is the earliest unix-nano time the next
|
||||
// request may leave; throttle() spaces requests 60s/RPM apart. Zero when
|
||||
// cfg.RPM <= 0 (unlimited).
|
||||
rateMu sync.Mutex
|
||||
nextOK int64
|
||||
rpmGap time.Duration
|
||||
lastProbe struct {
|
||||
ok bool
|
||||
err string
|
||||
@ -178,6 +203,9 @@ func New(cfg config.Source, vm *lua.VM) *Provider {
|
||||
if cfg.MaxConcurrent <= 0 {
|
||||
p.sem = nil
|
||||
}
|
||||
if cfg.RPM > 0 {
|
||||
p.rpmGap = time.Minute / time.Duration(cfg.RPM)
|
||||
}
|
||||
for _, m := range cfg.Models {
|
||||
p.states[m.ID] = &ModelState{}
|
||||
}
|
||||
@ -337,7 +365,7 @@ func (p *Provider) ModelAvailable(model string) bool {
|
||||
p.mu.Lock()
|
||||
defer p.mu.Unlock()
|
||||
if s, ok := p.states[model]; ok {
|
||||
return s.Available()
|
||||
return s.Available() && s.Pref() > prefMin
|
||||
}
|
||||
return true
|
||||
}
|
||||
@ -407,7 +435,7 @@ func (p *Provider) probeModels(ctx context.Context) (bool, string) {
|
||||
case resp.StatusCode == 404 || resp.StatusCode == 405:
|
||||
return false, ""
|
||||
default:
|
||||
return false, fmt.Sprintf("api error %d: %s", resp.StatusCode, truncate(string(raw), 300))
|
||||
return false, p.apiErrReason(resp.StatusCode, string(raw))
|
||||
}
|
||||
}
|
||||
|
||||
@ -444,7 +472,7 @@ func (p *Provider) probeChat(ctx context.Context) (bool, string) {
|
||||
} else if status == 200 {
|
||||
ok = true
|
||||
} else {
|
||||
msg = fmt.Sprintf("api error %d: %s", status, truncate(raw, 500))
|
||||
msg = p.apiErrReason(status, raw)
|
||||
}
|
||||
} else {
|
||||
msg = err.Error()
|
||||
@ -491,15 +519,20 @@ func (p *Provider) RecordSuccess(model string) {
|
||||
|
||||
// ReportStatus records an upstream HTTP status for the given model and drives
|
||||
// the (source, model) backoff state. 401/403 → capped self-healing cooldown
|
||||
// with doubled penalty; 429 and 5xx → normal exponential backoff. Other codes
|
||||
// (400 client schema errors, 402 billing errors) are not penalized here —
|
||||
// they surface via the status page / audit instead.
|
||||
// with doubled penalty; 5xx → normal exponential backoff; 429 → short fixed
|
||||
// cooldown (rateLimitCooldown): a rate rejection means "too fast", not
|
||||
// "broken", so the slot must re-enter rotation quickly instead of escalating
|
||||
// to a 30-minute lockout that wastes remaining quota.
|
||||
func (p *Provider) ReportStatus(model string, code int) {
|
||||
if code == 401 || code == 403 {
|
||||
p.RecordFailure(model, code)
|
||||
return
|
||||
}
|
||||
if code >= 500 || code == 429 {
|
||||
if code == 429 {
|
||||
p.state(model).RecordRateLimit()
|
||||
return
|
||||
}
|
||||
if code >= 500 {
|
||||
p.RecordFailure(model, code)
|
||||
}
|
||||
}
|
||||
@ -586,6 +619,37 @@ func (p *Provider) Release() {
|
||||
<-p.sem
|
||||
}
|
||||
|
||||
// ---- proactive rate limiting ----
|
||||
|
||||
// Throttle blocks until this source's rate-limit window allows the next
|
||||
// request, spacing outbound requests at least rpmGap apart (60s/RPM). It is a
|
||||
// no-op when no rpm cap is configured. The wait happens AFTER the concurrency
|
||||
// slot is taken, so a queued request still counts against max_concurrent —
|
||||
// bounded by the caller's ctx (queue timeout / client disconnect).
|
||||
func (p *Provider) Throttle(ctx context.Context) error {
|
||||
if p.rpmGap <= 0 {
|
||||
return nil
|
||||
}
|
||||
p.rateMu.Lock()
|
||||
now := time.Now().UnixNano()
|
||||
wait := p.nextOK - now
|
||||
if wait > 0 {
|
||||
p.nextOK += int64(p.rpmGap) // reserve the next window for the follower
|
||||
p.rateMu.Unlock()
|
||||
t := time.NewTimer(time.Duration(wait))
|
||||
defer t.Stop()
|
||||
select {
|
||||
case <-t.C:
|
||||
return nil
|
||||
case <-ctx.Done():
|
||||
return ctx.Err()
|
||||
}
|
||||
}
|
||||
p.nextOK = now + int64(p.rpmGap)
|
||||
p.rateMu.Unlock()
|
||||
return nil
|
||||
}
|
||||
|
||||
// ---- request construction ----
|
||||
|
||||
func (p *Provider) buildHeaders(body, url string) (http.Header, error) {
|
||||
@ -630,6 +694,9 @@ func (p *Provider) Chat(ctx context.Context, req *types.ChatRequest) (*types.Uni
|
||||
return nil, err
|
||||
}
|
||||
defer p.Release()
|
||||
if err := p.Throttle(ctx); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
model := p.ModelFor(req.Model)
|
||||
|
||||
body, err := marshalTransform(p.vm, p.adapter, "transform_request", req)
|
||||
@ -651,7 +718,7 @@ func (p *Provider) Chat(ctx context.Context, req *types.ChatRequest) (*types.Uni
|
||||
}
|
||||
if status != 200 {
|
||||
p.ReportStatus(model, status)
|
||||
return nil, fmt.Errorf("api error %d: %s", status, truncate(raw, 500))
|
||||
return nil, fmt.Errorf("%s", p.apiErrReason(status, raw))
|
||||
}
|
||||
unified, err := p.vm.Transform(p.adapter, "transform_response", raw)
|
||||
if err != nil {
|
||||
@ -684,6 +751,10 @@ func (p *Provider) ChatStream(ctx context.Context, req *types.ChatRequest) (<-ch
|
||||
if err := p.TryAcquire(ctx); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if err := p.Throttle(ctx); err != nil {
|
||||
p.Release()
|
||||
return nil, err
|
||||
}
|
||||
model := p.ModelFor(req.Model)
|
||||
req.Stream = true
|
||||
body, err := marshalTransform(p.vm, p.adapter, "transform_request", req)
|
||||
@ -707,7 +778,7 @@ func (p *Provider) ChatStream(ctx context.Context, req *types.ChatRequest) (<-ch
|
||||
rc <- respOrErr{resp, err}
|
||||
}()
|
||||
|
||||
ch := make(chan types.UnifiedChunk, 64)
|
||||
inner := make(chan types.UnifiedChunk, 64)
|
||||
sel := <-rc
|
||||
if sel.err != nil {
|
||||
// client disconnect/cancel before the first byte: not a scheduling
|
||||
@ -723,16 +794,24 @@ func (p *Provider) ChatStream(ctx context.Context, req *types.ChatRequest) (<-ch
|
||||
sel.resp.Body.Close()
|
||||
p.ReportStatus(model, sel.resp.StatusCode)
|
||||
p.Release()
|
||||
return nil, fmt.Errorf("api error %d: %s", sel.resp.StatusCode, truncate(string(raw), 500))
|
||||
return nil, fmt.Errorf("%s", p.apiErrReason(sel.resp.StatusCode, string(raw)))
|
||||
}
|
||||
go func() {
|
||||
defer p.Release()
|
||||
defer close(ch)
|
||||
defer close(inner)
|
||||
defer sel.resp.Body.Close()
|
||||
scanner := bufio.NewScanner(sel.resp.Body)
|
||||
scanner.Buffer(make([]byte, 0, 64*1024), 1024*1024)
|
||||
var chunks int
|
||||
var realChunks int
|
||||
var doneSeen bool
|
||||
var doneSent bool
|
||||
// toolIdx maps a per-stream tool fragment's upstream index (Anthropic
|
||||
// content-block ordinal, which is sparse when thinking/text blocks
|
||||
// precede tool_use) to a dense 0-based OpenAI tool_call ordinal so
|
||||
// clients that accumulate fragments by index reassemble multi-tool
|
||||
// responses correctly.
|
||||
toolIdx := newToolIndexRemapper()
|
||||
for scanner.Scan() {
|
||||
line := strings.TrimSpace(scanner.Text())
|
||||
if line == "" || !strings.HasPrefix(line, "data:") {
|
||||
@ -744,9 +823,15 @@ func (p *Provider) ChatStream(ctx context.Context, req *types.ChatRequest) (<-ch
|
||||
}
|
||||
if data == "[DONE]" {
|
||||
doneSeen = true
|
||||
select {
|
||||
case ch <- types.UnifiedChunk{Done: true}:
|
||||
case <-ctx.Done():
|
||||
// adapters that already emitted their terminating done chunk
|
||||
// (with the real finish reason) must not get a second,
|
||||
// reason-less done from [DONE] — it would override the true
|
||||
// finish_reason downstream.
|
||||
if !doneSent {
|
||||
select {
|
||||
case inner <- types.UnifiedChunk{Done: true}:
|
||||
case <-ctx.Done():
|
||||
}
|
||||
}
|
||||
continue
|
||||
}
|
||||
@ -764,9 +849,16 @@ func (p *Provider) ChatStream(ctx context.Context, req *types.ChatRequest) (<-ch
|
||||
if err := json.Unmarshal([]byte(unified), &ck); err != nil {
|
||||
continue
|
||||
}
|
||||
ck.ToolCalls = toolIdx.remap(ck.ToolCalls)
|
||||
if ck.Done {
|
||||
doneSent = true
|
||||
}
|
||||
if !errorOnlyChunk(ck) {
|
||||
realChunks++
|
||||
}
|
||||
chunks++
|
||||
select {
|
||||
case ch <- ck:
|
||||
case inner <- ck:
|
||||
case <-ctx.Done():
|
||||
return
|
||||
}
|
||||
@ -777,15 +869,142 @@ func (p *Provider) ChatStream(ctx context.Context, req *types.ChatRequest) (<-ch
|
||||
// success nor failure for scheduling purposes. A 200 that produced
|
||||
// zero chunks and no [DONE] is an empty stream, i.e. a failure
|
||||
// before the first chunk — record it so the slot can fall back.
|
||||
// A stream whose only payload was error-only done chunks (e.g.
|
||||
// finish_reason:"network_error") is likewise a failure, not a success.
|
||||
if ctx.Err() == nil && scanner.Err() == nil {
|
||||
if doneSeen || chunks > 0 {
|
||||
if realChunks > 0 || (doneSeen && chunks == 0) {
|
||||
p.RecordSuccess(model)
|
||||
} else {
|
||||
p.RecordFailure(model, 0)
|
||||
}
|
||||
}
|
||||
}()
|
||||
return ch, nil
|
||||
|
||||
// Hold back the first chunk to validate the stream actually carries
|
||||
// content: some upstreams answer HTTP 200 with a degenerate stream whose
|
||||
// only payload is an error finish reason (zen free pool sends
|
||||
// finish_reason:"network_error" with empty delta). Failing the candidate
|
||||
// here — before any byte reaches the gateway — lets the scheduler fall
|
||||
// through to the next source instead of serving the client an empty reply.
|
||||
type heldChunk struct {
|
||||
ck types.UnifiedChunk
|
||||
ok bool
|
||||
}
|
||||
var first heldChunk
|
||||
select {
|
||||
case ck, ok := <-inner:
|
||||
first = heldChunk{ck, ok}
|
||||
case <-ctx.Done():
|
||||
return nil, ctx.Err()
|
||||
}
|
||||
if !first.ok {
|
||||
return nil, fmt.Errorf("provider %s: empty stream", p.Name())
|
||||
}
|
||||
if errorOnlyChunk(first.ck) {
|
||||
go func() {
|
||||
for range inner {
|
||||
}
|
||||
}()
|
||||
return nil, fmt.Errorf("provider %s: upstream returned %q stream",
|
||||
p.Name(), first.ck.FinishReason)
|
||||
}
|
||||
out := make(chan types.UnifiedChunk, 64)
|
||||
go func() {
|
||||
defer close(out)
|
||||
out <- first.ck
|
||||
for ck := range inner {
|
||||
out <- ck
|
||||
}
|
||||
}()
|
||||
return out, nil
|
||||
}
|
||||
|
||||
// errorOnlyChunk reports whether ck carries nothing but an upstream error
|
||||
// signal: a done chunk with a non-standard finish reason and zero content,
|
||||
// tool calls, reasoning text or usage. Standard OpenAI finish reasons are
|
||||
// never classified as errors, so legitimate instant-empty completions
|
||||
// (finish_reason:"stop", no output) still reach the client.
|
||||
|
||||
// toolIndexRemapper maps per-stream tool_call fragment indices from a sparse
|
||||
// upstream ordinal (e.g. Anthropic content_block index, which skips over
|
||||
// thinking/text blocks) to a dense 0-based OpenAI tool_call ordinal.
|
||||
type toolIndexRemapper struct {
|
||||
seen map[int]int // upstream index -> dense ordinal
|
||||
next int // next dense ordinal to assign
|
||||
}
|
||||
|
||||
func newToolIndexRemapper() toolIndexRemapper {
|
||||
return toolIndexRemapper{seen: make(map[int]int)}
|
||||
}
|
||||
|
||||
func (t *toolIndexRemapper) remap(raw json.RawMessage) json.RawMessage {
|
||||
if len(raw) == 0 {
|
||||
return raw
|
||||
}
|
||||
// Decode into generic maps so every upstream field is preserved verbatim;
|
||||
// only the index value is rewritten.
|
||||
var frags []map[string]json.RawMessage
|
||||
if err := json.Unmarshal(raw, &frags); err != nil || len(frags) == 0 {
|
||||
return raw
|
||||
}
|
||||
changed := false
|
||||
for i := range frags {
|
||||
orig := 0
|
||||
if v, ok := frags[i]["index"]; ok {
|
||||
if err := json.Unmarshal(v, &orig); err != nil {
|
||||
continue
|
||||
}
|
||||
}
|
||||
mapped, ok := t.seen[orig]
|
||||
if !ok {
|
||||
mapped = t.next
|
||||
t.next++
|
||||
t.seen[orig] = mapped
|
||||
}
|
||||
if mapped != orig {
|
||||
b, err := json.Marshal(mapped)
|
||||
if err != nil {
|
||||
continue
|
||||
}
|
||||
frags[i]["index"] = b
|
||||
changed = true
|
||||
}
|
||||
}
|
||||
if !changed {
|
||||
return raw
|
||||
}
|
||||
out, err := json.Marshal(frags)
|
||||
if err != nil {
|
||||
return raw
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
func errorOnlyChunk(ck types.UnifiedChunk) bool {
|
||||
if !ck.Done || ck.FinishReason == "" {
|
||||
return false
|
||||
}
|
||||
switch ck.FinishReason {
|
||||
case "stop", "length", "tool_calls", "function_call", "content_filter":
|
||||
return false
|
||||
}
|
||||
return ck.Content == "" && len(ck.ToolCalls) == 0 &&
|
||||
ck.ReasoningContent == "" && ck.Usage == nil
|
||||
}
|
||||
|
||||
// apiErrReason builds the client-facing reason for a non-200 upstream
|
||||
// response: the adapter's optional transform_error hook wins (per-source
|
||||
// protocol knowledge lives in Lua), otherwise clients get a uniform
|
||||
// "unknown error" while the raw body stays in the server log for debugging.
|
||||
func (p *Provider) apiErrReason(status int, raw string) string {
|
||||
if reason, ok, err := p.vm.TransformError(p.adapter, status, raw); err == nil && ok {
|
||||
if trimmed := strings.TrimSpace(reason); trimmed != "" {
|
||||
return fmt.Sprintf("api error %d: %s", status, types.OneLine(trimmed, 200))
|
||||
}
|
||||
}
|
||||
log.Printf("[provider] unhandled upstream error body (adapter %q lacks transform_error): status=%d body=%.300s",
|
||||
p.adapter, status, types.OneLine(raw, 300))
|
||||
return fmt.Sprintf("api error %d: unknown error", status)
|
||||
}
|
||||
|
||||
// Image generates images via /v1/images/generations. Same scheduling-state
|
||||
@ -825,13 +1044,14 @@ func (p *Provider) Image(ctx context.Context, req *types.ImageGenRequest) (*type
|
||||
}
|
||||
if status != 200 {
|
||||
p.ReportStatus(model, status)
|
||||
return nil, fmt.Errorf("image api error %d: %s", status, truncate(raw, 500))
|
||||
return nil, fmt.Errorf("%s", p.apiErrReason(status, raw))
|
||||
}
|
||||
var out types.UnifiedResponse
|
||||
// try adapter transform_response; if missing, parse standard openai image format
|
||||
unified, terr := p.vm.Transform(p.adapter+"_image", "transform_response", raw)
|
||||
if terr == nil && unified != raw {
|
||||
if err := json.Unmarshal([]byte(unified), &out); err == nil {
|
||||
out.Model = model // actual model served (AUTO resolved to bestImageModel)
|
||||
p.RecordSuccess(model)
|
||||
return &out, nil
|
||||
}
|
||||
@ -844,6 +1064,7 @@ func (p *Provider) Image(ctx context.Context, req *types.ImageGenRequest) (*type
|
||||
return nil, fmt.Errorf("unmarshal image response: %w", err)
|
||||
}
|
||||
out.ImageData = img.Data
|
||||
out.Model = model // actual model served (AUTO resolved to bestImageModel)
|
||||
p.RecordSuccess(model)
|
||||
return &out, nil
|
||||
}
|
||||
@ -902,12 +1123,17 @@ func standardSSEChunk(data string) string {
|
||||
FinishReason *string `json:"finish_reason"`
|
||||
} `json:"choices"`
|
||||
UpstreamUsage struct {
|
||||
PromptTokens int `json:"prompt_tokens"`
|
||||
CompletionTokens int `json:"completion_tokens"`
|
||||
TotalTokens int `json:"total_tokens"`
|
||||
Prompt int `json:"prompt"`
|
||||
Completion int `json:"completion"`
|
||||
Total int `json:"total"`
|
||||
PromptTokens int `json:"prompt_tokens"`
|
||||
CompletionTokens int `json:"completion_tokens"`
|
||||
TotalTokens int `json:"total_tokens"`
|
||||
Prompt int `json:"prompt"`
|
||||
Completion int `json:"completion"`
|
||||
Total int `json:"total"`
|
||||
PromptCacheHit int `json:"prompt_cache_hit_tokens"`
|
||||
PromptCacheMiss int `json:"prompt_cache_miss_tokens"`
|
||||
PromptTokensDetails *struct {
|
||||
CachedTokens int `json:"cached_tokens"`
|
||||
} `json:"prompt_tokens_details"`
|
||||
} `json:"usage"`
|
||||
}
|
||||
if err := json.Unmarshal([]byte(data), &raw); err != nil {
|
||||
@ -920,9 +1146,25 @@ func standardSSEChunk(data string) string {
|
||||
pu := raw.UpstreamUsage
|
||||
if pu.Total > 0 || pu.TotalTokens > 0 {
|
||||
usage = &types.TokenUsage{
|
||||
Prompt: pickFirst(pu.PromptTokens, pu.Prompt),
|
||||
Completion: pickFirst(pu.CompletionTokens, pu.Completion),
|
||||
Total: pickFirst(pu.TotalTokens, pu.Total),
|
||||
Prompt: pickFirst(pu.PromptTokens, pu.Prompt),
|
||||
Completion: pickFirst(pu.CompletionTokens, pu.Completion),
|
||||
Total: pickFirst(pu.TotalTokens, pu.Total),
|
||||
PromptCacheHit: pu.PromptCacheHit,
|
||||
PromptCacheMiss: pu.PromptCacheMiss,
|
||||
}
|
||||
if pu.PromptTokensDetails != nil && pu.PromptTokensDetails.CachedTokens > 0 {
|
||||
usage.PromptTokensDetails = &types.PromptTokensDetails{
|
||||
CachedTokens: pu.PromptTokensDetails.CachedTokens,
|
||||
}
|
||||
}
|
||||
}
|
||||
finish := ""
|
||||
done := false
|
||||
if len(raw.Choices) > 0 && raw.Choices[0].FinishReason != nil {
|
||||
// empty-string finish reasons (sensenova sends "" on every chunk)
|
||||
// are not a finish signal
|
||||
if fr := *raw.Choices[0].FinishReason; fr != "" {
|
||||
finish, done = fr, true
|
||||
}
|
||||
}
|
||||
out, _ := json.Marshal(types.UnifiedChunk{
|
||||
@ -932,8 +1174,9 @@ func standardSSEChunk(data string) string {
|
||||
}
|
||||
return ""
|
||||
}(),
|
||||
Done: len(raw.Choices) > 0 && raw.Choices[0].FinishReason != nil,
|
||||
Usage: usage,
|
||||
Done: done,
|
||||
FinishReason: finish,
|
||||
Usage: usage,
|
||||
})
|
||||
return string(out)
|
||||
}
|
||||
@ -946,10 +1189,3 @@ func pickFirst(a, b int) int {
|
||||
}
|
||||
return b
|
||||
}
|
||||
|
||||
func truncate(s string, n int) string {
|
||||
if len(s) <= n {
|
||||
return s
|
||||
}
|
||||
return s[:n] + "..."
|
||||
}
|
||||
|
||||
@ -7,7 +7,9 @@ import (
|
||||
"fmt"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
"sync/atomic"
|
||||
"testing"
|
||||
"time"
|
||||
@ -245,6 +247,77 @@ func TestModelStateCooldownAndRecovery(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
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(prefFailStep) && st.Pref() > int64(prefMin) {
|
||||
t.Fatalf("pref = %d", st.Pref())
|
||||
}
|
||||
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))
|
||||
}
|
||||
}
|
||||
|
||||
func TestTryAcquire(t *testing.T) {
|
||||
p := newTestProvider(t, src("mock", "http://127.0.0.1:1", "openai", "m"))
|
||||
p.cfg.MaxConcurrent = 1
|
||||
@ -345,7 +418,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 +431,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 +444,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 +528,153 @@ 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)
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
@ -217,6 +217,12 @@ type SourceStatus struct {
|
||||
// probe was rate-limited by the upstream).
|
||||
RecentOK int64 `json:"recent_ok,omitempty"`
|
||||
RecentErr int64 `json:"recent_err,omitempty"`
|
||||
// AvgFirstByteMs is the mean time-to-first-byte (ms) for recently
|
||||
// successful requests served by this source (0 when unmeasured).
|
||||
AvgFirstByteMs int64 `json:"avg_first_byte_ms,omitempty"`
|
||||
// AvgTokPerS is the aggregate completion throughput (tokens/s) for
|
||||
// recently successful requests served by this source (0 when no data).
|
||||
AvgTokPerS int64 `json:"avg_tok_per_s,omitempty"`
|
||||
}
|
||||
|
||||
// ProbeAll runs a live reachability check for every provider (in parallel).
|
||||
|
||||
@ -248,10 +248,9 @@ func (s *Scheduler) chainDrive(ctx context.Context, chain *Chain, req *types.Cha
|
||||
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 ModelAvailable
|
||||
// (which checks Pref > prefMin).
|
||||
base := tn.NextStart()
|
||||
res := runTier(ctx, tn, cands, base, req, stream)
|
||||
if res.resp != nil || res.chunks != nil {
|
||||
@ -327,6 +326,61 @@ 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 {
|
||||
var cands []*Slot
|
||||
for _, sl := range tn.Slots {
|
||||
if !sl.Prov.ModelAvailable(sl.Model) {
|
||||
continue
|
||||
}
|
||||
cands = append(cands, sl)
|
||||
}
|
||||
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()
|
||||
var hard []TierError
|
||||
for i := 0; i < len(cands); i++ {
|
||||
sl := cands[(int(base)+i)%len(cands)]
|
||||
if !sl.Prov.ModelAvailable(sl.Model) {
|
||||
continue
|
||||
}
|
||||
r := *req
|
||||
r.Model = sl.Model
|
||||
resp, err := sl.Prov.Image(ctx, &r)
|
||||
if ctx.Err() != nil {
|
||||
return nil, "", "", ctx.Err()
|
||||
}
|
||||
if err == nil {
|
||||
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})
|
||||
}
|
||||
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
|
||||
|
||||
@ -177,14 +177,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)
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@ -5,6 +5,7 @@ package types
|
||||
import (
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"strings"
|
||||
"time"
|
||||
)
|
||||
|
||||
@ -70,6 +71,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"`
|
||||
@ -83,6 +85,24 @@ 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"`
|
||||
}
|
||||
|
||||
// PromptTokensDetails is the OpenAI v2 prompt_tokens_details object. Only
|
||||
// CachedTokens is emitted (omitempty drops the whole object when zero).
|
||||
type PromptTokensDetails struct {
|
||||
// CachedTokens is always emitted (even 0) so clients can distinguish
|
||||
// "upstream reports cache, this request missed" from "no cache data".
|
||||
CachedTokens int `json:"cached_tokens"`
|
||||
}
|
||||
|
||||
// MarshalJSON emits both the legacy short keys (prompt/completion/total, used
|
||||
@ -90,20 +110,34 @@ type TokenUsage struct {
|
||||
// standard keys (prompt_tokens/completion_tokens/total_tokens). Standard
|
||||
// clients such as DSH and DevEco Code read the *_tokens fields.
|
||||
func (t TokenUsage) MarshalJSON() ([]byte, error) {
|
||||
// Auto-generate prompt_tokens_details from legacy DeepSeek fields when
|
||||
// the upstream adapter only set the standalone hit count (openai.lua
|
||||
// does this in Lua, but other adapters or the standardSSEChunk fallback
|
||||
// may not). dsh reads prompt_tokens_details.cached_tokens first.
|
||||
pdetails := t.PromptTokensDetails
|
||||
if pdetails == nil && t.PromptCacheHit > 0 {
|
||||
pdetails = &PromptTokensDetails{CachedTokens: t.PromptCacheHit}
|
||||
}
|
||||
return json.Marshal(struct {
|
||||
PromptTokens int `json:"prompt_tokens"`
|
||||
CompletionTokens int `json:"completion_tokens"`
|
||||
TotalTokens int `json:"total_tokens"`
|
||||
Prompt int `json:"prompt"`
|
||||
Completion int `json:"completion"`
|
||||
Total int `json:"total"`
|
||||
PromptTokens int `json:"prompt_tokens"`
|
||||
CompletionTokens int `json:"completion_tokens"`
|
||||
TotalTokens int `json:"total_tokens"`
|
||||
Prompt int `json:"prompt"`
|
||||
Completion int `json:"completion"`
|
||||
Total int `json:"total"`
|
||||
PromptTokensDetails *PromptTokensDetails `json:"prompt_tokens_details,omitempty"`
|
||||
PromptCacheHitTokens int `json:"prompt_cache_hit_tokens,omitempty"`
|
||||
PromptCacheMissTokens int `json:"prompt_cache_miss_tokens,omitempty"`
|
||||
}{
|
||||
PromptTokens: t.Prompt,
|
||||
CompletionTokens: t.Completion,
|
||||
TotalTokens: t.Total,
|
||||
Prompt: t.Prompt,
|
||||
Completion: t.Completion,
|
||||
Total: t.Total,
|
||||
PromptTokens: t.Prompt,
|
||||
CompletionTokens: t.Completion,
|
||||
TotalTokens: t.Total,
|
||||
Prompt: t.Prompt,
|
||||
Completion: t.Completion,
|
||||
Total: t.Total,
|
||||
PromptTokensDetails: pdetails,
|
||||
PromptCacheHitTokens: t.PromptCacheHit,
|
||||
PromptCacheMissTokens: t.PromptCacheMiss,
|
||||
})
|
||||
}
|
||||
|
||||
@ -136,6 +170,10 @@ type ImageGenResponse struct {
|
||||
type UnifiedChunk struct {
|
||||
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
|
||||
@ -155,4 +193,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
|
||||
}
|
||||
70
internal/types/types_test.go
Normal file
70
internal/types/types_test.go
Normal file
@ -0,0 +1,70 @@
|
||||
package types
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"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))
|
||||
}
|
||||
Reference in New Issue
Block a user