openteam/server/internal/proxy/hold.go

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
}