openteam/server/internal/proxy/passthrough.go

package proxy

import (
	"bufio"
	"bytes"
	"crypto/rand"
	"encoding/hex"
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"strconv"
	"strings"
	"time"

	"github.com/gin-gonic/gin"
	"github.com/openteam/server/internal/channel"
	"github.com/openteam/server/internal/pkg/tokenizer"
	"github.com/openteam/server/internal/store"
)

// newTraceID 生成请求 trace（用于日志与记账幂等 ref）。
func newTraceID() string {
	b := make([]byte, 8)
	_, _ = rand.Read(b)
	return hex.EncodeToString(b)
}

// bodyReq 统一取出请求体并解析 model / stream 字段。
type bodyReq struct {
	Model  string `json:"model"`
	Stream bool   `json:"stream"`
}

// parseBody 读取并回填请求体，解析 model/stream。
func parseBody(c *gin.Context) (*bodyReq, []byte, error) {
	body, err := io.ReadAll(c.Request.Body)
	if err != nil {
		return nil, nil, err
	}
	c.Request.Body = io.NopCloser(bytes.NewReader(body))
	// 中断请求输入文本（token 估算用，仅存文本不立即分词）
	c.Set("est_input_text", requestText(body))
	br := &bodyReq{}
	_ = json.Unmarshal(body, br) // 解析失败按空处理，直通仍可转发
	return br, body, nil
}

// requestText 提取请求体中的用户输入文本（chat/messages 的 content、responses 的 input/instructions），
// 用于中断时估算输入 token。
func requestText(body []byte) string {
	var m map[string]any
	if json.Unmarshal(body, &m) != nil {
		return ""
	}
	var parts []string
	add := func(s string) {
		if s = strings.TrimSpace(s); s != "" {
			parts = append(parts, s)
		}
	}
	if s, ok := m["instructions"].(string); ok {
		add(s)
	}
	if s, ok := m["system"].(string); ok {
		add(s)
	}
	switch input := m["input"].(type) {
	case string:
		add(input)
	case []any:
		for _, it := range input {
			if im, ok := it.(map[string]any); ok {
				if s, ok := im["content"].(string); ok {
					add(s)
				}
			}
		}
	}
	if msgs, ok := m["messages"].([]any); ok {
		for _, msg := range msgs {
			mm, ok := msg.(map[string]any)
			if !ok {
				continue
			}
			switch c := mm["content"].(type) {
			case string:
				add(c)
			case []any:
				for _, b := range c {
					if bm, ok := b.(map[string]any); ok {
						if s, ok := bm["text"].(string); ok {
							add(s)
						}
					}
				}
			}
		}
	}
	return strings.Join(parts, "\n")
}

// sseDataPayload 提取一条 SSE 的 JSON 载荷（去掉 data: 前缀与空白）。
// 兼容三种写法：
//   - 单 data: 行：data: {...} 或 data:{...}（上游如火山方舟会省略 data: 后的空格）
//   - event:+data: 多行块：转换器 eventLine 产出的块（event: xxx\ndata: {...} 拼在一个 []byte）
func sseDataPayload(line []byte) (string, bool) {
	s := string(line)
	idx := strings.LastIndex(s, "\ndata:")
	if idx >= 0 {
		s = s[idx+len("\ndata:"):] // 跳过 event: 头，落在 data: 之后
	} else if strings.HasPrefix(s, "data:") {
		s = strings.TrimPrefix(s, "data:")
	} else {
		return "", false
	}
	s = strings.TrimLeft(s, " ") // data: 后的可选空格
	s = strings.TrimSpace(s)
	if s == "" || s == "[DONE]" {
		return "", false
	}
	return s, true
}

// sseContentText 提取一条 SSE 中的内容文本（chat delta.content / responses delta / messages delta.text）。
func sseContentText(line []byte) string {
	s, ok := sseDataPayload(line)
	if !ok {
		return ""
	}
	var m map[string]any
	if json.Unmarshal([]byte(s), &m) != nil {
		return ""
	}
	// responses output_text.delta: {"delta":"..."}
	if d, ok := m["delta"].(string); ok {
		return d
	}
	// messages content_block_delta: {"delta":{"text":"..."}}
	if dm, ok := m["delta"].(map[string]any); ok {
		if t, ok := dm["text"].(string); ok {
			return t
		}
	}
	// chat chunk: {"choices":[{"delta":{"content":"..."}}]}（GLM 思考模型另有 reasoning_content）
	if choices, ok := m["choices"].([]any); ok && len(choices) > 0 {
		if c0, ok := choices[0].(map[string]any); ok {
			if delta, ok := c0["delta"].(map[string]any); ok {
				if t, ok := delta["content"].(string); ok {
					return t
				}
				if t, ok := delta["reasoning_content"].(string); ok {
					return t
				}
			}
		}
	}
	return ""
}

// upstreamURL 组装上游地址：按协议选 base_url 再拼资源路径（见 store.Channel.UpstreamURL）。
func upstreamURL(ch *store.Channel, proto, path string) string {
	return ch.UpstreamURL(proto, path)
}

// doProxy 通用代理（M5）：加权随机选一个候选渠道；失败自动故障转移。
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")
		g.recordError(c, nil, nil, now(), "no_available_channel")
		return
	}
	// 加权随机选择起始渠道
	picked := g.ch.Pick(cands)
	startIdx := 0
	for i, cand := range cands {
		if cand.Channel.ID == picked.ID {
			startIdx = i
			break
		}
	}

	var lastStatus = http.StatusBadGateway
	var lastBody = []byte("all upstream channels failed")
	// 从选中的渠道开始遍历，到末尾后再从头遍历到选中渠道之前
	for offset := 0; offset < len(cands); offset++ {
		idx := (startIdx + offset) % len(cands)
		cand := cands[idx]
		ch := cand.Channel
		plan, err := prepareUpstream(ch, clientProto, body, cand.UpstreamModel)
		if err != nil {
			lastStatus, lastBody = http.StatusInternalServerError, []byte("conversion error: "+err.Error())
			continue
		}
		release, ok := g.ch.TryAcquire(ch)
		if !ok {
			continue // 渠道满载，溢出到下一个
		}
		written, retry, st, b := g.proxyOne(c, ch, plan, stream, sink)
		release()
		if written {
			return
		}
		if !retry {
			// 非重试性失败（如 400）：透传上游错误体
			c.Header("Content-Type", "application/json")
			c.Data(st, "application/json", b)
			g.recordError(c, ch, nil, now(), "upstream_http_"+strconv.Itoa(st))
			return
		}
		lastStatus, lastBody = st, b
	}
	// 全部候选重试性失败
	apiError(c, lastStatus, "upstream_error", string(lastBody))
	g.recordError(c, nil, nil, now(), "all_channels_failed")
}

// proxyOne 对单个渠道执行一次代理。
// 返回：written=是否已写客户端响应；retry=是否可安全换渠道重试；status+respBody=失败信息。
func (g *Gateway) proxyOne(c *gin.Context, ch *store.Channel, plan *upstreamPlan, stream bool, sink *usageSink) (written bool, retry bool, status int, respBody []byte) {
	upKey, err := g.ch.UpstreamKey(ch)
	if err != nil {
		return false, true, http.StatusInternalServerError, []byte("failed to decrypt channel key")
	}

	upBody := plan.body
	// 直通 chat 流式：注入 stream_options.include_usage，保证末块带 usage（OpenAI 行为）
	if stream && plan.path == "/chat/completions" && plan.lineConv == nil && !bytes.Contains(upBody, []byte(`"include_usage"`)) {
		var m map[string]any
		if json.Unmarshal(upBody, &m) == nil {
			m["stream_options"] = map[string]any{"include_usage": true}
			if b, err := json.Marshal(m); err == nil {
				upBody = b
			}
		}
	}

	req, err := http.NewRequest(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")
	}
	req.Header.Set("Content-Type", "application/json")
	req.Header.Set("Authorization", "Bearer "+upKey)
	req.Header.Set("Accept", c.GetHeader("Accept"))
	if ua := c.GetHeader("User-Agent"); ua != "" {
		req.Header.Set("User-Agent", ua)
	}
	if plan.path == "/messages" {
		req.Header.Set("anthropic-version", "2023-06-01")
	}
	for _, h := range []string{"OpenAI-Organization", "OpenAI-Project", "OpenAI-Beta"} {
		if v := c.GetHeader(h); v != "" {
			req.Header.Set(h, v)
		}
	}

	start := time.Now()
	resp, err := g.hc.Do(req)
	if err != nil {
		return false, true, http.StatusBadGateway, []byte("upstream request failed: " + err.Error())
	}
	defer resp.Body.Close()

	// 非 2xx：429/5xx 可重试；其余透传错误体
	if resp.StatusCode < 200 || resp.StatusCode >= 300 {
		errBody, _ := io.ReadAll(resp.Body)
		if resp.StatusCode == http.StatusTooManyRequests || resp.StatusCode >= 500 {
			return false, true, resp.StatusCode, errBody
		}
		return false, false, resp.StatusCode, errBody
	}

	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)
	}
	return true, false, http.StatusOK, nil
}

// 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)
	if err != nil {
		apiError(c, http.StatusBadGateway, "upstream_error", "failed reading upstream response")
		g.recordError(c, ch, nil, start, "read_error")
		return
	}
	if usageRaw := extractUsage(data); usageRaw != nil && sink != nil {
		sink.push(usageRaw)
	}
	out := data
	if bodyConv != nil {
		if converted, cerr := bodyConv(data); cerr == nil {
			out = converted
		}
	}
	_, _ = c.Writer.Write(out)
	if _, ok := c.Get("raw_request"); ok {
		c.Set("raw_response", string(data)) // 上游原始响应（未转换）
	}
	g.finishUsage(c, ch, start, store.UsageStatusSuccess, "")
}

// 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{}
	}

	// 原始响应捕获：仅管理员且开关开启（raw_request 已 set）时累积上游原始行
	_, capture := c.Get("raw_request")
	var rawResp strings.Builder

	// commitRaw 在记账前把已累积的原始响应写入 context
	commitRaw := func() {
		if capture {
			c.Set("raw_response", rawResp.String())
		}
	}

	scanner := newSSEScanner(r)
	for {
		line, err := scanner.Next()
		if line != nil {
			if capture {
				rawResp.Write(line)
			}
			out := line
			if lineConv != nil {
				out = lineConv(line)
			}
			if out != nil {
				if _, werr := w.Write(out); werr != nil {
					// 客户端意外断开：按已生成部分收费（canceled）
					commitRaw()
					g.finishUsage(c, ch, start, store.UsageStatusCanceled, "client_disconnect")
					return
				}
				flusher.Flush()
				if sink != nil {
					sink.outputText += sseContentText(out)
				}
			}
			if usageRaw := scanUsage(line); usageRaw != nil && sink != nil {
				sink.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() {}

// ---------------------------------------------------------------------------
// usage 提取

// usageShape 兼容 chat (prompt/completion)、responses (input/output)、messages (input/output) 命名。
type usageShape struct {
	PromptTokens             int64 `json:"prompt_tokens"`
	CompletionTokens         int64 `json:"completion_tokens"`
	InputTokens              int64 `json:"input_tokens"`
	OutputTokens             int64 `json:"output_tokens"`
	TotalTokens              int64 `json:"total_tokens"`
	CacheReadInputTokens     int64 `json:"cache_read_input_tokens"`
	CacheCreationInputTokens int64 `json:"cache_creation_input_tokens"`
}

// extractUsage 从完整响应体提取 usage 子对象（chat / responses / messages）。
func extractUsage(data []byte) json.RawMessage {
	var m map[string]json.RawMessage
	if json.Unmarshal(data, &m) != nil {
		return nil
	}
	if u := usageFromMap(m); u != nil {
		return u
	}
	if respRaw, ok := m["response"]; ok {
		var resp map[string]json.RawMessage
		if json.Unmarshal(respRaw, &resp) == nil {
			if u := usageFromMap(resp); u != nil {
				return u
			}
		}
	}
	if choices, ok := m["choices"]; ok {
		var cs []map[string]json.RawMessage
		if json.Unmarshal(choices, &cs) == nil {
			for _, ch := range cs {
				if msgRaw, ok := ch["message"]; ok {
					var msg map[string]json.RawMessage
					if json.Unmarshal(msgRaw, &msg) == nil {
						if u := usageFromMap(msg); u != nil {
							return u
						}
					}
				}
			}
		}
	}
	return nil
}

// scanUsage 从 SSE 一行中提取 usage（OpenAI 末块 / responses completed / messages message_delta 等）。
func scanUsage(line []byte) json.RawMessage {
	if !bytes.Contains(line, []byte(`"usage"`)) {
		return nil
	}
	s, ok := sseDataPayload(line)
	if !ok {
		return nil
	}
	var m map[string]json.RawMessage
	if json.Unmarshal([]byte(s), &m) != nil {
		return nil
	}
	if u := usageFromMap(m); u != nil {
		return u
	}
	if respRaw, ok := m["response"]; ok {
		var resp map[string]json.RawMessage
		if json.Unmarshal(respRaw, &resp) == nil {
			if u := usageFromMap(resp); u != nil {
				return u
			}
		}
	}
	return nil
}

// usageFromMap 从 map 顶层、message 或 delta 子对象中取 usage。
// 覆盖三种形态：
//   - OpenAI 流式末块顶层 usage
//   - Anthropic 非流式 / message_start 的 message.usage
//   - Anthropic 流式 message_delta 的 delta.usage（真实 token 计数所在）
func usageFromMap(m map[string]json.RawMessage) json.RawMessage {
	if u, ok := m["usage"]; ok && string(u) != "null" {
		return u
	}
	if msgRaw, ok := m["message"]; ok {
		var msg map[string]json.RawMessage
		if json.Unmarshal(msgRaw, &msg) == nil {
			if u, ok := msg["usage"]; ok && string(u) != "null" {
				return u
			}
		}
	}
	if deltaRaw, ok := m["delta"]; ok {
		var delta map[string]json.RawMessage
		if json.Unmarshal(deltaRaw, &delta) == nil {
			if u, ok := delta["usage"]; ok && string(u) != "null" {
				return u
			}
		}
	}
	return nil
}

// sseScanner 按 SSE 行边界读取（兼容 \n 与 \r\n），保留原始行内容。
type sseScanner struct {
	r *bufio.Reader
}

func newSSEScanner(r io.Reader) *sseScanner { return &sseScanner{r: bufio.NewReaderSize(r, 32*1024)} }

func (s *sseScanner) Next() ([]byte, error) {
	line, err := s.r.ReadBytes('\n')
	if len(line) > 0 {
		return line, nil
	}
	if err != nil {
		return nil, err
	}
	return nil, io.EOF
}

// ---------------------------------------------------------------------------
// 记账

// usageSink 累积多次 usage：合并各事件字段（message_start 给 input，message_delta 给 output）。
// outputText 累积已转发内容文本，用于流式中断时按 tiktoken 估算输出 token。
type usageSink struct {
	us         usageShape
	outputText string
}

func (u *usageSink) push(raw json.RawMessage) {
	if len(raw) == 0 {
		return
	}
	var t usageShape
	if json.Unmarshal(raw, &t) != nil {
		return
	}
	// messages 流式最终事件（message_delta 的 usage）带 cache_* 字段，是上游的最终计费口径，
	// 其中 input_tokens 仅指"非缓存输入"（与 message_start 的"总输入"语义不同）。
	// 整体替换而非字段合并，避免 delta 的非缓存 input 覆盖 start 的总 input 后语义错乱
	// （实际消耗由 finishUsage 按 input + cache_read + cache_creation 汇总）。
	if t.CacheReadInputTokens > 0 || t.CacheCreationInputTokens > 0 {
		u.us = t
		return
	}
	// 零值不覆盖：不同事件携带不同字段
	if t.PromptTokens > 0 {
		u.us.PromptTokens = t.PromptTokens
	}
	if t.CompletionTokens > 0 {
		u.us.CompletionTokens = t.CompletionTokens
	}
	if t.InputTokens > 0 {
		u.us.InputTokens = t.InputTokens
	}
	if t.OutputTokens > 0 {
		u.us.OutputTokens = t.OutputTokens
	}
	if t.CacheReadInputTokens > 0 {
		u.us.CacheReadInputTokens = t.CacheReadInputTokens
	}
	if t.CacheCreationInputTokens > 0 {
		u.us.CacheCreationInputTokens = t.CacheCreationInputTokens
	}
}

// Shape 返回合并后的用量。
func (u *usageSink) Shape() usageShape { return u.us }

// finishUsage 落账：计算成本并异步写入。
func (g *Gateway) finishUsage(c *gin.Context, ch *store.Channel, start time.Time, status, errCode string) {
	uid, _ := c.Get(CtxUserID)
	kid, _ := c.Get(CtxKeyID)
	trace, _ := c.Get(CtxTrace)

	var us usageShape
	var sink *usageSink
	if h, ok := c.Get("usage_raw"); ok {
		if holder, ok := h.(*sinkHolder); ok && holder.sink != nil {
			us = holder.sink.Shape()
			sink = holder.sink
		}
	}

	in := us.PromptTokens + us.InputTokens
	out := us.CompletionTokens + us.OutputTokens
	cacheRead := us.CacheReadInputTokens
	cacheCreate := us.CacheCreationInputTokens

	modelName, _ := c.Get("model_name")
	mn, _ := modelName.(string)

	// 上游未返回真实 usage 时估算（tiktoken）：
	//   - 流式中断（canceled）：上游最终 usage 可能未返回
	//   - 流式正常结束（success）：部分上游（如火山方舟 Anthropic 流式）message_start.usage 恒为 0、
	//     message_delta 不带 usage，只能按已收发内容估算，否则记账为 0 消耗
	// 非流式上游必返回 usage，此处 in/out 非 0 不受影响。
	if status == store.UsageStatusCanceled || status == store.UsageStatusSuccess {
		if in == 0 && cacheRead == 0 && cacheCreate == 0 {
			if est, ok := c.Get("est_input_text"); ok {
				if v, ok2 := est.(string); ok2 && v != "" {
					in = int64(tokenizer.Count(v, mn))
				}
			}
		}
		if out == 0 && sink != nil && sink.outputText != "" {
			out = int64(tokenizer.Count(sink.outputText, mn))
		}
	}

	var model store.Model
	var cost float64
	var modelID uint64
	_ = g.db.Where("name = ?", mn).First(&model).Error
	if model.ID > 0 {
		modelID = model.ID
		// 计价口径：in=非缓存输入、cacheRead=缓存读、cacheCreate=缓存写（Anthropic 语义，
		// messages 流式 message_delta 的 input_tokens 即非缓存部分）。
		// 缓存写按 1.25× 输入价（Anthropic 5m 口径）。
		cost = float64(in)/1e6*model.InputPrice +
			float64(out)/1e6*model.OutputPrice +
			float64(cacheRead)/1e6*model.CacheReadPrice +
			float64(cacheCreate)/1e6*model.InputPrice*1.25
	} else {
		cost = float64(in+cacheRead+cacheCreate)/1e6*0.15 + float64(out)/1e6*0.60 // 无定价模型时按示例价
	}
	// 落库的 input_tokens 存输入总量（含缓存读/写），与上游 message_start 口径一致，便于对账展示。
	in += cacheRead + cacheCreate

	proto, _ := c.Get("protocol")
	p, _ := proto.(string)
	if p == "" {
		p = "chat"
	}
	traceStr, _ := trace.(string)
	latency := int(time.Since(start).Milliseconds())

	if status == store.UsageStatusSuccess && c.Writer.Status() >= 400 {
		status = store.UsageStatusError
	}

	var errCodePtr *string
	if errCode != "" {
		errCodePtr = &errCode
	}

	var uidVal, kidVal uint64
	if u, ok := uid.(uint64); ok {
		uidVal = u
	}
	if k, ok := kid.(uint64); ok {
		kidVal = k
	}
	var chID uint64
	if ch != nil {
		chID = ch.ID
	}

	var rawReq, rawResp string
	if v, ok := c.Get("raw_request"); ok {
		rawReq, _ = v.(string)
	}
	if v, ok := c.Get("raw_response"); ok {
		rawResp, _ = v.(string)
	}

	// 密钥今日 token 用量累计（配额检查用）
	if g.lim != nil && kidVal > 0 {
		g.lim.AddTokens(kidVal, in+out)
	}

	g.rec.Record(&store.UsageLog{
		RequestID:           fmt.Sprintf("trace-%s", traceStr),
		TraceID:             traceStr,
		UserID:              uidVal,
		KeyID:               kidVal,
		ChannelID:           chID,
		ModelID:             modelID,
		ModelName:           mn,
		Protocol:            p,
		InputTokens:         in,
		OutputTokens:        out,
		CacheReadTokens:     cacheRead,
		CacheCreationTokens: cacheCreate,
		InputPrice:          model.InputPrice,
		OutputPrice:         model.OutputPrice,
		CacheReadPrice:      model.CacheReadPrice,
		Cost:                cost,
		LatencyMS:           latency,
		Status:              status,
		ErrorCode:           errCodePtr,
		RawRequest:          rawReq,
		RawResponse:         rawResp,
		CreatedAt:           time.Now().UTC(),
	})
}

// recordError 失败请求的记账（不产生扣费，status=error）。
func (g *Gateway) recordError(c *gin.Context, ch *store.Channel, resp *http.Response, start time.Time, code string) {
	_ = resp
	g.finishUsage(c, ch, start, store.UsageStatusError, code)
}

func now() time.Time { return time.Now() }