diff --git a/.env.example b/.env.example index 246b5b6..58017ca 100644 --- a/.env.example +++ b/.env.example @@ -44,6 +44,9 @@ OT_PROXY_STREAM_HEARTBEAT=15s # 流式在放行给客户端之前最多等内容多久（0=不限）：超时未出内容就取消本次尝试、 # 换下一个渠道。上游挂起（连接正常但一直不吐字）时的兜底。 OT_PROXY_FIRST_TOKEN_WAIT=0s +# 会话亲和 TTL（0=关闭，默认 5m）：同一会话粘住上次应答的渠道，命中上游的提示词 +# 缓存（整段上下文每轮重发），而不是被加权随机打散后全价重算。 +OT_PROXY_AFFINITY_TTL=5m # 渠道健康检查 OT_PROXY_HEALTH_INTERVAL=60s diff --git a/README.md b/README.md index 336de32..8923045 100644 --- a/README.md +++ b/README.md @@ -13,12 +13,12 @@ - `GET /v1/models` — 可用模型列表：只列调用方真能用的模型（叠加用户与密钥白名单），带上下文长度与能力位；Anthropic 客户端（带 `x-api-key`/`anthropic-version`）返回 Anthropic 形状 - **三协议互转**：客户端协议 × 渠道协议不匹配时自动转换（如 Chat 调用 Claude、Messages 调用 OpenAI、Responses 调用 Claude），流式逐事件转换；协议匹配时直通 - 错误按客户端协议返回（OpenAI 格式 / Anthropic 格式） -- **渠道体系**：按模型绑定选渠道 + 加权负载均衡；模型已有绑定时只用绑定渠道（都不可用则 503，不会把请求发给不认识它的渠道）；每渠道并发信号量（满载溢出）；后台健康检查（探测 `/models`，连续失败进 cooldown、恢复放回）+ 运行期健康（按失败原因分类冷却：网络错误/5xx 连续 3 次进冷却且指数退避，余额/配额类第一次即冷却 30 分钟，成功清零退避计数）；可安全重试的失败自动故障转移（网络错误/429/5xx/超时/流式首个内容事件之前的任何失败） +- **渠道体系**：按模型绑定选渠道 + 加权负载均衡；模型已有绑定时只用绑定渠道（都不可用则 503，不会把请求发给不认识它的渠道）；每渠道并发信号量（满载溢出）；后台健康检查（探测 `/models`，连续失败进 cooldown、恢复放回）+ 运行期健康（按失败原因分类冷却：网络错误/5xx 连续 3 次进冷却且指数退避，余额/配额类第一次即冷却 30 分钟，成功清零退避计数）；可安全重试的失败自动故障转移（网络错误/429/5xx/超时/流式首个内容事件之前的任何失败）；**会话亲和**（`OT_PROXY_AFFINITY_TTL`，默认 5m）：同一会话粘住上次应答的渠道，命中上游提示词缓存而不是全价重算，冷却/下线自动回退加权随机 - **上游超时**：渠道 `timeout_ms` 生效——非流式为整段请求上限，流式仅约束首包（正文不限时长）；上游请求绑定客户端 context，断连即取消 - **故障转移与流式细节**（借鉴 magpie 的 holdWriter）：流式请求在**首个内容事件**到达前由网关缓冲——上游的错误事件、空流、断连、协议形态不符（要流给了整包，didNotStream）都在放行前换渠道消化，客户端无感知；`OT_PROXY_FIRST_TOKEN_WAIT` 限定放行前等内容的最长时间（0=不限）；429 不一定是限流——错误体带余额/配额特征时按配额类长冷却；上游 `Retry-After` 会延长该渠道冷却（夹在 10 分钟内）；重试受 `OT_PROXY_FAILOVER_BUDGET` 时间预算约束（首个渠道不受限，避免坏渠道把客户端拖到 N×timeout）；流式空闲超过 `OT_PROXY_STREAM_HEARTBEAT` 补一行 SSE 注释心跳（扣住期间也照发），防止 nginx/Cloudflare 的空闲超时掐掉长思考；响应头已被心跳先行写出时，全部候选失败以 SSE 错误事件收尾 - **渠道观测**：`GET /api/v1/admin/channels/stats?window=5m|1h|6h|24h|7d` 按渠道聚合请求量、成功率、错误码分布与 P50/P95 延迟；渠道页直接展示近况与运行期冷却剩余时间，可一键恢复 - **限流/配额**（内存计数）：密钥级每日请求数 / Token 数配额、用户级每秒速率（`OT_RATELIMIT_USER_RPS`）；超限返回 429 -- **用量计费**：请求级 `usage_logs` 异步批量落库，按模型价格扣减余额，日粒度预聚合（`usage_daily`）；模型可单独配置缓存写价（留空按输入价 ×1.25） +- **用量计费**：请求级 `usage_logs` 异步批量落库，按模型价格扣减余额，日粒度预聚合（`usage_daily`）；模型可单独配置缓存写价（留空按输入价 ×1.25）；记录思考 token（`reasoning_tokens`，展示口径不重复计费）、Anthropic 1h 缓存写（计价 2× 输入）与上游实际应答的模型名（`served_model`，映射/聚合商场景对账用） - **认证面防护**：登录按「单 IP」+「单 IP×账号」双桶滑动窗口限速、注册按 IP 限速（429 + `Retry-After`）；`/v1` 请求体上限 `OT_PROXY_MAX_BODY_MB`（超限按客户端协议返回 413）；CORS 改为显式 origin 白名单（默认不发跨域头）；`env=production` 启动时断言 JWT secret / 主密钥 / 管理员密码不能是仓库默认值，否则拒绝启动 - **前端**：Dashboard / 管理总览骨架屏加载态 - **用户体系**：注册（开放/邀请码可切换，管理后台可改）、登录（JWT access + HttpOnly refresh cookie）、argon2id 密码 diff --git a/server/internal/cli/serve.go b/server/internal/cli/serve.go index 440aff8..d0818bf 100644 --- a/server/internal/cli/serve.go +++ b/server/internal/cli/serve.go @@ -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) diff --git a/server/internal/config/config.go b/server/internal/config/config.go index 59ea9ab..b59b2df 100644 --- a/server/internal/config/config.go +++ b/server/internal/config/config.go @@ -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"), diff --git a/server/internal/proxy/affinity.go b/server/internal/proxy/affinity.go new file mode 100644 index 0000000..0594f8d --- /dev/null +++ b/server/internal/proxy/affinity.go @@ -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])) +} diff --git a/server/internal/proxy/affinity_test.go b/server/internal/proxy/affinity_test.go new file mode 100644 index 0000000..e0ed42a --- /dev/null +++ b/server/internal/proxy/affinity_test.go @@ -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()) + } +} diff --git a/server/internal/proxy/convert/convert_test.go b/server/internal/proxy/convert/convert_test.go index fdfe812..dbc6cac 100644 --- a/server/internal/proxy/convert/convert_test.go +++ b/server/internal/proxy/convert/convert_test.go @@ -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"]) + } +} diff --git a/server/internal/proxy/convert/json_chat.go b/server/internal/proxy/convert/json_chat.go index eed3e69..45216e4 100644 --- a/server/internal/proxy/convert/json_chat.go +++ b/server/internal/proxy/convert/json_chat.go @@ -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 请求。 diff --git a/server/internal/proxy/convert/json_responses.go b/server/internal/proxy/convert/json_responses.go index 032b9af..c55b499 100644 --- a/server/internal/proxy/convert/json_responses.go +++ b/server/internal/proxy/convert/json_responses.go @@ -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 diff --git a/server/internal/proxy/gateway.go b/server/internal/proxy/gateway.go index f3b3864..4e38185 100644 --- a/server/internal/proxy/gateway.go +++ b/server/internal/proxy/gateway.go @@ -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{}, } } diff --git a/server/internal/proxy/passthrough.go b/server/internal/proxy/passthrough.go index a43cfc2..bbd9602 100644 --- a/server/internal/proxy/passthrough.go +++ b/server/internal/proxy/passthrough.go @@ -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) } diff --git a/server/internal/proxy/usage_refine_test.go b/server/internal/proxy/usage_refine_test.go new file mode 100644 index 0000000..48e28c5 --- /dev/null +++ b/server/internal/proxy/usage_refine_test.go @@ -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()) + } +} diff --git a/server/internal/store/db.go b/server/internal/store/db.go index 82e01d0..d21663c 100644 --- a/server/internal/store/db.go +++ b/server/internal/store/db.go @@ -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。 diff --git a/server/internal/store/models.go b/server/internal/store/models.go index ba63eb6..f59c93a 100644 --- a/server/internal/store/models.go +++ b/server/internal/store/models.go @@ -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"` // 快照