Files
openteam/server/internal/proxy/passthrough.go
T
Sakurasan c39c3e7da4 feat: 补齐上游细节、渠道观测与模型定价元数据
- 上游 429 的 Retry-After 转成渠道冷却时长(夹在 10 分钟内;比默认值短则不缩短,
  否则 1s 就把坏渠道放回轮转)。失败计数改由 proxyOne 就地报——只有它读得到响应头。
- 故障转移时间预算 OT_PROXY_FAILOVER_BUDGET(默认 30s):只约束第二个及以后的渠道,
  首个不受限(长回答本就要跑几分钟),但坏渠道不再把客户端拖到 N×timeout。
- 流式空闲心跳 OT_PROXY_STREAM_HEARTBEAT(默认 15s):补一行 SSE 注释 `: ping`,
  防止 nginx/Cloudflare 的空闲超时掐掉长思考。心跳与正文共用一把写的锁,
  gin 的 ResponseWriter 不是并发安全的;注释行不进转换也不计费。
- 渠道观测:GET /api/v1/admin/channels/stats?window=5m|1h|6h|24h|7d 聚合请求量、
  成功率、P50/P95/max 延迟、成本与 top_errors;分位只取成功样本,canceled(客户端
  主动断开)不计入渠道错误率。渠道页桌面表格加「近N小时」列、移动端加同内容条,
  顶部可切窗口。
- 修一个让统计瞎掉的归因问题:全部渠道失败时 recordError 传的是 nil,
  usage_logs 里 channel_id=0,这类失败在渠道维度完全看不见。改为归到最后一个
  真正试过的渠道。
- 模型定价与元数据:Model 新增 cache_write_price(0 则回退输入价×1.25,可显式配
  Anthropic 1h 的 ×2)、context_length 与 supports_tools/vision/reasoning;
  finishUsage 走 CacheWritePriceOr()。/v1/models 叠加用户与密钥白名单过滤
  (原来会把限制外的模型名直接泄漏在列表里)并返回元数据,价格不外泄。
- /v1/models 的 Anthropic 分支原是死代码:协议只按路径判,而两种客户端共用这个路径。
  改为按 x-api-key / anthropic-version 头判,OpenAI SDK 从不发这两个头,互不影响。
- NewGateway 的十个位置参数换成 Options 结构体(三个 Duration 两个 int,传错编译期发现不了)。

新增 6 个用例:Retry-After 两种写法与夹顶、重试预算截止、空闲心跳出帧、
/v1/models 过滤与双形状、缓存写价两条路径、失败归因。
2026-10-17 04:09:39 +08:00

805 lines
24 KiB
Go
Raw Blame History

This file contains ambiguous Unicode characters
This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.
package proxy
import (
"bufio"
"bytes"
"context"
"crypto/rand"
"encoding/hex"
"encoding/json"
"errors"
"fmt"
"io"
"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/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 ""
}
// retryAfter 解析上游 Retry-After(delta-秒或 HTTP 日期),用于按对端意愿延长渠道冷却。
// 解析不出或为非正返回 0,交给 channel.Service 用默认时长(并夹到上限)。
func retryAfter(h http.Header) time.Duration {
v := strings.TrimSpace(h.Get("Retry-After"))
if v == "" {
return 0
}
if secs, err := strconv.Atoi(v); err == nil {
if secs <= 0 {
return 0
}
return time.Duration(secs) * time.Second
}
if t, err := http.ParseTime(v); err == nil {
if d := time.Until(t); d > 0 {
return d
}
}
return 0
}
// 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")
// 最后真正试过的渠道:全部失败时把错误归到它头上,否则 usage_logs 里
// channel_id=0,渠道维度的统计与排障就看不到这类失败。
var lastCh *store.Channel
startAll := time.Now()
// 从选中的渠道开始遍历,到末尾后再从头遍历到选中渠道之前
for offset := 0; offset < len(cands); offset++ {
// 重试预算只约束第二个及以后的渠道:首个渠道不受限(长回答本就要跑几分钟),
// 但坏渠道一个接一个地等下去会让客户端挂到 N×timeout。
if offset > 0 && g.failoverBudget > 0 && time.Since(startAll) > g.failoverBudget {
lastBody = []byte("failover budget exhausted: " + string(lastBody))
break
}
idx := (startIdx + offset) % len(cands)
cand := cands[idx]
ch := cand.Channel
// 客户端已断开:上游请求会随 context 立刻失败,换渠道也没有意义
if c.Request.Context().Err() != nil {
g.finishUsage(c, ch, now(), store.UsageStatusCanceled, "client_disconnect")
return
}
if g.ch.InCooldown(ch.ID) {
continue // 本次请求期间已被并发请求判入运行期冷却
}
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 // 渠道满载,溢出到下一个
}
// 失败/成功计数由 proxyOne 就地报:只有它读得到上游响应头(429 的 Retry-After)
lastCh = ch
done, retry, st, b := g.proxyOne(c, ch, plan, stream, sink)
release()
if done {
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, lastCh, nil, now(), "all_channels_failed")
}
// 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) {
upKey, err := g.ch.UpstreamKey(ch)
if err != nil {
g.ch.ReportFailure(ch, 0)
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
}
}
}
// 上游请求绑定客户端 context:客户端断开即取消上游,不再空转生成、继续付费。
// 非流式再叠加渠道超时作为整段硬上限;流式只靠 Transport.ResponseHeaderTimeout
// 约束首包,正文时长不设限(长回答不能被超时砍断)。
limit := g.channelTimeout(ch)
ctx := c.Request.Context()
if !stream && limit > 0 {
var cancel context.CancelFunc
ctx, cancel = context.WithTimeout(ctx, limit)
defer cancel()
}
req, err := http.NewRequestWithContext(ctx, 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.clientFor(limit).Do(req)
if err != nil {
// 客户端断开导致的取消:就地记账并终结本次请求(不计入渠道失败,也不换渠道)
if c.Request.Context().Err() != nil {
g.finishUsage(c, ch, start, store.UsageStatusCanceled, "client_disconnect")
return true, false, 0, nil
}
g.ch.ReportFailure(ch, 0) // 网络错误/超时
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 {
g.ch.ReportFailure(ch, 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)
}
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 {
// 走到这里响应头已按 200 写出,无法再回错误体,只按中断原因记账
switch {
case c.Request.Context().Err() != nil:
g.finishUsage(c, ch, start, store.UsageStatusCanceled, "client_disconnect")
case errors.Is(err, context.DeadlineExceeded):
g.finishUsage(c, ch, start, store.UsageStatusError, "upstream_timeout")
default:
g.finishUsage(c, ch, start, store.UsageStatusError, "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())
}
}
lw := &lockedWriter{w: w, f: flusher}
lw.lastN.Store(time.Now().UnixNano())
stopHB := lw.heartbeat(g.streamHeartbeat)
defer stopHB()
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 := lw.write(out); werr != nil {
// 客户端意外断开:按已生成部分收费(canceled)
commitRaw()
g.finishUsage(c, ch, start, store.UsageStatusCanceled, "client_disconnect")
return
}
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() {}
// 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
}
// 心跳写失败说明客户端已断开,交给主循环下一次写出时统一记账
_, _ = lw.write([]byte(": ping\n\n"))
}
}
}()
return func() { once.Do(func() { close(done) }) }
}
// ---------------------------------------------------------------------------
// 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.CacheWritePriceOr()
} 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() }