Sakurasan
04dfb4285d
会话亲和（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 提取（非流式与流式）。
88 lines
3.7 KiB
Go
88 lines
3.7 KiB
Go
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())
|
||
}
|
||
}
|