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) { // 廉价的子串预过滤:绝大多数事件与错误无关,先排除再谈解析。 // 放行前每条事件都会过这里,长思考阶段的几十上百条事件不值得逐条反序列化。 if !bytes.Contains(line, []byte(`"error"`)) && !bytes.Contains(line, []byte("response.failed")) { return "", false } 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 }