feat: 会话亲和粘住缓存渠道;透传 opaque 字段;用量记思考/1h 缓存写/实际模型

会话亲和(proxy/affinity.go):同一会话粘住上次应答的渠道,命中上游提示词
缓存而非被加权随机打散后全价重算。会话身份 = 密钥 + 模型 + 请求文本稳定前缀
摘要,不依赖各客户端会话头;TTL 默认 5m(OT_PROXY_AFFINITY_TTL,0=关闭),
成功应答刷新。亲和渠道必须在候选集内(健康/未冷却/绑定),否则回退加权随机。
proxyOne 增加 served 返回值区分「应答成功」与「已终结但失败」。

opaque 透传(convert):chat↔responses 互转原样携带 metadata、service_tier、
parallel_tool_calls。此前逐字段重建请求体会丢它们,而中转上游会校验/使用。

用量细化(proxy+store):usage_logs 新增 reasoning_tokens(展示口径,不重复
计费)、cache_write_1h_tokens(Anthropic 1h 缓存写按 2× 输入价,此前统一
1.25× 少算 60%)、served_model(上游实际应答模型,映射/聚合商场景对账)。
schema v6→v7。

验证:go vet ./... 干净;go test ./... 与 go test -race ./... 全绿;
新增测试覆盖亲和粘性与冷却回退、opaque 双向透传、reasoning 不重复计费、
1h/5m 缓存写分开计价、served_model 提取(非流式与流式)。
This commit is contained in:
Sakurasan committed 2026-10-30 11:41:51 +08:00
1 parent e7f8a5e5ad
commit 04dfb4285d
14 files changed
+541 -38

No files matched your search

+1
View File
@@ -39,6 +39,7 @@ func serve() error {
FailoverBudget: cfg.Proxy.FailoverBudget,
StreamHeartbeat: cfg.Proxy.StreamHeartbeat,
FirstTokenWait: cfg.Proxy.FirstTokenWait,
AffinityTTL: cfg.Proxy.AffinityTTL,
})
router := api.NewRouter(a, gw)
+3
View File
@@ -82,6 +82,7 @@ type ProxyConfig struct {
FailoverBudget time.Duration // 故障转移给"第二个及以后的渠道"留的时间预算,0=不限
StreamHeartbeat time.Duration // 流式空闲多久发一次 SSE 注释心跳,0=关闭
FirstTokenWait time.Duration // 流式放行前等内容的最长时间,0=不限;超时未出内容换下一个渠道
AffinityTTL time.Duration // 会话亲和 TTL,0=关闭;同一会话粘住上次应答的渠道以命中上游缓存
HealthInterval time.Duration // 渠道健康检查周期
HealthFailThreshold int // 连续失败 N 次进 cooldown
LogRaw bool // 记录管理员原始请求体+响应到 usage_logs(调试用,默认关)
@@ -150,6 +151,7 @@ func Load() (*Config, error) {
v.SetDefault("proxy.failover_budget", "30s")
v.SetDefault("proxy.stream_heartbeat", "15s")
v.SetDefault("proxy.first_token_wait", "0s")
v.SetDefault("proxy.affinity_ttl", "5m")
v.SetDefault("proxy.health_interval", "60s")
v.SetDefault("proxy.health_fail_threshold", 2)
v.SetDefault("proxy.log_raw", false)
@@ -200,6 +202,7 @@ func Load() (*Config, error) {
FailoverBudget: v.GetDuration("proxy.failover_budget"),
StreamHeartbeat: v.GetDuration("proxy.stream_heartbeat"),
FirstTokenWait: v.GetDuration("proxy.first_token_wait"),
AffinityTTL: v.GetDuration("proxy.affinity_ttl"),
HealthInterval: v.GetDuration("proxy.health_interval"),
HealthFailThreshold: v.GetInt("proxy.health_fail_threshold"),
LogRaw: v.GetBool("proxy.log_raw"),
+84
View File
@@ -0,0 +1,84 @@
package proxy
import (
"crypto/sha256"
"encoding/hex"
"fmt"
"sync"
"time"
)
// 会话亲和(借鉴 magpie 的 cache-economics 思路,做服务端可用的简化版):
// 同一会话粘住上次应答的渠道,让上游的提示词缓存(整段上下文每轮重发)被再次
// 命中,而不是被加权随机打散后全价重算。
//
// 会话身份 = 密钥 + 模型 + 请求文本稳定前缀的摘要。编码代理每轮重发全量上下文,
// 系统提示与最早的用户消息不变,取前缀即可区分会话;不需要解析各家客户端的
// 会话头。命中亲和的渠道仍必须在候选集里(健康、未冷却、认识该模型)——
// 冷却或下线后自然回退到加权随机。
//
// TTL 默认 5 分钟(对齐 Anthropic 提示词缓存窗口),每次成功应答刷新;
// 上游不报缓存的(OpenAI 自动前缀缓存)同样受益:粘住就是省。
const (
// affinityMax 亲和表上限:超过先清过期,仍超则不再新增(防御性,正常流量到不了)。
affinityMax = 8192
// affinityPrefix 会话身份摘要取请求文本的前缀字节数。
affinityPrefix = 512
)
type affinityTable struct {
mu sync.Mutex
m map[string]affinityEntry
}
type affinityEntry struct {
channelID uint64
until time.Time
}
func newAffinityTable() *affinityTable {
return &affinityTable{m: map[string]affinityEntry{}}
}
// get 返回该会话上次应答的渠道;过期条目顺手清理。
func (t *affinityTable) get(key string) (uint64, bool) {
t.mu.Lock()
defer t.mu.Unlock()
e, ok := t.m[key]
if !ok {
return 0, false
}
if time.Now().After(e.until) {
delete(t.m, key)
return 0, false
}
return e.channelID, true
}
// put 记录(或刷新)会话→渠道的绑定。
func (t *affinityTable) put(key string, channelID uint64, ttl time.Duration) {
t.mu.Lock()
defer t.mu.Unlock()
if len(t.m) >= affinityMax {
now := time.Now()
for k, e := range t.m {
if now.After(e.until) {
delete(t.m, k)
}
}
if len(t.m) >= affinityMax {
return // 表满且全是活的:放弃本次记录,旧绑定继续服务
}
}
t.m[key] = affinityEntry{channelID: channelID, until: time.Now().Add(ttl)}
}
// affinityKey 会话身份:密钥 + 模型 + 请求文本稳定前缀的摘要。
func affinityKey(keyID uint64, model, firstText string) string {
if len(firstText) > affinityPrefix {
firstText = firstText[:affinityPrefix]
}
h := sha256.Sum256([]byte(firstText))
return fmt.Sprintf("%d|%s|%s", keyID, model, hex.EncodeToString(h[:8]))
}
+87
View File
@@ -0,0 +1,87 @@
package proxy
import (
"io"
"net/http"
"net/http/httptest"
"strings"
"sync/atomic"
"testing"
"time"
"github.com/openteam/server/internal/channel"
"github.com/openteam/server/internal/store"
)
// 会话亲和:同一会话(同密钥+模型+同首条消息)粘住上次应答的渠道——
// 即使加权随机此时会选别的渠道;不同会话不受影响,照常负载均衡。
func TestAffinitySticksToAnsweringChannel(t *testing.T) {
var ch1Hits, ch2Hits atomic.Int64
ch1 := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
ch1Hits.Add(1)
_, _ = io.WriteString(w, `{"id":"ok","choices":[{"message":{"role":"assistant","content":"from-1"},"finish_reason":"stop"}],"usage":{"prompt_tokens":1,"completion_tokens":1}}`)
}))
defer ch1.Close()
ch2 := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
ch2Hits.Add(1)
_, _ = io.WriteString(w, `{"id":"ok","choices":[{"message":{"role":"assistant","content":"from-2"},"finish_reason":"stop"}],"usage":{"prompt_tokens":1,"completion_tokens":1}}`)
}))
defer ch2.Close()
h := newHarness(t, ch1.URL, 0, true, nil)
h.gw.affinityTTL = 5 * time.Minute
addChannel(t, h, "c2", ch2.URL, 1)
pushWeightTo(t, h, "c1", 1000)
sameConv := `{"model":"m1","messages":[{"role":"user","content":"会话的第一条消息"}]}`
// 第一轮:权重压在渠道1,应答来自渠道1,亲和绑定建立
if w := h.run(t, nil, sameConv); w.Code != http.StatusOK || !strings.Contains(w.Body.String(), "from-1") {
t.Fatalf("第一轮应来自渠道1:%d %s", w.Code, w.Body.String())
}
// 反转权重:加权随机此刻会选渠道2
pushWeightTo(t, h, "c1", 1)
pushWeightTo(t, h, "c2", 1000)
// 同一会话第二轮:仍应粘在渠道1(缓存值得保留)
if w := h.run(t, nil, sameConv); !strings.Contains(w.Body.String(), "from-1") {
t.Fatalf("同一会话应粘住渠道1:%s", w.Body.String())
}
// 不同会话(不同首条消息):走加权随机,落到渠道2
otherConv := `{"model":"m1","messages":[{"role":"user","content":"另一个会话的开场"}]}`
if w := h.run(t, nil, otherConv); !strings.Contains(w.Body.String(), "from-2") {
t.Fatalf("不同会话应走加权随机落到渠道2:%s", w.Body.String())
}
if ch1Hits.Load() != 2 || ch2Hits.Load() != 1 {
t.Fatalf("命中次数不符:ch1=%d ch2=%d", ch1Hits.Load(), ch2Hits.Load())
}
}
// 亲和渠道进入运行期冷却后自动失效,回退加权随机;不会把请求发给冷却中的渠道。
func TestAffinityRespectsCooldown(t *testing.T) {
ch1 := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
t.Error("冷却中的渠道1 不应收到请求")
_, _ = io.WriteString(w, `{"id":"bad"}`)
}))
defer ch1.Close()
ch2 := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
_, _ = io.WriteString(w, `{"id":"ok","choices":[{"message":{"role":"assistant","content":"from-2"},"finish_reason":"stop"}],"usage":{"prompt_tokens":1,"completion_tokens":1}}`)
}))
defer ch2.Close()
h := newHarness(t, ch1.URL, 0, true, nil)
h.gw.affinityTTL = 5 * time.Minute
addChannel(t, h, "c2", ch2.URL, 1)
// 预置亲和绑定指向渠道1,随后让渠道1 进运行期冷却
var ch1row store.Channel
h.db.Where("name = ?", "c1").First(&ch1row)
h.gw.aff.put(affinityKey(1, "m1", "会话的第一条消息"), ch1row.ID, 5*time.Minute)
for i := 0; i < 3; i++ {
h.gw.ch.ReportFailure(&ch1row, channel.FailRate, 0)
}
w := h.run(t, nil, `{"model":"m1","messages":[{"role":"user","content":"会话的第一条消息"}]}`)
if !strings.Contains(w.Body.String(), "from-2") {
t.Fatalf("冷却后亲和应失效并切换渠道2:%d %s", w.Code, w.Body.String())
}
}
@@ -1136,3 +1136,49 @@ func TestStreamResponsesToChatTextStillWorks(t *testing.T) {
t.Fatalf("missing finish chunk: %s", out)
}
}
// OpenAI 两协议同形的 opaque 字段(metadata / service_tier / parallel_tool_calls)
// 在 chat↔responses 互转时原样透传——上游可能是另一层中转,会校验/使用它们。
func TestChatToResponsesCarriesOpaque(t *testing.T) {
body := `{"model":"m1","messages":[{"role":"user","content":"hi"}],"service_tier":"priority","metadata":{"session_id":"s1"},"parallel_tool_calls":false}`
out, err := ConvertRequest([]byte(body), ProtoChat, ProtoResponses)
if err != nil {
t.Fatal(err)
}
var m map[string]any
if json.Unmarshal(out, &m) != nil {
t.Fatal("invalid json")
}
if m["service_tier"] != "priority" {
t.Fatalf("service_tier 丢失:%v", m["service_tier"])
}
md, _ := m["metadata"].(map[string]any)
if md == nil || md["session_id"] != "s1" {
t.Fatalf("metadata 丢失:%v", m["metadata"])
}
if m["parallel_tool_calls"] != false {
t.Fatalf("parallel_tool_calls 丢失:%v", m["parallel_tool_calls"])
}
}
func TestResponsesToChatCarriesOpaque(t *testing.T) {
body := `{"model":"m1","input":"hi","service_tier":"flex","metadata":{"installation_id":"i9"},"parallel_tool_calls":true}`
out, err := ConvertRequest([]byte(body), ProtoResponses, ProtoChat)
if err != nil {
t.Fatal(err)
}
var m map[string]any
if json.Unmarshal(out, &m) != nil {
t.Fatal("invalid json")
}
if m["service_tier"] != "flex" {
t.Fatalf("service_tier 丢失:%v", m["service_tier"])
}
md, _ := m["metadata"].(map[string]any)
if md == nil || md["installation_id"] != "i9" {
t.Fatalf("metadata 丢失:%v", m["metadata"])
}
if m["parallel_tool_calls"] != true {
t.Fatalf("parallel_tool_calls 丢失:%v", m["parallel_tool_calls"])
}
}
@@ -39,6 +39,12 @@ type chatReq struct {
MaxTokens *int `json:"max_tokens"`
Stop []string `json:"stop"`
Stream bool `json:"stream"`
// OpenAI 两协议同形的 opaque 字段:转换到 Responses 时原样透传。
// 上游可能是另一层中转,会校验/使用它们(service_tier 影响优先级,
// metadata 承载 Codex 等客户端的会话/安装标识)。
ServiceTier string `json:"service_tier"`
Metadata map[string]string `json:"metadata"`
ParallelToolCalls *bool `json:"parallel_tool_calls"`
}
// chatToMessagesReq 将 OpenAI Chat 请求转为 Anthropic Messages 请求。
@@ -27,6 +27,13 @@ func responsesToChatReq(body []byte) ([]byte, error) {
if v, ok := m["max_output_tokens"]; ok {
out["max_tokens"] = v
}
// OpenAI 两协议同形的 opaque 字段原样透传:metadata(map[string]string)、
// service_tier、parallel_tool_calls——丢了它们,中转上游的校验与路由会失效
for _, k := range []string{"metadata", "service_tier", "parallel_tool_calls"} {
if v := rawJSON(m, k); v != nil {
out[k] = json.RawMessage(v)
}
}
var msgs []any
if ins := str(rawJSON(m, "instructions")); ins != "" {
@@ -199,6 +206,16 @@ func chatToResponsesReq(body []byte) ([]byte, error) {
if req.MaxTokens != nil {
out["max_output_tokens"] = *req.MaxTokens
}
// OpenAI 两协议同形的 opaque 字段原样透传(见 chatReq 注释)
if req.ServiceTier != "" {
out["service_tier"] = req.ServiceTier
}
if len(req.Metadata) > 0 {
out["metadata"] = req.Metadata
}
if req.ParallelToolCalls != nil {
out["parallel_tool_calls"] = *req.ParallelToolCalls
}
var system []string
var input []any
+8 -2
View File
@@ -50,8 +50,11 @@ type Gateway struct {
streamHeartbeat time.Duration
// 流式放行前等多久内容(0=不限制)。超时未出内容取消本次尝试换下一个渠道。
firstTokenWait time.Duration
tpMu sync.Mutex
tps map[time.Duration]*http.Transport
// 会话亲和 TTL(0=关闭):同一会话粘住上次应答的渠道,保住上游提示词缓存。
affinityTTL time.Duration
aff *affinityTable
tpMu sync.Mutex
tps map[time.Duration]*http.Transport
policyMu sync.Mutex
policy modelPolicy
@@ -167,6 +170,7 @@ type Options struct {
FailoverBudget time.Duration // 故障转移重试的时间预算
StreamHeartbeat time.Duration // 流式空闲心跳间隔
FirstTokenWait time.Duration // 流式放行前等内容的最长时间(0=不限)
AffinityTTL time.Duration // 会话亲和 TTL(0=关闭)
}
func NewGateway(db *gorm.DB, enc *crypto.Encryptor, rec *usage.Recorder, lim *ratelimit.Limiter, o Options) *Gateway {
@@ -183,6 +187,8 @@ func NewGateway(db *gorm.DB, enc *crypto.Encryptor, rec *usage.Recorder, lim *ra
failoverBudget: o.FailoverBudget,
streamHeartbeat: o.StreamHeartbeat,
firstTokenWait: o.FirstTokenWait,
affinityTTL: o.AffinityTTL,
aff: newAffinityTable(),
tps: map[time.Duration]*http.Transport{},
}
}
+165 -33
View File
@@ -197,13 +197,28 @@ func (g *Gateway) doProxy(c *gin.Context, cands []channel.Candidate, clientProto
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
// 会话亲和:上次应答过该会话的渠道优先(候选集已过滤健康/冷却/绑定),
// 让上游提示词缓存被再次命中;命中不了(首次/冷却/下线)走加权随机。
affKey := g.affinityKeyFor(c)
startIdx := -1
if affKey != "" {
if chID, ok := g.aff.get(affKey); ok {
for i, cand := range cands {
if cand.Channel.ID == chID {
startIdx = i
break
}
}
}
}
if startIdx < 0 {
// 加权随机选择起始渠道
picked := g.ch.Pick(cands)
for i, cand := range cands {
if cand.Channel.ID == picked.ID {
startIdx = i
break
}
}
}
@@ -253,9 +268,12 @@ func (g *Gateway) doProxy(c *gin.Context, cands []channel.Candidate, clientProto
}
// 失败/成功计数由 proxyOne 就地报:只有它读得到上游响应头(429 的 Retry-After)
lastCh = ch
done, retry, st, b := g.proxyOne(c, ch, plan, stream, sink, sh, clientProto)
done, served, retry, st, b := g.proxyOne(c, ch, plan, stream, sink, sh, clientProto)
release()
if done {
if served && affKey != "" {
g.aff.put(affKey, ch.ID, g.affinityTTL)
}
return
}
if !retry {
@@ -283,6 +301,23 @@ func (g *Gateway) doProxy(c *gin.Context, cands []channel.Candidate, clientProto
g.recordError(c, lastCh, nil, now(), "all_channels_failed")
}
// affinityKeyFor 计算本请求的会话亲和键;亲和关闭或缺少身份要素时返回空。
func (g *Gateway) affinityKeyFor(c *gin.Context) string {
if g.affinityTTL <= 0 {
return ""
}
kid, ok := c.Get(CtxKeyID)
kID, ok2 := kid.(uint64)
if !ok || !ok2 {
return ""
}
mn, _ := c.Get("model_name")
model, _ := mn.(string)
text, _ := c.Get("est_input_text")
firstText, _ := text.(string)
return affinityKey(kID, model, firstText)
}
// streamErrorEvent 把错误构造成客户端协议的 SSE 错误事件(响应头已写出时的唯一交付方式)。
func streamErrorEvent(proto string, status int, msg string) []byte {
errObj := map[string]any{"message": msg, "type": "api_error", "code": "upstream_error", "status": status}
@@ -301,13 +336,14 @@ func streamErrorEvent(proto string, status int, msg string) []byte {
// proxyOne 对单个渠道执行一次代理。
// 返回:done=本次请求已终结(客户端响应已写出或已记账,停止故障转移);
// retry=是否可安全换渠道重试;status+respBody=失败信息。
// served=该渠道成功应答(会话亲和据此绑定);retry=是否可安全换渠道重试;
// status+respBody=失败信息。
// sh 为流式容错缓冲(非流式为 nil);sink 为请求级记账 sink。
func (g *Gateway) proxyOne(c *gin.Context, ch *store.Channel, plan *upstreamPlan, stream bool, sink *usageSink, sh *streamHold, clientProto string) (done bool, retry bool, status int, respBody []byte) {
func (g *Gateway) proxyOne(c *gin.Context, ch *store.Channel, plan *upstreamPlan, stream bool, sink *usageSink, sh *streamHold, clientProto string) (done bool, served bool, retry bool, status int, respBody []byte) {
upKey, err := g.ch.UpstreamKey(ch)
if err != nil {
g.ch.ReportFailure(ch, channel.FailNet, 0)
return false, true, http.StatusInternalServerError, []byte("failed to decrypt channel key")
return false, false, true, http.StatusInternalServerError, []byte("failed to decrypt channel key")
}
upBody := plan.body
@@ -351,7 +387,7 @@ func (g *Gateway) proxyOne(c *gin.Context, ch *store.Channel, plan *upstreamPlan
req, err := http.NewRequestWithContext(tryCtx, 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")
return false, false, false, http.StatusInternalServerError, []byte("failed to build upstream request")
}
req.Header.Set("Content-Type", "application/json")
req.Header.Set("Authorization", "Bearer "+upKey)
@@ -380,10 +416,10 @@ func (g *Gateway) proxyOne(c *gin.Context, ch *store.Channel, plan *upstreamPlan
// 客户端断开导致的取消:就地记账并终结本次请求(不计入渠道失败,也不换渠道)
if c.Request.Context().Err() != nil {
g.finishUsage(c, ch, start, store.UsageStatusCanceled, "client_disconnect")
return true, false, 0, nil
return true, false, false, 0, nil
}
g.ch.ReportFailure(ch, channel.FailNet, 0) // 网络错误/超时
return false, true, http.StatusBadGateway, []byte("upstream request failed: " + err.Error())
return false, false, true, http.StatusBadGateway, []byte("upstream request failed: " + err.Error())
}
defer resp.Body.Close()
@@ -394,9 +430,9 @@ func (g *Gateway) proxyOne(c *gin.Context, ch *store.Channel, plan *upstreamPlan
kind, retryable := classifyUpstreamError(resp.StatusCode, errBody)
if retryable {
g.ch.ReportFailure(ch, kind, retryAfter(resp.Header))
return false, true, resp.StatusCode, errBody
return false, false, true, resp.StatusCode, errBody
}
return false, false, resp.StatusCode, errBody
return false, false, false, resp.StatusCode, errBody
}
g.ch.ReportSuccess(ch)
@@ -411,7 +447,7 @@ func (g *Gateway) proxyOne(c *gin.Context, ch *store.Channel, plan *upstreamPlan
if stream {
asked, got = "流式", "整包响应"
}
return false, true, http.StatusBadGateway,
return false, false, true, http.StatusBadGateway,
[]byte(ch.Name + " answered a " + asked + " request with " + got + " (didNotStream)")
}
@@ -425,25 +461,28 @@ func (g *Gateway) proxyOne(c *gin.Context, ch *store.Channel, plan *upstreamPlan
// copyAndCapture 非流式:先读完整上游响应再写客户端——响应头与正文都在"确定成功"后
// 才写出,读失败(网络错误/超时)发生在写出之前,可换渠道重试。
// 此前实现先写 200 头再读,超时只能给客户端留一个 200 空响应。
func (g *Gateway) copyAndCapture(c *gin.Context, ch *store.Channel, resp *http.Response, start time.Time, bodyConv func([]byte) ([]byte, error), sink *usageSink) (done bool, retry bool, status int, respBody []byte) {
func (g *Gateway) copyAndCapture(c *gin.Context, ch *store.Channel, resp *http.Response, start time.Time, bodyConv func([]byte) ([]byte, error), sink *usageSink) (done bool, served bool, retry bool, status int, respBody []byte) {
data, err := io.ReadAll(resp.Body)
if err != nil {
switch {
case c.Request.Context().Err() != nil:
// 客户端断开:就地记账终结
g.finishUsage(c, ch, start, store.UsageStatusCanceled, "client_disconnect")
return true, false, 0, nil
return true, false, false, 0, nil
case errors.Is(err, context.DeadlineExceeded):
g.ch.ReportFailure(ch, channel.FailNet, 0)
return false, true, http.StatusBadGateway, []byte(ch.Name + " upstream timeout")
return false, false, true, http.StatusBadGateway, []byte(ch.Name + " upstream timeout")
default:
g.ch.ReportFailure(ch, channel.FailNet, 0)
return false, true, http.StatusBadGateway, []byte(ch.Name + " read error: " + err.Error())
return false, false, true, http.StatusBadGateway, []byte(ch.Name + " read error: " + err.Error())
}
}
if usageRaw := extractUsage(data); usageRaw != nil && sink != nil {
sink.push(usageRaw)
}
if sink != nil && sink.servedModel == "" {
sink.servedModel = extractServedModel(data)
}
out := data
if bodyConv != nil {
if converted, cerr := bodyConv(data); cerr == nil {
@@ -457,7 +496,7 @@ func (g *Gateway) copyAndCapture(c *gin.Context, ch *store.Channel, resp *http.R
c.Set("raw_response", string(data)) // 上游原始响应(未转换)
}
g.finishUsage(c, ch, start, store.UsageStatusSuccess, "")
return true, false, 0, nil
return true, true, false, 0, nil
}
// streamHoldCopy 流式转发:边读上游 SSE 边经 hold 写客户端(放行前缓冲)。
@@ -467,7 +506,7 @@ func (g *Gateway) copyAndCapture(c *gin.Context, ch *store.Channel, resp *http.R
// 失败渠道的 message_start 用量不能计入最终账单。
func (g *Gateway) streamHoldCopy(c *gin.Context, sh *streamHold, ch *store.Channel, r io.Reader, start time.Time,
lineConv func([]byte) []byte, clientProto string, trySink, reqSink *usageSink, onRelease func(), tryCtx context.Context,
) (done bool, retry bool, status int, respBody []byte) {
) (done bool, served bool, retry bool, status int, respBody []byte) {
// 原始响应捕获:仅管理员且开关开启(raw_request 已 set)时累积上游原始行
_, capture := c.Get("raw_request")
var rawResp strings.Builder
@@ -479,9 +518,9 @@ func (g *Gateway) streamHoldCopy(c *gin.Context, sh *streamHold, ch *store.Chann
scanner := newSSEScanner(r)
releaseSeen := false
fail := func(msg string, kind channel.FailKind) (bool, bool, int, []byte) {
fail := func(msg string, kind channel.FailKind) (bool, bool, bool, int, []byte) {
g.ch.ReportFailure(ch, kind, 0)
return false, true, http.StatusBadGateway, []byte(ch.Name + ": " + msg)
return false, false, true, http.StatusBadGateway, []byte(ch.Name + ": " + msg)
}
for {
line, err := scanner.Next()
@@ -504,7 +543,7 @@ func (g *Gateway) streamHoldCopy(c *gin.Context, sh *streamHold, ch *store.Chann
// 客户端意外断开:按已生成部分收费(canceled)
commitRaw()
g.finishUsage(c, ch, start, store.UsageStatusCanceled, "client_disconnect")
return true, false, 0, nil
return true, false, false, 0, nil
}
if trySink != nil {
trySink.outputText += sseContentText(out)
@@ -520,6 +559,11 @@ func (g *Gateway) streamHoldCopy(c *gin.Context, sh *streamHold, ch *store.Chann
if usageRaw := scanUsage(line); usageRaw != nil && trySink != nil {
trySink.push(usageRaw)
}
if trySink != nil && trySink.servedModel == "" && bytes.Contains(line, []byte(`"model"`)) {
if payload, ok := sseDataPayload(line); ok {
trySink.servedModel = extractServedModel([]byte(payload))
}
}
}
if err != nil {
if releaseSeen {
@@ -533,12 +577,12 @@ func (g *Gateway) streamHoldCopy(c *gin.Context, sh *streamHold, ch *store.Chann
} else {
g.recordError(c, ch, nil, start, "stream_read_error")
}
return true, false, 0, nil
return true, err == io.EOF, false, 0, nil
}
// 放行前失败:换渠道重试
if c.Request.Context().Err() != nil {
g.finishUsage(c, ch, start, store.UsageStatusCanceled, "client_disconnect")
return true, false, 0, nil
return true, false, false, 0, nil
}
if tryCtx != nil && tryCtx.Err() != nil {
return fail(fmt.Sprintf("first token timeout (no content in %s)", g.firstTokenWait), channel.FailNet)
@@ -563,6 +607,22 @@ type usageShape struct {
TotalTokens int64 `json:"total_tokens"`
CacheReadInputTokens int64 `json:"cache_read_input_tokens"`
CacheCreationInputTokens int64 `json:"cache_creation_input_tokens"`
// Reasoning:思考 token(chat 的 completion_tokens_details / responses 的
// output_tokens_details)。上游已把它计入 completion/output,只做展示口径,不重复计费。
CompletionDetails *usageTokenDetails `json:"completion_tokens_details,omitempty"`
OutputDetails *usageTokenDetails `json:"output_tokens_details,omitempty"`
// Anthropic 的 cache_creation 细分:1h 写价是输入价 2×、5m 是 1.25×,必须拆开算钱
CacheCreation *usageCacheCreation `json:"cache_creation,omitempty"`
}
type usageTokenDetails struct {
ReasoningTokens int64 `json:"reasoning_tokens"`
}
type usageCacheCreation struct {
Ephemeral5m int64 `json:"ephemeral_5m_input_tokens"`
Ephemeral1h int64 `json:"ephemeral_1h_input_tokens"`
}
// extractUsage 从完整响应体提取 usage 子对象(chat / responses / messages)。
@@ -676,11 +736,45 @@ func (s *sseScanner) Next() ([]byte, error) {
// ---------------------------------------------------------------------------
// 记账
// extractServedModel 从响应体(或流式事件载荷)提取上游实际应答的模型名:
// chat/responses/messages 三协议都在顶层带 model;流式事件在 message.model /
// response.model。模型映射、聚合商场景下它与请求的模型名不同,落库便于对账。
func extractServedModel(data []byte) string {
var m map[string]json.RawMessage
if json.Unmarshal(data, &m) != nil {
return ""
}
if raw, ok := m["model"]; ok {
var s string
if json.Unmarshal(raw, &s) == nil && s != "" {
return s
}
}
for _, key := range []string{"response", "message"} {
raw, ok := m[key]
if !ok {
continue
}
var sub map[string]json.RawMessage
if json.Unmarshal(raw, &sub) != nil {
continue
}
if mraw, ok := sub["model"]; ok {
var s string
if json.Unmarshal(mraw, &s) == nil && s != "" {
return s
}
}
}
return ""
}
// usageSink 累积多次 usage:合并各事件字段(message_start 给 input,message_delta 给 output)。
// outputText 累积已转发内容文本,用于流式中断时按 tiktoken 估算输出 token。
type usageSink struct {
us usageShape
outputText string
us usageShape
outputText string
servedModel string // 上游实际应答的模型名(模型映射/聚合商场景下与请求的不同)
}
func (u *usageSink) push(raw json.RawMessage) {
@@ -718,6 +812,15 @@ func (u *usageSink) push(raw json.RawMessage) {
if t.CacheCreationInputTokens > 0 {
u.us.CacheCreationInputTokens = t.CacheCreationInputTokens
}
if t.CompletionDetails != nil {
u.us.CompletionDetails = t.CompletionDetails
}
if t.OutputDetails != nil {
u.us.OutputDetails = t.OutputDetails
}
if t.CacheCreation != nil {
u.us.CacheCreation = t.CacheCreation
}
}
// Shape 返回合并后的用量。
@@ -748,6 +851,9 @@ func (u *usageSink) merge(o *usageSink) {
u.us.CacheCreationInputTokens = o.us.CacheCreationInputTokens
}
u.outputText += o.outputText
if o.servedModel != "" && u.servedModel == "" {
u.servedModel = o.servedModel
}
}
// finishUsage 落账:计算成本并异步写入。
@@ -769,9 +875,30 @@ func (g *Gateway) finishUsage(c *gin.Context, ch *store.Channel, start time.Time
out := us.CompletionTokens + us.OutputTokens
cacheRead := us.CacheReadInputTokens
cacheCreate := us.CacheCreationInputTokens
reasoning := int64(0)
if us.CompletionDetails != nil {
reasoning += us.CompletionDetails.ReasoningTokens
}
if us.OutputDetails != nil {
reasoning += us.OutputDetails.ReasoningTokens
}
// Anthropic 缓存写细分:1h 是输入价 2×、5m 是 1.25×,倍率 1.6。
// CacheWritePriceOr() 是 5m 口径(缺省输入价 ×1.25),1h 按同比例放大。
cache1h := int64(0)
if us.CacheCreation != nil && us.CacheCreation.Ephemeral1h > 0 {
cache1h = us.CacheCreation.Ephemeral1h
if cache1h > cacheCreate {
cache1h = cacheCreate
}
}
cache5m := cacheCreate - cache1h
modelName, _ := c.Get("model_name")
mn, _ := modelName.(string)
servedModel := ""
if sink != nil {
servedModel = sink.servedModel
}
// 上游未返回真实 usage 时估算(tiktoken):
// - 流式中断(canceled):上游最终 usage 可能未返回
@@ -800,10 +927,12 @@ func (g *Gateway) finishUsage(c *gin.Context, ch *store.Channel, start time.Time
// 计价口径:in=非缓存输入、cacheRead=缓存读、cacheCreate=缓存写(Anthropic 语义,
// messages 流式 message_delta 的 input_tokens 即非缓存部分)。
// 缓存写价未单独配置时按输入价 1.25×(Anthropic 5m 口径)。
cwp := model.CacheWritePriceOr()
cost = float64(in)/1e6*model.InputPrice +
float64(out)/1e6*model.OutputPrice +
float64(cacheRead)/1e6*model.CacheReadPrice +
float64(cacheCreate)/1e6*model.CacheWritePriceOr()
float64(cache5m)/1e6*cwp +
float64(cache1h)/1e6*cwp*1.6
} else {
// 模型在记账时已被删除(请求进行中被下架):无法定价,按 0 记账留痕。
// 虚构单价扣用户的钱既不可审计也无法对账。
@@ -871,6 +1000,9 @@ func (g *Gateway) finishUsage(c *gin.Context, ch *store.Channel, start time.Time
OutputPrice: model.OutputPrice,
CacheReadPrice: model.CacheReadPrice,
Cost: cost,
ReasoningTokens: reasoning,
CacheWrite1hTokens: cache1h,
ServedModel: servedModel,
LatencyMS: latency,
Status: status,
ErrorCode: errCodePtr,
@@ -880,8 +1012,8 @@ func (g *Gateway) finishUsage(c *gin.Context, ch *store.Channel, start time.Time
})
// 请求级一行日志:usage_logs 落库是异步的且可能丢,排障时日志是唯一即时依据
log.Printf("proxy: trace=%s user=%d key=%d model=%q channel=%d proto=%s in=%d out=%d cache_r=%d cache_w=%d cost=%.6f status=%s http=%d latency=%dms err=%s",
traceStr, uidVal, kidVal, mn, chID, p, in, out, cacheRead, cacheCreate, cost, status,
log.Printf("proxy: trace=%s user=%d key=%d model=%q served=%q channel=%d proto=%s in=%d out=%d r=%d cache_r=%d cache_w=%d(1h=%d) cost=%.6f status=%s http=%d latency=%dms err=%s",
traceStr, uidVal, kidVal, mn, servedModel, chID, p, in, out, reasoning, cacheRead, cacheCreate, cache1h, cost, status,
c.Writer.Status(), latency, errCode)
}
+114
View File
@@ -0,0 +1,114 @@
package proxy
import (
"math"
"net/http"
"net/http/httptest"
"strings"
"testing"
"github.com/gin-gonic/gin"
"github.com/openteam/server/internal/store"
)
// runPath 走完整链路(Auth → Handle),路径与请求体由调用方给。
func runPath(t *testing.T, h *harness, path, body string) *httptest.ResponseRecorder {
t.Helper()
w := httptest.NewRecorder()
c, _ := gin.CreateTestContext(w)
req := httptest.NewRequest(http.MethodPost, path, strings.NewReader(body))
req.Header.Set("Authorization", "Bearer "+h.key)
c.Request = req
h.gw.Auth(c)
if !c.IsAborted() {
h.gw.Handle(c)
}
return w
}
// usage 细化:reasoning token 与上游实际应答模型落库;计费口径不变
// (reasoning 已含在 completion_tokens 内,不重复计费)。
func TestUsageReasoningAndServedModelRecorded(t *testing.T) {
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
w.Header().Set("Content-Type", "application/json")
_, _ = w.Write([]byte(`{"model":"gpt-x-actual","choices":[{"message":{"role":"assistant","content":"hi"},"finish_reason":"stop"}],"usage":{"prompt_tokens":10,"completion_tokens":5,"total_tokens":15,"completion_tokens_details":{"reasoning_tokens":3}}}`))
}))
defer srv.Close()
h := newHarness(t, srv.URL, 0, true, nil)
w := h.run(t, nil, `{"model":"m1","messages":[{"role":"user","content":"hi"}]}`)
if w.Code != http.StatusOK {
t.Fatalf("code=%d body=%s", w.Code, w.Body.String())
}
h.rec.Close()
var logs []store.UsageLog
h.db.Find(&logs)
if len(logs) != 1 {
t.Fatalf("应有一条记账,实际 %d", len(logs))
}
l := logs[0]
if l.ServedModel != "gpt-x-actual" {
t.Fatalf("served_model 应记录上游应答的模型名,实际 %q", l.ServedModel)
}
if l.ReasoningTokens != 3 {
t.Fatalf("reasoning_tokens 应为 3,实际 %d", l.ReasoningTokens)
}
// 计费不含 reasoning 加成:10×0.15/1e6 + 5×0.6/1e6 = 4.5e-6
if math.Abs(l.Cost-4.5e-6) > 1e-12 {
t.Fatalf("cost 应为 4.5e-6(reasoning 不重复计费),实际 %g", l.Cost)
}
}
// Anthropic 缓存写细分:1h 与 5m 分开计价(1h = 2× 输入价,5m = 1.25×)。
func TestAnthropicCache1hBillingSplit(t *testing.T) {
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
w.Header().Set("Content-Type", "application/json")
_, _ = w.Write([]byte(`{"id":"msg_1","model":"claude-x","role":"assistant","content":[{"type":"text","text":"hi"}],"stop_reason":"end_turn","usage":{"input_tokens":100,"output_tokens":10,"cache_read_input_tokens":200,"cache_creation_input_tokens":300,"cache_creation":{"ephemeral_5m_input_tokens":100,"ephemeral_1h_input_tokens":200}}}`))
}))
defer srv.Close()
h := newHarness(t, srv.URL, 0, true, nil)
// 渠道改为 Anthropic 形态,模型定价:输入 1、输出 2、缓存读 0.1、缓存写未配置(1.25× 输入)
var ch1 store.Channel
if err := h.db.Where("name = ?", "c1").First(&ch1).Error; err != nil {
t.Fatal(err)
}
ch1.Formats = []string{store.FormatMessages}
if err := h.db.Save(&ch1).Error; err != nil {
t.Fatal(err)
}
if err := h.db.Model(&store.Model{}).Where("name = ?", "m1").
Updates(map[string]any{"input_price": 1.0, "output_price": 2.0, "cache_read_price": 0.1}).Error; err != nil {
t.Fatal(err)
}
w := runPath(t, h, "/v1/messages", `{"model":"m1","max_tokens":16,"messages":[{"role":"user","content":"hi"}]}`)
if w.Code != http.StatusOK {
t.Fatalf("code=%d body=%s", w.Code, w.Body.String())
}
h.rec.Close()
var logs []store.UsageLog
h.db.Find(&logs)
if len(logs) != 1 {
t.Fatalf("应有一条记账,实际 %d", len(logs))
}
l := logs[0]
if l.CacheWrite1hTokens != 200 {
t.Fatalf("cache_write_1h_tokens 应为 200,实际 %d", l.CacheWrite1hTokens)
}
if l.ServedModel != "claude-x" {
t.Fatalf("served_model 应为 claude-x,实际 %q", l.ServedModel)
}
// 100/1e6*1 + 10/1e6*2 + 200/1e6*0.1 + 100/1e6*1.25 + 200/1e6*(1.25*1.6) = 6.65e-4
want := 100.0/1e6*1.0 + 10.0/1e6*2.0 + 200.0/1e6*0.1 + 100.0/1e6*1.25 + 200.0/1e6*2.0
if math.Abs(l.Cost-want) > 1e-9 {
t.Fatalf("cost 应为 %g(1h/5m 分开计价),实际 %g", want, l.Cost)
}
// 落库 input 为含缓存的总量口径
if l.InputTokens != 600 || l.OutputTokens != 10 {
t.Fatalf("in=%d out=%d,期望 600/10", l.InputTokens, l.OutputTokens)
}
if !strings.Contains(w.Body.String(), `"text":"hi"`) {
t.Fatalf("响应应原样透传:%s", w.Body.String())
}
}
+2 -1
View File
@@ -25,7 +25,8 @@ import (
// v5:usage_logs 补 key_id/channel_id 索引与 (user,created)/(channel,created) 复合索引。
// v6:修正 v5 复合索引列序——AutoMigrate 按名跳过已存在的索引,旧库上的
// idx_user_created 是 user_id 单列,光改 struct tag 不会重建,需删掉重迁一次。
const schemaVersion = 6
// v7:usage_logs 加 reasoning_tokens / cache_write_1h_tokens / served_model 三列。
const schemaVersion = 7
// Open 打开数据库连接并自动迁移。
// 开发默认 SQLite(自动追加 WAL 与 busy_timeout,见 withSQLitePragmas),生产可切 postgres。
+3
View File
@@ -196,6 +196,9 @@ type UsageLog struct {
OutputTokens int64 `json:"output_tokens"`
CacheReadTokens int64 `json:"cache_read_tokens"`
CacheCreationTokens int64 `json:"cache_creation_tokens"`
ReasoningTokens int64 `json:"reasoning_tokens"` // 思考 token(上游已计入 output,展示口径不重复计费)
CacheWrite1hTokens int64 `json:"cache_write_1h_tokens"` // Anthropic 1h 缓存写 token(计价 2× 输入)
ServedModel string `gorm:"size:128" json:"served_model"` // 上游实际应答的模型名(映射/聚合商场景与请求名不同)
InputPrice float64 `gorm:"type:numeric(20,8)" json:"input_price"` // 快照
OutputPrice float64 `gorm:"type:numeric(20,8)" json:"output_price"` // 快照
CacheReadPrice float64 `gorm:"type:numeric(20,8)" json:"cache_read_price"` // 快照