mirror of
https://gitcode.com/JianFeeeee/ModelRouter.git
synced 2026-11-07 09:47:23 +00:00
Compare commits
43 Commits
| Author | SHA1 | Date | |
|---|---|---|---|
| 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 | |||
| 3069cfce4e | |||
| cb9c025952 | |||
| 83e6d88813 | |||
| 6f1c806591 | |||
| 98c08d51c4 | |||
| ca81048637 | |||
| b9f2486037 | |||
| 8c5279b416 | |||
| 23bd9e5552 | |||
| d75d8411c0 | |||
| f67831c3c2 | |||
| e356885130 |
4
.gitignore
vendored
4
.gitignore
vendored
@ -17,3 +17,7 @@ cmd/gui/*.log
|
||||
.pi-glla/
|
||||
.codegraph/
|
||||
/build/
|
||||
|
||||
# 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"
|
||||
}
|
||||
}
|
||||
}
|
||||
18
Makefile
18
Makefile
@ -9,6 +9,12 @@ export GUI_DIR
|
||||
GUI_DIR=cmd/gui
|
||||
VERSION ?= $(shell git describe --tags --dirty 2>/dev/null || echo "0.1.0")
|
||||
|
||||
# Clean env for GUI packaging: the host may have LD_LIBRARY_PATH polluted by
|
||||
# third-party runtimes (e.g. /opt/cangjie) which electron-builder bakes into
|
||||
# the produced binary's DT_RPATH/dep resolution -> the packaged app then
|
||||
# SIGSEGVs inside ld.so on clean machines. Strip it for every pack step.
|
||||
CLEAN_PACK_ENV = env -u LD_LIBRARY_PATH
|
||||
|
||||
all: build
|
||||
|
||||
# ---- core binary (server users / embedded core) ----
|
||||
@ -26,23 +32,25 @@ gui-dev:
|
||||
|
||||
# build the embedded core into the GUI, then run electron-builder
|
||||
# targets: deb + AppImage (linux), rpm needs system rpmbuild
|
||||
# NOTE: run in a clean env (no LD_LIBRARY_PATH pollution) so the packaged
|
||||
# binary does not link against host-only libs and crash on user machines.
|
||||
gui-dist:
|
||||
@cd $(GUI_DIR) && npm run dist:linux
|
||||
@cd $(GUI_DIR) && $(CLEAN_PACK_ENV) npm run dist:linux
|
||||
|
||||
# deb only (for local sharing), rpm needs rpmbuild
|
||||
gui-deb:
|
||||
@cd $(GUI_DIR) && npm run dist:debian
|
||||
@cd $(GUI_DIR) && $(CLEAN_PACK_ENV) npm run dist:debian
|
||||
|
||||
# windows nsis installer: host-built core (needs apt install gcc-mingw-w64-x86-64 + wine)
|
||||
gui-win:
|
||||
@cd $(GUI_DIR) && npm run dist:win
|
||||
@cd $(GUI_DIR) && $(CLEAN_PACK_ENV) npm run dist:win
|
||||
|
||||
# windows nsis installer: full cross-build in docker (self-sufficient, no host mingw)
|
||||
gui-win-docker:
|
||||
@cd $(GUI_DIR) && ./scripts/dist-win-docker.sh
|
||||
@cd $(GUI_DIR) && $(CLEAN_PACK_ENV) ./scripts/dist-win-docker.sh
|
||||
|
||||
gui-dist-dir:
|
||||
@cd $(GUI_DIR) && npm run dist:dir
|
||||
@cd $(GUI_DIR) && $(CLEAN_PACK_ENV) npm run dist:dir
|
||||
|
||||
test:
|
||||
go test ./...
|
||||
|
||||
12
README.md
12
README.md
@ -89,8 +89,7 @@ GOMODCACHE=... GOPROXY=off go build -tags luajit -o llmsproxy ./cmd/llmsproxy
|
||||
|
||||
> **配置不入库**:仓库不携带任何 config 文件(配置文件含密钥)。二进制首次运行时
|
||||
> 会在 `-config` 指定的位置生成默认配置:随机 admin key(打印在启动日志中)、
|
||||
> 仅绑定 `127.0.0.1:8080`、内置一个免 key 的 opencode zen 源可直接对话。
|
||||
> 首次登录后请在 WebUI「密钥」页更换管理员密钥。
|
||||
> 仅绑定 `127.0.0.1:8080`。首次登录后请在 WebUI「密钥」页更换管理员密钥。
|
||||
|
||||
```bash
|
||||
# 无 key -> 401
|
||||
@ -281,19 +280,12 @@ return {
|
||||
|
||||
### 内置适配器
|
||||
|
||||
`openai` `deepseek` `anthropic` `gemini` `github` `groq` `mistral` `ollama` `kimicode`
|
||||
`opencode`。
|
||||
`openai` `deepseek` `anthropic` `gemini` `github` `groq` `mistral` `ollama` `kimicode`。
|
||||
|
||||
`anthropic`/`gemini`/`ollama` 适配器内置多模态转换(`image_url` → 各自上游格式);若
|
||||
源启用 `disable_thinking`,`deepseek` 适配器会把 `extra_body.thinking.type` 置为
|
||||
`disabled`。
|
||||
|
||||
**opencode** 适配器面向 opencode.ai zen 免费池(`https://opencode.ai/zen/v1`):
|
||||
zen 按 User-Agent 指纹识别官方客户端并分流——非官方 UA 的请求(curl、Go 默认 UA)
|
||||
会被丢进匿名池触发 `FreeUsageLimitError`。该适配器固定发送 opencode 客户端 UA,
|
||||
配合 `api_key: "public"`(官方无 key 客户端实际发送 `Bearer public`)即可走免费池,
|
||||
零成本接入 `deepseek-v4-flash-free` 等免费模型。
|
||||
|
||||
**kimicode** 是展示 `build_headers` 的样例:云端校验调用方 app,需要按
|
||||
`meta.app_secret` 对时间戳+URL+请求体哈希做 HMAC 签名并附 `X-App-Sign` 等头。
|
||||
配好 `sources[].meta.{app_id, app_secret, app_agent}` 即可。
|
||||
|
||||
11
README_EN.md
11
README_EN.md
@ -77,7 +77,7 @@ GOMODCACHE=... GOPROXY=off go build -tags luajit -o llmsproxy ./cmd/llmsproxy
|
||||
> **No config is shipped in the repo** (config files carry real keys). On
|
||||
> first run the binary generates a default config at the `-config` path: a
|
||||
> random admin key (printed to the startup log), loopback-only `127.0.0.1:8080`,
|
||||
> and a keyless opencode zen source ready to chat. Rotate the admin key in the
|
||||
> and a keyless source ready to chat. Rotate the admin key in the
|
||||
> WebUI after first login.
|
||||
|
||||
```bash
|
||||
@ -289,19 +289,12 @@ Shared helpers: `hmac_sha256_hex(key, data)`, `sha256_hex(data)`,
|
||||
### Built-in adapters
|
||||
|
||||
`openai` `deepseek` `anthropic` `gemini` `github` `groq` `mistral` `ollama`
|
||||
`kimicode` `opencode`.
|
||||
`kimicode`.
|
||||
|
||||
`anthropic`/`gemini`/`ollama` include multimodal conversion
|
||||
(`image_url` → their native format); with `disable_thinking` the `deepseek`
|
||||
adapter sets `extra_body.thinking.type` to `disabled`.
|
||||
|
||||
**opencode** targets the opencode.ai zen free pool
|
||||
(`https://opencode.ai/zen/v1`): zen fingerprints clients by User-Agent and
|
||||
routes non-official UAs (curl, Go's default) into an anonymous pool that
|
||||
hits `FreeUsageLimitError`. The adapter always sends the opencode client UA;
|
||||
combined with `api_key: "public"` (the keyless official client actually sends
|
||||
`Bearer public`) it gets the free pool, e.g. `deepseek-v4-flash-free`.
|
||||
|
||||
**kimicode** demonstrates `build_headers`: the cloud validates the calling
|
||||
app, so you HMAC-sign timestamp+URL+body with `meta.app_secret` and add
|
||||
`X-App-Sign`-style headers. Configure `sources[].meta.{app_id, app_secret,
|
||||
|
||||
@ -10,6 +10,7 @@ const {
|
||||
Tray,
|
||||
shell,
|
||||
nativeImage,
|
||||
nativeTheme,
|
||||
} = require("electron");
|
||||
const path = require("path");
|
||||
const fs = require("fs");
|
||||
@ -46,6 +47,15 @@ const CORE_EXE = path.join(
|
||||
"bin",
|
||||
process.platform === "win32" ? "llmsproxy.exe" : "llmsproxy",
|
||||
);
|
||||
// Packaged window icon: must be a REAL file on disk (extraResource), not an
|
||||
// asar path — Electron/GTK pulls BrowserWindow icons from the filesystem and
|
||||
// crashes (SIGSEGV) when given an asar virtual path on Linux.
|
||||
// Dev: cmd/gui/build/icon.png.
|
||||
const WINDOW_ICON = path.join(
|
||||
app.isPackaged ? process.resourcesPath : __dirname,
|
||||
"build",
|
||||
"icon.png",
|
||||
);
|
||||
// default port for the embedded core (configurable via settings)
|
||||
const DEFAULT_PORT = 8787;
|
||||
|
||||
@ -53,6 +63,7 @@ let mainWindow = null;
|
||||
let tray = null;
|
||||
let coreProc = null;
|
||||
let coreStartedAt = 0;
|
||||
let coreReady = false; // true only after the embedded core actually answers HTTP
|
||||
let settings = {}; // { port, autoStart, minimizeToTray }
|
||||
let authRules = [];
|
||||
|
||||
@ -120,9 +131,9 @@ function coreStarted() {
|
||||
|
||||
function writeDefaultEmbeddedConfig() {
|
||||
// Generate a self-contained embedded profile with a fresh random admin key.
|
||||
// (Like Clash Verge generates its profile; the core's EnsureDefault default
|
||||
// is 127.0.0.1:8080 — we want our own port + key on first run.) If a config
|
||||
// already exists we do NOT overwrite it (keep user's edits).
|
||||
// No upstream sources are preconfigured — the desktop user adds them via
|
||||
// the WebUI's sources page (like Clash Verge's empty-by-default profile).
|
||||
// If a config already exists we do NOT overwrite it (keep user's edits).
|
||||
if (fs.existsSync(CONFIG_FILE)) return;
|
||||
const key = "sk-gw-" + crypto.randomBytes(16).toString("hex");
|
||||
const yaml = [
|
||||
@ -132,16 +143,6 @@ function writeDefaultEmbeddedConfig() {
|
||||
"adapter_dir: " + JSON.stringify(path.join(PROFILE_DIR, "adapters")),
|
||||
"runtime_file: " + JSON.stringify(path.join(PROFILE_DIR, "runtime.json")),
|
||||
"",
|
||||
"sources:",
|
||||
" - name: zen",
|
||||
" base_url: https://opencode.ai/zen/v1",
|
||||
" api_key: public",
|
||||
" adapter: opencode",
|
||||
" models:",
|
||||
" - id: deepseek-v4-flash-free",
|
||||
" priority: 100",
|
||||
" kind: chat",
|
||||
"",
|
||||
// Admin key as a plain `keys` entry (no seed flag) so the web UI never
|
||||
// shows the replace-the-initial-key warning. gateway_keys entries are
|
||||
// marked seed:true by the core.
|
||||
@ -211,6 +212,7 @@ function startCore() {
|
||||
coreProc.on("exit", (code) => {
|
||||
const wasManaged = coreStartedAt > 0;
|
||||
coreStartedAt = 0;
|
||||
coreReady = false;
|
||||
coreProc = null;
|
||||
console.log("core exited:", code);
|
||||
// auto-restart if still running GUI (crash recovery, max every ~5s)
|
||||
@ -223,6 +225,12 @@ function startCore() {
|
||||
}, 1500);
|
||||
}
|
||||
});
|
||||
// Broadcast starting (NOT ready) immediately. The renderer must not
|
||||
// attempt to load the iframe until the core actually answers HTTP —
|
||||
// otherwise the frame load races the listener and fails permanently.
|
||||
// (Fix for blank-screen: running=true was previously broadcast while the
|
||||
// core was still booting, so the iframe ERR_CONNECTION_REFUSED and the
|
||||
// renderer never retried.)
|
||||
notifyCoreState();
|
||||
// pump until the core answers, then inject the admin key and tell the
|
||||
// renderer the web UI is ready (the iframe must load only after auth is set)
|
||||
@ -230,6 +238,7 @@ function startCore() {
|
||||
.then((ok) => {
|
||||
if (ok) {
|
||||
setEmbeddedAuth();
|
||||
coreReady = true;
|
||||
notifyCoreState();
|
||||
}
|
||||
})
|
||||
@ -254,7 +263,7 @@ function waitForCoreReady(timeoutMs = 10000) {
|
||||
let coreLastExit = 0;
|
||||
|
||||
function stopCore() {
|
||||
app.isQuitting = true;
|
||||
coreReady = false;
|
||||
if (coreProc && coreProc.exitCode === null) {
|
||||
try {
|
||||
coreProc.kill();
|
||||
@ -357,17 +366,27 @@ function createTray() {
|
||||
);
|
||||
}
|
||||
tray = new Tray(img);
|
||||
// Rebuild the context menu when the OS theme flips (dark/light) so menu
|
||||
// text colors follow the desktop theme; nativeTheme is the Electron API
|
||||
// that tracks GTK/appindicator dark mode on Linux.
|
||||
const themeRebuild = () => {
|
||||
try {
|
||||
if (global.__rebuildTray) global.__rebuildTray();
|
||||
} catch (e) {}
|
||||
};
|
||||
try {
|
||||
nativeTheme.on("updated", themeRebuild);
|
||||
} catch (e) {}
|
||||
// Status rows used to be enabled:false, which GTK renders in a fixed light
|
||||
// grey that never flips with the system dark/light theme — invisible on
|
||||
// light desktops. They are now enabled (normal theme-aware fg) with a
|
||||
// no-op click so they stay display-only but inherit the menu text color.
|
||||
const statusRow = (text) => ({ label: text, enabled: true, click: () => {} });
|
||||
const rebuild = () => {
|
||||
const running = coreStarted();
|
||||
const menu = Menu.buildFromTemplate([
|
||||
{
|
||||
label: "ModelRouter Desktop",
|
||||
enabled: false,
|
||||
},
|
||||
{
|
||||
label: running ? "● 内核运行中" : "○ 内核已停止",
|
||||
enabled: false,
|
||||
},
|
||||
statusRow("ModelRouter Desktop"),
|
||||
statusRow(running ? "● 内核运行中" : "○ 内核已停止"),
|
||||
{ type: "separator" },
|
||||
{
|
||||
label: "显示主窗口",
|
||||
@ -438,7 +457,8 @@ function createTray() {
|
||||
function notifyCoreState() {
|
||||
if (mainWindow && !mainWindow.isDestroyed()) {
|
||||
mainWindow.webContents.send("core:state", {
|
||||
running: coreStarted(),
|
||||
running: coreStarted() && coreReady,
|
||||
ready: coreReady,
|
||||
startedAt: coreStartedAt,
|
||||
});
|
||||
}
|
||||
@ -496,6 +516,12 @@ function createWindow() {
|
||||
const menu = Menu.buildFromTemplate([]);
|
||||
Menu.setApplicationMenu(menu);
|
||||
mainWindow = new BrowserWindow({
|
||||
// Window icon (taskbar/dock). Packaged: real file under resourcesPath (see
|
||||
// WINDOW_ICON) — an asar path here makes GTK segfault on Linux.
|
||||
icon: WINDOW_ICON,
|
||||
// GNOME taskbar icons match the window's WM_CLASS (lowercase
|
||||
// "modelrouter-gui") against StartupWMClass in the .desktop file.
|
||||
backgroundColor: "#1f2937",
|
||||
width: 1280,
|
||||
height: 860,
|
||||
minWidth: 900,
|
||||
@ -560,7 +586,8 @@ ipcMain.handle(
|
||||
);
|
||||
|
||||
ipcMain.handle("core:state", () => ({
|
||||
running: coreStarted(),
|
||||
running: coreStarted() && coreReady,
|
||||
ready: coreReady,
|
||||
startedAt: coreStartedAt,
|
||||
baseUrl: embeddedBaseUrl(),
|
||||
port: settings.port || DEFAULT_PORT,
|
||||
|
||||
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",
|
||||
@ -22,7 +22,7 @@
|
||||
"build": {
|
||||
"appId": "com.modelrouter.gui",
|
||||
"productName": "ModelRouter",
|
||||
"linux": {
|
||||
"linux": {
|
||||
"target": [
|
||||
"deb",
|
||||
"rpm",
|
||||
@ -32,10 +32,19 @@
|
||||
"icon": "build/icon.png",
|
||||
"maintainer": "ModelRouter",
|
||||
"synopsis": "ModelRouter Desktop",
|
||||
"desktop": {
|
||||
"entry": {
|
||||
"StartupWMClass": "modelrouter-gui"
|
||||
}
|
||||
},
|
||||
"extraResources": [
|
||||
{
|
||||
"from": "bin/llmsproxy",
|
||||
"to": "bin/llmsproxy"
|
||||
},
|
||||
{
|
||||
"from": "build/icon.png",
|
||||
"to": "build/icon.png"
|
||||
}
|
||||
]
|
||||
},
|
||||
@ -44,7 +53,13 @@
|
||||
},
|
||||
"mac": {
|
||||
"target": "dmg",
|
||||
"icon": "build/icon.png"
|
||||
"icon": "build/icon.png",
|
||||
"extraResources": [
|
||||
{
|
||||
"from": "build/icon.png",
|
||||
"to": "build/icon.png"
|
||||
}
|
||||
]
|
||||
},
|
||||
"win": {
|
||||
"target": "nsis",
|
||||
@ -57,6 +72,10 @@
|
||||
{
|
||||
"from": "bin/lua51.dll",
|
||||
"to": "bin/lua51.dll"
|
||||
},
|
||||
{
|
||||
"from": "build/icon.png",
|
||||
"to": "build/icon.png"
|
||||
}
|
||||
]
|
||||
},
|
||||
|
||||
@ -5,6 +5,7 @@
|
||||
const state = {
|
||||
core: {
|
||||
running: false,
|
||||
ready: false,
|
||||
baseUrl: "",
|
||||
port: 8787,
|
||||
profileDir: "",
|
||||
@ -18,21 +19,13 @@ const state = {
|
||||
minimizeToTray: true,
|
||||
},
|
||||
theme: localStorage.getItem("mr-theme") || "light",
|
||||
loadedOnce: false,
|
||||
loadedOnce: false, // iframe has been pointed at the web UI at least once
|
||||
frameLoaded: false, // last loadFrame actually finished loading (did-finish-load)
|
||||
loadTries: 0,
|
||||
};
|
||||
|
||||
const $ = (s) => document.querySelector(s);
|
||||
|
||||
function esc(s) {
|
||||
return String(s == null ? "" : s).replace(
|
||||
/[&<>"']/g,
|
||||
(c) =>
|
||||
({ "&": "&", "<": "<", ">": ">", '"': """, "'": "'" })[
|
||||
c
|
||||
],
|
||||
);
|
||||
}
|
||||
|
||||
function toast(msg, isErr, ms) {
|
||||
const t = $("#toast");
|
||||
t.textContent = msg;
|
||||
@ -55,7 +48,11 @@ function renderCoreStatus() {
|
||||
const core = state.core;
|
||||
const dot = $("#conn-dot");
|
||||
dot.className = "status-dot" + (core.running ? " ok" : " bad");
|
||||
dot.title = core.running ? "内核运行中" : "内核未运行";
|
||||
dot.title = !core.running
|
||||
? "内核未运行"
|
||||
: core.ready
|
||||
? "内核运行中"
|
||||
: "内核启动中…";
|
||||
}
|
||||
|
||||
function showOffline() {
|
||||
@ -70,11 +67,47 @@ function showLoading() {
|
||||
$("#console-loading").style.display = "flex";
|
||||
}
|
||||
|
||||
// probe the core's web root directly (fetch from the renderer). Using
|
||||
// fetch() instead of relying on the iframe's own load means the core's
|
||||
// readiness is confirmed by an actual HTTP round-trip before we point the
|
||||
// iframe at it — this is what prevents the blank-screen race where the
|
||||
// frame load happens while the listener is still coming up.
|
||||
function coreReachable() {
|
||||
return fetch(
|
||||
state.core.baseUrl || "http://127.0.0.1:" + (state.core.port || 8787),
|
||||
{ method: "GET", redirect: "follow", cache: "no-store" },
|
||||
)
|
||||
.then(() => true)
|
||||
.catch(() => false);
|
||||
}
|
||||
|
||||
// point the iframe at the WebUI, but only after the core is reachable.
|
||||
// If the load fails (core still booting / raced), we retry with backoff
|
||||
// instead of giving up forever (the old code set loadedOnce blindly and
|
||||
// never retried -> permanent blank screen).
|
||||
let frameRetryTimer = null;
|
||||
function scheduleFrameLoad() {
|
||||
if (frameRetryTimer) return; // already scheduled
|
||||
frameRetryTimer = setTimeout(async () => {
|
||||
frameRetryTimer = null;
|
||||
if (!state.core.running) return; // core stopped meanwhile
|
||||
if (!(await coreReachable())) {
|
||||
// core not reachable yet: wait a bit and try again
|
||||
state.loadTries += 1;
|
||||
showLoading();
|
||||
scheduleFrameLoad();
|
||||
return;
|
||||
}
|
||||
loadFrame(state.core.baseUrl || "http://127.0.0.1:" + state.core.port);
|
||||
}, state.loadTries === 0 ? 800 : Math.min(3000, 400 * state.loadTries));
|
||||
}
|
||||
|
||||
function loadFrame(url) {
|
||||
const f = frame();
|
||||
f.style.display = "block";
|
||||
$("#offline").style.display = "none";
|
||||
$("#console-loading").style.display = "none";
|
||||
if (state.loadedOnce && f.src === url) return; // already pointing there
|
||||
f.src = url;
|
||||
state.loadedOnce = true;
|
||||
}
|
||||
@ -82,8 +115,9 @@ function loadFrame(url) {
|
||||
async function ensureCore() {
|
||||
await refreshCore();
|
||||
if (state.core.running) {
|
||||
loadFrame(state.core.baseUrl);
|
||||
scheduleFrameLoad();
|
||||
} else {
|
||||
state.loadedOnce = false;
|
||||
showOffline();
|
||||
}
|
||||
}
|
||||
@ -141,10 +175,21 @@ async function saveSettings() {
|
||||
toast("设置已保存");
|
||||
await refreshCore();
|
||||
if (state.core.running) {
|
||||
loadFrame(state.core.baseUrl);
|
||||
scheduleFrameLoad();
|
||||
} else {
|
||||
state.loadedOnce = false;
|
||||
showOffline();
|
||||
}
|
||||
}
|
||||
|
||||
// ===== theme =====
|
||||
function applyTheme() {
|
||||
document.documentElement.dataset.theme = state.theme;
|
||||
try {
|
||||
localStorage.setItem("mr-theme", state.theme);
|
||||
} catch (e) {}
|
||||
}
|
||||
|
||||
// ===== init =====
|
||||
function init() {
|
||||
$("#tb-min").onclick = () => window.modelrouter.win.minimize();
|
||||
@ -171,15 +216,30 @@ function init() {
|
||||
showLoading();
|
||||
await window.modelrouter.core.start();
|
||||
await refreshCore();
|
||||
if (state.core.running) loadFrame(state.core.baseUrl);
|
||||
if (state.core.running) scheduleFrameLoad();
|
||||
else showOffline();
|
||||
};
|
||||
frame().onload = () => refreshCore();
|
||||
// when the iframe finishes loading, mark loaded and refresh core status so
|
||||
// the conn-dot reflects the actual WebUI (not just the core reachability)
|
||||
frame().onload = () => {
|
||||
state.frameLoaded = true;
|
||||
refreshCore();
|
||||
};
|
||||
// load failure (e.g. the frame was pointed at the core just before it
|
||||
// finished booting) -> clear loadedOnce and let scheduleFrameLoad retry
|
||||
frame().onerror = () => {
|
||||
state.loadedOnce = false;
|
||||
state.frameLoaded = false;
|
||||
scheduleFrameLoad();
|
||||
};
|
||||
window.modelrouter.core.onState((d) => {
|
||||
state.core = Object.assign({}, state.core, d);
|
||||
renderCoreStatus();
|
||||
if (d.running && !state.loadedOnce) loadFrame(state.core.baseUrl);
|
||||
if (!d.running) showOffline();
|
||||
if (d.running && !state.loadedOnce) scheduleFrameLoad();
|
||||
if (!d.running) {
|
||||
state.loadedOnce = false;
|
||||
showOffline();
|
||||
}
|
||||
});
|
||||
// Escape closes overlay
|
||||
document.addEventListener("keydown", (e) => {
|
||||
@ -189,6 +249,7 @@ function init() {
|
||||
}
|
||||
|
||||
async function boot() {
|
||||
applyTheme();
|
||||
init();
|
||||
state.settings = await window.modelrouter.settings.get();
|
||||
renderRail();
|
||||
|
||||
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
|
||||
@ -103,9 +104,10 @@ func EnsureDefault(path string) (bool, error) {
|
||||
}
|
||||
|
||||
// writeDefaultConfig writes a minimal, safe-by-default config: loopback-only
|
||||
// listen, a fresh random admin key, and a no-key zen source that works out of
|
||||
// the box. adapter_dir / runtime_file live next to the config file so the
|
||||
// binary works regardless of the working directory it is started from.
|
||||
// listen and a fresh random admin key. No upstream sources are preconfigured
|
||||
// (the repo ships none — add them via WebUI or config.yaml). adapter_dir /
|
||||
// runtime_file live next to the config file so the binary works regardless of
|
||||
// the working directory it is started from.
|
||||
func writeDefaultConfig(path string) error {
|
||||
key, err := NewGatewayKey()
|
||||
if err != nil {
|
||||
@ -122,15 +124,6 @@ func writeDefaultConfig(path string) error {
|
||||
DefaultModel: "AUTO",
|
||||
AdapterDir: filepath.Join(abs, "adapters"),
|
||||
RuntimeFile: filepath.Join(abs, "runtime.json"),
|
||||
Sources: []Source{
|
||||
{
|
||||
Name: "zen",
|
||||
BaseURL: "https://opencode.ai/zen/v1",
|
||||
APIKey: "public", // zen 免费池:官方无 key 客户端实际发送 Bearer public
|
||||
Adapter: "opencode",
|
||||
Models: []Model{{ID: "deepseek-v4-flash-free", Priority: 100, Kind: "chat"}},
|
||||
},
|
||||
},
|
||||
}
|
||||
out, err := yaml.Marshal(&cfg)
|
||||
if err != nil {
|
||||
@ -321,11 +314,31 @@ 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"`
|
||||
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"
|
||||
|
||||
@ -91,14 +91,8 @@ func TestEnsureDefaultGeneratesOnMissingFile(t *testing.T) {
|
||||
if cfg.Listen != "127.0.0.1:8080" {
|
||||
t.Fatalf("listen = %q, want loopback-only", cfg.Listen)
|
||||
}
|
||||
if len(cfg.Sources) != 1 || cfg.Sources[0].Name != "zen" || cfg.Sources[0].Adapter != "opencode" {
|
||||
t.Fatalf("default sources = %+v", cfg.Sources)
|
||||
}
|
||||
if cfg.Sources[0].APIKey != "public" {
|
||||
t.Fatalf("zen api_key = %q", cfg.Sources[0].APIKey)
|
||||
}
|
||||
if cfg.Sources[0].Models[0].ID != "deepseek-v4-flash-free" {
|
||||
t.Fatalf("default model = %+v", cfg.Sources[0].Models)
|
||||
if len(cfg.Sources) != 0 {
|
||||
t.Fatalf("default config must ship no sources (user adds them), got %+v", cfg.Sources)
|
||||
}
|
||||
// adapter_dir / runtime_file resolve next to the config file
|
||||
if !strings.HasPrefix(cfg.AdapterDir, dir) || !strings.HasPrefix(cfg.RuntimeFile, dir) {
|
||||
|
||||
@ -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,49 @@ func (s *Store) DeletedAdapters() map[string]bool {
|
||||
return out
|
||||
}
|
||||
|
||||
// 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 +229,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.
|
||||
@ -181,6 +182,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 +298,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 +338,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 +452,7 @@ func (c *Core) rebuildRegistry() error {
|
||||
c.registry.Replace(providers)
|
||||
}
|
||||
c.buildAutoChain()
|
||||
c.buildAutoImageChain()
|
||||
return nil
|
||||
}
|
||||
|
||||
@ -466,6 +499,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 +675,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"
|
||||
@ -53,11 +54,14 @@ type RespMessage struct {
|
||||
}
|
||||
|
||||
type ChatChunk struct {
|
||||
ID string `json:"id"`
|
||||
Object string `json:"object"`
|
||||
Created int64 `json:"created"`
|
||||
Model string `json:"model"`
|
||||
Choices []ChunkChoice `json:"choices"`
|
||||
ID string `json:"id"`
|
||||
Object string `json:"object"`
|
||||
Created int64 `json:"created"`
|
||||
Model string `json:"model"`
|
||||
Choices []ChunkChoice `json:"choices"`
|
||||
// Usage is sent in the final chunk of a stream (empty choices) so
|
||||
// OpenAI-compatible clients can read token usage.
|
||||
Usage *types.TokenUsage `json:"usage,omitempty"`
|
||||
}
|
||||
|
||||
type ChunkChoice struct {
|
||||
@ -384,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 {
|
||||
@ -520,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)
|
||||
@ -543,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
|
||||
@ -557,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 {
|
||||
@ -566,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 {
|
||||
@ -577,27 +648,47 @@ func (g *Gateway) writeRec(rec *Req) {
|
||||
g.stats.Record(*rec)
|
||||
}
|
||||
|
||||
func (g *Gateway) streamChat(w http.ResponseWriter, ctx context.Context, cands []*provider.Provider, req *types.ChatRequest, effective string, rec *Req) {
|
||||
rec.LatMs = 0
|
||||
t0 := time.Now()
|
||||
rec.OK = true
|
||||
rec.Status = http.StatusOK
|
||||
defer func() {
|
||||
rec.LatMs = time.Since(t0).Milliseconds()
|
||||
g.writeRec(rec)
|
||||
}()
|
||||
chunks, _, usedModel, err := g.core.Scheduler().ChatStream(ctx, toScheduler(cands), req)
|
||||
if err != nil {
|
||||
rec.OK = false
|
||||
rec.Status = upstreamErrStatus(err)
|
||||
rec.Err = err.Error()
|
||||
writeError(w, rec.Status, "upstream_error", err.Error())
|
||||
return
|
||||
// mergeUsage combines token usage across stream chunks additively. Some
|
||||
// providers split usage across chunks (e.g. Anthropic reports prompt tokens
|
||||
// in message_start and the final completion tokens in message_delta); a plain
|
||||
// "last non-nil wins" would discard the prompt half. Non-zero fields from cur
|
||||
// override prev; total is recomputed from the merged parts so a partial later
|
||||
// chunk can't shrink it. For the common single-chunk case (OpenAI's terminal
|
||||
// empty-choices+usage chunk) upstream totals are preserved exactly.
|
||||
func mergeUsage(prev, cur *types.TokenUsage) *types.TokenUsage {
|
||||
if prev == nil {
|
||||
u := *cur
|
||||
if u.Total == 0 && (u.Prompt > 0 || u.Completion > 0) {
|
||||
u.Total = u.Prompt + u.Completion
|
||||
}
|
||||
return &u
|
||||
}
|
||||
if usedModel != "" {
|
||||
rec.Model = usedModel
|
||||
out := *prev
|
||||
if cur.Prompt > 0 {
|
||||
out.Prompt = cur.Prompt
|
||||
}
|
||||
rec.Prompt = estimatePromptTokens(req)
|
||||
if cur.Completion > 0 {
|
||||
out.Completion = cur.Completion
|
||||
}
|
||||
out.Total = out.Prompt + out.Completion
|
||||
if cur.PromptTokensDetails != nil && cur.PromptTokensDetails.CachedTokens > 0 {
|
||||
out.PromptTokensDetails = cur.PromptTokensDetails
|
||||
}
|
||||
if cur.PromptCacheHit > 0 {
|
||||
out.PromptCacheHit = cur.PromptCacheHit
|
||||
}
|
||||
if cur.PromptCacheMiss > 0 {
|
||||
out.PromptCacheMiss = cur.PromptCacheMiss
|
||||
}
|
||||
return &out
|
||||
}
|
||||
|
||||
// pumpStream writes the full SSE sequence for a started stream: role
|
||||
// preamble, one chunk per unified delta, the terminating finish_reason, the
|
||||
// OpenAI-standard final usage chunk (empty choices) and [DONE]. modelName
|
||||
// follows writeChatCompletion's rule (requested id for direct routes, exact
|
||||
// slot model for AUTO).
|
||||
func (g *Gateway) pumpStream(w http.ResponseWriter, rec *Req, chunks <-chan types.UnifiedChunk, modelName string, t0 time.Time) {
|
||||
w.Header().Set("Content-Type", "text/event-stream")
|
||||
w.Header().Set("Cache-Control", "no-cache")
|
||||
w.Header().Set("Connection", "keep-alive")
|
||||
@ -621,14 +712,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 != "" {
|
||||
@ -639,8 +738,12 @@ func (g *Gateway) streamChat(w http.ResponseWriter, ctx context.Context, cands [
|
||||
}
|
||||
choice := ChunkChoice{Index: 0, Delta: delta}
|
||||
if ck.Done {
|
||||
stop := "stop"
|
||||
choice.FinishReason = &stop
|
||||
fin := ck.FinishReason
|
||||
if fin == "" {
|
||||
fin = "stop"
|
||||
}
|
||||
lastFinish = fin
|
||||
choice.FinishReason = &fin
|
||||
}
|
||||
chunk.Choices = []ChunkChoice{choice}
|
||||
rec.Compl += int64(len(ck.Content)+len(ck.ReasoningContent)+len(ck.ToolCalls)) / 3
|
||||
@ -648,17 +751,82 @@ func (g *Gateway) streamChat(w http.ResponseWriter, ctx context.Context, cands [
|
||||
return
|
||||
}
|
||||
}
|
||||
stop := "stop"
|
||||
finalFinish := lastFinish
|
||||
if finalFinish == "" {
|
||||
finalFinish = "stop"
|
||||
}
|
||||
send(ChatChunk{
|
||||
ID: id, Object: "chat.completion.chunk", Created: created, Model: effective,
|
||||
Choices: []ChunkChoice{{Index: 0, Delta: RespMessage{}, FinishReason: &stop}},
|
||||
ID: id, Object: "chat.completion.chunk", Created: created, Model: modelName,
|
||||
Choices: []ChunkChoice{{Index: 0, Delta: RespMessage{}, FinishReason: &finalFinish}},
|
||||
})
|
||||
// Final usage chunk (OpenAI standard: empty choices + usage before [DONE]).
|
||||
// Prefer the upstream's exact usage if the stream carried it; fall back to
|
||||
// the gateway's estimate otherwise.
|
||||
var tut *types.TokenUsage
|
||||
if lastUsage != nil {
|
||||
tut = lastUsage
|
||||
// Write the upstream's cache accounting back onto the request record
|
||||
// so the audit trail carries the hit/miss split for streaming too.
|
||||
if d := tut.PromptTokensDetails; d != nil {
|
||||
rec.CacheReported = true
|
||||
if d.CachedTokens > 0 {
|
||||
rec.CacheHit = int64(d.CachedTokens)
|
||||
}
|
||||
} else if tut.PromptCacheHit > 0 {
|
||||
rec.CacheHit = int64(tut.PromptCacheHit)
|
||||
rec.CacheReported = true
|
||||
}
|
||||
if tut.PromptCacheMiss > 0 {
|
||||
rec.CacheMiss = int64(tut.PromptCacheMiss)
|
||||
}
|
||||
} else if rec.Prompt+rec.Compl > 0 {
|
||||
u := types.TokenUsage{
|
||||
Prompt: int(rec.Prompt),
|
||||
Completion: int(rec.Compl),
|
||||
Total: int(rec.Prompt + rec.Compl),
|
||||
}
|
||||
tut = &u
|
||||
}
|
||||
if tut != nil {
|
||||
send(ChatChunk{
|
||||
ID: id, Object: "chat.completion.chunk", Created: created, Model: modelName,
|
||||
Choices: []ChunkChoice{},
|
||||
Usage: tut,
|
||||
})
|
||||
}
|
||||
fmt.Fprintf(w, "data: [DONE]\n\n")
|
||||
if flusher != nil {
|
||||
flusher.Flush()
|
||||
}
|
||||
}
|
||||
|
||||
// streamChat runs a direct (model-pinned) streaming request across the
|
||||
// candidate list, falling back early on connect errors.
|
||||
func (g *Gateway) streamChat(w http.ResponseWriter, ctx context.Context, cands []*provider.Provider, req *types.ChatRequest, effective string, rec *Req) {
|
||||
rec.LatMs = 0
|
||||
t0 := time.Now()
|
||||
rec.OK = true
|
||||
rec.Status = http.StatusOK
|
||||
defer func() {
|
||||
rec.LatMs = time.Since(t0).Milliseconds()
|
||||
g.writeRec(rec)
|
||||
}()
|
||||
chunks, usedSrc, usedModel, err := g.core.Scheduler().ChatStream(ctx, toScheduler(cands), req)
|
||||
if err != nil {
|
||||
g.failChat(w, rec, err)
|
||||
return
|
||||
}
|
||||
if usedModel != "" {
|
||||
rec.Model = usedModel
|
||||
}
|
||||
// Audit accuracy: pin the source that actually served the stream (after a
|
||||
// failover it differs from the first candidate). Direct streams previously
|
||||
// discarded it.
|
||||
rec.Source = usedSrc
|
||||
rec.Prompt = estimatePromptTokens(req)
|
||||
g.pumpStream(w, rec, chunks, effective, t0)
|
||||
}
|
||||
|
||||
// singleChatAuto runs a non-streaming AUTO request down the chain (see
|
||||
// scheduler.ChainChat): tiers ascending (tier 1 = highest priority first),
|
||||
// per-tier round-robin ordered by preference, cooldown as the only hard skip,
|
||||
@ -671,48 +839,18 @@ func (g *Gateway) singleChatAuto(w http.ResponseWriter, ctx context.Context, cha
|
||||
resp, usedSrc, usedModel, err := g.core.Scheduler().ChainChat(ctx, chain, req, quotaExhausted)
|
||||
rec.LatMs = time.Since(t0).Milliseconds()
|
||||
if err != nil {
|
||||
rec.OK = false
|
||||
rec.Status = upstreamErrStatus(err)
|
||||
rec.Err = err.Error()
|
||||
if ce, ok := err.(*scheduler.ChainErr); ok && len(ce.Tiers) > 0 {
|
||||
rec.Source = ce.Tiers[0].Source
|
||||
rec.Model = ce.Tiers[0].Model
|
||||
}
|
||||
g.failChat(w, rec, err)
|
||||
g.writeRec(rec)
|
||||
writeError(w, rec.Status, "upstream_error", err.Error())
|
||||
return
|
||||
}
|
||||
rec.OK = true
|
||||
rec.Status = http.StatusOK
|
||||
rec.Prompt = int64(resp.TokenUsage.Prompt)
|
||||
if rec.Prompt == 0 {
|
||||
rec.Prompt = estimatePromptTokens(req)
|
||||
}
|
||||
rec.Compl = int64(resp.TokenUsage.Completion)
|
||||
if rec.Compl == 0 {
|
||||
rec.Compl = estimateTextTokens(resp.Content, resp.ReasoningContent, resp.ToolCalls)
|
||||
}
|
||||
recordChatUsage(rec, req, resp)
|
||||
rec.Source = usedSrc
|
||||
rec.Model = usedModel
|
||||
rec.FirstByteMs = rec.LatMs
|
||||
g.writeRec(rec)
|
||||
msg := RespMessage{Role: "assistant", Content: resp.Content}
|
||||
if resp.ReasoningContent != "" {
|
||||
msg.ReasoningContent = resp.ReasoningContent
|
||||
}
|
||||
if len(resp.ToolCalls) > 0 {
|
||||
msg.ToolCalls = toolCallsWire(resp.ToolCalls)
|
||||
}
|
||||
out := ChatCompletion{
|
||||
ID: newID(),
|
||||
Object: "chat.completion",
|
||||
Created: time.Now().Unix(),
|
||||
Model: usedModel,
|
||||
Choices: []ChatChoice{{Index: 0, Message: msg, FinishReason: resp.FinishReason}},
|
||||
}
|
||||
if resp.TokenUsage.Total > 0 || resp.TokenUsage.Prompt > 0 || resp.TokenUsage.Completion > 0 {
|
||||
out.Usage = &resp.TokenUsage
|
||||
}
|
||||
writeJSON(w, http.StatusOK, out)
|
||||
writeChatCompletion(w, resp, usedModel)
|
||||
}
|
||||
|
||||
// streamChatAuto streams an AUTO request down the chain. A slot is abandoned
|
||||
@ -730,14 +868,7 @@ func (g *Gateway) streamChatAuto(w http.ResponseWriter, ctx context.Context, cha
|
||||
}()
|
||||
chunks, usedSrc, usedModel, err := g.core.Scheduler().ChainChatStream(ctx, chain, req, quotaExhausted)
|
||||
if err != nil {
|
||||
rec.OK = false
|
||||
rec.Status = upstreamErrStatus(err)
|
||||
rec.Err = err.Error()
|
||||
if ce, ok := err.(*scheduler.ChainErr); ok && len(ce.Tiers) > 0 {
|
||||
rec.Source = ce.Tiers[0].Source
|
||||
rec.Model = ce.Tiers[0].Model
|
||||
}
|
||||
writeError(w, rec.Status, "upstream_error", err.Error())
|
||||
g.failChat(w, rec, err)
|
||||
return
|
||||
}
|
||||
if usedModel != "" {
|
||||
@ -745,63 +876,7 @@ func (g *Gateway) streamChatAuto(w http.ResponseWriter, ctx context.Context, cha
|
||||
}
|
||||
rec.Source = usedSrc
|
||||
rec.Prompt = estimatePromptTokens(req)
|
||||
w.Header().Set("Content-Type", "text/event-stream")
|
||||
w.Header().Set("Cache-Control", "no-cache")
|
||||
w.Header().Set("Connection", "keep-alive")
|
||||
w.WriteHeader(http.StatusOK)
|
||||
flusher, _ := w.(http.Flusher)
|
||||
id := newID()
|
||||
created := time.Now().Unix()
|
||||
send := func(obj interface{}) bool {
|
||||
b, err := json.Marshal(obj)
|
||||
if err != nil {
|
||||
return false
|
||||
}
|
||||
if _, err := fmt.Fprintf(w, "data: %s\n\n", b); err != nil {
|
||||
return false
|
||||
}
|
||||
if flusher != nil {
|
||||
flusher.Flush()
|
||||
}
|
||||
return true
|
||||
}
|
||||
if !send(ChatChunk{
|
||||
ID: id, Object: "chat.completion.chunk", Created: created, Model: rec.Model,
|
||||
Choices: []ChunkChoice{{Index: 0, Delta: RespMessage{Role: "assistant"}}},
|
||||
}) {
|
||||
return
|
||||
}
|
||||
for ck := range chunks {
|
||||
chunk := ChatChunk{
|
||||
ID: id, Object: "chat.completion.chunk", Created: created, Model: rec.Model,
|
||||
}
|
||||
delta := RespMessage{Role: "assistant", Content: ck.Content}
|
||||
if ck.ReasoningContent != "" {
|
||||
delta.ReasoningContent = ck.ReasoningContent
|
||||
}
|
||||
if len(ck.ToolCalls) > 0 {
|
||||
delta.ToolCalls = ck.ToolCalls
|
||||
}
|
||||
choice := ChunkChoice{Index: 0, Delta: delta}
|
||||
if ck.Done {
|
||||
stop := "stop"
|
||||
choice.FinishReason = &stop
|
||||
}
|
||||
chunk.Choices = []ChunkChoice{choice}
|
||||
rec.Compl += int64(len(ck.Content)+len(ck.ReasoningContent)+len(ck.ToolCalls)) / 3
|
||||
if !send(chunk) {
|
||||
return
|
||||
}
|
||||
}
|
||||
stop := "stop"
|
||||
send(ChatChunk{
|
||||
ID: id, Object: "chat.completion.chunk", Created: created, Model: rec.Model,
|
||||
Choices: []ChunkChoice{{Index: 0, Delta: RespMessage{}, FinishReason: &stop}},
|
||||
})
|
||||
fmt.Fprintf(w, "data: [DONE]\n\n")
|
||||
if flusher != nil {
|
||||
flusher.Flush()
|
||||
}
|
||||
g.pumpStream(w, rec, chunks, usedModel, t0)
|
||||
}
|
||||
|
||||
func (g *Gateway) handleImage(w http.ResponseWriter, r *http.Request) {
|
||||
@ -822,11 +897,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)
|
||||
@ -851,12 +960,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))
|
||||
|
||||
@ -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,34 +420,59 @@ document.getElementById('key').addEventListener('keydown',function(e){if(e.key==
|
||||
apply();
|
||||
</script></body></html>`
|
||||
|
||||
// handleLoginAPI validates the gateway key and issues a session cookie.
|
||||
// handleLoginAPI validates the gateway key and issues a session cookie,
|
||||
// guarded by per-IP + global brute-force throttling (see loginGuard).
|
||||
func (g *Gateway) handleLoginAPI(w http.ResponseWriter, r *http.Request) {
|
||||
if r.Method != http.MethodPost {
|
||||
writeError(w, http.StatusMethodNotAllowed, "method_not_allowed", "use POST")
|
||||
return
|
||||
}
|
||||
now := time.Now().Unix()
|
||||
if !g.loginAllow(clientIP(r), now) {
|
||||
// Don't reveal whether the key was right: answer the same 401 the
|
||||
// failure path uses, but skip the expensive key lookup.
|
||||
writeError(w, http.StatusUnauthorized, "invalid_api_key", "invalid gateway api key")
|
||||
return
|
||||
}
|
||||
var body struct {
|
||||
Key string `json:"key"`
|
||||
}
|
||||
if err := json.NewDecoder(r.Body).Decode(&body); err != nil {
|
||||
g.loginRecord(clientIP(r), false, now)
|
||||
writeError(w, http.StatusBadRequest, "invalid_request", "invalid json")
|
||||
return
|
||||
}
|
||||
if _, ok := g.core.FindKey(body.Key); !ok {
|
||||
g.loginRecord(clientIP(r), false, now)
|
||||
writeError(w, http.StatusUnauthorized, "invalid_api_key", "invalid gateway api key")
|
||||
return
|
||||
}
|
||||
g.loginRecord(clientIP(r), true, now)
|
||||
http.SetCookie(w, &http.Cookie{
|
||||
Name: "gw_key",
|
||||
Value: body.Key,
|
||||
Path: "/",
|
||||
MaxAge: 86400 * 30,
|
||||
Name: "gw_key",
|
||||
Value: body.Key,
|
||||
Path: "/",
|
||||
MaxAge: 86400 * 30,
|
||||
HttpOnly: true,
|
||||
SameSite: http.SameSiteLaxMode,
|
||||
})
|
||||
writeJSON(w, http.StatusOK, map[string]interface{}{"ok": true})
|
||||
}
|
||||
|
||||
// clientIP returns the caller's IP for rate limiting. Behind a trusted
|
||||
// reverse proxy (nginx) the X-Real-IP header is preferred; otherwise the
|
||||
// remote address is used with the port stripped.
|
||||
func clientIP(r *http.Request) string {
|
||||
if x := r.Header.Get("X-Real-IP"); x != "" {
|
||||
return x
|
||||
}
|
||||
host, _, err := net.SplitHostPort(r.RemoteAddr)
|
||||
if err != nil {
|
||||
return r.RemoteAddr
|
||||
}
|
||||
return host
|
||||
}
|
||||
|
||||
func (g *Gateway) handleModels(w http.ResponseWriter, r *http.Request) {
|
||||
if r.Method != http.MethodGet {
|
||||
writeError(w, http.StatusMethodNotAllowed, "method_not_allowed", "use GET")
|
||||
@ -456,7 +576,20 @@ func (g *Gateway) handleStatusAPI(w http.ResponseWriter, r *http.Request) {
|
||||
"gateway_keys": ks,
|
||||
}
|
||||
if reqRole(r.Context()) == "admin" {
|
||||
resp["sources"] = g.core.Registry().Status()
|
||||
sts := g.core.Registry().Status()
|
||||
recent := g.stats.SourceRecent(300)
|
||||
avgs := g.stats.SourceAverages(300)
|
||||
for i := range sts {
|
||||
if v, ok := recent[sts[i].Name]; ok {
|
||||
sts[i].RecentOK = v[0]
|
||||
sts[i].RecentErr = v[1]
|
||||
}
|
||||
if a, ok := avgs[sts[i].Name]; ok {
|
||||
sts[i].AvgFirstByteMs = a.AvgFirstByteMs
|
||||
sts[i].AvgTokPerS = a.AvgTokPerS
|
||||
}
|
||||
}
|
||||
resp["sources"] = sts
|
||||
resp["adapters"] = g.core.ListAdapters()
|
||||
}
|
||||
writeJSON(w, http.StatusOK, resp)
|
||||
@ -491,4 +624,4 @@ func writeJSON(w http.ResponseWriter, code int, v interface{}) {
|
||||
if err := json.NewEncoder(w).Encode(v); err != nil {
|
||||
log.Printf("[gateway] write json: %v", err)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@ -26,7 +26,20 @@ type Req struct {
|
||||
Compl int64 `json:"completion_tokens"`
|
||||
// LatMs total handling time ms
|
||||
LatMs int64 `json:"latency_ms"`
|
||||
OK bool `json:"ok"`
|
||||
// FirstByteMs time-to-first-byte for streaming (ms from request start
|
||||
// to the first SSE chunk sent to the client); for non-streaming it
|
||||
// equals LatMs. 0 when unmeasured (legacy records).
|
||||
FirstByteMs int64 `json:"first_byte_ms,omitempty"`
|
||||
// CacheHit / CacheMiss carry the upstream prompt-cache accounting
|
||||
// (DeepSeek-style hit/miss tokens) when the upstream reports it.
|
||||
// Both 0 = upstream gave no cache data.
|
||||
CacheHit int64 `json:"cache_hit_tokens,omitempty"`
|
||||
CacheMiss int64 `json:"cache_miss_tokens,omitempty"`
|
||||
// CacheReported marks that the upstream usage reported cache
|
||||
// accounting at all (even when the hit count is 0). The WebUI shows
|
||||
// "0%" instead of "—" for such rows.
|
||||
CacheReported bool `json:"cache_reported,omitempty"`
|
||||
OK bool `json:"ok"`
|
||||
// Status http status code
|
||||
Status int `json:"status"`
|
||||
// Err short error message
|
||||
@ -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
|
||||
@ -445,6 +442,92 @@ func (s *Stats) AuditRecords(from, to int64, key string) []Req {
|
||||
return out
|
||||
}
|
||||
|
||||
// SourceRecent counts real gateway requests per source within the last
|
||||
// window (unix seconds). It reads only the in-memory ring, so it is cheap and
|
||||
// reflects live traffic — used by /api/status so the source status column is
|
||||
// driven by what actually happens, not just a probe.
|
||||
func (s *Stats) SourceRecent(windowSec int64) map[string][2]int64 {
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
cut := time.Now().Unix() - windowSec
|
||||
out := map[string][2]int64{}
|
||||
for _, r := range s.recs {
|
||||
if r.Time/1000 < cut {
|
||||
continue
|
||||
}
|
||||
if r.Source == "" {
|
||||
continue
|
||||
}
|
||||
v := out[r.Source]
|
||||
if r.OK {
|
||||
v[0]++
|
||||
} else {
|
||||
v[1]++
|
||||
}
|
||||
out[r.Source] = v
|
||||
}
|
||||
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{} {
|
||||
@ -506,4 +589,4 @@ func (s *Stats) Snapshot(limit int, key string) map[string]interface{} {
|
||||
"by_status": bs,
|
||||
"records": append([]Req(nil), recs...),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@ -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
|
||||
@ -82,6 +82,13 @@ function adapter.transform_response(raw_body)
|
||||
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)
|
||||
-- Anthropic reports cache_read_input_tokens; normalize into
|
||||
-- OpenAI-standard prompt_tokens_details.cached_tokens so clients
|
||||
-- (dsh) see the cache hit count.
|
||||
local cacheRead = resp.usage.cache_read_input_tokens or 0
|
||||
if cacheRead > 0 then
|
||||
unified.token_usage.prompt_tokens_details = { cached_tokens = cacheRead }
|
||||
end
|
||||
end
|
||||
|
||||
if resp.content and #resp.content > 0 then
|
||||
@ -99,9 +106,52 @@ end
|
||||
function adapter.transform_stream_chunk(raw_chunk)
|
||||
local ok, chunk = pcall(json.decode, raw_chunk)
|
||||
if not ok then return "" end
|
||||
if chunk.type == "message_start" then return "" end
|
||||
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 cacheRead = u.cache_read_input_tokens or 0
|
||||
if cacheRead > 0 then
|
||||
uses.prompt_tokens_details = { cached_tokens = cacheRead }
|
||||
end
|
||||
end
|
||||
end
|
||||
if uses ~= nil then
|
||||
return json.encode({ usage = uses, done = false })
|
||||
end
|
||||
return ""
|
||||
end
|
||||
if chunk.type == "message_delta" then
|
||||
return json.encode({ content = "", done = (chunk.delta and chunk.delta.stop_reason ~= nil) })
|
||||
local uses = nil
|
||||
if type(chunk.usage) == "table" then
|
||||
local u = chunk.usage
|
||||
local p = u.input_tokens or 0
|
||||
local c = u.output_tokens or 0
|
||||
if p > 0 or c > 0 then
|
||||
uses = { prompt = p, completion = c, total = p + c }
|
||||
end
|
||||
end
|
||||
local finish = nil
|
||||
if chunk.delta and chunk.delta.stop_reason ~= nil then
|
||||
-- Anthropic stop_reason -> OpenAI finish_reason
|
||||
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
|
||||
if uses ~= nil then
|
||||
-- completion is final here; prompt is merged from message_start
|
||||
return json.encode({ content = "", done = (finish ~= nil), finish_reason = finish, usage = uses })
|
||||
end
|
||||
return json.encode({ content = "", done = (finish ~= nil), finish_reason = finish })
|
||||
end
|
||||
if chunk.type == "content_block_start" and chunk.content_block
|
||||
and chunk.content_block.type == "tool_use" then
|
||||
@ -141,4 +191,15 @@ function adapter.transform_stream_chunk(raw_chunk)
|
||||
return ""
|
||||
end
|
||||
|
||||
-- 错误收敛:Anthropic 信封 {type:"error", error:{type, 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
|
||||
if resp.type == "error" and 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
|
||||
@ -82,14 +87,47 @@ end
|
||||
function adapter.transform_stream_chunk(raw_chunk)
|
||||
local ok, chunk = pcall(json.decode, raw_chunk)
|
||||
if not ok then return "" end
|
||||
if not chunk.choices or #chunk.choices == 0 then return "" end
|
||||
|
||||
-- OpenAI-style streams may attach usage to a chunk with empty choices
|
||||
-- (the final usage chunk). Keys must match Go's TokenUsage json tags
|
||||
-- (prompt/completion/total); the gateway re-emits standard *_tokens.
|
||||
local uses = nil
|
||||
if type(chunk.usage) == "table" then
|
||||
uses = {
|
||||
prompt = chunk.usage.prompt_tokens or chunk.usage.prompt or 0,
|
||||
completion = chunk.usage.completion_tokens or chunk.usage.completion or 0,
|
||||
total = chunk.usage.total_tokens or chunk.usage.total or 0,
|
||||
prompt_cache_hit_tokens = chunk.usage.prompt_cache_hit_tokens or 0,
|
||||
prompt_cache_miss_tokens = chunk.usage.prompt_cache_miss_tokens or 0,
|
||||
}
|
||||
if chunk.usage.prompt_tokens_details and chunk.usage.prompt_tokens_details.cached_tokens ~= nil then
|
||||
uses.prompt_tokens_details = { cached_tokens = chunk.usage.prompt_tokens_details.cached_tokens }
|
||||
end
|
||||
end
|
||||
|
||||
if not chunk.choices or #chunk.choices == 0 then
|
||||
if uses ~= nil then
|
||||
-- usage-only chunk is not a content/finish signal; the gateway
|
||||
-- emits its own terminal stop chunk and merges this usage.
|
||||
return json.encode({ usage = uses, done = false })
|
||||
end
|
||||
return ""
|
||||
end
|
||||
local delta = chunk.choices[1].delta or {}
|
||||
local fr = chunk.choices[1].finish_reason
|
||||
|
||||
local finish = (type(fr) == "string" and fr ~= "") and fr or nil
|
||||
|
||||
local unified = {
|
||||
content = delta.content or "",
|
||||
done = (fr ~= nil)
|
||||
done = (finish ~= nil)
|
||||
}
|
||||
if finish then
|
||||
unified.finish_reason = finish
|
||||
end
|
||||
if uses ~= nil then
|
||||
unified.usage = uses
|
||||
end
|
||||
if delta.reasoning_content then
|
||||
unified.reasoning_content = delta.reasoning_content
|
||||
end
|
||||
@ -99,4 +137,15 @@ function adapter.transform_stream_chunk(raw_chunk)
|
||||
return json.encode(unified)
|
||||
end
|
||||
|
||||
-- 错误收敛:DeepSeek 走标准 OpenAI 信封 {error:{message,...}}
|
||||
function adapter.transform_error(status, body)
|
||||
local ok, resp = pcall(json.decode, body)
|
||||
if not ok or type(resp) ~= "table" then return nil end
|
||||
local e = resp.error
|
||||
if type(e) == "table" and type(e.message) == "string" then
|
||||
return e.message
|
||||
end
|
||||
return nil
|
||||
end
|
||||
|
||||
return adapter
|
||||
|
||||
@ -91,9 +91,42 @@ function adapter.transform_stream_chunk(raw_chunk)
|
||||
local ok, chunk = pcall(json.decode, raw_chunk)
|
||||
if not ok then return "" end
|
||||
|
||||
if not chunk.candidates or #chunk.candidates == 0 then return "" end
|
||||
-- Gemini attaches usageMetadata to the final chunk (alongside or after
|
||||
-- candidates). Keys map to Go's TokenUsage json tags (prompt/...).
|
||||
local uses = nil
|
||||
if type(chunk.usageMetadata) == "table" then
|
||||
local p = chunk.usageMetadata.promptTokenCount or 0
|
||||
local c = chunk.usageMetadata.candidatesTokenCount or 0
|
||||
local t = chunk.usageMetadata.totalTokenCount or 0
|
||||
if p > 0 or c > 0 or t > 0 then
|
||||
uses = { prompt = p, completion = c, total = t }
|
||||
local cacheRead = chunk.usageMetadata.cachedContentTokenCount or 0
|
||||
if cacheRead > 0 then
|
||||
uses.prompt_tokens_details = { cached_tokens = cacheRead }
|
||||
end
|
||||
end
|
||||
end
|
||||
|
||||
if not chunk.candidates or #chunk.candidates == 0 then
|
||||
if uses ~= nil then
|
||||
return json.encode({ usage = uses, done = false })
|
||||
end
|
||||
return ""
|
||||
end
|
||||
local cand = chunk.candidates[1]
|
||||
local unified = { content = "", done = (cand.finishReason ~= nil) }
|
||||
-- Gemini finishReason -> OpenAI finish_reason
|
||||
local finish = nil
|
||||
if cand.finishReason ~= nil then
|
||||
if cand.finishReason == "MAX_TOKENS" then
|
||||
finish = "length"
|
||||
elseif cand.finishReason == "SAFETY" or cand.finishReason == "RECITATION"
|
||||
or cand.finishReason == "BLOCKLIST" then
|
||||
finish = "content_filter"
|
||||
else
|
||||
finish = "stop"
|
||||
end
|
||||
end
|
||||
local unified = { content = "", done = (finish ~= nil), finish_reason = finish }
|
||||
local reasoning = ""
|
||||
local tools = {}
|
||||
if cand.content and cand.content.parts then
|
||||
@ -117,7 +150,26 @@ function adapter.transform_stream_chunk(raw_chunk)
|
||||
end
|
||||
if reasoning ~= "" then unified.reasoning_content = reasoning end
|
||||
if #tools > 0 then unified.tool_calls = tools end
|
||||
if uses ~= nil then
|
||||
unified.usage = uses
|
||||
end
|
||||
return json.encode(unified)
|
||||
end
|
||||
|
||||
-- 错误收敛:Gemini REST 信封 {error:{code, message, status}}
|
||||
function adapter.transform_error(status, body)
|
||||
local ok, resp = pcall(json.decode, body)
|
||||
if not ok or type(resp) ~= "table" then return nil end
|
||||
local e = resp.error
|
||||
if type(e) == "table" then
|
||||
if type(e.message) == "string" then
|
||||
if type(e.status) == "string" then
|
||||
return e.status .. ": " .. e.message
|
||||
end
|
||||
return e.message
|
||||
end
|
||||
end
|
||||
return nil
|
||||
end
|
||||
|
||||
return adapter
|
||||
|
||||
@ -33,6 +33,18 @@ function adapter.transform_response(raw_body)
|
||||
unified.token_usage.prompt = resp.usage.prompt_tokens or 0
|
||||
unified.token_usage.completion = resp.usage.completion_tokens or 0
|
||||
unified.token_usage.total = resp.usage.total_tokens or 0
|
||||
local hit = 0
|
||||
if type(resp.usage.prompt_tokens_details) == "table" and resp.usage.prompt_tokens_details.cached_tokens ~= nil then
|
||||
hit = resp.usage.prompt_tokens_details.cached_tokens
|
||||
unified.token_usage.prompt_tokens_details = { cached_tokens = hit }
|
||||
end
|
||||
if (resp.usage.prompt_cache_hit_tokens or 0) > 0 then
|
||||
unified.token_usage.prompt_cache_hit_tokens = resp.usage.prompt_cache_hit_tokens
|
||||
unified.token_usage.prompt_cache_miss_tokens = resp.usage.prompt_cache_miss_tokens or 0
|
||||
if hit == 0 then
|
||||
unified.token_usage.prompt_tokens_details = { cached_tokens = resp.usage.prompt_cache_hit_tokens }
|
||||
end
|
||||
end
|
||||
end
|
||||
|
||||
if type(resp.choices) == "table" and #resp.choices > 0 then
|
||||
@ -63,14 +75,49 @@ end
|
||||
function adapter.transform_stream_chunk(raw_chunk)
|
||||
local ok, chunk = pcall(json.decode, raw_chunk)
|
||||
if not ok then return "" end
|
||||
if not chunk.choices or #chunk.choices == 0 then return "" end
|
||||
|
||||
-- OpenAI-style streams may attach usage to a chunk with empty choices
|
||||
-- (the final usage chunk). Keys must match Go's TokenUsage json tags
|
||||
-- (prompt/completion/total); the gateway re-emits standard *_tokens.
|
||||
local uses = nil
|
||||
if type(chunk.usage) == "table" then
|
||||
uses = {
|
||||
prompt = chunk.usage.prompt_tokens or chunk.usage.prompt or 0,
|
||||
completion = chunk.usage.completion_tokens or chunk.usage.completion or 0,
|
||||
total = chunk.usage.total_tokens or chunk.usage.total or 0,
|
||||
}
|
||||
if type(chunk.usage.prompt_tokens_details) == "table" and chunk.usage.prompt_tokens_details.cached_tokens ~= nil then
|
||||
uses.prompt_tokens_details = { cached_tokens = chunk.usage.prompt_tokens_details.cached_tokens }
|
||||
elseif (chunk.usage.prompt_cache_hit_tokens or 0) > 0 then
|
||||
uses.prompt_cache_hit_tokens = chunk.usage.prompt_cache_hit_tokens
|
||||
uses.prompt_cache_miss_tokens = chunk.usage.prompt_cache_miss_tokens or 0
|
||||
uses.prompt_tokens_details = { cached_tokens = chunk.usage.prompt_cache_hit_tokens }
|
||||
end
|
||||
end
|
||||
|
||||
if not chunk.choices or #chunk.choices == 0 then
|
||||
if uses ~= nil then
|
||||
-- usage-only chunk is not a content/finish signal; the gateway
|
||||
-- emits its own terminal stop chunk and merges this usage.
|
||||
return json.encode({ usage = uses, done = false })
|
||||
end
|
||||
return ""
|
||||
end
|
||||
local delta = chunk.choices[1].delta or {}
|
||||
local fr = chunk.choices[1].finish_reason
|
||||
|
||||
local finish = (type(fr) == "string" and fr ~= "") and fr or nil
|
||||
|
||||
local unified = {
|
||||
content = delta.content or "",
|
||||
done = (fr ~= nil)
|
||||
done = (finish ~= nil)
|
||||
}
|
||||
if finish then
|
||||
unified.finish_reason = finish
|
||||
end
|
||||
if uses ~= nil then
|
||||
unified.usage = uses
|
||||
end
|
||||
if delta.reasoning_content then
|
||||
unified.reasoning_content = delta.reasoning_content
|
||||
end
|
||||
@ -80,4 +127,15 @@ function adapter.transform_stream_chunk(raw_chunk)
|
||||
return json.encode(unified)
|
||||
end
|
||||
|
||||
-- 错误收敛:GitHub Models 走标准 OpenAI 信封 {error:{message,...}}
|
||||
function adapter.transform_error(status, body)
|
||||
local ok, resp = pcall(json.decode, body)
|
||||
if not ok or type(resp) ~= "table" then return nil end
|
||||
local e = resp.error
|
||||
if type(e) == "table" and type(e.message) == "string" then
|
||||
return e.message
|
||||
end
|
||||
return nil
|
||||
end
|
||||
|
||||
return adapter
|
||||
|
||||
@ -32,6 +32,18 @@ function adapter.transform_response(raw_body)
|
||||
unified.token_usage.prompt = resp.usage.prompt_tokens or 0
|
||||
unified.token_usage.completion = resp.usage.completion_tokens or 0
|
||||
unified.token_usage.total = resp.usage.total_tokens or 0
|
||||
local hit = 0
|
||||
if type(resp.usage.prompt_tokens_details) == "table" and resp.usage.prompt_tokens_details.cached_tokens ~= nil then
|
||||
hit = resp.usage.prompt_tokens_details.cached_tokens
|
||||
unified.token_usage.prompt_tokens_details = { cached_tokens = hit }
|
||||
end
|
||||
if (resp.usage.prompt_cache_hit_tokens or 0) > 0 then
|
||||
unified.token_usage.prompt_cache_hit_tokens = resp.usage.prompt_cache_hit_tokens
|
||||
unified.token_usage.prompt_cache_miss_tokens = resp.usage.prompt_cache_miss_tokens or 0
|
||||
if hit == 0 then
|
||||
unified.token_usage.prompt_tokens_details = { cached_tokens = resp.usage.prompt_cache_hit_tokens }
|
||||
end
|
||||
end
|
||||
end
|
||||
|
||||
if type(resp.choices) == "table" and #resp.choices > 0 then
|
||||
@ -62,14 +74,49 @@ end
|
||||
function adapter.transform_stream_chunk(raw_chunk)
|
||||
local ok, chunk = pcall(json.decode, raw_chunk)
|
||||
if not ok then return "" end
|
||||
if not chunk.choices or #chunk.choices == 0 then return "" end
|
||||
|
||||
-- OpenAI-style streams may attach usage to a chunk with empty choices
|
||||
-- (the final usage chunk). Keys must match Go's TokenUsage json tags
|
||||
-- (prompt/completion/total); the gateway re-emits standard *_tokens.
|
||||
local uses = nil
|
||||
if type(chunk.usage) == "table" then
|
||||
uses = {
|
||||
prompt = chunk.usage.prompt_tokens or chunk.usage.prompt or 0,
|
||||
completion = chunk.usage.completion_tokens or chunk.usage.completion or 0,
|
||||
total = chunk.usage.total_tokens or chunk.usage.total or 0,
|
||||
}
|
||||
if type(chunk.usage.prompt_tokens_details) == "table" and chunk.usage.prompt_tokens_details.cached_tokens ~= nil then
|
||||
uses.prompt_tokens_details = { cached_tokens = chunk.usage.prompt_tokens_details.cached_tokens }
|
||||
elseif (chunk.usage.prompt_cache_hit_tokens or 0) > 0 then
|
||||
uses.prompt_cache_hit_tokens = chunk.usage.prompt_cache_hit_tokens
|
||||
uses.prompt_cache_miss_tokens = chunk.usage.prompt_cache_miss_tokens or 0
|
||||
uses.prompt_tokens_details = { cached_tokens = chunk.usage.prompt_cache_hit_tokens }
|
||||
end
|
||||
end
|
||||
|
||||
if not chunk.choices or #chunk.choices == 0 then
|
||||
if uses ~= nil then
|
||||
-- usage-only chunk is not a content/finish signal; the gateway
|
||||
-- emits its own terminal stop chunk and merges this usage.
|
||||
return json.encode({ usage = uses, done = false })
|
||||
end
|
||||
return ""
|
||||
end
|
||||
local delta = chunk.choices[1].delta or {}
|
||||
local fr = chunk.choices[1].finish_reason
|
||||
|
||||
local finish = (type(fr) == "string" and fr ~= "") and fr or nil
|
||||
|
||||
local unified = {
|
||||
content = delta.content or "",
|
||||
done = (fr ~= nil)
|
||||
done = (finish ~= nil)
|
||||
}
|
||||
if finish then
|
||||
unified.finish_reason = finish
|
||||
end
|
||||
if uses ~= nil then
|
||||
unified.usage = uses
|
||||
end
|
||||
if delta.reasoning_content then
|
||||
unified.reasoning_content = delta.reasoning_content
|
||||
end
|
||||
@ -79,4 +126,15 @@ function adapter.transform_stream_chunk(raw_chunk)
|
||||
return json.encode(unified)
|
||||
end
|
||||
|
||||
-- 错误收敛:Groq 走标准 OpenAI 信封 {error:{message,...}}
|
||||
function adapter.transform_error(status, body)
|
||||
local ok, resp = pcall(json.decode, body)
|
||||
if not ok or type(resp) ~= "table" then return nil end
|
||||
local e = resp.error
|
||||
if type(e) == "table" and type(e.message) == "string" then
|
||||
return e.message
|
||||
end
|
||||
return nil
|
||||
end
|
||||
|
||||
return adapter
|
||||
|
||||
@ -64,6 +64,18 @@ function adapter.transform_response(raw_body)
|
||||
unified.token_usage.prompt = resp.usage.prompt_tokens or 0
|
||||
unified.token_usage.completion = resp.usage.completion_tokens or 0
|
||||
unified.token_usage.total = resp.usage.total_tokens or 0
|
||||
local hit = 0
|
||||
if type(resp.usage.prompt_tokens_details) == "table" and resp.usage.prompt_tokens_details.cached_tokens ~= nil then
|
||||
hit = resp.usage.prompt_tokens_details.cached_tokens
|
||||
unified.token_usage.prompt_tokens_details = { cached_tokens = hit }
|
||||
end
|
||||
if (resp.usage.prompt_cache_hit_tokens or 0) > 0 then
|
||||
unified.token_usage.prompt_cache_hit_tokens = resp.usage.prompt_cache_hit_tokens
|
||||
unified.token_usage.prompt_cache_miss_tokens = resp.usage.prompt_cache_miss_tokens or 0
|
||||
if hit == 0 then
|
||||
unified.token_usage.prompt_tokens_details = { cached_tokens = resp.usage.prompt_cache_hit_tokens }
|
||||
end
|
||||
end
|
||||
end
|
||||
if type(resp.choices) == "table" and #resp.choices > 0 then
|
||||
local ch = resp.choices[1]
|
||||
@ -95,14 +107,49 @@ end
|
||||
function adapter.transform_stream_chunk(raw_chunk)
|
||||
local ok, chunk = pcall(json.decode, raw_chunk)
|
||||
if not ok then return "" end
|
||||
if not chunk.choices or #chunk.choices == 0 then return "" end
|
||||
|
||||
-- OpenAI-style streams may attach usage to a chunk with empty choices
|
||||
-- (the final usage chunk). Keys must match Go's TokenUsage json tags
|
||||
-- (prompt/completion/total); the gateway re-emits standard *_tokens.
|
||||
local uses = nil
|
||||
if type(chunk.usage) == "table" then
|
||||
uses = {
|
||||
prompt = chunk.usage.prompt_tokens or chunk.usage.prompt or 0,
|
||||
completion = chunk.usage.completion_tokens or chunk.usage.completion or 0,
|
||||
total = chunk.usage.total_tokens or chunk.usage.total or 0,
|
||||
}
|
||||
if type(chunk.usage.prompt_tokens_details) == "table" and chunk.usage.prompt_tokens_details.cached_tokens ~= nil then
|
||||
uses.prompt_tokens_details = { cached_tokens = chunk.usage.prompt_tokens_details.cached_tokens }
|
||||
elseif (chunk.usage.prompt_cache_hit_tokens or 0) > 0 then
|
||||
uses.prompt_cache_hit_tokens = chunk.usage.prompt_cache_hit_tokens
|
||||
uses.prompt_cache_miss_tokens = chunk.usage.prompt_cache_miss_tokens or 0
|
||||
uses.prompt_tokens_details = { cached_tokens = chunk.usage.prompt_cache_hit_tokens }
|
||||
end
|
||||
end
|
||||
|
||||
if not chunk.choices or #chunk.choices == 0 then
|
||||
if uses ~= nil then
|
||||
-- usage-only chunk is not a content/finish signal; the gateway
|
||||
-- emits its own terminal stop chunk and merges this usage.
|
||||
return json.encode({ usage = uses, done = false })
|
||||
end
|
||||
return ""
|
||||
end
|
||||
local delta = chunk.choices[1].delta or {}
|
||||
local fr = chunk.choices[1].finish_reason
|
||||
|
||||
local finish = (type(fr) == "string" and fr ~= "") and fr or nil
|
||||
|
||||
local unified = {
|
||||
content = delta.content or "",
|
||||
done = (fr ~= nil)
|
||||
done = (finish ~= nil)
|
||||
}
|
||||
if finish then
|
||||
unified.finish_reason = finish
|
||||
end
|
||||
if uses ~= nil then
|
||||
unified.usage = uses
|
||||
end
|
||||
if delta.reasoning_content then
|
||||
unified.reasoning_content = delta.reasoning_content
|
||||
end
|
||||
@ -112,4 +159,15 @@ function adapter.transform_stream_chunk(raw_chunk)
|
||||
return json.encode(unified)
|
||||
end
|
||||
|
||||
return adapter
|
||||
-- 错误收敛:kimi 网关为 OpenAI 风格 {error:{message,...}}
|
||||
function adapter.transform_error(status, body)
|
||||
local ok, resp = pcall(json.decode, body)
|
||||
if not ok or type(resp) ~= "table" then return nil end
|
||||
local e = resp.error
|
||||
if type(e) == "table" and type(e.message) == "string" then
|
||||
return e.message
|
||||
end
|
||||
return nil
|
||||
end
|
||||
|
||||
return adapter
|
||||
|
||||
@ -32,6 +32,18 @@ function adapter.transform_response(raw_body)
|
||||
unified.token_usage.prompt = resp.usage.prompt_tokens or 0
|
||||
unified.token_usage.completion = resp.usage.completion_tokens or 0
|
||||
unified.token_usage.total = resp.usage.total_tokens or 0
|
||||
local hit = 0
|
||||
if type(resp.usage.prompt_tokens_details) == "table" and resp.usage.prompt_tokens_details.cached_tokens ~= nil then
|
||||
hit = resp.usage.prompt_tokens_details.cached_tokens
|
||||
unified.token_usage.prompt_tokens_details = { cached_tokens = hit }
|
||||
end
|
||||
if (resp.usage.prompt_cache_hit_tokens or 0) > 0 then
|
||||
unified.token_usage.prompt_cache_hit_tokens = resp.usage.prompt_cache_hit_tokens
|
||||
unified.token_usage.prompt_cache_miss_tokens = resp.usage.prompt_cache_miss_tokens or 0
|
||||
if hit == 0 then
|
||||
unified.token_usage.prompt_tokens_details = { cached_tokens = resp.usage.prompt_cache_hit_tokens }
|
||||
end
|
||||
end
|
||||
end
|
||||
|
||||
if type(resp.choices) == "table" and #resp.choices > 0 then
|
||||
@ -62,14 +74,49 @@ end
|
||||
function adapter.transform_stream_chunk(raw_chunk)
|
||||
local ok, chunk = pcall(json.decode, raw_chunk)
|
||||
if not ok then return "" end
|
||||
if not chunk.choices or #chunk.choices == 0 then return "" end
|
||||
|
||||
-- OpenAI-style streams may attach usage to a chunk with empty choices
|
||||
-- (the final usage chunk). Keys must match Go's TokenUsage json tags
|
||||
-- (prompt/completion/total); the gateway re-emits standard *_tokens.
|
||||
local uses = nil
|
||||
if type(chunk.usage) == "table" then
|
||||
uses = {
|
||||
prompt = chunk.usage.prompt_tokens or chunk.usage.prompt or 0,
|
||||
completion = chunk.usage.completion_tokens or chunk.usage.completion or 0,
|
||||
total = chunk.usage.total_tokens or chunk.usage.total or 0,
|
||||
}
|
||||
if type(chunk.usage.prompt_tokens_details) == "table" and chunk.usage.prompt_tokens_details.cached_tokens ~= nil then
|
||||
uses.prompt_tokens_details = { cached_tokens = chunk.usage.prompt_tokens_details.cached_tokens }
|
||||
elseif (chunk.usage.prompt_cache_hit_tokens or 0) > 0 then
|
||||
uses.prompt_cache_hit_tokens = chunk.usage.prompt_cache_hit_tokens
|
||||
uses.prompt_cache_miss_tokens = chunk.usage.prompt_cache_miss_tokens or 0
|
||||
uses.prompt_tokens_details = { cached_tokens = chunk.usage.prompt_cache_hit_tokens }
|
||||
end
|
||||
end
|
||||
|
||||
if not chunk.choices or #chunk.choices == 0 then
|
||||
if uses ~= nil then
|
||||
-- usage-only chunk is not a content/finish signal; the gateway
|
||||
-- emits its own terminal stop chunk and merges this usage.
|
||||
return json.encode({ usage = uses, done = false })
|
||||
end
|
||||
return ""
|
||||
end
|
||||
local delta = chunk.choices[1].delta or {}
|
||||
local fr = chunk.choices[1].finish_reason
|
||||
|
||||
local finish = (type(fr) == "string" and fr ~= "") and fr or nil
|
||||
|
||||
local unified = {
|
||||
content = delta.content or "",
|
||||
done = (fr ~= nil)
|
||||
done = (finish ~= nil)
|
||||
}
|
||||
if finish then
|
||||
unified.finish_reason = finish
|
||||
end
|
||||
if uses ~= nil then
|
||||
unified.usage = uses
|
||||
end
|
||||
if delta.reasoning_content then
|
||||
unified.reasoning_content = delta.reasoning_content
|
||||
end
|
||||
@ -79,4 +126,14 @@ function adapter.transform_stream_chunk(raw_chunk)
|
||||
return json.encode(unified)
|
||||
end
|
||||
|
||||
-- 错误收敛:Mistral 用顶层 {message="...", type="..."}(非嵌套 error)
|
||||
function adapter.transform_error(status, body)
|
||||
local ok, resp = pcall(json.decode, body)
|
||||
if not ok or type(resp) ~= "table" then return nil end
|
||||
if type(resp.message) == "string" then
|
||||
return resp.message
|
||||
end
|
||||
return nil
|
||||
end
|
||||
|
||||
return adapter
|
||||
|
||||
@ -53,11 +53,14 @@ function adapter.transform_response(raw_body)
|
||||
local ok, resp = pcall(json.decode, raw_body)
|
||||
if not ok then return raw_body end
|
||||
|
||||
local p = resp.prompt_eval_count or 0
|
||||
local c = resp.eval_count or 0
|
||||
local unified = {
|
||||
content = "",
|
||||
finish_reason = resp.done_reason or "",
|
||||
tool_calls = {},
|
||||
usage = { prompt = 0, completion = 0, total = 0 }
|
||||
-- key must be token_usage to match Go's UnifiedResponse json tag
|
||||
token_usage = { prompt = p, completion = c, total = p + c }
|
||||
}
|
||||
|
||||
if resp.message then
|
||||
@ -70,12 +73,43 @@ end
|
||||
function adapter.transform_stream_chunk(raw_chunk)
|
||||
local ok, chunk = pcall(json.decode, raw_chunk)
|
||||
if not ok then return "" end
|
||||
if not chunk.message then return "" end
|
||||
|
||||
-- Ollama's terminal chunk (done=true) carries token counts but may omit
|
||||
-- message; pass them through so the gateway emits real usage.
|
||||
local uses = nil
|
||||
if chunk.done then
|
||||
local p = chunk.prompt_eval_count or 0
|
||||
local c = chunk.eval_count or 0
|
||||
if p > 0 or c > 0 then
|
||||
uses = { prompt = p, completion = c, total = p + c }
|
||||
end
|
||||
end
|
||||
|
||||
-- Ollama done_reason -> OpenAI finish_reason ("length" 透传,其余归一 stop)
|
||||
local finish = nil
|
||||
if chunk.done then
|
||||
if chunk.done_reason == "length" then
|
||||
finish = "length"
|
||||
else
|
||||
finish = "stop"
|
||||
end
|
||||
end
|
||||
|
||||
if not chunk.message then
|
||||
if uses ~= nil then
|
||||
return json.encode({ content = "", done = true, finish_reason = finish, usage = uses })
|
||||
end
|
||||
return json.encode({ content = "", done = true, finish_reason = finish })
|
||||
end
|
||||
|
||||
local unified = {
|
||||
content = chunk.message.content or "",
|
||||
done = chunk.done or false
|
||||
done = chunk.done or false,
|
||||
finish_reason = finish
|
||||
}
|
||||
if uses ~= nil then
|
||||
unified.usage = uses
|
||||
end
|
||||
if chunk.message.reasoning_content then
|
||||
unified.reasoning_content = chunk.message.reasoning_content
|
||||
end
|
||||
@ -97,4 +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
|
||||
@ -67,22 +83,67 @@ function adapter.transform_stream_chunk(raw_chunk)
|
||||
local ok, chunk = pcall(json.decode, raw_chunk)
|
||||
if not ok then return "" end
|
||||
|
||||
if not chunk.choices or #chunk.choices == 0 then return "" end
|
||||
-- OpenAI-style streams may attach usage to a chunk with empty choices
|
||||
-- (the final usage chunk). Keys must match Go's TokenUsage json tags
|
||||
-- (prompt/completion/total); the gateway re-emits standard *_tokens.
|
||||
local uses = nil
|
||||
if type(chunk.usage) == "table" then
|
||||
uses = {
|
||||
prompt = chunk.usage.prompt_tokens or chunk.usage.prompt or 0,
|
||||
completion = chunk.usage.completion_tokens or chunk.usage.completion or 0,
|
||||
total = chunk.usage.total_tokens or chunk.usage.total or 0,
|
||||
}
|
||||
if type(chunk.usage.prompt_tokens_details) == "table" and chunk.usage.prompt_tokens_details.cached_tokens ~= nil then
|
||||
uses.prompt_tokens_details = { cached_tokens = chunk.usage.prompt_tokens_details.cached_tokens }
|
||||
elseif (chunk.usage.prompt_cache_hit_tokens or 0) > 0 then
|
||||
uses.prompt_cache_hit_tokens = chunk.usage.prompt_cache_hit_tokens
|
||||
uses.prompt_cache_miss_tokens = chunk.usage.prompt_cache_miss_tokens or 0
|
||||
uses.prompt_tokens_details = { cached_tokens = chunk.usage.prompt_cache_hit_tokens }
|
||||
end
|
||||
end
|
||||
|
||||
if not chunk.choices or #chunk.choices == 0 then
|
||||
if uses ~= nil then
|
||||
-- usage-only chunk is not a content/finish signal; the gateway
|
||||
-- emits its own terminal stop chunk and merges this usage.
|
||||
return json.encode({ usage = uses, done = false })
|
||||
end
|
||||
return ""
|
||||
end
|
||||
local delta = chunk.choices[1].delta or {}
|
||||
local fr = chunk.choices[1].finish_reason
|
||||
|
||||
local finish = (type(fr) == "string" and fr ~= "") and fr or nil
|
||||
|
||||
local unified = {
|
||||
content = delta.content or "",
|
||||
done = (fr ~= nil)
|
||||
done = (finish ~= nil)
|
||||
}
|
||||
if finish then
|
||||
unified.finish_reason = finish
|
||||
end
|
||||
if uses ~= nil then
|
||||
unified.usage = uses
|
||||
end
|
||||
if delta.reasoning_content then
|
||||
unified.reasoning_content = delta.reasoning_content
|
||||
end
|
||||
if delta.tool_calls then
|
||||
-- pass raw streaming fragments through; OpenAI clients accumulate index+id+name+arguments
|
||||
unified.tool_calls = delta.tool_calls
|
||||
end
|
||||
return json.encode(unified)
|
||||
end
|
||||
|
||||
-- 错误收敛:标准 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
|
||||
if type(e) == "string" then return e 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 = {}
|
||||
@ -112,16 +147,53 @@ function adapter.transform_stream_chunk(raw_chunk)
|
||||
local ok, chunk = pcall(json.decode, raw_chunk)
|
||||
if not ok then return "" end
|
||||
|
||||
if not chunk.choices or #chunk.choices == 0 then return "" end
|
||||
-- OpenAI-style streams may attach usage to a chunk with empty choices
|
||||
-- (the final usage chunk). Preserve it; the gateway emits it as the
|
||||
-- terminal usage chunk. Note: keys must match Go's TokenUsage json tags
|
||||
-- (prompt/completion/total); the gateway re-emits standard *_tokens.
|
||||
local uses = nil
|
||||
if type(chunk.usage) == "table" then
|
||||
uses = {
|
||||
prompt = chunk.usage.prompt_tokens or chunk.usage.prompt or 0,
|
||||
completion = chunk.usage.completion_tokens or chunk.usage.completion or 0,
|
||||
total = chunk.usage.total_tokens or chunk.usage.total or 0,
|
||||
}
|
||||
if type(chunk.usage.prompt_tokens_details) == "table" and chunk.usage.prompt_tokens_details.cached_tokens ~= nil then
|
||||
uses.prompt_tokens_details = { cached_tokens = chunk.usage.prompt_tokens_details.cached_tokens }
|
||||
elseif (chunk.usage.prompt_cache_hit_tokens or 0) > 0 then
|
||||
uses.prompt_cache_hit_tokens = chunk.usage.prompt_cache_hit_tokens
|
||||
uses.prompt_cache_miss_tokens = chunk.usage.prompt_cache_miss_tokens or 0
|
||||
uses.prompt_tokens_details = { cached_tokens = chunk.usage.prompt_cache_hit_tokens }
|
||||
end
|
||||
end
|
||||
|
||||
if not chunk.choices or #chunk.choices == 0 then
|
||||
if uses ~= nil then
|
||||
-- usage-only chunk is not a content/finish signal; the gateway
|
||||
-- emits its own terminal stop chunk and merges this usage.
|
||||
return json.encode({ usage = uses, done = false })
|
||||
end
|
||||
return ""
|
||||
end
|
||||
local delta = chunk.choices[1].delta or {}
|
||||
local fr = chunk.choices[1].finish_reason
|
||||
|
||||
local finish = (type(fr) == "string" and fr ~= "") and fr or nil
|
||||
|
||||
local unified = {
|
||||
content = delta.content or "",
|
||||
done = (fr ~= nil)
|
||||
done = (finish ~= nil)
|
||||
}
|
||||
if delta.reasoning_content then
|
||||
unified.reasoning_content = delta.reasoning_content
|
||||
if finish then
|
||||
unified.finish_reason = finish
|
||||
end
|
||||
if uses ~= nil then
|
||||
unified.usage = uses
|
||||
end
|
||||
-- zen 用 reasoning 字段承载推理文本(OpenAI 惯例是 reasoning_content)
|
||||
local reasoning = delta.reasoning_content or delta.reasoning
|
||||
if reasoning then
|
||||
unified.reasoning_content = reasoning
|
||||
end
|
||||
if delta.tool_calls then
|
||||
-- pass raw streaming fragments through; OpenAI clients accumulate index+id+name+arguments
|
||||
@ -130,4 +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"])
|
||||
@ -135,6 +154,39 @@ func TestOpenCodeAdapterFingerprint(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
func TestOpenCodeNormalizesRoles(t *testing.T) {
|
||||
vm := NewVM(freshAdapterDir(t))
|
||||
if err := vm.Start(); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer vm.Stop()
|
||||
// zen 上游只接受 system/user/assistant/tool/latest_reminder;
|
||||
// OpenAI 系客户端发 `developer`(Claude Code 还会用 `function`)必须归一化成 system。
|
||||
body := `{"model":"x","messages":[
|
||||
{"role":"developer","content":"you are helpful"},
|
||||
{"role":"system","content":"sys"},
|
||||
{"role":"user","content":"hi"},
|
||||
{"role":"assistant","content":"ok"},
|
||||
{"role":"tool","tool_call_id":"t1","content":"{\"ok\":true}"}
|
||||
]}`
|
||||
out, err := vm.Transform("opencode", "transform_request", body)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if strings.Contains(out, `"developer"`) {
|
||||
t.Fatalf("developer role must be normalized: %s", out)
|
||||
}
|
||||
if !strings.Contains(out, `"system"`) || !strings.Contains(out, `"user"`) {
|
||||
t.Fatalf("allowed roles must survive: %s", out)
|
||||
}
|
||||
if !strings.Contains(out, `"assistant"`) {
|
||||
t.Fatalf("assistant role must survive: %s", out)
|
||||
}
|
||||
if !strings.Contains(out, `"tool"`) {
|
||||
t.Fatalf("tool role must survive: %s", out)
|
||||
}
|
||||
}
|
||||
|
||||
func TestOpenCodeStripsMultiModalParts(t *testing.T) {
|
||||
vm := NewVM(freshAdapterDir(t))
|
||||
if err := vm.Start(); err != nil {
|
||||
@ -370,3 +422,35 @@ 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)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@ -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{}
|
||||
}
|
||||
@ -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,18 @@ 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
|
||||
for scanner.Scan() {
|
||||
line := strings.TrimSpace(scanner.Text())
|
||||
if line == "" || !strings.HasPrefix(line, "data:") {
|
||||
@ -744,9 +817,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 +843,15 @@ func (p *Provider) ChatStream(ctx context.Context, req *types.ChatRequest) (<-ch
|
||||
if err := json.Unmarshal([]byte(unified), &ck); err != nil {
|
||||
continue
|
||||
}
|
||||
if ck.Done {
|
||||
doneSent = true
|
||||
}
|
||||
if !errorOnlyChunk(ck) {
|
||||
realChunks++
|
||||
}
|
||||
chunks++
|
||||
select {
|
||||
case ch <- ck:
|
||||
case inner <- ck:
|
||||
case <-ctx.Done():
|
||||
return
|
||||
}
|
||||
@ -777,15 +862,86 @@ 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.
|
||||
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 +981,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 +1001,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
|
||||
}
|
||||
@ -901,23 +1059,70 @@ func standardSSEChunk(data string) string {
|
||||
} `json:"delta"`
|
||||
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"`
|
||||
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 {
|
||||
return ""
|
||||
}
|
||||
if len(raw.Choices) == 0 {
|
||||
if len(raw.Choices) == 0 && raw.UpstreamUsage.Total == 0 && raw.UpstreamUsage.TotalTokens == 0 {
|
||||
return ""
|
||||
}
|
||||
var usage *types.TokenUsage
|
||||
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),
|
||||
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{
|
||||
Content: raw.Choices[0].Delta.Content,
|
||||
Done: raw.Choices[0].FinishReason != nil,
|
||||
Content: func() string {
|
||||
if len(raw.Choices) > 0 {
|
||||
return raw.Choices[0].Delta.Content
|
||||
}
|
||||
return ""
|
||||
}(),
|
||||
Done: done,
|
||||
FinishReason: finish,
|
||||
Usage: usage,
|
||||
})
|
||||
return string(out)
|
||||
}
|
||||
|
||||
func truncate(s string, n int) string {
|
||||
if len(s) <= n {
|
||||
return s
|
||||
// pickFirst returns a if non-zero, else b (for usage keys that may appear in
|
||||
// either standard *_tokens or legacy short form).
|
||||
func pickFirst(a, b int) int {
|
||||
if a != 0 {
|
||||
return a
|
||||
}
|
||||
return s[:n] + "..."
|
||||
return b
|
||||
}
|
||||
|
||||
@ -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,73 @@ 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)
|
||||
}
|
||||
}
|
||||
|
||||
@ -211,6 +211,18 @@ type SourceStatus struct {
|
||||
FailCount int `json:"fail_count,omitempty"`
|
||||
BackoffUntil int64 `json:"backoff_until,omitempty"`
|
||||
Permanent bool `json:"permanent,omitempty"`
|
||||
// RecentOK / RecentErr count real gateway requests served by this source
|
||||
// within the last 300s (drives the status column so a source that is
|
||||
// actually serving traffic can never be shown as down just because a
|
||||
// probe was rate-limited by the upstream).
|
||||
RecentOK int64 `json:"recent_ok,omitempty"`
|
||||
RecentErr int64 `json:"recent_err,omitempty"`
|
||||
// AvgFirstByteMs is the mean time-to-first-byte (ms) for recently
|
||||
// successful requests served by this source (0 when unmeasured).
|
||||
AvgFirstByteMs int64 `json:"avg_first_byte_ms,omitempty"`
|
||||
// AvgTokPerS is the aggregate completion throughput (tokens/s) for
|
||||
// recently successful requests served by this source (0 when no data).
|
||||
AvgTokPerS int64 `json:"avg_tok_per_s,omitempty"`
|
||||
}
|
||||
|
||||
// ProbeAll runs a live reachability check for every provider (in parallel).
|
||||
@ -260,4 +272,4 @@ func (r *Registry) Status() []SourceStatus {
|
||||
out = append(out, s)
|
||||
}
|
||||
return out
|
||||
}
|
||||
}
|
||||
|
||||
@ -327,6 +327,63 @@ 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
|
||||
}
|
||||
sort.SliceStable(cands, func(i, j int) bool {
|
||||
return cands[i].Prov.Pref(cands[i].Model) > cands[j].Prov.Pref(cands[j].Model)
|
||||
})
|
||||
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
|
||||
|
||||
@ -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,60 @@ 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
|
||||
// by the internal unified representation and older clients) and the OpenAI
|
||||
// standard keys (prompt_tokens/completion_tokens/total_tokens). Standard
|
||||
// clients such as DSH and DevEco Code read the *_tokens fields.
|
||||
func (t TokenUsage) MarshalJSON() ([]byte, error) {
|
||||
// Auto-generate prompt_tokens_details from legacy DeepSeek fields when
|
||||
// the upstream adapter only set the standalone hit count (openai.lua
|
||||
// does this in Lua, but other adapters or the standardSSEChunk fallback
|
||||
// may not). dsh reads prompt_tokens_details.cached_tokens first.
|
||||
pdetails := t.PromptTokensDetails
|
||||
if pdetails == nil && t.PromptCacheHit > 0 {
|
||||
pdetails = &PromptTokensDetails{CachedTokens: t.PromptCacheHit}
|
||||
}
|
||||
return json.Marshal(struct {
|
||||
PromptTokens int `json:"prompt_tokens"`
|
||||
CompletionTokens int `json:"completion_tokens"`
|
||||
TotalTokens int `json:"total_tokens"`
|
||||
Prompt int `json:"prompt"`
|
||||
Completion int `json:"completion"`
|
||||
Total int `json:"total"`
|
||||
PromptTokensDetails *PromptTokensDetails `json:"prompt_tokens_details,omitempty"`
|
||||
PromptCacheHitTokens int `json:"prompt_cache_hit_tokens,omitempty"`
|
||||
PromptCacheMissTokens int `json:"prompt_cache_miss_tokens,omitempty"`
|
||||
}{
|
||||
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,
|
||||
})
|
||||
}
|
||||
|
||||
type ImageData struct {
|
||||
@ -114,8 +170,17 @@ 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
|
||||
// it (OpenAI-style streams attach usage to some chunks; the final one
|
||||
// often has empty choices). Gateway uses it to emit exact usage in the
|
||||
// final stream chunk instead of estimates.
|
||||
Usage *TokenUsage `json:"usage,omitempty"`
|
||||
}
|
||||
|
||||
// Meta passed to Lua build_headers hook
|
||||
@ -128,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