- 上游 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 过滤与双形状、缓存写价两条路径、失败归因。
608 lines
21 KiB
Go
608 lines
21 KiB
Go
package proxy
|
||
|
||
import (
|
||
"context"
|
||
"encoding/json"
|
||
"errors"
|
||
"io"
|
||
"net/http"
|
||
"net/http/httptest"
|
||
"os"
|
||
"strings"
|
||
"sync/atomic"
|
||
"testing"
|
||
"time"
|
||
|
||
"github.com/gin-gonic/gin"
|
||
"github.com/openteam/server/internal/pkg/apikey"
|
||
"github.com/openteam/server/internal/pkg/crypto"
|
||
"github.com/openteam/server/internal/pkg/ratelimit"
|
||
"github.com/openteam/server/internal/store"
|
||
"github.com/openteam/server/internal/usage"
|
||
"gorm.io/gorm"
|
||
)
|
||
|
||
type harness struct {
|
||
gw *Gateway
|
||
db *gorm.DB
|
||
rec *usage.Recorder
|
||
key string
|
||
}
|
||
|
||
// newHarness 起一个内存 SQLite 网关:1 用户 + 1 密钥 + 1 模型 + 1 渠道(可选绑定)。
|
||
func newHarness(t *testing.T, baseURL string, timeoutMS int, bind bool, keyModels []string) *harness {
|
||
t.Helper()
|
||
dsn := "/tmp/proxy_fix_test_" + time.Now().Format("150405.000000000") + ".db"
|
||
db, err := store.Open("sqlite", dsn)
|
||
if err != nil {
|
||
t.Fatalf("open db: %v", err)
|
||
}
|
||
t.Cleanup(func() {
|
||
sqlDB, _ := db.DB()
|
||
sqlDB.Close()
|
||
os.Remove(dsn)
|
||
})
|
||
|
||
enc := crypto.NewEncryptor("test-master-key-1234567890")
|
||
encKey, _ := enc.Encrypt("upstream-key")
|
||
u := store.User{Username: "u1", Email: "u1@x.com", PasswordHash: "x", Role: store.RoleUser,
|
||
Balance: 100, Status: store.UserStatusActive}
|
||
if err := db.Create(&u).Error; err != nil {
|
||
t.Fatal(err)
|
||
}
|
||
plain, hash, prefix, _ := apikey.Generate()
|
||
k := store.APIKey{UserID: u.ID, Name: "t", KeyHash: hash, KeyPrefix: prefix,
|
||
Status: store.KeyStatusActive, AllowedModels: keyModels}
|
||
if err := db.Create(&k).Error; err != nil {
|
||
t.Fatal(err)
|
||
}
|
||
m := store.Model{Name: "m1", InputPrice: 0.15, OutputPrice: 0.6, Enabled: true}
|
||
if err := db.Create(&m).Error; err != nil {
|
||
t.Fatal(err)
|
||
}
|
||
ch := store.Channel{Name: "c1", Provider: store.ChannelProviderOpenAI,
|
||
Formats: []string{store.FormatChat}, BaseURL: baseURL, APIKeyEnc: encKey,
|
||
Weight: 1, TimeoutMS: timeoutMS, MaxConcurrency: 4,
|
||
HealthStatus: store.ChannelHealthHealthy, Enabled: true}
|
||
if err := db.Create(&ch).Error; err != nil {
|
||
t.Fatal(err)
|
||
}
|
||
if bind {
|
||
if err := db.Create(&store.ChannelModelBinding{ChannelID: ch.ID, ModelID: m.ID, UpstreamModel: "m1"}).Error; err != nil {
|
||
t.Fatal(err)
|
||
}
|
||
}
|
||
|
||
rec := usage.NewRecorder(db)
|
||
t.Cleanup(rec.Close)
|
||
return &harness{
|
||
gw: NewGateway(db, enc, rec, ratelimit.New(), Options{
|
||
GWTimeout: 300 * time.Second,
|
||
MaxBodyMB: 32,
|
||
FailoverBudget: 30 * time.Second,
|
||
StreamHeartbeat: 15 * time.Second,
|
||
}),
|
||
db: db, rec: rec, key: plain,
|
||
}
|
||
}
|
||
|
||
// run 走完整链路:Auth(含新的密钥白名单注入)→ Handle。
|
||
func (h *harness) run(t *testing.T, ctx context.Context, body string) *httptest.ResponseRecorder {
|
||
t.Helper()
|
||
w := httptest.NewRecorder()
|
||
c, _ := gin.CreateTestContext(w)
|
||
req := httptest.NewRequest(http.MethodPost, "/v1/chat/completions", strings.NewReader(body))
|
||
req.Header.Set("Authorization", "Bearer "+h.key)
|
||
if ctx != nil {
|
||
req = req.WithContext(ctx)
|
||
}
|
||
c.Request = req
|
||
h.gw.Auth(c)
|
||
if !c.IsAborted() {
|
||
h.gw.Handle(c)
|
||
}
|
||
return w
|
||
}
|
||
|
||
// 修复①:渠道 timeout_ms 真正生效(此前被 int64 溢出成"永不超时")。
|
||
func TestChannelTimeoutIsEnforced(t *testing.T) {
|
||
var hits atomic.Int64
|
||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||
hits.Add(1)
|
||
select {
|
||
case <-r.Context().Done():
|
||
case <-time.After(1 * time.Second):
|
||
}
|
||
_, _ = io.WriteString(w, `{"id":"late"}`)
|
||
}))
|
||
defer srv.Close()
|
||
|
||
h := newHarness(t, srv.URL, 200, true, nil)
|
||
start := time.Now()
|
||
w := h.run(t, nil, `{"model":"m1","messages":[{"role":"user","content":"hi"}]}`)
|
||
elapsed := time.Since(start)
|
||
|
||
if elapsed > 600*time.Millisecond {
|
||
t.Fatalf("渠道 timeout_ms=200 没生效(上游 1s 后才返回),实际耗时 %s", elapsed)
|
||
}
|
||
if w.Code != http.StatusBadGateway {
|
||
t.Fatalf("超时后应报 upstream_error,实际 %d: %s", w.Code, w.Body.String())
|
||
}
|
||
if hits.Load() == 0 {
|
||
t.Fatal("上游没被调用")
|
||
}
|
||
}
|
||
|
||
// 修复①延伸:流式不受整段超时限制,但首包超时由 transport 兜住。
|
||
func TestStreamHasNoWholeBodyDeadline(t *testing.T) {
|
||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||
w.Header().Set("Content-Type", "text/event-stream")
|
||
w.WriteHeader(200)
|
||
f := w.(http.Flusher)
|
||
for i := 0; i < 3; i++ {
|
||
_, _ = io.WriteString(w, `data: {"choices":[{"delta":{"content":"x"}}]}`+"\n\n")
|
||
f.Flush()
|
||
time.Sleep(150 * time.Millisecond) // 总计 450ms > 渠道 timeout_ms
|
||
}
|
||
_, _ = io.WriteString(w, "data: [DONE]\n\n")
|
||
}))
|
||
defer srv.Close()
|
||
|
||
h := newHarness(t, srv.URL, 200, true, nil)
|
||
w := h.run(t, nil, `{"model":"m1","stream":true,"messages":[{"role":"user","content":"hi"}]}`)
|
||
if w.Code != http.StatusOK {
|
||
t.Fatalf("stream code=%d body=%s", w.Code, w.Body.String())
|
||
}
|
||
if n := strings.Count(w.Body.String(), `"content":"x"`); n != 3 {
|
||
t.Fatalf("流被超时截断,只收到 %d 块,期望 3", n)
|
||
}
|
||
}
|
||
|
||
// 修复②:客户端断连必须把取消传到上游,两侧都停止,而不是让上游继续生成。
|
||
func TestClientDisconnectCancelsUpstream(t *testing.T) {
|
||
started := make(chan struct{})
|
||
upstreamStop := make(chan error, 1)
|
||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||
w.Header().Set("Content-Type", "text/event-stream")
|
||
w.WriteHeader(http.StatusOK)
|
||
f := w.(http.Flusher)
|
||
close(started)
|
||
// 模拟持续生成:客户端断开后写会失败、request context 会被取消
|
||
for i := 0; i < 40; i++ {
|
||
if _, err := io.WriteString(w, `data: {"choices":[{"delta":{"content":"x"}}]}`+"\n\n"); err != nil {
|
||
upstreamStop <- err
|
||
return
|
||
}
|
||
f.Flush()
|
||
if err := r.Context().Err(); err != nil {
|
||
upstreamStop <- err
|
||
return
|
||
}
|
||
time.Sleep(50 * time.Millisecond)
|
||
}
|
||
upstreamStop <- errors.New("上游跑完了 2s 都没察觉客户端断开")
|
||
}))
|
||
defer srv.Close()
|
||
|
||
h := newHarness(t, srv.URL, 0, true, nil)
|
||
ctx, cancel := context.WithCancel(context.Background())
|
||
done := make(chan *httptest.ResponseRecorder, 1)
|
||
go func() {
|
||
done <- h.run(t, ctx, `{"model":"m1","stream":true,"messages":[{"role":"user","content":"hi"}]}`)
|
||
}()
|
||
|
||
<-started
|
||
time.Sleep(120 * time.Millisecond) // 先收几块,保证已进入流式转发
|
||
cancel()
|
||
|
||
select {
|
||
case err := <-upstreamStop:
|
||
t.Logf("上游因客户端断开而停止:%v", err)
|
||
case <-time.After(2 * time.Second):
|
||
t.Fatal("客户端断开后上游仍在生成:context 没有传播")
|
||
}
|
||
<-done
|
||
h.rec.Close()
|
||
|
||
var n int64
|
||
h.db.Model(&store.UsageLog{}).Where("status = ? AND error_code = ?",
|
||
store.UsageStatusCanceled, "client_disconnect").Count(&n)
|
||
if n != 1 {
|
||
t.Fatalf("canceled 记账条数 = %d,期望 1", n)
|
||
}
|
||
}
|
||
|
||
// 修复③:APIKey.AllowedModels 现在真的是道闸门。
|
||
func TestKeyModelWhitelistIsEnforced(t *testing.T) {
|
||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||
_, _ = io.WriteString(w, `{"id":"ok","choices":[{"message":{"role":"assistant","content":"hi"},"finish_reason":"stop"}],"usage":{"prompt_tokens":1,"completion_tokens":1}}`)
|
||
}))
|
||
defer srv.Close()
|
||
|
||
h := newHarness(t, srv.URL, 0, true, []string{"gpt-4o"})
|
||
w := h.run(t, nil, `{"model":"m1","messages":[{"role":"user","content":"hi"}]}`)
|
||
if w.Code != http.StatusForbidden {
|
||
t.Fatalf("密钥白名单外的模型应 403,实际 %d: %s", w.Code, w.Body.String())
|
||
}
|
||
var body map[string]any
|
||
_ = json.Unmarshal(w.Body.Bytes(), &body)
|
||
errObj, _ := body["error"].(map[string]any)
|
||
if errObj == nil || errObj["code"] != "model_not_allowed" {
|
||
t.Fatalf("错误体不符:%v", body)
|
||
}
|
||
}
|
||
|
||
// 修复④(渠道侧):运行期连续失败让坏渠道从候选里消失。
|
||
func TestRuntimeCooldownEvictsFailingChannel(t *testing.T) {
|
||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||
w.WriteHeader(http.StatusTooManyRequests)
|
||
_, _ = io.WriteString(w, `{"error":{"message":"rate limited"}}`)
|
||
}))
|
||
defer srv.Close()
|
||
|
||
h := newHarness(t, srv.URL, 0, true, nil)
|
||
var ch store.Channel
|
||
h.db.First(&ch)
|
||
|
||
// 3 次 429:每次换不到别的渠道,最终 502,同时累计运行期失败
|
||
body := `{"model":"m1","messages":[{"role":"user","content":"hi"}]}`
|
||
for i := 0; i < 3; i++ {
|
||
if got := len(h.gw.candidateChannels("m1")); got != 1 {
|
||
t.Fatalf("第 %d 次失败前候选应为 1,实际 %d", i+1, got)
|
||
}
|
||
if w := h.run(t, nil, body); w.Code != http.StatusTooManyRequests {
|
||
t.Fatalf("429 应作为最后状态返回,实际 %d: %s", w.Code, w.Body.String())
|
||
}
|
||
}
|
||
if !h.gw.ch.InCooldown(ch.ID) {
|
||
t.Fatal("3 次连续 429 后渠道应进入运行期冷却")
|
||
}
|
||
if got := h.gw.candidateChannels("m1"); len(got) != 0 {
|
||
t.Fatalf("冷却中不应出现在候选集,实际 %d", len(got))
|
||
}
|
||
if w := h.run(t, nil, body); w.Code != http.StatusServiceUnavailable {
|
||
t.Fatalf("冷却后应 503 no_channel,实际 %d: %s", w.Code, w.Body.String())
|
||
}
|
||
}
|
||
|
||
// 修复④延伸:上游恢复正常后清空失败计数,不会带着旧账被误踢。
|
||
func TestRuntimeFailureCounterResetsOnSuccess(t *testing.T) {
|
||
fail := true
|
||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||
if fail {
|
||
w.WriteHeader(http.StatusBadGateway)
|
||
_, _ = io.WriteString(w, `{"error":{"message":"boom"}}`)
|
||
return
|
||
}
|
||
_, _ = io.WriteString(w, `{"id":"ok","choices":[{"message":{"role":"assistant","content":"hi"},"finish_reason":"stop"}],"usage":{"prompt_tokens":1,"completion_tokens":1}}`)
|
||
}))
|
||
defer srv.Close()
|
||
|
||
h := newHarness(t, srv.URL, 0, true, nil)
|
||
body := `{"model":"m1","messages":[{"role":"user","content":"hi"}]}`
|
||
h.run(t, nil, body)
|
||
h.run(t, nil, body)
|
||
fail = false
|
||
if w := h.run(t, nil, body); w.Code != http.StatusOK {
|
||
t.Fatalf("应成功,实际 %d: %s", w.Code, w.Body.String())
|
||
}
|
||
var ch store.Channel
|
||
h.db.First(&ch)
|
||
// 成功清零后再失败两次仍不该进冷却(阈值是"连续"3 次)
|
||
fail = true
|
||
h.run(t, nil, body)
|
||
h.run(t, nil, body)
|
||
if h.gw.ch.InCooldown(ch.ID) {
|
||
t.Fatal("失败计数没有被成功请求清零")
|
||
}
|
||
}
|
||
|
||
// 回归:全局兜底超时值不再被二次换算。
|
||
func TestGatewayTimeoutIsNotMultiplied(t *testing.T) {
|
||
gw := NewGateway(nil, nil, nil, nil, Options{GWTimeout: 30 * time.Second, MaxBodyMB: 8})
|
||
if got := gw.channelTimeout(&store.Channel{TimeoutMS: 0}); got != 30*time.Second {
|
||
t.Fatalf("兜底超时 = %s,期望 30s", got)
|
||
}
|
||
if got := gw.channelTimeout(&store.Channel{TimeoutMS: 1500}); got != 1500*time.Millisecond {
|
||
t.Fatalf("渠道超时 = %s,期望 1.5s", got)
|
||
}
|
||
if c := gw.clientFor(30 * time.Second); c.Timeout != 0 {
|
||
t.Fatalf("client.Timeout 应为 0(由 context 控制整段时长),实际 %s", c.Timeout)
|
||
}
|
||
}
|
||
|
||
// 请求体超过上限时应按客户端协议返回 413(上限由 /v1 组的 BodyLimit 中间件把住)。
|
||
func TestRequestBodyTooLarge(t *testing.T) {
|
||
h := newHarness(t, "http://127.0.0.1:1", 0, true, nil)
|
||
|
||
w := httptest.NewRecorder()
|
||
c, _ := gin.CreateTestContext(w)
|
||
req := httptest.NewRequest(http.MethodPost, "/v1/chat/completions",
|
||
strings.NewReader(`{"model":"m1","messages":[{"role":"user","content":"`+strings.Repeat("x", 4096)+`"}]}`))
|
||
req.Header.Set("Authorization", "Bearer "+h.key)
|
||
// 模拟 middleware.BodyLimit:超限后 ReadAll 拿到 *http.MaxBytesError
|
||
req.Body = http.MaxBytesReader(w, req.Body, 64)
|
||
c.Request = req
|
||
|
||
h.gw.Auth(c)
|
||
if !c.IsAborted() {
|
||
h.gw.Handle(c)
|
||
}
|
||
if w.Code != http.StatusRequestEntityTooLarge {
|
||
t.Fatalf("应返回 413,实际 %d: %s", w.Code, w.Body.String())
|
||
}
|
||
if !strings.Contains(w.Body.String(), "request_too_large") {
|
||
t.Fatalf("错误码不符:%s", w.Body.String())
|
||
}
|
||
if !strings.Contains(w.Body.String(), "32 MB") {
|
||
t.Fatalf("报错应点名上限,实际:%s", w.Body.String())
|
||
}
|
||
}
|
||
|
||
// 429 的 Retry-After 要能读出来(数字与 HTTP 日期两种写法)。
|
||
func TestRetryAfterParsing(t *testing.T) {
|
||
cases := []struct {
|
||
val string
|
||
want time.Duration
|
||
}{
|
||
{"", 0},
|
||
{"7", 7 * time.Second},
|
||
{" 12 ", 12 * time.Second},
|
||
{"0", 0},
|
||
{"-5", 0},
|
||
{"not-a-number", 0},
|
||
}
|
||
for _, tc := range cases {
|
||
h := http.Header{}
|
||
if tc.val != "" {
|
||
h.Set("Retry-After", tc.val)
|
||
}
|
||
if got := retryAfter(h); got != tc.want {
|
||
t.Fatalf("Retry-After=%q 解析成 %s,期望 %s", tc.val, got, tc.want)
|
||
}
|
||
}
|
||
h := http.Header{}
|
||
h.Set("Retry-After", time.Now().UTC().Add(20*time.Second).Format(http.TimeFormat))
|
||
if got := retryAfter(h); got < 15*time.Second || got > 25*time.Second {
|
||
t.Fatalf("HTTP 日期写法解析异常:%s", got)
|
||
}
|
||
}
|
||
|
||
// 故障转移的时间预算只约束第二个及以后的渠道:坏渠道把预算吃光后不再拖着客户端等下一条。
|
||
func TestFailoverBudgetStopsRetrying(t *testing.T) {
|
||
var hitFast atomic.Int64
|
||
slow := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||
time.Sleep(1200 * time.Millisecond) // 慢且最终失败,专门吃预算
|
||
w.WriteHeader(http.StatusBadGateway)
|
||
_, _ = io.WriteString(w, `{"error":{"message":"boom"}}`)
|
||
}))
|
||
defer slow.Close()
|
||
fast := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||
hitFast.Add(1)
|
||
_, _ = io.WriteString(w, `{"id":"ok","choices":[{"message":{"role":"assistant","content":"hi"},"finish_reason":"stop"}],"usage":{"prompt_tokens":1,"completion_tokens":1}}`)
|
||
}))
|
||
defer fast.Close()
|
||
|
||
h := newHarness(t, slow.URL, 0, true, nil) // 渠道1:慢坏
|
||
// 加权随机决定起点,把权重压到一边,否则起点可能直接落在健康渠道上、测不到重试预算
|
||
if err := h.db.Model(&store.Channel{}).Where("name = ?", "c1").Update("weight", 1000).Error; err != nil {
|
||
t.Fatal(err)
|
||
}
|
||
// 渠道2:健康,绑定同一个模型
|
||
encKey, _ := crypto.NewEncryptor("test-master-key-1234567890").Encrypt("upstream-key")
|
||
var m store.Model
|
||
h.db.Where("name = ?", "m1").First(&m)
|
||
var ch2 store.Channel
|
||
ch2 = store.Channel{Name: "c2", Provider: store.ChannelProviderOpenAI,
|
||
Formats: []string{store.FormatChat}, BaseURL: fast.URL, APIKeyEnc: encKey,
|
||
Weight: 1, TimeoutMS: 0, MaxConcurrency: 4,
|
||
HealthStatus: store.ChannelHealthHealthy, Enabled: true}
|
||
if err := h.db.Create(&ch2).Error; err != nil {
|
||
t.Fatal(err)
|
||
}
|
||
if err := h.db.Create(&store.ChannelModelBinding{ChannelID: ch2.ID, ModelID: m.ID, UpstreamModel: "m1"}).Error; err != nil {
|
||
t.Fatal(err)
|
||
}
|
||
body := `{"model":"m1","messages":[{"role":"user","content":"hi"}]}`
|
||
|
||
// 预算 300ms:渠道1 耗掉 1.2s 后就不再试渠道2,直接回最后一个错误
|
||
h.gw.failoverBudget = 300 * time.Millisecond
|
||
w := h.run(t, nil, body)
|
||
if hitFast.Load() != 0 {
|
||
t.Fatal("重试预算耗尽后不应再去试第二条渠道")
|
||
}
|
||
if w.Code != http.StatusBadGateway {
|
||
t.Fatalf("应返回 502,实际 %d: %s", w.Code, w.Body.String())
|
||
}
|
||
|
||
// 关掉预算:正常故障转移到渠道2 并成功
|
||
h.gw.failoverBudget = 0
|
||
w2 := h.run(t, nil, body)
|
||
if hitFast.Load() == 0 {
|
||
t.Fatalf("不限预算时应故障转移到渠道2,响应 %d: %s", w2.Code, w2.Body.String())
|
||
}
|
||
if w2.Code != http.StatusOK {
|
||
t.Fatalf("转移到健康渠道应成功,实际 %d: %s", w2.Code, w2.Body.String())
|
||
}
|
||
}
|
||
|
||
// 流式空闲时补 SSE 注释心跳:上游长时间不出字(长思考)不能被反代空闲超时掐掉。
|
||
func TestStreamHeartbeatWhileIdle(t *testing.T) {
|
||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||
w.Header().Set("Content-Type", "text/event-stream")
|
||
w.WriteHeader(http.StatusOK)
|
||
w.(http.Flusher).Flush()
|
||
time.Sleep(700 * time.Millisecond) // 首包之后上游沉默
|
||
_, _ = io.WriteString(w, `data: {"choices":[{"delta":{"content":"久等了"}}]}`+"\n\n")
|
||
_, _ = io.WriteString(w, "data: [DONE]\n\n")
|
||
}))
|
||
defer srv.Close()
|
||
|
||
h := newHarness(t, srv.URL, 0, true, nil)
|
||
h.gw.streamHeartbeat = 120 * time.Millisecond
|
||
w := h.run(t, nil, `{"model":"m1","stream":true,"messages":[{"role":"user","content":"hi"}]}`)
|
||
|
||
if w.Code != http.StatusOK {
|
||
t.Fatalf("code=%d body=%s", w.Code, w.Body.String())
|
||
}
|
||
out := w.Body.String()
|
||
if !strings.Contains(out, ": ping") {
|
||
t.Fatalf("空闲期没有发出心跳:%q", out)
|
||
}
|
||
if !strings.Contains(out, "久等了") {
|
||
t.Fatalf("心跳把正文挤掉了:%q", out)
|
||
}
|
||
// 心跳是注释行,不该被当成输出去计费
|
||
var logs []store.UsageLog
|
||
h.db.Find(&logs)
|
||
for _, l := range logs {
|
||
if strings.Contains(l.RawResponse, ": ping") && l.Status == store.UsageStatusSuccess && l.OutputTokens == 0 {
|
||
t.Fatalf("心跳被当作输出了:%+v", l)
|
||
}
|
||
}
|
||
}
|
||
|
||
// GET /v1/models 只返回调用方真能用的模型,并带目录元数据。
|
||
func TestModelsListRespectsCallerAllowlist(t *testing.T) {
|
||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { _, _ = io.WriteString(w, `{}`) }))
|
||
defer srv.Close()
|
||
h := newHarness(t, srv.URL, 0, true, nil)
|
||
|
||
// 给 m1 补上元数据
|
||
if err := h.db.Model(&store.Model{}).Where("name = ?", "m1").
|
||
Updates(map[string]any{"context_length": 200000, "supports_tools": true, "supports_vision": true}).Error; err != nil {
|
||
t.Fatal(err)
|
||
}
|
||
|
||
listAs := func(key string) (int, string) {
|
||
w := httptest.NewRecorder()
|
||
c, _ := gin.CreateTestContext(w)
|
||
req := httptest.NewRequest(http.MethodGet, "/v1/models", nil)
|
||
req.Header.Set("Authorization", "Bearer "+key)
|
||
c.Request = req
|
||
h.gw.Auth(c)
|
||
if !c.IsAborted() {
|
||
h.gw.Handle(c)
|
||
}
|
||
var got struct {
|
||
Data []map[string]any `json:"data"`
|
||
}
|
||
_ = json.Unmarshal(w.Body.Bytes(), &got)
|
||
n := 0
|
||
if len(got.Data) > 0 {
|
||
n = len(got.Data)
|
||
}
|
||
return n, w.Body.String()
|
||
}
|
||
|
||
// unrestricted key:m1 可见且带元数据
|
||
n, body := listAs(h.key)
|
||
if n != 1 || !strings.Contains(body, `"context_length":200000`) || !strings.Contains(body, `"tools":true`) {
|
||
t.Fatalf(" unrestricted 列表不对:%d %s", n, body)
|
||
}
|
||
|
||
// Anthropic 客户端(只有 x-api-key)应拿到它自己的形状
|
||
{
|
||
w := httptest.NewRecorder()
|
||
c, _ := gin.CreateTestContext(w)
|
||
req := httptest.NewRequest(http.MethodGet, "/v1/models", nil)
|
||
req.Header.Set("x-api-key", h.key)
|
||
c.Request = req
|
||
h.gw.Auth(c)
|
||
h.gw.Handle(c)
|
||
if !strings.Contains(w.Body.String(), `"type":"model"`) || !strings.Contains(w.Body.String(), "display_name") {
|
||
t.Fatalf("x-api-key 客户端应得到 Anthropic 形状:%s", w.Body.String())
|
||
}
|
||
if strings.Contains(w.Body.String(), `"owned_by"`) {
|
||
t.Fatalf("Anthropic 形状里不该有 OpenAI 字段:%s", w.Body.String())
|
||
}
|
||
}
|
||
|
||
// 受限密钥:模型名都不该泄漏出去
|
||
token := "sk-ot-" + strings.Repeat("a", 44)
|
||
var uid uint64
|
||
h.db.Model(&store.User{}).Select("id").Row().Scan(&uid)
|
||
if err := h.db.Create(&store.APIKey{UserID: uid, Name: "limited",
|
||
KeyHash: apikey.Hash(token), KeyPrefix: apikey.Prefix(token),
|
||
Status: store.KeyStatusActive, AllowedModels: []string{"another-model"}}).Error; err != nil {
|
||
t.Fatal(err)
|
||
}
|
||
if n, body := listAs(token); n != 0 {
|
||
t.Fatalf("密钥白名单外的模型不应出现在 /v1/models:%s", body)
|
||
}
|
||
}
|
||
|
||
// 缓存写价单独配置后必须按它计费,而不是继续按输入价×1.25。
|
||
func TestCacheWritePriceBilledSeparately(t *testing.T) {
|
||
body := `{"id":"x","choices":[{"message":{"role":"assistant","content":"hi"},"finish_reason":"stop"}],` +
|
||
`"usage":{"prompt_tokens":10,"completion_tokens":5,"cache_creation_input_tokens":1000000}}`
|
||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||
_, _ = io.WriteString(w, body)
|
||
}))
|
||
defer srv.Close()
|
||
|
||
reqBody := `{"model":"m1","messages":[{"role":"user","content":"hi"}]}`
|
||
costOf := func(h *harness) float64 {
|
||
if w := h.run(t, nil, reqBody); w.Code != http.StatusOK {
|
||
t.Fatalf("code=%d body=%s", w.Code, w.Body.String())
|
||
}
|
||
h.rec.Close()
|
||
var l store.UsageLog
|
||
if err := h.db.Order("id DESC").First(&l).Error; err != nil {
|
||
t.Fatal(err)
|
||
}
|
||
if l.CacheCreationTokens != 1000000 {
|
||
t.Fatalf("缓存写 token 没记上:%+v", l)
|
||
}
|
||
return l.Cost
|
||
}
|
||
|
||
// 默认未配置:回退输入价 0.15 × 1.25 = 0.1875,加 in/out 的零头
|
||
h1 := newHarness(t, srv.URL, 0, true, nil)
|
||
fallback := costOf(h1)
|
||
if fallback < 0.18 || fallback > 0.20 {
|
||
t.Fatalf("未配置缓存写价应按输入价×1.25,实际 cost=%v", fallback)
|
||
}
|
||
|
||
// 显式配置 2.0/百万:应显著高于回退值
|
||
h2 := newHarness(t, srv.URL, 0, true, nil)
|
||
if err := h2.db.Model(&store.Model{}).Where("name = ?", "m1").
|
||
Update("cache_write_price", 2.0).Error; err != nil {
|
||
t.Fatal(err)
|
||
}
|
||
explicit := costOf(h2)
|
||
if explicit < 1.99 || explicit > 2.02 {
|
||
t.Fatalf("按 cache_write_price=2.0 计费应为 ~2.0,实际 %v", explicit)
|
||
}
|
||
if !(explicit > fallback) {
|
||
t.Fatalf("显式价应高于回退价:%v vs %v", explicit, fallback)
|
||
}
|
||
}
|
||
|
||
// 全部渠道都失败时,错误行要归到最后一个试过的渠道上,
|
||
// 否则 channel_id=0,渠道维度的成功率统计永远看不到这类失败。
|
||
func TestAllChannelsFailedAttributesLastChannel(t *testing.T) {
|
||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||
w.WriteHeader(http.StatusInternalServerError)
|
||
_, _ = io.WriteString(w, `{"error":{"message":"boom"}}`)
|
||
}))
|
||
defer srv.Close()
|
||
|
||
h := newHarness(t, srv.URL, 0, true, nil)
|
||
if w := h.run(t, nil, `{"model":"m1","messages":[{"role":"user","content":"hi"}]}`); w.Code != http.StatusInternalServerError {
|
||
t.Fatalf("上游 500 应透出 500,实际 %d: %s", w.Code, w.Body.String())
|
||
}
|
||
h.rec.Close()
|
||
|
||
var l store.UsageLog
|
||
if err := h.db.Where("error_code = ?", "all_channels_failed").First(&l).Error; err != nil {
|
||
t.Fatalf("没有 all_channels_failed 记录:%v", err)
|
||
}
|
||
if l.ChannelID == 0 {
|
||
t.Fatalf("失败记录没归到渠道上:%+v", l)
|
||
}
|
||
if l.Status != store.UsageStatusError || l.Cost != 0 {
|
||
t.Fatalf("失败不应产生扣费:%+v", l)
|
||
}
|
||
}
|