goaiapi/src/go/sse.go

package main

import (
	"bytes"
	"encoding/json"
	"io"
	"strings"
)

// streamRecorder 包在 resp.Body 外层：数据被客户端读取时同步流过，
// 一边原样转发（客户端实时收到每个 chunk），一边把全量原文拼回 + 解析 SSE usage。
// 读到 EOF（流结束）时回调 onDone，由它落库。
//
// 这是流式代理的标准做法：用一个 io.ReadCloser 包装上游 body，
// 不做 io.ReadAll，避免缓冲整条流导致客户端要等所有 chunk 到齐才收到第一个字。
type streamRecorder struct {
	src    io.ReadCloser
	raw    bytes.Buffer // 全量原文（拼回所有 chunk）
	parser *sseParser
	onDone func(raw []byte, in, out int, model string)
	done   bool
}

func newStreamRecorder(src io.ReadCloser, format string, onDone func(raw []byte, in, out int, model string)) *streamRecorder {
	return &streamRecorder{
		src:    src,
		parser: newSSEParser(format),
		onDone: onDone,
	}
}

func (s *streamRecorder) Read(p []byte) (int, error) {
	n, err := s.src.Read(p)
	if n > 0 {
		// 原样留一份，并喂给 SSE 解析器累积 usage
		s.raw.Write(p[:n])
		s.parser.feed(p[:n])
	}
	if err == io.EOF && !s.done {
		s.done = true
		s.parser.flush()
		in, out, model := s.parser.result()
		s.onDone(s.raw.Bytes(), in, out, model)
	}
	return n, err
}

func (s *streamRecorder) Close() error {
	// 客户端提前断开等情况：若还没触发 onDone，这里兜底落库一次。
	if !s.done {
		s.done = true
		s.parser.flush()
		in, out, model := s.parser.result()
		s.onDone(s.raw.Bytes(), in, out, model)
	}
	return s.src.Close()
}

// sseParser 按行累积 SSE 数据，跨 chunk 解析 OpenAI / Claude 两家流式 usage。
//
// OpenAI 流式：usage 出现在末尾带 stream_options.include_usage 的 chunk 里，
//   形如 data: {...,"usage":{"prompt_tokens":..,"completion_tokens":..}}；
//   没开 include_usage 则 usage 落空（如实记 0）。
// Claude 流式：message_start 事件给 input_tokens（usage.input_tokens），
//   message_delta 事件累加 output_tokens（usage.output_tokens），需跨事件攒。
type sseParser struct {
	format string
	buf    bytes.Buffer // 未成行的残留字节

	in    int
	out   int
	model string
}

func newSSEParser(format string) *sseParser {
	return &sseParser{format: format}
}

// feed 接收任意切片的字节，按行切出完整的 SSE 行交给 parseLine。
func (p *sseParser) feed(b []byte) {
	p.buf.Write(b)
	data := p.buf.Bytes()
	for {
		i := bytes.IndexByte(data, '\n')
		if i < 0 {
			break
		}
		line := data[:i]
		p.parseLine(line)
		data = data[i+1:]
	}
	// 保留未成行的尾部
	rest := make([]byte, len(data))
	copy(rest, data)
	p.buf.Reset()
	p.buf.Write(rest)
}

// flush 处理流结束时缓冲里残留的最后一行（可能没有结尾换行）。
func (p *sseParser) flush() {
	if p.buf.Len() > 0 {
		p.parseLine(p.buf.Bytes())
		p.buf.Reset()
	}
}

// parseLine 解析单行 SSE。只关心 "data:" 行里的 JSON。
func (p *sseParser) parseLine(line []byte) {
	s := strings.TrimRight(string(line), "\r")
	s = strings.TrimSpace(s)
	if !strings.HasPrefix(s, "data:") {
		return
	}
	payload := strings.TrimSpace(strings.TrimPrefix(s, "data:"))
	if payload == "" || payload == "[DONE]" {
		return
	}

	var m map[string]json.RawMessage
	if err := json.Unmarshal([]byte(payload), &m); err != nil {
		return // 非 JSON 行，跳过
	}

	// model 字段（两家都有，取到就记）
	if raw, ok := m["model"]; ok {
		var mod string
		if json.Unmarshal(raw, &mod) == nil && mod != "" {
			p.model = mod
		}
	}

	// usage 可能直接在顶层（OpenAI 末尾 chunk），
	// 也可能在 message.usage 里（Claude message_start）。
	if raw, ok := m["usage"]; ok {
		p.applyUsage(raw)
	}
	if raw, ok := m["message"]; ok {
		var msg map[string]json.RawMessage
		if json.Unmarshal(raw, &msg) == nil {
			if u, ok := msg["usage"]; ok {
				p.applyUsage(u)
			}
			if mraw, ok := msg["model"]; ok {
				var mod string
				if json.Unmarshal(mraw, &mod) == nil && mod != "" {
					p.model = mod
				}
			}
		}
	}
}

// applyUsage 从一段 usage JSON 里抽 token 数，兼容两家命名。
// Claude 的 output_tokens 在 message_delta 里是累计值（取最大即可），
// input_tokens 在 message_start 给定后不再变；OpenAI 末尾一次性给全量。
func (p *sseParser) applyUsage(raw json.RawMessage) {
	var u struct {
		// OpenAI
		PromptTokens     int `json:"prompt_tokens"`
		CompletionTokens int `json:"completion_tokens"`
		// Claude
		InputTokens  int `json:"input_tokens"`
		OutputTokens int `json:"output_tokens"`
	}
	if json.Unmarshal(raw, &u) != nil {
		return
	}
	// 取较大值，避免后到的事件覆盖成 0（Claude 各事件字段不全）。
	if u.PromptTokens > p.in {
		p.in = u.PromptTokens
	}
	if u.InputTokens > p.in {
		p.in = u.InputTokens
	}
	if u.CompletionTokens > p.out {
		p.out = u.CompletionTokens
	}
	if u.OutputTokens > p.out {
		p.out = u.OutputTokens
	}
}

func (p *sseParser) result() (in, out int, model string) {
	return p.in, p.out, p.model
}