feat: 流式首包缓冲容错(hold)与失败分类冷却
流式请求在首个内容事件到达前由网关缓冲(借鉴 magpie 的 holdWriter):上游的错误 事件、空流、断连、协议形态不符(要流给了整包)都在放行前换渠道消化,客户端无感知。 扣住期间照常发 SSE 注释心跳;缓冲超过 1MB 直接放行退化为普通转发,判定失效不吞正文。 - 新增 proxy/hold.go:streamHold(缓冲/放行/心跳/首 token 计时与放行互斥)、 sseEventHasContent(chat/messages/responses 三种协议的内容判定)、 sseUpstreamError(200+流式形式的错误事件识别)、classifyUpstreamError + quotaSignals(余额/配额特征识别) - 非流式改为先读完整响应再写头:读失败(网络错误/超时)发生在写出之前,可换渠道 重试;此前先写 200 头再读,超时只能给客户端留一个 200 空响应 - 协议形态校验:要流却给了整包(didNotStream)按失败换渠道,反复如此的渠道会被 健康计数剔除 - 记账:每次尝试独立 trySink,只有成功走到 EOF 才并入请求级 sink——失败渠道的 message_start 用量不计入账单 - 冷却按原因分类(channel.FailKind):余额/配额类第一次失败即冷却 30 分钟(429 不 一定是限流,智谱 429「余额不足」也在此列);其余按连续进冷却次数指数退避 30s→60s→120s,封顶 10 分钟;短于退避值的 Retry-After 不再缩短冷却 - 新增 OT_PROXY_FIRST_TOKEN_WAIT(默认 0=不限):放行前等内容的最长时间,超时即 取消本次尝试换下一个渠道 - 渠道列表接口返回 cooldown_reason,后台可区分「限流」与「余额耗尽」 - 测试:hold 缓冲/放行/心跳、内容与错误事件判定、失败分类与指数退避、容量热生效 - docs/magpie-plugin-mechanism.md:参考实现的插件机制调研
This commit is contained in:
1 parent
d357134e97
commit
979973ddde
12 files changed
+1377
-168
No files matched your search
@@ -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{},
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
@@ -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)
|
||||
|
||||
Reference in new issue
Block a user