diff --git a/.env.example b/.env.example index c6032d2..246b5b6 100644 --- a/.env.example +++ b/.env.example @@ -41,6 +41,9 @@ OT_PROXY_FAILOVER_BUDGET=30s # 流式空闲多久补一行 SSE 注释心跳(0=关闭)。挂在 nginx/Cloudflare 后面时, # 上游长时间不出字(长思考)会被对端的空闲超时掐掉,心跳用来续住这条流。 OT_PROXY_STREAM_HEARTBEAT=15s +# 流式在放行给客户端之前最多等内容多久(0=不限):超时未出内容就取消本次尝试、 +# 换下一个渠道。上游挂起(连接正常但一直不吐字)时的兜底。 +OT_PROXY_FIRST_TOKEN_WAIT=0s # 渠道健康检查 OT_PROXY_HEALTH_INTERVAL=60s diff --git a/README.md b/README.md index 7885a1d..336de32 100644 --- a/README.md +++ b/README.md @@ -13,9 +13,9 @@ - `GET /v1/models` — 可用模型列表:只列调用方真能用的模型(叠加用户与密钥白名单),带上下文长度与能力位;Anthropic 客户端(带 `x-api-key`/`anthropic-version`)返回 Anthropic 形状 - **三协议互转**:客户端协议 × 渠道协议不匹配时自动转换(如 Chat 调用 Claude、Messages 调用 OpenAI、Responses 调用 Claude),流式逐事件转换;协议匹配时直通 - 错误按客户端协议返回(OpenAI 格式 / Anthropic 格式) -- **渠道体系**:按模型绑定选渠道 + 加权负载均衡;模型已有绑定时只用绑定渠道(都不可用则 503,不会把请求发给不认识它的渠道);每渠道并发信号量(满载溢出);后台健康检查(探测 `/models`,连续失败进 cooldown、恢复放回)+ 运行期健康(连续 3 次网络错误/超时/429/5xx 进 30s 冷却,成功即清零);可安全重试的失败自动故障转移(网络错误/429/5xx/超时且未写出响应头) +- **渠道体系**:按模型绑定选渠道 + 加权负载均衡;模型已有绑定时只用绑定渠道(都不可用则 503,不会把请求发给不认识它的渠道);每渠道并发信号量(满载溢出);后台健康检查(探测 `/models`,连续失败进 cooldown、恢复放回)+ 运行期健康(按失败原因分类冷却:网络错误/5xx 连续 3 次进冷却且指数退避,余额/配额类第一次即冷却 30 分钟,成功清零退避计数);可安全重试的失败自动故障转移(网络错误/429/5xx/超时/流式首个内容事件之前的任何失败) - **上游超时**:渠道 `timeout_ms` 生效——非流式为整段请求上限,流式仅约束首包(正文不限时长);上游请求绑定客户端 context,断连即取消 -- **故障转移与流式细节**:上游 429 的 `Retry-After` 会延长该渠道冷却(夹在 10 分钟内);重试受 `OT_PROXY_FAILOVER_BUDGET` 时间预算约束(首个渠道不受限,避免坏渠道把客户端拖到 N×timeout);流式空闲超过 `OT_PROXY_STREAM_HEARTBEAT` 补一行 SSE 注释心跳,防止 nginx/Cloudflare 的空闲超时掐掉长思考 +- **故障转移与流式细节**(借鉴 magpie 的 holdWriter):流式请求在**首个内容事件**到达前由网关缓冲——上游的错误事件、空流、断连、协议形态不符(要流给了整包,didNotStream)都在放行前换渠道消化,客户端无感知;`OT_PROXY_FIRST_TOKEN_WAIT` 限定放行前等内容的最长时间(0=不限);429 不一定是限流——错误体带余额/配额特征时按配额类长冷却;上游 `Retry-After` 会延长该渠道冷却(夹在 10 分钟内);重试受 `OT_PROXY_FAILOVER_BUDGET` 时间预算约束(首个渠道不受限,避免坏渠道把客户端拖到 N×timeout);流式空闲超过 `OT_PROXY_STREAM_HEARTBEAT` 补一行 SSE 注释心跳(扣住期间也照发),防止 nginx/Cloudflare 的空闲超时掐掉长思考;响应头已被心跳先行写出时,全部候选失败以 SSE 错误事件收尾 - **渠道观测**:`GET /api/v1/admin/channels/stats?window=5m|1h|6h|24h|7d` 按渠道聚合请求量、成功率、错误码分布与 P50/P95 延迟;渠道页直接展示近况与运行期冷却剩余时间,可一键恢复 - **限流/配额**(内存计数):密钥级每日请求数 / Token 数配额、用户级每秒速率(`OT_RATELIMIT_USER_RPS`);超限返回 429 - **用量计费**:请求级 `usage_logs` 异步批量落库,按模型价格扣减余额,日粒度预聚合(`usage_daily`);模型可单独配置缓存写价(留空按输入价 ×1.25) diff --git a/docs/magpie-plugin-mechanism.md b/docs/magpie-plugin-mechanism.md new file mode 100644 index 0000000..854b860 --- /dev/null +++ b/docs/magpie-plugin-mechanism.md @@ -0,0 +1,273 @@ +# magpie 插件机制深度调研 + +> 调研对象:[yetone/magpie](https://github.com/yetone/magpie)(本地 LLM 网关,Go,~855K 行)。 +> 本文聚焦其插件机制本身,目标是把可移植的设计提炼出来供其他项目参考。所有行号基于 2026-10-30 的 main 分支。 + +## 0. 总览:三层并存的「插件」 + +magpie 里「插件」一词覆盖三个完全不同的机制,看懂分层是看懂代码的前提: + +| 层 | 位置 | 形态 | 运行环境 | 典型用途 | +| --- | --- | --- | --- | --- | +| 内置插件(built-in) | `plugins/` 包 + `internal/pluginrt` | 编译进二进制的 Go 包,用户可开关 | 网关主进程 | routing、usage、redact、sessions 等功能开关 | +| JS 网关中间件 | `internal/middleware` | 用户安装的 JS 模块 | 进程内 moejs 引擎(Go 写的 JS 引擎) | 改写请求/响应、路由策略、用量观察 | +| Provider 插件 | `internal/plugin`(`host.js`) | OpenCode v1 兼容的 npm/git/本地包 | Bun 子进程,stdin/stdout JSON-RPC | 供应商接入:登录、模型列表、代发请求 | + +三层的调用频率完全不同,这决定了它们的技术选型(见 §4.1)。 + +--- + +## 1. 内置插件系统(pluginrt) + +### 1.1 设计动机 + +包注释(`internal/pluginrt/pluginrt.go:1-10`)说得直白:magpie 的功能「不是所有人都需要」(Lxxyx on X: "功能太多了"),所以把非核心功能拆成可开关的包。关键约束写在注释里并**由测试强制**: + +- 运行时是 core,**永不 import 任何插件**; +- 插件之间**互相不 import**,只能通过彼此 Provide 的服务通信。 + +这两条不是 review 纪律,而是 `TestPluginsImportOnlyCore`(`internal/pluginrt/pluginrt_test.go:94`):用 AST 扫描全部 `.go` 文件的 import,违反即编译期红。架构边界测试化,比文档和口头约定可靠得多。 + +### 1.2 Manifest:插件的全部声明 + +`plugins/` 包在 `init()` 里调 `pluginrt.Register(Manifest)`(`pluginrt.go:128`,重复 ID 直接 panic)。Manifest 字段(`pluginrt.go:35-72`): + +```go +type Manifest struct { + ID string + API string // 宿主 API 版本区间,caret 语义 "^1" + Requires, Provides []string // 服务名:插件间唯一的依赖/通信表达 + Essential bool // essentials 预设是否保留 + Name, Sub string // 展示名(窗口负责翻译) + One, Many string // "1 个分组在用" / "{n} 个分组在用" + InUse func() (int, error) // 运行时清点"有多少东西在用它" + Apply func(*Ctx) error // 注册钩子;nil 表示纯 UI 插件 + Parts []Part // 前端 UI 插槽:nav.page / settings.section / part + Defaults map[string]any + Settings any // 设置结构体 → 反射生成 JSON Schema → 合并进 Defaults + Client []string // 窗口侧资源(脚本/i18n/css) +} +``` + +三个值得单独说的设计: + +**InUse 运行时清点**(`pluginrt.go:48-51`)。插件能否被关掉不靠用户记忆,而是运行时数「还有多少用户创建的东西依赖它」(分组、预热任务……)。InUse > 0 时无视 plugins.yml 强制开启;**读取出错也算在用**——"unknown is never unused",宁可多开不少开,避免用户的东西失去宿主却还在运行。 + +**Settings → JSON Schema**(`pluginrt.go:137-150` + `schema.go`)。Register 时反射设置结构体生成 schema 并把零值合并进 Defaults;窗口用 schema 画表单,`SetConfig` 只接受 schema 允许的字段——用户手改的非法值被拒收但保留在文件里,不会静默丢失。 + +**API 版本 caret 区间**(`pluginrt.go:391-402`)。`Compatible("^1.2")` = 主版本相同且 minor ≤ 宿主。不匹配的插件不 apply,状态记为 Failed 并说明原因。 + +### 1.3 生命周期:Reconcile 不动点状态机 + +唯一的入口是 `Runtime.Reconcile(on, config)`(`pluginrt.go:211-278`),每次配置变化都整体收敛一遍: + +1. **先关**:迭代求出「不再 wanted」∪「Requires 的服务被某个要关插件 Provides」的闭包(`needsGoing`,:281-292),然后**按 apply 的逆序** dispose——保证服务在它的使用者之后才被撤回。每个 dispose 闭包单独 recover(`undo`,:342-354),一个插件清理时 panic 不影响其他插件。 +2. **后开**:循环「能应用就应用」直到不动点——Requires 全部满足才 apply(core 提供 `settings/providers/gateway/window` 四个基础服务,:31),不满足则挂 Pending 并记录在等谁。这个不动点循环天然实现了**按服务依赖图的拓扑序 apply、逆拓扑序 dispose**,不需要声明式 priority。 +3. 状态集:`off / pending / active / failed / disposed`(:83-89),GUI 可见每个插件为什么没起来(`Why` 字段)。 + +Apply 的 panic 被 recover 成该插件 FAILED(:308-330),**它注册了一半的东西也全部 undo**——不存在「半应用」状态。 + +### 1.4 Ctx:注册即登记撤销 + +`Apply` 拿到的 `Ctx`(`pluginrt.go:93-119`)只有四个能力,全部自带撤销: + +- `OnDispose(f)`:登记清理函数,关插件时逆序执行; +- `Contribute(slot, key)`:往窗口放一块 UI; +- `Provide(name, v)`:提供 named service,撤销时从 services map 删除; +- `Service(name)`:取别的插件 Provide 的服务。 + +配套的核心模式是「**注册函数返回 remove 闭包**」:每个扩展点的 `Use/Serve/Register` 都返回 `remove func()`,插件把它交给 `Ctx.OnDispose`,一行完成即插即拔。remove 闭包统一用 `sync.Once` 包裹防重复(`gwmw.go:141-154` 的泛型 `add[T]` 是模板实现)。 + +### 1.5 配置系统 + +**plugins.yml 三层合并**(`pluginrt/config.go:126-175` `Compose`):内置 Defaults → preset(essentials/all)→ 用户行。每字段记录来源 `Source`(default/preset/user/in use),`magpie plugin dump` 能解释「这个值从哪来」——排查配置不生效的利器。 + +**热更新三件套**(`config.go:363-395` + `gateway.go:649-650`): + +1. `filememo`:按 (mtime, size, 原始字节) 记忆解析结果,本进程写文件后经 channel 立即失效; +2. fsnotify 看其他进程的写入; +3. 轮询兜底; +4. 再加一条「**请求路径按需拉起**」:请求到达时发现「需要 masker/limiter/中间件但没有」,节流(≤1 次/秒)触发 `pluginrt.Refresh` 立刻重新收敛(`gateway.go:631-644`)。效果是配置改了不用重启,稳态开销近乎为零。 + +**settings 读取的健壮性**(`internal/settings/settings.go` + `internal/filememo`):解析结果按原始字节缓存 + `Shared()` 免拷贝只读实例;损坏文件回退 lastgood(最后一次完好的代);「读不了」与「没有文件」语义分开(`LoadChecked`);写入走全量校验 + temp+rename 原子替换 + `.bak`。 + +--- + +## 2. 网关扩展点(core 侧) + +core 为插件预留了五个「点」,每个是一个独立小包,统一模式:**接口 + 注册返回 remove 闭包 + 堆叠时后注册者生效 + 显式的失败策略**。 + +### 2.1 中间件点(gwmw)——同一包内两种失败策略 + +`internal/gwmw/gwmw.go` 包注释(1-16 行)开门见山:这个包有两个位置,**各有各的失败策略**: + +- **Agent 中间件**(fail-open):挂在客户端请求进网关路由之前。接口(:45-59): + + ```go + type Agent interface { Begin(Info) Run } + type Run interface { + Request(body []byte) ([]byte, *Reject) // 改写或拒绝 + Wrap(w http.ResponseWriter) (http.ResponseWriter, func()) // 包响应写手 + End() + } + ``` + + 中间件挂了/超时了 → 请求原样放行。可选接口 `Tried(body, To) []byte`(:72-74)让网关**每次换上游重试前重新变形 body**,每次 try 从 onRequest 后的 body 出发、互不污染。 +- **Masker**(fail-closed):挂在发往供应商之前。`Masking(wanted)` 在「设置要求遮罩但没有 masker 在服务」时返回 `ErrUnmasked` 拒发(:269-277),错误信息直接告诉用户去哪里打开插件——**错误写到用户能行动的层面**。`Safely`(:281-288)把 masker 的 panic 也转成错误,绝不把未遮罩的请求发出去。 + +另有一个精巧的「拉起」机制:`Served()` 发现没有 masker 时节流地回调 refresh(最多 1 秒一次、不重入,`throttle` :194-212),于是「设置打开了遮罩」会在下一个请求时自动把 redact 插件拉起来。 + +### 2.2 路由策略点(routes)——只排序,不决策 + +`internal/routes/routes.go`:`Strategy` 只做一件事——给出候选池的**排列**(`Order func(pool string, seats []Seat) []int`)。resting、亲和、故障转移全部留在网关核心,策略碰不到。`Apply`(routes.go 尾部)对策略输出做校验:panic、超时、或者不是合法排列,就 fail-open 回退到列表顺序并打日志。`Seat`(:29-72)携带策略可能需要的全部事实:使用率、各窗口重置时间、衰减 token 数、权重、是否在冷却、实测 TTFT/TPS、钱包余额……**策略是无状态函数,事实由核心喂给它**。 + +### 2.3 预算点(gwbudget)——fail-closed 的限流 + +`Limiter.Reserve(who, body, model, now) (release, Refusal)`(`gwbudget.go:33-35`)。key 配了 limit 但没有 limiter 在服务时返回 `ErrUnlimited` 拒绝(:106-124)——计费限额和遮罩一样 fail-closed:**该拦的没拦住,比误拦严重**。 + +### 2.4 观察点(usage.Observer)——按需声明开销 + +`internal/usage/observe.go:19-31`: + +```go +type Observer interface { + Serve() (stop func()) // 常驻 goroutine 的生命周期 + Bodies() Bodies // None / 256KiB / Whole + Begin(Request) Watch // Attempt(每此尝试) + End(整请求) + Ledger(Record) // 账本每一行 +} +``` + +精髓在 `Bodies()`:**观察者自己声明要多少数据**,网关据此决定要不要为它缓冲完整流(`ObservedBodies` :146-152 取最大值);没有观察者注册时,网关连整包缓冲都不做,「不做任何观察工作,连回复的 body 都不留」——零观察者零开销。 + +### 2.5 工具点(gwtool) + +极简的按名开关:`Gate()` 网关启动时全局关闸、`Serve(name)`/`Served(name)` 按名放行(`gwtool.go:34-72`)。search 插件供 `web_search`、tools 插件供 `plugin_tools`,就这么点东西——**能用一个 gate 表达的就不用接口**。 + +--- + +## 3. JS 中间件引擎(moejs,进程内) + +### 3.1 为什么不用 Bun:按调用频率选引擎 + +包注释(`internal/middleware/middleware.go:1-12`)给了决策依据:一个流式回复有数千个事件,**进程内 moejs 跑一个 hook 约 1µs**,Bun 子进程一个来回要数十 µs——每事件都走子进程是不可接受的。于是: + +- **高频钩子**(每请求、每 SSE 事件)→ 进程内 moejs; +- **低频异步钩子**(可 await I/O、观察完整往返)→ Bun 子进程(见 §4)。 + +这是「插件热路径留在进程内、重 I/O 放子进程」的教科书分层。 + +### 3.2 运行时管理 + +- 每个中间件**编译一次**成模块;每个请求从 `sync.Pool` 取独立 runtime(`middleware.go:44-71`),请求的多个 hook 共享 `ctx.state`; +- 超时硬限:请求级 hook 250ms、事件级 50ms、加载 1s(:35-39),超时用 `runtime.Interrupt` 打断,输入保持原样; +- 热加载用 stamp(plugins.json mtime + 各中间件文件 stat),**每请求最多重查一次/秒**(`current()` :95-116),热路径几乎零开销。 + +### 3.3 四个 hook 与两个扩展导出 + +| Hook | 签名 | 语义 | +| --- | --- | --- | +| `onRequest` | `(body, ctx) → body \| undefined` | 改写入站请求;`ctx.reject()` 可拒;undefined 原样放行 | +| `onProvider` | `(body, ctx) → body` | 每次换上游重试前调用,`ctx.provider` 告诉你是谁;返回 null 报错(已路由不能丢) | +| `onEvent` | `(data, ctx)` | SSE 逐事件;`export const events` 可白名单事件名 | +| `onResponse` | `(body, ctx)` | 整段 JSON 响应(缓冲后调用) | + +两个额外导出(`points.go:1-10`,都跑在模块自己的单线程 runtime 里,互相可见): + +- **`strategies`**:导出路由策略函数 `(pool, seats, options) → 顺序`,注册成 routes.Strategy。在请求路径上跑,20ms 超时、必须返回座位下标的排列,否则 fail-open 回退原顺序; +- **`onUsage`**:账本行观察者,**跑在离线队列**——容量 1024,满了丢弃并计数,绝不阻塞请求路径(`Ledger` 的 select/default,points.go:376-398);row 是白名单化的 `usageRow`,**没有 body、prompt、错误文本、会话 id**——第三方代码拿不到敏感数据。 + +防第三方代码拖垮网关的完整清单:超时 Interrupt、日志限速 60 行/分钟(`middleware.go:383-404`)、onUsage 满队丢弃并计数、钩子 context 白名单化(无鉴权头)、策略必须返回合法排列。 + +--- + +## 4. Provider 插件宿主(Bun 子进程) + +### 4.1 线协议(`internal/plugin/host.js:1-20`) + +stdin/stdout,一行一个 JSON: + +``` +magpie → host: {id, method, params} host 应答 {id, result} 或 {id, error} +``` + +fetch 类请求的流式应答分三段:`{id, event:"head", status, headers}` → 多个 `{id, event:"chunk", data(base64)}` → `{id, result:null}`。`{method:"abort"}` 取消单个请求。host 主动上报:`log`/`toast`(插件想展示的)、`auth`(插件保存或刷新了登录态)。 + +两个值得抄的细节: + +- **credit 窗口背压**(`host.js:22-38` 一带):每个 fetch 有字节窗口,读端消费后用 `{method:"credit", id, n}` 归还额度——**慢读端会暂停上游,而不是让宿主无界缓冲**。子进程 IPC 背压的标准解法。 +- **多账号槽位**:第一个账号存 provider id 下,其余存 `id#`;请求内的一切(getAuth、auth.set、loader)通过 `AsyncLocalStorage` 只见自己的账号作用域。 + +### 4.2 入口声明与异步请求钩子 + +Provider 插件走 OpenCode v1 约定;网关异步钩子另有两个入口(`internal/plugin/gateway.go:15-45`):package.json 的 `magpie.hooks` 字段,或 `*.hooks.{js,mjs,ts,mts}` 文件(Bun 直接跑 TS)。与中间件的 `*.gateway.js` **入口名刻意不重叠**,一个文件永远不会被当两种东西加载。 + +钩子协议是 begin/finish/close 三段(`gateway.go:174-260`):Begin 返回改写后的 body 和 Run 句柄,Finish 异步通知回复,Close 释放;总预算 11 秒;**传给插件的 context 刻意剥离鉴权头**(:82-92)。 + +### 4.3 供应链 + +Bun 本体按平台下载,SHA-256 钉死在代码里(对照 Bun 官方 SHASUMS256.txt 校验);npm 没有的包从社区仓库 `dist` 分支拉,sha512 与代码同源——信任模型是「谁有权限推 dist 分支」而非哈希本身,这个判断写在 LESSONS.md 里。插件更新有 `min` 版本机制,内置代码里记录着「要求插件 ≥ 某版」。 + +--- + +## 5. 设计原则提炼(可移植清单) + +1. **架构边界测试化**。core 不依赖插件、插件互不依赖,写成 AST 扫描测试。一条测试顶一百条 review 约定。 +2. **注册表 + 服务依赖 + 逆序 dispose**。依赖用服务名表达(`Requires/Provides`),Reconcile 不动点自动拓扑排序,不需要 priority 字段;关停永远逆序,清理函数各自 recover。 +3. **remove 闭包 + sync.Once** 是唯一的注册/注销协议。所有扩展点长一个样,学一次就会用。 +4. **每个扩展点显式声明 fail-open 还是 fail-closed**,并把错误信息写到用户能行动的层面。判断标准:功能性增强(改写、路由、观察)fail-open;安全/计费/隐私(遮罩、限额)fail-closed。 +5. **按调用频率选执行引擎**。热路径(每事件 µs 级)留进程内,重 I/O 放子进程,两者入口命名隔离。 +6. **给第三方代码上全套缰绳**:超时 Interrupt、日志限速、队列满丢弃计数、context 白名单、输出结构校验。宁可功能降级,不让插件拖垮宿主。 +7. **观察者按需声明数据需求**(Bodies 分级),零注册零开销。 +8. **InUse 运行时清点**决定插件可否关闭,错误也算在用。 +9. **配置四层合并 + 来源追踪 + lastgood 兜底 + temp+rename 原子写**;热更新 = 文件记忆缓存 + fsnotify + 轮询兜底 + 请求路径按需拉起。 +10. **策略无状态化**:策略只给排列,事实(冷却、用量、延迟)由核心以 `Seat` 喂入,输出非法就 fail-open。 + +## 6. 最小移植骨架 + +如果要在别的 Go 项目里复刻这套机制,最小核心约 600 行: + +```go +// 1. 运行时(~200 行) +type Manifest struct { + ID string + Requires, Provides []string + Apply func(*Ctx) error +} +type Ctx struct { /* id + disposes + services */ } +func (c *Ctx) OnDispose(f func()) { c.disposes = append(c.disposes, f) } +func (rt *Runtime) Reconcile(on map[string]bool) { + // 先关:wanted=false ∪ Requires 命中正在关闭者 Provides 的闭包,逆序 dispose + // 后开:不动点循环,Requires 满足才 apply,apply panic → undo + Failed +} + +// 2. 扩展点(每个 ~30 行) +func Use(h Hook) (remove func()) { // 泛型 add[T] + // append 到注册表;返回 sync.Once 包裹的 remove +} + +// 3. 边界测试(~50 行) +// AST 扫描:core 包不得 import plugins/*;plugins/* 互不 import +``` + +再按需加:配置三层合并(yaml + Source 追踪)、一个 fail-open 的请求中间件点、一个 fail-closed 的计费点。Bun/JS 宿主只在确实需要第三方生态时再引入。 + +## 7. 关键文件索引 + +| 文件 | 内容 | +| --- | --- | +| `internal/pluginrt/pluginrt.go` | 注册表、Manifest、Reconcile 状态机、API 兼容 | +| `internal/pluginrt/config.go` | plugins.yml 三层合并、热更新、原子写 | +| `internal/pluginrt/schema.go` | Settings → JSON Schema | +| `internal/pluginrt/pluginrt_test.go` | `TestPluginsImportOnlyCore`(AST 边界测试) | +| `internal/gwmw/gwmw.go` | 中间件点:Agent(fail-open)/ Masker(fail-closed) | +| `internal/routes/routes.go` | 路由策略点:Seat 与 Strategy | +| `internal/gwbudget/gwbudget.go` | 预算点(fail-closed) | +| `internal/usage/observe.go` | 观察点:Bodies 分级、Observer 生命周期 | +| `internal/middleware/middleware.go` | moejs 引擎、编译缓存、热加载 stamp | +| `internal/middleware/run.go` | 每请求 runtime、hook 派发 | +| `internal/middleware/points.go` | strategies / onUsage 扩展导出、离线队列 | +| `internal/plugin/host.js` | Bun 宿主:JSON 行协议、credit 背压、账号作用域 | +| `internal/plugin/gateway.go` | 异步请求钩子入口与三段协议 | +| `docs/subsystems/plugin-runtime.md` | 官方设计文档(P3/P4 扩展点清单) | +| `docs/subsystems/gateway-middleware.md` | 中间件点官方文档 | diff --git a/server/internal/api/admin_channels.go b/server/internal/api/admin_channels.go index 6dd1647..f150a22 100644 --- a/server/internal/api/admin_channels.go +++ b/server/internal/api/admin_channels.go @@ -30,28 +30,28 @@ func (h *Handler) AdminChannels(c *gin.Context) { } else if err == nil { masked = "****" } - fails, cooldownSec := h.channelRuntime(ch.ID) + fails, cooldownSec, reason := h.channelRuntime(ch.ID) out = append(out, gin.H{ "id": ch.ID, "name": ch.Name, "provider": ch.Provider, "formats": ch.FormatsEffective(), "base_url": ch.BaseURL, "base_urls": ch.BaseURLs, "api_key_masked": masked, "weight": ch.Weight, "priority": ch.Priority, "timeout_ms": ch.TimeoutMS, "max_concurrency": ch.MaxConcurrency, "health_status": ch.HealthStatus, "enabled": ch.Enabled, - "runtime_fails": fails, "cooldown_seconds": cooldownSec, + "runtime_fails": fails, "cooldown_seconds": cooldownSec, "cooldown_reason": reason, "created_at": ch.CreatedAt, }) } resp.OK(c, gin.H{"items": out}) } -// channelRuntime 渠道运行期健康:连续失败次数与冷却剩余秒数(0 表示无冷却)。 +// channelRuntime 渠道运行期健康:连续失败次数、冷却剩余秒数(0 表示无冷却)与最近失败原因。 // 运行期冷却只存在于进程内存,DB 的 health_status 看不到它,故列表单独返回。 -func (h *Handler) channelRuntime(id uint64) (fails int, cooldownSeconds int64) { +func (h *Handler) channelRuntime(id uint64) (fails int, cooldownSeconds int64, reason string) { if h.gw == nil { - return 0, 0 + return 0, 0, "" } - fails, remaining := h.gw.ChannelRuntime(id) - return fails, int64((remaining + time.Second - time.Nanosecond) / time.Second) // 向上取整 + fails, remaining, reason := h.gw.ChannelRuntime(id) + return fails, int64((remaining + time.Second - time.Nanosecond) / time.Second), reason // 向上取整 } type channelBody struct { diff --git a/server/internal/channel/channel.go b/server/internal/channel/channel.go index ea95101..3bf4e81 100644 --- a/server/internal/channel/channel.go +++ b/server/internal/channel/channel.go @@ -17,8 +17,38 @@ import ( var ErrNoChannel = errors.New("no available channel") -// maxCooldown 运行期冷却上限:上游 Retry-After 再长也不让渠道缺席超过这个时间。 -const maxCooldown = 10 * time.Minute +const ( + // maxCooldown 运行期冷却上限:上游 Retry-After 再长也不让渠道缺席超过这个时间。 + maxCooldown = 10 * time.Minute + // quotaCooldown 余额/配额耗尽的一次性冷却时长:这类失败不会自己恢复, + // 冷却到上限也没意义,半小时后渠道大概率已被管理员处理或配额已重置。 + quotaCooldown = 30 * time.Minute +) + +// FailKind 上游失败原因分类。冷却策略随原因不同: +// 余额/配额类第一次失败就长冷却(重试无意义);其余按指数退避。 +type FailKind uint8 + +const ( + FailNet FailKind = iota // 网络错误、超时、连接失败 + FailServer // 5xx / 529 服务端错误 + FailRate // 429 限流 + FailQuota // 余额/配额耗尽(402、带配额特征的 429/403) +) + +func (k FailKind) String() string { + switch k { + case FailNet: + return "network" + case FailServer: + return "server_error" + case FailRate: + return "rate_limited" + case FailQuota: + return "quota_exhausted" + } + return "unknown" +} type Service struct { db *gorm.DB @@ -34,11 +64,14 @@ type Service struct { } // runtimeHealth 运行期健康:连续可重试失败(网络错误/超时/429/5xx)累计到阈值进入冷却。 +// trips 记录连续进冷却的次数,退避时长按它翻倍;一次成功即清零。 // 与 HealthMonitor 的探测结果相互独立——探测只看 /models 是否 200, // 覆盖不了"models 正常但对话接口一直报错"的渠道。 type runtimeHealth struct { fails int until time.Time + kind FailKind + trips int } func NewService(db *gorm.DB, enc *crypto.Encryptor) *Service { @@ -221,62 +254,78 @@ func (s *Service) UpstreamKey(ch *store.Channel) (string, error) { // --------------------------------------------------------------------------- // 运行期健康(进程内,与 /v1/models 探测独立) -// ReportFailure 记一次可重试的上游失败;连续达 failThreshold 次则该渠道进入冷却, -// 期间不出现在候选集里。冷却到期后需再攒满阈值才会被剔除。 -// cooldown>0 时用它覆盖默认时长(上游 429 带的 Retry-After),并夹到 maxCooldown -// ——不能让一个写错的重试头把渠道钉死几小时。 -func (s *Service) ReportFailure(ch *store.Channel, cooldown time.Duration) { - if cooldown > s.cooldownFor { - if cooldown > maxCooldown { - cooldown = maxCooldown - } - } else { - cooldown = s.cooldownFor - } +// ReportFailure 记一次可重试的上游失败,按原因分类冷却: +// - FailQuota(余额/配额耗尽):渠道在充值或配额重置前不会恢复,第一次失败就冷却 +// quotaCooldown,不等阈值——429 不一定是限流(智谱 429 "余额不足" 也在此列); +// - 其余原因:连续 failThreshold 次进入冷却,时长为基准冷却按连续进冷却次数的 +// 指数退避(30s→60s→120s…),坏渠道不会以固定节拍被反复打。 +// +// retryAfter>0(上游 429 带的 Retry-After)且长于退避值时覆盖退避时长,仍夹到 +// maxCooldown——不能让一个写错的重试头把渠道钉死几小时;比退避值短的 Retry-After +// 不缩短冷却,否则 1s 就把坏渠道放回轮转。 +func (s *Service) ReportFailure(ch *store.Channel, kind FailKind, retryAfter time.Duration) { s.mu.Lock() st := s.runt[ch.ID] if st == nil { st = &runtimeHealth{} s.runt[ch.ID] = st } - st.fails++ - tripped := false - if st.fails >= s.failThreshold { + st.kind = kind + if kind == FailQuota { st.fails = 0 - st.until = time.Now().Add(cooldown) - tripped = true + st.until = time.Now().Add(quotaCooldown) + s.mu.Unlock() + log.Printf("channel: %q cooled down for %s: %s", ch.Name, quotaCooldown, kind) + return } + d := s.cooldownFor << st.trips // 指数退避 + if d > maxCooldown || d <= 0 { + d = maxCooldown + } + if retryAfter > d { + d = retryAfter + if d > maxCooldown { + d = maxCooldown + } + } + st.fails++ + if st.fails < s.failThreshold { + s.mu.Unlock() + return + } + st.fails = 0 + st.trips++ + st.until = time.Now().Add(d) s.mu.Unlock() - if tripped { - log.Printf("channel: %q entered runtime cooldown (%s)", ch.Name, cooldown) - } + log.Printf("channel: %q entered runtime cooldown (%s, %s)", ch.Name, d, kind) } -// ReportSuccess 上游正常返回响应头,清空该渠道的连续失败计数。 +// ReportSuccess 上游正常返回响应,清空该渠道的连续失败与冷却退避计数。 func (s *Service) ReportSuccess(ch *store.Channel) { s.mu.Lock() if st, ok := s.runt[ch.ID]; ok { st.fails = 0 + st.trips = 0 } s.mu.Unlock() } -// RuntimeState 渠道运行期健康快照:连续失败次数与冷却剩余时间(管理后台展示用)。 +// RuntimeState 渠道运行期健康快照:连续失败次数、冷却剩余时间与最近失败原因(管理后台展示用)。 // 只读,不做过期清理——清理仍由 InCooldown 在选路时惰性完成。 -func (s *Service) RuntimeState(id uint64) (fails int, remaining time.Duration) { +func (s *Service) RuntimeState(id uint64) (fails int, remaining time.Duration, reason string) { s.mu.Lock() defer s.mu.Unlock() st, ok := s.runt[id] if !ok { - return 0, 0 + return 0, 0, "" } if !st.until.IsZero() { if d := time.Until(st.until); d > 0 { - return 0, d // 冷却中:计数在写入时已归零 + return 0, d, st.kind.String() // 冷却中:计数在写入时已归零 } - return 0, 0 // 已到期,等下次选路清理 + return 0, 0, st.kind.String() // 已到期,等下次选路清理 } - return st.fails, 0 + return st.fails, 0, st.kind.String() } // ClearRuntime 清除运行期状态(管理后台"立即恢复")。 diff --git a/server/internal/channel/channel_test.go b/server/internal/channel/channel_test.go index 6911d4c..5644bf2 100644 --- a/server/internal/channel/channel_test.go +++ b/server/internal/channel/channel_test.go @@ -174,12 +174,12 @@ func TestRuntimeCooldownThresholdAndExpiry(t *testing.T) { s.cooldownFor = 60 * time.Millisecond // 默认阈值 3:前两次失败不该踢,第三次才进冷却 - s.ReportFailure(&ch, 0) - s.ReportFailure(&ch, 0) + s.ReportFailure(&ch, FailNet, 0) + s.ReportFailure(&ch, FailNet, 0) if s.InCooldown(id) { t.Fatal("未达阈值(3 次连续)不应进冷却") } - s.ReportFailure(&ch, 0) + s.ReportFailure(&ch, FailNet, 0) if !s.InCooldown(id) { t.Fatal("连续达阈值应进冷却") } @@ -198,9 +198,9 @@ func TestRuntimeCooldownThresholdAndExpiry(t *testing.T) { t.Fatalf("恢复后应重新进候选,实际 %d", len(got)) } - s.ReportFailure(&ch, 0) + s.ReportFailure(&ch, FailNet, 0) s.ReportSuccess(&ch) - s.ReportFailure(&ch, 0) + s.ReportFailure(&ch, FailNet, 0) if s.InCooldown(id) { t.Fatal("非连续失败不应触发冷却") } @@ -214,16 +214,16 @@ func TestRuntimeStateAndClear(t *testing.T) { db.First(&ch, id) s := NewService(db, crypto.NewEncryptor("test-master-key-1234567890")) - if fails, rem := s.RuntimeState(id); fails != 0 || rem != 0 { + if fails, rem, _ := s.RuntimeState(id); fails != 0 || rem != 0 { t.Fatalf("初始状态 = %d/%s,期望 0/0", fails, rem) } - s.ReportFailure(&ch, 0) - s.ReportFailure(&ch, 0) - if fails, rem := s.RuntimeState(id); fails != 2 || rem != 0 { + s.ReportFailure(&ch, FailNet, 0) + s.ReportFailure(&ch, FailNet, 0) + if fails, rem, _ := s.RuntimeState(id); fails != 2 || rem != 0 { t.Fatalf("两次失败后 = %d/%s,期望 2/0", fails, rem) } - s.ReportFailure(&ch, 0) - fails, rem := s.RuntimeState(id) + s.ReportFailure(&ch, FailNet, 0) + fails, rem, _ := s.RuntimeState(id) if fails != 0 || rem <= 0 || rem > s.cooldownFor { t.Fatalf("冷却中 = %d/%s,期望 0/(0,%s]", fails, rem, s.cooldownFor) } @@ -232,7 +232,7 @@ func TestRuntimeStateAndClear(t *testing.T) { } s.ClearRuntime(id) - if fails, rem := s.RuntimeState(id); fails != 0 || rem != 0 { + if fails, rem, _ := s.RuntimeState(id); fails != 0 || rem != 0 { t.Fatalf("清除后 = %d/%s,期望 0/0", fails, rem) } if s.InCooldown(id) { @@ -254,26 +254,26 @@ func TestReportFailureCooldownOverrideAndClamp(t *testing.T) { s.cooldownFor = 30 * time.Second for i := 0; i < 3; i++ { - s.ReportFailure(&ch, 2*time.Minute) + s.ReportFailure(&ch, FailRate, 2*time.Minute) } - if _, rem := s.RuntimeState(id); rem < 115*time.Second || rem > 2*time.Minute { + if _, rem, _ := s.RuntimeState(id); rem < 115*time.Second || rem > 2*time.Minute { t.Fatalf("应按 Retry-After 冷却 2min,实际剩余 %s", rem) } s.ClearRuntime(id) for i := 0; i < 3; i++ { - s.ReportFailure(&ch, time.Hour) + s.ReportFailure(&ch, FailRate, time.Hour) } - if _, rem := s.RuntimeState(id); rem > maxCooldown { + if _, rem, _ := s.RuntimeState(id); rem > maxCooldown { t.Fatalf("冷却应夹到上限 %s,实际 %s", maxCooldown, rem) } // 比默认值短的 Retry-After 不缩短冷却(否则 1s 就把坏渠道放回轮转) s.ClearRuntime(id) for i := 0; i < 3; i++ { - s.ReportFailure(&ch, 5*time.Second) + s.ReportFailure(&ch, FailRate, 5*time.Second) } - if _, rem := s.RuntimeState(id); rem < 25*time.Second { + if _, rem, _ := s.RuntimeState(id); rem < 25*time.Second { t.Fatalf("短 Retry-After 不应缩短默认冷却,实际 %s", rem) } } @@ -354,3 +354,88 @@ func TestTryAcquireCapacityChange(t *testing.T) { r1() r2() } + +// 失败分类冷却:配额类第一次失败即长冷却;普通类指数退避;成功清零退避计数。 +func TestFailKindCooldownPolicy(t *testing.T) { + db := newTestDB(t) + id := seedChannel(t, db, "kind", 1, 0, 4) + var ch store.Channel + db.First(&ch, id) + + s := NewService(db, crypto.NewEncryptor("test-master-key-1234567890")) + s.cooldownFor = 30 * time.Second + + // 配额耗尽:不攒阈值,第一次失败立即冷却 quotaCooldown + s.ReportFailure(&ch, FailQuota, 0) + if !s.InCooldown(id) { + t.Fatal("配额类失败应立即进冷却") + } + if _, rem, reason := s.RuntimeState(id); reason != FailQuota.String() || rem <= 0 { + t.Fatalf("冷却原因应为 quota_exhausted 且剩余>0,实际 %s/%s", reason, rem) + } + s.ClearRuntime(id) + + // 指数退避:连续进冷却时长翻倍,封顶 maxCooldown + s.ReportFailure(&ch, FailServer, 0) + s.ReportFailure(&ch, FailServer, 0) + s.ReportFailure(&ch, FailServer, 0) // 第 1 次冷却:30s + s.mu.Lock() // 只看 trips,把 until 抹掉以便观察下一次退避 + s.runt[id].until = time.Time{} + trips1 := s.runt[id].trips + s.mu.Unlock() + if trips1 != 1 { + t.Fatalf("第 1 次冷却后 trips=1,实际 %d", trips1) + } + s.ReportFailure(&ch, FailServer, 0) + s.ReportFailure(&ch, FailServer, 0) + s.ReportFailure(&ch, FailServer, 0) // 第 2 次冷却:60s + s.mu.Lock() + d := s.runt[id].until + s.mu.Unlock() + if time.Until(d) <= 45*time.Second || time.Until(d) > 61*time.Second { + t.Fatalf("第 2 次冷却应为 60s 左右,实际 %s", time.Until(d)) + } + + // Retry-After 比退避值长时覆盖(夹上限) + s.ClearRuntime(id) + s.ReportFailure(&ch, FailRate, 2*time.Minute) + s.ReportFailure(&ch, FailRate, 2*time.Minute) + s.ReportFailure(&ch, FailRate, 2*time.Minute) + if _, rem, _ := s.RuntimeState(id); rem <= 0 { + t.Fatal("Retry-After 覆盖后应有冷却剩余") + } + + // 一次成功清零退避计数 + s.ClearRuntime(id) + s.ReportFailure(&ch, FailNet, 0) + s.ReportFailure(&ch, FailNet, 0) + s.ReportFailure(&ch, FailNet, 0) // trips=1 + s.ReportSuccess(&ch) + s.mu.Lock() + rt := s.runt[id] + fails, trips := 0, 0 + if rt != nil { + fails, trips = rt.fails, rt.trips + } + s.mu.Unlock() + if fails != 0 || trips != 0 { + t.Fatalf("成功后失败计数与退避计数都应清零,实际 fails=%d trips=%d", fails, trips) + } +} + +// 短 Retry-After 不缩短默认冷却(1s 不能把坏渠道立刻放回轮转)。 +func TestShortRetryAfterDoesNotShortenCooldown(t *testing.T) { + db := newTestDB(t) + id := seedChannel(t, db, "short-ra", 1, 0, 4) + var ch store.Channel + db.First(&ch, id) + + s := NewService(db, crypto.NewEncryptor("test-master-key-1234567890")) + s.cooldownFor = 30 * time.Second + s.ReportFailure(&ch, FailRate, time.Second) + s.ReportFailure(&ch, FailRate, time.Second) + s.ReportFailure(&ch, FailRate, time.Second) + if _, rem, _ := s.RuntimeState(id); rem < 25*time.Second { + t.Fatalf("短 Retry-After 不应缩短冷却,实际剩余 %s", rem) + } +} diff --git a/server/internal/cli/serve.go b/server/internal/cli/serve.go index 661dfa4..440aff8 100644 --- a/server/internal/cli/serve.go +++ b/server/internal/cli/serve.go @@ -38,6 +38,7 @@ func serve() error { MaxBodyMB: cfg.Proxy.MaxBodyMB, FailoverBudget: cfg.Proxy.FailoverBudget, StreamHeartbeat: cfg.Proxy.StreamHeartbeat, + FirstTokenWait: cfg.Proxy.FirstTokenWait, }) router := api.NewRouter(a, gw) diff --git a/server/internal/config/config.go b/server/internal/config/config.go index 7e34568..59ea9ab 100644 --- a/server/internal/config/config.go +++ b/server/internal/config/config.go @@ -81,6 +81,7 @@ type ProxyConfig struct { MaxBodyMB int // /v1 请求体上限(MB),超限按客户端协议返回 413 FailoverBudget time.Duration // 故障转移给"第二个及以后的渠道"留的时间预算,0=不限 StreamHeartbeat time.Duration // 流式空闲多久发一次 SSE 注释心跳,0=关闭 + FirstTokenWait time.Duration // 流式放行前等内容的最长时间,0=不限;超时未出内容换下一个渠道 HealthInterval time.Duration // 渠道健康检查周期 HealthFailThreshold int // 连续失败 N 次进 cooldown LogRaw bool // 记录管理员原始请求体+响应到 usage_logs(调试用,默认关) @@ -148,6 +149,7 @@ func Load() (*Config, error) { v.SetDefault("proxy.max_body_mb", 32) v.SetDefault("proxy.failover_budget", "30s") v.SetDefault("proxy.stream_heartbeat", "15s") + v.SetDefault("proxy.first_token_wait", "0s") v.SetDefault("proxy.health_interval", "60s") v.SetDefault("proxy.health_fail_threshold", 2) v.SetDefault("proxy.log_raw", false) @@ -197,6 +199,7 @@ func Load() (*Config, error) { MaxBodyMB: v.GetInt("proxy.max_body_mb"), FailoverBudget: v.GetDuration("proxy.failover_budget"), StreamHeartbeat: v.GetDuration("proxy.stream_heartbeat"), + FirstTokenWait: v.GetDuration("proxy.first_token_wait"), HealthInterval: v.GetDuration("proxy.health_interval"), HealthFailThreshold: v.GetInt("proxy.health_fail_threshold"), LogRaw: v.GetBool("proxy.log_raw"), diff --git a/server/internal/proxy/gateway.go b/server/internal/proxy/gateway.go index 731aefb..f3b3864 100644 --- a/server/internal/proxy/gateway.go +++ b/server/internal/proxy/gateway.go @@ -48,8 +48,10 @@ type Gateway struct { failoverBudget time.Duration // 流式空闲多久补一行 SSE 注释心跳(0=关闭)。反向代理的空闲超时会掐掉长思考的流。 streamHeartbeat time.Duration - tpMu sync.Mutex - tps map[time.Duration]*http.Transport + // 流式放行前等多久内容(0=不限制)。超时未出内容取消本次尝试换下一个渠道。 + firstTokenWait time.Duration + tpMu sync.Mutex + tps map[time.Duration]*http.Transport policyMu sync.Mutex policy modelPolicy @@ -71,8 +73,8 @@ func (g *Gateway) ResetModelPolicy() { g.policyMu.Unlock() } -// ChannelRuntime 渠道运行期健康快照:连续失败次数 + 冷却剩余时间(管理后台展示)。 -func (g *Gateway) ChannelRuntime(id uint64) (fails int, remaining time.Duration) { +// ChannelRuntime 渠道运行期健康快照:连续失败次数 + 冷却剩余时间 + 最近失败原因(管理后台展示)。 +func (g *Gateway) ChannelRuntime(id uint64) (fails int, remaining time.Duration, reason string) { return g.ch.RuntimeState(id) } @@ -164,6 +166,7 @@ type Options struct { MaxBodyMB int // /v1 请求体上限(MB) FailoverBudget time.Duration // 故障转移重试的时间预算 StreamHeartbeat time.Duration // 流式空闲心跳间隔 + FirstTokenWait time.Duration // 流式放行前等内容的最长时间(0=不限) } func NewGateway(db *gorm.DB, enc *crypto.Encryptor, rec *usage.Recorder, lim *ratelimit.Limiter, o Options) *Gateway { @@ -179,6 +182,7 @@ func NewGateway(db *gorm.DB, enc *crypto.Encryptor, rec *usage.Recorder, lim *ra maxBodyMB: o.MaxBodyMB, failoverBudget: o.FailoverBudget, streamHeartbeat: o.StreamHeartbeat, + firstTokenWait: o.FirstTokenWait, tps: map[time.Duration]*http.Transport{}, } } diff --git a/server/internal/proxy/hold.go b/server/internal/proxy/hold.go new file mode 100644 index 0000000..f67ae57 --- /dev/null +++ b/server/internal/proxy/hold.go @@ -0,0 +1,320 @@ +package proxy + +import ( + "bytes" + "encoding/json" + "net/http" + "strings" + "sync" + "sync/atomic" + "time" + + "github.com/gin-gonic/gin" + "github.com/openteam/server/internal/channel" + "github.com/openteam/server/internal/proxy/convert" +) + +// streamHold 流式容错缓冲(借鉴 magpie 的 holdWriter):首个内容事件到达之前, +// 上游的错误事件、空流、断连都还能换渠道重试——客户端只在"确定有货"之后才看到正文。 +// 扣住期间照常发 SSE 注释心跳,防反代的空闲超时掐掉长思考。 +// +// SSE 只增不删:即使响应头已因心跳先行写出(writeHead),换渠道重试也只是在同一条 +// 流上重头写事件,客户端看到的是"注释、注释、…、正确内容",无感知切换。 +// 缓冲超限则直接放行退化为普通转发——判定失效不能吞正文。 +type streamHold struct { + c *gin.Context + w http.ResponseWriter + f http.Flusher + mu sync.Mutex + head bool // 响应头已写给客户端 + passed bool // 已放行:事件直写,不再缓冲 + buf bytes.Buffer + lastN atomic.Int64 // 最近一次向客户端写出时间(unix nano),心跳据此判断空闲 +} + +// maxHoldBytes 放行前最多缓冲多少字节:超过说明"内容判定"没认出上游的数据形态, +// 再缓冲只会吞正文,直接放行。 +const maxHoldBytes = 1 << 20 + +func newStreamHold(c *gin.Context) *streamHold { + f, _ := c.Writer.(http.Flusher) + h := &streamHold{c: c, w: c.Writer, f: f} + h.lastN.Store(time.Now().UnixNano()) + return h +} + +// writeHead 把 200 + event-stream 响应头写给客户端(幂等)。写出后状态码不可再改, +// 但流本身还能继续:之后的失败以 SSE 错误事件收尾。 +func (h *streamHold) writeHead() { + h.mu.Lock() + defer h.mu.Unlock() + h.writeHeadLocked() +} + +func (h *streamHold) writeHeadLocked() { + if h.head { + return + } + h.head = true + h.c.Header("Content-Type", "text/event-stream") + h.c.Header("Cache-Control", "no-cache") + h.w.WriteHeader(http.StatusOK) +} + +// writeEvent 写一条已转换的 SSE 输出:放行前缓冲,放行后直写。 +func (h *streamHold) writeEvent(p []byte) error { + h.mu.Lock() + defer h.mu.Unlock() + if h.passed { + _, err := h.writeLocked(p) + return err + } + h.buf.Write(p) + if h.buf.Len() > maxHoldBytes { + h.releaseLocked() + } + return nil +} + +// release 放行:把缓冲写给客户端,后续事件直写。首个内容事件到达时调用。 +func (h *streamHold) release() { + h.mu.Lock() + defer h.mu.Unlock() + h.releaseLocked() +} + +func (h *streamHold) releaseLocked() { + if h.passed { + return + } + h.writeHeadLocked() + if h.buf.Len() > 0 { + _, _ = h.w.Write(h.buf.Bytes()) + h.buf.Reset() + } + if h.f != nil { + h.f.Flush() + } + h.passed = true + h.lastN.Store(time.Now().UnixNano()) +} + +func (h *streamHold) released() bool { + h.mu.Lock() + defer h.mu.Unlock() + return h.passed +} + +// headSent 响应头是否已写出(心跳先行或已放行)。此后错误只能以 SSE 事件交付。 +func (h *streamHold) headSent() bool { + h.mu.Lock() + defer h.mu.Unlock() + return h.head +} + +// writeDirect 注释/错误事件直写(需要响应头先行)。与心跳 goroutine 共用一把锁, +// gin 的 ResponseWriter 不是并发安全的。 +func (h *streamHold) writeDirect(p []byte) error { + h.mu.Lock() + defer h.mu.Unlock() + h.writeHeadLocked() + _, err := h.writeLocked(p) + return err +} + +func (h *streamHold) writeLocked(p []byte) (int, error) { + n, err := h.w.Write(p) + if err == nil && h.f != nil { + h.f.Flush() + } + h.lastN.Store(time.Now().UnixNano()) + return n, err +} + +// cancelIfHeld 只在尚未放行时调用 cancel:与 release 互斥, +// 杜绝首 token 计时器把一条已经放行的健康流掐死。 +func (h *streamHold) cancelIfHeld(cancel func()) { + h.mu.Lock() + defer h.mu.Unlock() + if !h.passed { + cancel() + } +} + +// heartbeat 距上次写出超过 interval 就补一行 ": ping"(SSE 注释行,客户端解析器忽略, +// 不影响计费与转换)。挂在 nginx/Cloudflare 后面时长思考会被空闲超时掐掉。 +func (h *streamHold) heartbeat(interval time.Duration) func() { + if interval <= 0 { + return func() {} + } + done := make(chan struct{}) + var once sync.Once + go func() { + tk := time.NewTicker(interval) + defer tk.Stop() + for { + select { + case <-done: + return + case <-tk.C: + if time.Since(time.Unix(0, h.lastN.Load())) < interval { + continue + } + // 心跳写失败说明客户端已断开;主循环下一次写出会以断连记账 + _ = h.writeDirect([]byte(": ping\n\n")) + } + } + }() + return func() { once.Do(func() { close(done) }) } +} + +// --------------------------------------------------------------------------- +// SSE 事件判定 + +// sseEventHasContent 判定一条(客户端协议形状的)SSE 数据是否携带实质内容: +// 文本、思考、工具调用都算;role/usage/ping/message_start 这类元事件不算。 +// hold 容错以"首个内容"为放行点——内容没到,上游的错误就还能换渠道消化。 +func sseEventHasContent(line []byte, proto string) bool { + s, ok := sseDataPayload(line) + if !ok { + return false + } + var m map[string]any + if json.Unmarshal([]byte(s), &m) != nil { + return false + } + switch proto { + case convert.ProtoChat: + choices, _ := m["choices"].([]any) + if len(choices) == 0 { + return false + } + c0, _ := choices[0].(map[string]any) + delta, _ := c0["delta"].(map[string]any) + if delta == nil { + return false + } + if t, _ := delta["content"].(string); t != "" { + return true + } + if t, _ := delta["reasoning_content"].(string); t != "" { + return true + } + if tc, ok := delta["tool_calls"].([]any); ok && len(tc) > 0 { + return true + } + return false + case convert.ProtoMessages: + typ, _ := m["type"].(string) + switch typ { + case "content_block_start": + cb, _ := m["content_block"].(map[string]any) + t, _ := cb["type"].(string) + // 任何块的开启都算内容(text/tool_use/thinking/server_tool_use/…) + return t != "" + case "content_block_delta": + delta, _ := m["delta"].(map[string]any) + if t, _ := delta["text"].(string); t != "" { + return true + } + if t, _ := delta["thinking"].(string); t != "" { + return true + } + if t, _ := delta["partial_json"].(string); t != "" { + return true + } + return false + } + return false + case convert.ProtoResponses: + typ, _ := m["type"].(string) + return strings.HasSuffix(typ, ".delta") || typ == "response.output_item.added" + } + // 未知协议:任何数据事件都算内容(退化为旧语义,不吞正文) + return true +} + +// sseUpstreamError 识别上游以 200+流式形式回的错误事件: +// - Anthropic:{"type":"error","error":{...}}(event: error 的 data 行) +// - OpenAI 系:data: {"error":{...}} +// - Responses:{"type":"response.failed","response":{...,"error":{...}}} +// +// 返回错误消息文本;不是错误事件时 ok=false。只在放行前判定——放行后错误 +// 属于客户端可见的回复内容,不能拿去换渠道。 +func sseUpstreamError(line []byte) (msg string, ok bool) { + s, ok := sseDataPayload(line) + if !ok { + return "", false + } + var m map[string]json.RawMessage + if json.Unmarshal([]byte(s), &m) != nil { + return "", false + } + if raw, exists := m["error"]; exists && string(raw) != "null" && len(raw) > 0 { + return streamErrMessage(raw), true + } + typ := strings.Trim(strings.TrimSpace(string(m["type"])), `"`) + switch typ { + case "error": + return "upstream stream error", true + case "response.failed": + var resp struct { + Error json.RawMessage `json:"error"` + } + if json.Unmarshal(m["response"], &resp) == nil && len(resp.Error) > 0 && string(resp.Error) != "null" { + return streamErrMessage(resp.Error), true + } + return "upstream response failed", true + } + return "", false +} + +func streamErrMessage(raw json.RawMessage) string { + var obj struct { + Message string `json:"message"` + } + if json.Unmarshal(raw, &obj) == nil && obj.Message != "" { + return obj.Message + } + var s string + if json.Unmarshal(raw, &s) == nil && s != "" { + return s + } + return "upstream stream error" +} + +// classifyUpstreamError 把一次上游失败归类为 (原因, 是否可换渠道重试)。 +// 余额/配额耗尽重试无意义,要和普通限流区分开——429 不一定是限流 +// (智谱 429 1113 "余额不足");402 原本直接透传,现在也换渠道: +// 渠道 A 没余额不代表渠道 B 没有。402 恒为配额;429/403/401 带配额字样的也是配额。 +func classifyUpstreamError(status int, body []byte) (channel.FailKind, bool) { + quota := status == http.StatusPaymentRequired || quotaSignalIn(body) + switch { + case quota && (status == http.StatusPaymentRequired || status == http.StatusTooManyRequests || + status == http.StatusForbidden || status == http.StatusUnauthorized): + return channel.FailQuota, true + case status == http.StatusTooManyRequests: + return channel.FailRate, true + case status >= 500: + return channel.FailServer, true + default: + return channel.FailServer, false + } +} + +// quotaSignalIn 从上游错误体识别"余额/配额耗尽"特征。 +var quotaSignals = []string{ + "insufficient", "balance", "billing", "arrears", + "exceeded your current quota", "quota exceeded", + "额度不足", "余额不足", "欠费", +} + +func quotaSignalIn(body []byte) bool { + s := strings.ToLower(string(body)) + for _, sig := range quotaSignals { + if strings.Contains(s, sig) { + return true + } + } + return false +} diff --git a/server/internal/proxy/hold_test.go b/server/internal/proxy/hold_test.go new file mode 100644 index 0000000..1ba37b2 --- /dev/null +++ b/server/internal/proxy/hold_test.go @@ -0,0 +1,395 @@ +package proxy + +import ( + "encoding/json" + "io" + "net/http" + "net/http/httptest" + "strings" + "sync/atomic" + "testing" + "time" + + "github.com/openteam/server/internal/pkg/crypto" + "github.com/openteam/server/internal/store" +) + +// addChannel 给 harness 加一个绑定同一模型的渠道,返回渠道名。 +// weight 传大数可以把加权随机的起点压到该渠道上(与既有测试同一手法)。 +func addChannel(t *testing.T, h *harness, name, baseURL string, weight int) { + t.Helper() + encKey, err := crypto.NewEncryptor("test-master-key-1234567890").Encrypt("upstream-key") + if err != nil { + t.Fatal(err) + } + var m store.Model + h.db.Where("name = ?", "m1").First(&m) + ch := store.Channel{Name: name, Provider: store.ChannelProviderOpenAI, + Formats: []string{store.FormatChat}, BaseURL: baseURL, APIKeyEnc: encKey, + Weight: weight, TimeoutMS: 0, MaxConcurrency: 4, + HealthStatus: store.ChannelHealthHealthy, Enabled: true} + if err := h.db.Create(&ch).Error; err != nil { + t.Fatal(err) + } + if err := h.db.Create(&store.ChannelModelBinding{ChannelID: ch.ID, ModelID: m.ID, UpstreamModel: "m1"}).Error; err != nil { + t.Fatal(err) + } +} + +// pushWeightTo 把指定名字渠道的权重改大,让加权随机几乎必然从它开始。 +func pushWeightTo(t *testing.T, h *harness, name string, weight int) { + t.Helper() + if err := h.db.Model(&store.Channel{}).Where("name = ?", name).Update("weight", weight).Error; err != nil { + t.Fatal(err) + } +} + +// 流式容错:渠道1 在首个内容事件之前回错误事件(200+流式形式的错误),应换渠道2 +// 重试——客户端只看到渠道2 的内容,全程 200,无感知切换。 +func TestStreamFailoverBeforeFirstContent(t *testing.T) { + var ch1Hits, ch2Hits atomic.Int64 + bad := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + ch1Hits.Add(1) + w.Header().Set("Content-Type", "text/event-stream") + w.WriteHeader(http.StatusOK) + _, _ = io.WriteString(w, `data: {"error":{"message":"overloaded","type":"server_error"}}`+"\n\n") + })) + defer bad.Close() + good := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + ch2Hits.Add(1) + w.Header().Set("Content-Type", "text/event-stream") + w.WriteHeader(http.StatusOK) + _, _ = io.WriteString(w, `data: {"choices":[{"delta":{"content":"来自健康渠道"}}]}`+"\n\n") + _, _ = io.WriteString(w, "data: [DONE]\n\n") + })) + defer good.Close() + + h := newHarness(t, bad.URL, 0, true, nil) + addChannel(t, h, "c2", good.URL, 1) + pushWeightTo(t, h, "c1", 1000) + + w := h.run(t, nil, `{"model":"m1","stream":true,"messages":[{"role":"user","content":"hi"}]}`) + if w.Code != http.StatusOK { + t.Fatalf("code=%d body=%s", w.Code, w.Body.String()) + } + out := w.Body.String() + if ch1Hits.Load() != 1 || ch2Hits.Load() != 1 { + t.Fatalf("渠道命中次数不符:ch1=%d ch2=%d", ch1Hits.Load(), ch2Hits.Load()) + } + if strings.Contains(out, "overloaded") { + t.Fatalf("渠道1 的错误不应到达客户端:%q", out) + } + if !strings.Contains(out, "来自健康渠道") { + t.Fatalf("应收到渠道2 的内容:%q", out) + } +} + +// Anthropic 形状的错误事件({"type":"error","error":{...}})同样在放行前被识别换渠道。 +func TestAnthropicStyleStreamErrorFailsOver(t *testing.T) { + bad := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "text/event-stream") + w.WriteHeader(http.StatusOK) + _, _ = io.WriteString(w, `event: error`+"\n"+`data: {"type":"error","error":{"type":"overloaded_error","message":"Overloaded"}}`+"\n\n") + })) + defer bad.Close() + good := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "text/event-stream") + w.WriteHeader(http.StatusOK) + _, _ = io.WriteString(w, `data: {"choices":[{"delta":{"content":"ok"}}]}`+"\n\n") + _, _ = io.WriteString(w, "data: [DONE]\n\n") + })) + defer good.Close() + + h := newHarness(t, bad.URL, 0, true, nil) + addChannel(t, h, "c2", good.URL, 1) + pushWeightTo(t, h, "c1", 1000) + + w := h.run(t, nil, `{"model":"m1","stream":true,"messages":[{"role":"user","content":"hi"}]}`) + if w.Code != http.StatusOK || !strings.Contains(w.Body.String(), `"content":"ok"`) { + t.Fatalf("应无感知切换到渠道2:code=%d body=%s", w.Code, w.Body.String()) + } +} + +// 渠道回了 200+event-stream 却一个字没写就断流:放行前视为失败,换渠道重试。 +func TestEmptyStreamBeforeContentFailsOver(t *testing.T) { + bad := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "text/event-stream") + w.WriteHeader(http.StatusOK) + w.(http.Flusher).Flush() + // 直接结束:EOF 且无内容 + })) + defer bad.Close() + good := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "text/event-stream") + w.WriteHeader(http.StatusOK) + _, _ = io.WriteString(w, `data: {"choices":[{"delta":{"content":"补位渠道"}}]}`+"\n\n") + _, _ = io.WriteString(w, "data: [DONE]\n\n") + })) + defer good.Close() + + h := newHarness(t, bad.URL, 0, true, nil) + addChannel(t, h, "c2", good.URL, 1) + pushWeightTo(t, h, "c1", 1000) + + w := h.run(t, nil, `{"model":"m1","stream":true,"messages":[{"role":"user","content":"hi"}]}`) + if w.Code != http.StatusOK || !strings.Contains(w.Body.String(), "补位渠道") { + t.Fatalf("空流应触发故障转移:code=%d body=%s", w.Code, w.Body.String()) + } +} + +// didNotStream:要流式,渠道却回了整包 JSON——视为失败换渠道,而不是把 JSON 塞给等 SSE 的客户端。 +func TestDidNotStreamFailsOver(t *testing.T) { + bad := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(http.StatusOK) + _, _ = io.WriteString(w, `{"id":"x","choices":[{"message":{"role":"assistant","content":"整包"}}]}`) + })) + defer bad.Close() + good := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "text/event-stream") + w.WriteHeader(http.StatusOK) + _, _ = io.WriteString(w, `data: {"choices":[{"delta":{"content":"真流式"}}]}`+"\n\n") + _, _ = io.WriteString(w, "data: [DONE]\n\n") + })) + defer good.Close() + + h := newHarness(t, bad.URL, 0, true, nil) + addChannel(t, h, "c2", good.URL, 1) + pushWeightTo(t, h, "c1", 1000) + + w := h.run(t, nil, `{"model":"m1","stream":true,"messages":[{"role":"user","content":"hi"}]}`) + if w.Code != http.StatusOK { + t.Fatalf("code=%d body=%s", w.Code, w.Body.String()) + } + out := w.Body.String() + if !strings.HasPrefix(out, "data:") || !strings.Contains(out, "真流式") { + t.Fatalf("应从渠道2 收到 SSE 流:%q", out) + } +} + +// 反向 didNotStream:要整包,渠道回了 event-stream,同样换渠道。 +func TestNonStreamReceivingStreamFailsOver(t *testing.T) { + bad := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "text/event-stream") + w.WriteHeader(http.StatusOK) + _, _ = io.WriteString(w, "data: [DONE]\n\n") + })) + defer bad.Close() + good := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + _, _ = io.WriteString(w, `{"id":"ok","choices":[{"message":{"role":"assistant","content":"hi"},"finish_reason":"stop"}],"usage":{"prompt_tokens":1,"completion_tokens":1}}`) + })) + defer good.Close() + + h := newHarness(t, bad.URL, 0, true, nil) + addChannel(t, h, "c2", good.URL, 1) + pushWeightTo(t, h, "c1", 1000) + + w := h.run(t, nil, `{"model":"m1","messages":[{"role":"user","content":"hi"}]}`) + if w.Code != http.StatusOK || !strings.Contains(w.Body.String(), `"content":"hi"`) { + t.Fatalf("应从渠道2 收到整包 JSON:code=%d body=%s", w.Code, w.Body.String()) + } +} + +// 首 token 计时:流式放行前迟迟不出内容,超时取消本次尝试、换下一个渠道。 +func TestFirstTokenTimeoutFailsOver(t *testing.T) { + bad := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "text/event-stream") + w.WriteHeader(http.StatusOK) + w.(http.Flusher).Flush() + select { + case <-r.Context().Done(): // 被首 token 计时器取消 + case <-time.After(5 * time.Second): + } + })) + defer bad.Close() + good := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "text/event-stream") + w.WriteHeader(http.StatusOK) + _, _ = io.WriteString(w, `data: {"choices":[{"delta":{"content":"及时的内容"}}]}`+"\n\n") + _, _ = io.WriteString(w, "data: [DONE]\n\n") + })) + defer good.Close() + + h := newHarness(t, bad.URL, 0, true, nil) + h.gw.firstTokenWait = 150 * time.Millisecond + addChannel(t, h, "c2", good.URL, 1) + pushWeightTo(t, h, "c1", 1000) + + w := h.run(t, nil, `{"model":"m1","stream":true,"messages":[{"role":"user","content":"hi"}]}`) + if w.Code != http.StatusOK || !strings.Contains(w.Body.String(), "及时的内容") { + t.Fatalf("首 token 超时应换渠道:code=%d body=%s", w.Code, w.Body.String()) + } +} + +// 失败尝试扫到的 usage 不得计入最终账单:渠道1 的 message_start 式用量随失败被丢弃。 +func TestUsageFromFailedTryNotBilled(t *testing.T) { + bad := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "text/event-stream") + w.WriteHeader(http.StatusOK) + // role delta + usage 但没有内容:usage 会被扫进 trySink,随后错误事件导致整次尝试废弃 + _, _ = io.WriteString(w, `data: {"choices":[{"delta":{"role":"assistant"}}],"usage":{"prompt_tokens":500,"completion_tokens":0}}`+"\n\n") + _, _ = io.WriteString(w, `data: {"error":{"message":"boom"}}`+"\n\n") + })) + defer bad.Close() + good := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "text/event-stream") + w.WriteHeader(http.StatusOK) + _, _ = io.WriteString(w, `data: {"choices":[{"delta":{"content":"答案"}}]}`+"\n\n") + _, _ = io.WriteString(w, `data: {"choices":[],"usage":{"prompt_tokens":7,"completion_tokens":3}}`+"\n\n") + _, _ = io.WriteString(w, "data: [DONE]\n\n") + })) + defer good.Close() + + h := newHarness(t, bad.URL, 0, true, nil) + addChannel(t, h, "c2", good.URL, 1) + pushWeightTo(t, h, "c1", 1000) + + w := h.run(t, nil, `{"model":"m1","stream":true,"messages":[{"role":"user","content":"hi"}]}`) + if w.Code != http.StatusOK { + t.Fatalf("code=%d body=%s", w.Code, w.Body.String()) + } + h.rec.Close() // usage 异步批量落库,先冲刷再查 + var logs []store.UsageLog + h.db.Where("status = ?", store.UsageStatusSuccess).Find(&logs) + if len(logs) != 1 { + t.Fatalf("成功记账应只有一条,实际 %d", len(logs)) + } + if logs[0].InputTokens != 7 || logs[0].OutputTokens != 3 { + t.Fatalf("应只计渠道2 的用量 (7/3),实际 in=%d out=%d", logs[0].InputTokens, logs[0].OutputTokens) + } +} + +// 响应头已被心跳先行写出后,全部候选失败只能以 SSE 错误事件收尾(状态码已不可改)。 +func TestStreamErrorEventWhenHeadAlreadySent(t *testing.T) { + bad := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "text/event-stream") + w.WriteHeader(http.StatusOK) + time.Sleep(200 * time.Millisecond) // 期间心跳先把 200 头写给客户端 + _, _ = io.WriteString(w, `data: {"error":{"message":"boom"}}`+"\n\n") + })) + defer bad.Close() + + h := newHarness(t, bad.URL, 0, true, nil) + h.gw.streamHeartbeat = 40 * time.Millisecond + + w := h.run(t, nil, `{"model":"m1","stream":true,"messages":[{"role":"user","content":"hi"}]}`) + if w.Code != http.StatusOK { + t.Fatalf("头已写出,状态应为 200:code=%d body=%s", w.Code, w.Body.String()) + } + out := w.Body.String() + if !strings.Contains(out, ": ping") { + t.Fatalf("扣住期间应有心跳:%q", out) + } + if !strings.Contains(out, `"upstream_error"`) { + t.Fatalf("应以 SSE 错误事件收尾:%q", out) + } +} + +// --------------------------------------------------------------------------- +// 判定函数单测 + +func TestSSEEventHasContent(t *testing.T) { + cases := []struct { + proto string + line string + want bool + }{ + // chat:role-only 与空 delta 不算内容,content/reasoning/tool_calls 算 + {"chat", `data: {"choices":[{"delta":{"role":"assistant"}}]}`, false}, + {"chat", `data: {"choices":[{"delta":{}}]}`, false}, + {"chat", `data: {"choices":[{"delta":{"content":"x"}}]}`, true}, + {"chat", `data: {"choices":[{"delta":{"reasoning_content":"想"}}]}`, true}, + {"chat", `data: {"choices":[{"delta":{"tool_calls":[{"id":"1"}]}}]}`, true}, + {"chat", "data: [DONE]", false}, + // messages:message_start / ping 不算,块开启与 delta 算 + {"messages", `data: {"type":"message_start","message":{"usage":{"input_tokens":1}}}`, false}, + {"messages", `data: {"type":"ping"}`, false}, + {"messages", `data: {"type":"content_block_start","index":0,"content_block":{"type":"text","text":""}}`, true}, + {"messages", `data: {"type":"content_block_start","index":0,"content_block":{"type":"tool_use","id":"t","name":"f"}}`, true}, + {"messages", `data: {"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"x"}}`, true}, + {"messages", `data: {"type":"content_block_delta","index":0,"delta":{"type":"thinking_delta","thinking":"想"}}`, true}, + // responses:.delta 与 output_item.added 算,created/completed 不算 + {"responses", `data: {"type":"response.created","response":{}}`, false}, + {"responses", `data: {"type":"response.output_text.delta","delta":"x"}`, true}, + {"responses", `data: {"type":"response.output_item.added","item":{"type":"function_call"}}`, true}, + } + for i, tc := range cases { + if got := sseEventHasContent([]byte(tc.line), tc.proto); got != tc.want { + t.Fatalf("case %d (%s): sseEventHasContent(%s) = %v, want %v", i, tc.proto, tc.line, got, tc.want) + } + } + // 未知协议:任何数据事件都算内容(退化放行,不吞正文) + if !sseEventHasContent([]byte(`data: {"weird":true}`), "other") { + t.Fatal("未知协议应放行") + } +} + +func TestSSEUpstreamError(t *testing.T) { + cases := []struct { + line string + want bool + contain string + }{ + {`data: {"error":{"message":"overloaded","type":"server_error"}}`, true, "overloaded"}, + {`data: {"type":"error","error":{"type":"overloaded_error","message":"Overloaded"}}`, true, "Overloaded"}, + {`data: {"type":"error","code":500}`, true, "upstream stream error"}, + {`data: {"type":"response.failed","response":{"error":{"message":"bad tool"}}}`, true, "bad tool"}, + {`data: {"choices":[{"delta":{"content":"hi"}}]}`, false, ""}, + {`data: {"type":"content_block_start","content_block":{"type":"text"}}`, false, ""}, + {`data: {"error":null}`, false, ""}, + {`event: message_start`, false, ""}, + } + for i, tc := range cases { + msg, ok := sseUpstreamError([]byte(tc.line)) + if ok != tc.want { + t.Fatalf("case %d: sseUpstreamError(%s) ok=%v, want %v", i, tc.line, ok, tc.want) + } + if ok && !strings.Contains(msg, tc.contain) { + t.Fatalf("case %d: msg=%q, want contains %q", i, msg, tc.contain) + } + } +} + +func TestClassifyUpstreamError(t *testing.T) { + cases := []struct { + status int + body string + kind string + retryable bool + }{ + {http.StatusTooManyRequests, `{"error":{"message":"Rate limit reached"}}`, "rate_limited", true}, + {http.StatusTooManyRequests, `{"error":{"message":"You exceeded your current quota, check your plan and billing details"}}`, "quota_exhausted", true}, + {http.StatusTooManyRequests, `{"error":{"code":"1113","message":"余额不足"}}`, "quota_exhausted", true}, + {http.StatusPaymentRequired, `{"error":{"message":"insufficient balance"}}`, "quota_exhausted", true}, + {http.StatusForbidden, `{"error":{"message":"insufficient quota"}}`, "quota_exhausted", true}, + {http.StatusForbidden, `{"error":{"message":"model not allowed"}}`, "server_error", false}, + {http.StatusBadGateway, `{"error":{"message":"upstream gone"}}`, "server_error", true}, + {http.StatusBadRequest, `{"error":{"message":"bad request"}}`, "server_error", false}, + } + for i, tc := range cases { + kind, retryable := classifyUpstreamError(tc.status, []byte(tc.body)) + if kind.String() != tc.kind || retryable != tc.retryable { + t.Fatalf("case %d: classify(%d) = (%s, %v), want (%s, %v)", i, tc.status, kind, retryable, tc.kind, tc.retryable) + } + } +} + +func TestStreamErrorEventShapes(t *testing.T) { + chat := string(streamErrorEvent("chat", 502, "boom")) + if !strings.HasPrefix(chat, "data: ") || !strings.Contains(chat, `"code":"upstream_error"`) || !strings.Contains(chat, "data: [DONE]") { + t.Fatalf("chat 错误事件形状不符:%q", chat) + } + msg := string(streamErrorEvent("messages", 502, "boom")) + if !strings.HasPrefix(msg, "event: error\ndata: ") { + t.Fatalf("messages 错误事件形状不符:%q", msg) + } + var m map[string]any + if err := json.Unmarshal([]byte(strings.TrimPrefix(strings.SplitN(msg, "data: ", 2)[1], "\n\n")), &m); err != nil { + t.Fatalf("messages 错误事件不是合法 JSON: %v (%s)", err, msg) + } + resp := string(streamErrorEvent("responses", 502, "boom")) + if !strings.HasPrefix(resp, "data: ") || strings.Contains(resp, "[DONE]") { + t.Fatalf("responses 错误事件形状不符:%q", resp) + } +} diff --git a/server/internal/proxy/passthrough.go b/server/internal/proxy/passthrough.go index 1e7e86e..a43cfc2 100644 --- a/server/internal/proxy/passthrough.go +++ b/server/internal/proxy/passthrough.go @@ -14,13 +14,12 @@ import ( "net/http" "strconv" "strings" - "sync" - "sync/atomic" "time" "github.com/gin-gonic/gin" "github.com/openteam/server/internal/channel" "github.com/openteam/server/internal/pkg/tokenizer" + "github.com/openteam/server/internal/proxy/convert" "github.com/openteam/server/internal/store" ) @@ -190,6 +189,8 @@ func upstreamURL(ch *store.Channel, proto, path string) string { } // doProxy 通用代理(M5):加权随机选一个候选渠道;失败自动故障转移。 +// 流式请求挂一条 streamHold:首个内容事件到达之前,上游的错误/空流/断连都还能 +// 换渠道消化,客户端无感知(借鉴 magpie 的 holdWriter)。 func (g *Gateway) doProxy(c *gin.Context, cands []channel.Candidate, clientProto string, body []byte, stream bool, sink *usageSink) { if len(cands) == 0 { apiError(c, http.StatusBadGateway, "upstream_error", "no available channel") @@ -206,6 +207,15 @@ func (g *Gateway) doProxy(c *gin.Context, cands []channel.Candidate, clientProto } } + // 流式容错缓冲 + 心跳,整条请求一条,跨渠道重试共用 + var sh *streamHold + stopHB := func() {} + if stream { + sh = newStreamHold(c) + stopHB = sh.heartbeat(g.streamHeartbeat) + defer stopHB() + } + var lastStatus = http.StatusBadGateway var lastBody = []byte("all upstream channels failed") // 最后真正试过的渠道:全部失败时把错误归到它头上,否则 usage_logs 里 @@ -243,13 +253,19 @@ func (g *Gateway) doProxy(c *gin.Context, cands []channel.Candidate, clientProto } // 失败/成功计数由 proxyOne 就地报:只有它读得到上游响应头(429 的 Retry-After) lastCh = ch - done, retry, st, b := g.proxyOne(c, ch, plan, stream, sink) + done, retry, st, b := g.proxyOne(c, ch, plan, stream, sink, sh, clientProto) release() if done { return } if !retry { - // 非重试性失败(如 400):透传上游错误体 + // 非重试性失败(如 400):透传上游错误体。 + // 响应头若已被心跳先行写出,状态码不可再改,只能以 SSE 错误事件收尾。 + if sh != nil && sh.headSent() { + _ = sh.writeDirect(streamErrorEvent(clientProto, st, string(b))) + g.recordError(c, ch, nil, now(), "upstream_http_"+strconv.Itoa(st)) + return + } c.Header("Content-Type", "application/json") c.Data(st, "application/json", b) g.recordError(c, ch, nil, now(), "upstream_http_"+strconv.Itoa(st)) @@ -258,17 +274,39 @@ func (g *Gateway) doProxy(c *gin.Context, cands []channel.Candidate, clientProto lastStatus, lastBody = st, b } // 全部候选重试性失败 + if sh != nil && sh.headSent() { + _ = sh.writeDirect(streamErrorEvent(clientProto, lastStatus, string(lastBody))) + g.recordError(c, lastCh, nil, now(), "all_channels_failed") + return + } apiError(c, lastStatus, "upstream_error", string(lastBody)) g.recordError(c, lastCh, nil, now(), "all_channels_failed") } +// streamErrorEvent 把错误构造成客户端协议的 SSE 错误事件(响应头已写出时的唯一交付方式)。 +func streamErrorEvent(proto string, status int, msg string) []byte { + errObj := map[string]any{"message": msg, "type": "api_error", "code": "upstream_error", "status": status} + switch proto { + case convert.ProtoMessages: + body, _ := json.Marshal(map[string]any{"type": "error", "error": errObj}) + return append([]byte("event: error\ndata: "), append(body, '\n', '\n')...) + case convert.ProtoResponses: + body, _ := json.Marshal(map[string]any{"type": "error", "code": "upstream_error", "message": msg}) + return append([]byte("data: "), append(body, '\n', '\n')...) + default: + body, _ := json.Marshal(map[string]any{"error": errObj}) + return append(append([]byte("data: "), append(body, '\n', '\n')...), []byte("data: [DONE]\n\n")...) + } +} + // proxyOne 对单个渠道执行一次代理。 // 返回:done=本次请求已终结(客户端响应已写出或已记账,停止故障转移); // retry=是否可安全换渠道重试;status+respBody=失败信息。 -func (g *Gateway) proxyOne(c *gin.Context, ch *store.Channel, plan *upstreamPlan, stream bool, sink *usageSink) (done bool, retry bool, status int, respBody []byte) { +// sh 为流式容错缓冲(非流式为 nil);sink 为请求级记账 sink。 +func (g *Gateway) proxyOne(c *gin.Context, ch *store.Channel, plan *upstreamPlan, stream bool, sink *usageSink, sh *streamHold, clientProto string) (done bool, retry bool, status int, respBody []byte) { upKey, err := g.ch.UpstreamKey(ch) if err != nil { - g.ch.ReportFailure(ch, 0) + g.ch.ReportFailure(ch, channel.FailNet, 0) return false, true, http.StatusInternalServerError, []byte("failed to decrypt channel key") } @@ -294,8 +332,24 @@ func (g *Gateway) proxyOne(c *gin.Context, ch *store.Channel, plan *upstreamPlan ctx, cancel = context.WithTimeout(ctx, limit) defer cancel() } + // 首 token 计时(可选):流式在放行前迟迟不出内容就取消本次尝试、换下一个渠道。 + // 计时器与放行互斥(cancelIfHeld),已放行的健康流不会被掐死。 + tryCtx := ctx + var tryCancel context.CancelFunc + var ftTimer *time.Timer + if stream && g.firstTokenWait > 0 { + tryCtx, tryCancel = context.WithCancel(ctx) + defer tryCancel() + ftTimer = time.AfterFunc(g.firstTokenWait, func() { sh.cancelIfHeld(tryCancel) }) + defer ftTimer.Stop() + } + onRelease := func() { + if ftTimer != nil { + ftTimer.Stop() + } + } - req, err := http.NewRequestWithContext(ctx, http.MethodPost, upstreamURL(ch, plan.proto, plan.path), bytes.NewReader(upBody)) + req, err := http.NewRequestWithContext(tryCtx, http.MethodPost, upstreamURL(ch, plan.proto, plan.path), bytes.NewReader(upBody)) if err != nil { return false, false, http.StatusInternalServerError, []byte("failed to build upstream request") } @@ -328,46 +382,64 @@ func (g *Gateway) proxyOne(c *gin.Context, ch *store.Channel, plan *upstreamPlan g.finishUsage(c, ch, start, store.UsageStatusCanceled, "client_disconnect") return true, false, 0, nil } - g.ch.ReportFailure(ch, 0) // 网络错误/超时 + g.ch.ReportFailure(ch, channel.FailNet, 0) // 网络错误/超时 return false, true, http.StatusBadGateway, []byte("upstream request failed: " + err.Error()) } defer resp.Body.Close() - // 非 2xx:429/5xx 可重试;其余透传错误体 + // 非 2xx:按原因分类记账——429 不一定是限流(可能是余额/配额),402 也值得 + // 换渠道(渠道 A 没余额不代表渠道 B 没有);其余透传错误体。 if resp.StatusCode < 200 || resp.StatusCode >= 300 { - errBody, _ := io.ReadAll(resp.Body) - if resp.StatusCode == http.StatusTooManyRequests || resp.StatusCode >= 500 { - g.ch.ReportFailure(ch, retryAfter(resp.Header)) + errBody, _ := io.ReadAll(io.LimitReader(resp.Body, 1<<20)) + kind, retryable := classifyUpstreamError(resp.StatusCode, errBody) + if retryable { + g.ch.ReportFailure(ch, kind, retryAfter(resp.Header)) return false, true, resp.StatusCode, errBody } return false, false, resp.StatusCode, errBody } g.ch.ReportSuccess(ch) - c.Header("Content-Type", resp.Header.Get("Content-Type")) - c.Status(http.StatusOK) - if stream { - g.streamCopy(c, ch, resp.Body, start, plan.lineConv, sink) - } else { - g.copyAndCapture(c, ch, resp.Body, start, plan.bodyConv, sink) + // 协议形态校验:要流给了整包(或反过来)是劣质中转的常见毛病,按失败换渠道; + // 反复如此的渠道会被健康计数剔除(didNotStream 不会因重试而变好)。 + ct := resp.Header.Get("Content-Type") + isSSE := strings.HasPrefix(ct, "text/event-stream") + if stream != isSSE { + _, _ = io.Copy(io.Discard, io.LimitReader(resp.Body, 4<<20)) // 排干以复用连接 + g.ch.ReportFailure(ch, channel.FailServer, 0) + asked, got := "整包", "流式响应" + if stream { + asked, got = "流式", "整包响应" + } + return false, true, http.StatusBadGateway, + []byte(ch.Name + " answered a " + asked + " request with " + got + " (didNotStream)") } - return true, false, http.StatusOK, nil + + if stream { + trySink := &usageSink{} + return g.streamHoldCopy(c, sh, ch, resp.Body, start, plan.lineConv, clientProto, trySink, sink, onRelease, tryCtx) + } + return g.copyAndCapture(c, ch, resp, start, plan.bodyConv, sink) } -// copyAndCapture 非流式:整体转发(可转换)+ 解析 usage + 记账。 -func (g *Gateway) copyAndCapture(c *gin.Context, ch *store.Channel, r io.Reader, start time.Time, bodyConv func([]byte) ([]byte, error), sink *usageSink) { - data, err := io.ReadAll(r) +// copyAndCapture 非流式:先读完整上游响应再写客户端——响应头与正文都在"确定成功"后 +// 才写出,读失败(网络错误/超时)发生在写出之前,可换渠道重试。 +// 此前实现先写 200 头再读,超时只能给客户端留一个 200 空响应。 +func (g *Gateway) copyAndCapture(c *gin.Context, ch *store.Channel, resp *http.Response, start time.Time, bodyConv func([]byte) ([]byte, error), sink *usageSink) (done bool, retry bool, status int, respBody []byte) { + data, err := io.ReadAll(resp.Body) if err != nil { - // 走到这里响应头已按 200 写出,无法再回错误体,只按中断原因记账 switch { case c.Request.Context().Err() != nil: + // 客户端断开:就地记账终结 g.finishUsage(c, ch, start, store.UsageStatusCanceled, "client_disconnect") + return true, false, 0, nil case errors.Is(err, context.DeadlineExceeded): - g.finishUsage(c, ch, start, store.UsageStatusError, "upstream_timeout") + g.ch.ReportFailure(ch, channel.FailNet, 0) + return false, true, http.StatusBadGateway, []byte(ch.Name + " upstream timeout") default: - g.finishUsage(c, ch, start, store.UsageStatusError, "read_error") + g.ch.ReportFailure(ch, channel.FailNet, 0) + return false, true, http.StatusBadGateway, []byte(ch.Name + " read error: " + err.Error()) } - return } if usageRaw := extractUsage(data); usageRaw != nil && sink != nil { sink.push(usageRaw) @@ -378,128 +450,105 @@ func (g *Gateway) copyAndCapture(c *gin.Context, ch *store.Channel, r io.Reader, out = converted } } + c.Header("Content-Type", resp.Header.Get("Content-Type")) + c.Status(http.StatusOK) _, _ = c.Writer.Write(out) if _, ok := c.Get("raw_request"); ok { c.Set("raw_response", string(data)) // 上游原始响应(未转换) } g.finishUsage(c, ch, start, store.UsageStatusSuccess, "") + return true, false, 0, nil } -// streamCopy 流式:边读上游 SSE 边写客户端,零缓冲转发;按 lineConv 转换;扫描 usage 记账。 -func (g *Gateway) streamCopy(c *gin.Context, ch *store.Channel, r io.Reader, start time.Time, lineConv func([]byte) []byte, sink *usageSink) { - w := c.Writer - flusher, ok := w.(http.Flusher) - if !ok { - flusher = nopFlusher{} - } - +// streamHoldCopy 流式转发:边读上游 SSE 边经 hold 写客户端(放行前缓冲)。 +// 放行前的一切失败(上游错误事件、空流、断连、首 token 超时)都可换渠道重试, +// 健康计数就地上报;放行后的失败与旧语义一致,就地记账终结。 +// trySink 收集本次尝试扫到的 usage,只有成功走到 EOF 才并入请求级 reqSink—— +// 失败渠道的 message_start 用量不能计入最终账单。 +func (g *Gateway) streamHoldCopy(c *gin.Context, sh *streamHold, ch *store.Channel, r io.Reader, start time.Time, + lineConv func([]byte) []byte, clientProto string, trySink, reqSink *usageSink, onRelease func(), tryCtx context.Context, +) (done bool, retry bool, status int, respBody []byte) { // 原始响应捕获:仅管理员且开关开启(raw_request 已 set)时累积上游原始行 _, capture := c.Get("raw_request") var rawResp strings.Builder - - // commitRaw 在记账前把已累积的原始响应写入 context commitRaw := func() { if capture { c.Set("raw_response", rawResp.String()) } } - lw := &lockedWriter{w: w, f: flusher} - lw.lastN.Store(time.Now().UnixNano()) - stopHB := lw.heartbeat(g.streamHeartbeat) - defer stopHB() - scanner := newSSEScanner(r) + releaseSeen := false + fail := func(msg string, kind channel.FailKind) (bool, bool, int, []byte) { + g.ch.ReportFailure(ch, kind, 0) + return false, true, http.StatusBadGateway, []byte(ch.Name + ": " + msg) + } for { line, err := scanner.Next() if line != nil { if capture { rawResp.Write(line) } + // 上游错误事件(200+流式形式的错误):放行前换渠道消化 + if !releaseSeen { + if msg, isErr := sseUpstreamError(line); isErr { + return fail("stream error: "+msg, channel.FailServer) + } + } out := line if lineConv != nil { out = lineConv(line) } if out != nil { - if _, werr := lw.write(out); werr != nil { + if werr := sh.writeEvent(out); werr != nil { // 客户端意外断开:按已生成部分收费(canceled) commitRaw() g.finishUsage(c, ch, start, store.UsageStatusCanceled, "client_disconnect") - return + return true, false, 0, nil } - if sink != nil { - sink.outputText += sseContentText(out) + if trySink != nil { + trySink.outputText += sseContentText(out) } } - if usageRaw := scanUsage(line); usageRaw != nil && sink != nil { - sink.push(usageRaw) + if !releaseSeen && sseEventHasContent(out, clientProto) { + releaseSeen = true + sh.release() + if onRelease != nil { + onRelease() + } + } + if usageRaw := scanUsage(line); usageRaw != nil && trySink != nil { + trySink.push(usageRaw) } } if err != nil { - commitRaw() - if err == io.EOF { - g.finishUsage(c, ch, start, store.UsageStatusSuccess, "") - } else if c.Request.Context().Err() != nil { - // 客户端意外断开:按已生成部分收费(canceled) - g.finishUsage(c, ch, start, store.UsageStatusCanceled, "client_disconnect") - } else { - g.recordError(c, ch, nil, start, "stream_read_error") - } - return - } - } -} - -type nopFlusher struct{} - -func (nopFlusher) Flush() {} - -// lockedWriter 给流式写出加锁:心跳在独立 goroutine 里写,而 gin 的 ResponseWriter -// 不是并发安全的,两路同时写会坏帧。 -type lockedWriter struct { - mu sync.Mutex - w io.Writer - f http.Flusher - lastN atomic.Int64 // 最近一次写出时间(unix nano),心跳据此判断空闲 -} - -func (lw *lockedWriter) write(b []byte) (int, error) { - lw.mu.Lock() - defer lw.mu.Unlock() - n, err := lw.w.Write(b) - lw.lastN.Store(time.Now().UnixNano()) - if err == nil { - lw.f.Flush() - } - return n, err -} - -// heartbeat 距上次写出超过 interval 时补一行 SSE 注释(`: ping`),返回停止函数。 -// 挂在 nginx/Cloudflare 后面时长思考(首包后上游安静几十秒)会被空闲超时掐掉, -// 客户端只看到流断在半路。注释行按 SSE 规范被解析器忽略,不影响计费与转换。 -func (lw *lockedWriter) heartbeat(interval time.Duration) func() { - if interval <= 0 { - return func() {} - } - done := make(chan struct{}) - var once sync.Once - go func() { - tk := time.NewTicker(interval) - defer tk.Stop() - for { - select { - case <-done: - return - case <-tk.C: - if time.Since(time.Unix(0, lw.lastN.Load())) < interval { - continue + if releaseSeen { + // 已放行:客户端正在看内容,失败就地记账终结,与旧语义一致 + commitRaw() + if err == io.EOF { + reqSink.merge(trySink) + g.finishUsage(c, ch, start, store.UsageStatusSuccess, "") + } else if c.Request.Context().Err() != nil { + g.finishUsage(c, ch, start, store.UsageStatusCanceled, "client_disconnect") + } else { + g.recordError(c, ch, nil, start, "stream_read_error") } - // 心跳写失败说明客户端已断开,交给主循环下一次写出时统一记账 - _, _ = lw.write([]byte(": ping\n\n")) + return true, false, 0, nil } + // 放行前失败:换渠道重试 + if c.Request.Context().Err() != nil { + g.finishUsage(c, ch, start, store.UsageStatusCanceled, "client_disconnect") + return true, false, 0, nil + } + if tryCtx != nil && tryCtx.Err() != nil { + return fail(fmt.Sprintf("first token timeout (no content in %s)", g.firstTokenWait), channel.FailNet) + } + if err == io.EOF { + return fail("ended the stream before any content", channel.FailServer) + } + return fail("stream failed before any content: "+err.Error(), channel.FailNet) } - }() - return func() { once.Do(func() { close(done) }) } + } } // --------------------------------------------------------------------------- @@ -674,6 +723,33 @@ func (u *usageSink) push(raw json.RawMessage) { // Shape 返回合并后的用量。 func (u *usageSink) Shape() usageShape { return u.us } +// merge 并入一次成功尝试的用量:失败渠道的用量整体丢弃(换渠道重试时), +// 成功渠道的 usage 事件(含放行前缓冲期扫到的 message_start)在这里落账。 +func (u *usageSink) merge(o *usageSink) { + if o == nil { + return + } + if o.us.PromptTokens > 0 { + u.us.PromptTokens = o.us.PromptTokens + } + if o.us.CompletionTokens > 0 { + u.us.CompletionTokens = o.us.CompletionTokens + } + if o.us.InputTokens > 0 { + u.us.InputTokens = o.us.InputTokens + } + if o.us.OutputTokens > 0 { + u.us.OutputTokens = o.us.OutputTokens + } + if o.us.CacheReadInputTokens > 0 { + u.us.CacheReadInputTokens = o.us.CacheReadInputTokens + } + if o.us.CacheCreationInputTokens > 0 { + u.us.CacheCreationInputTokens = o.us.CacheCreationInputTokens + } + u.outputText += o.outputText +} + // finishUsage 落账:计算成本并异步写入。 func (g *Gateway) finishUsage(c *gin.Context, ch *store.Channel, start time.Time, status, errCode string) { uid, _ := c.Get(CtxUserID)