Sakurasan
b2e500ff51
- 计费：19 个金额字段 float64→decimal.Decimal，DB 列保持 numeric(20,8) 零迁移； API 金额仍为 JSON number（MarshalJSONWithoutQuotes），前端零改动 - failover：保留时间预算，新增 OT_PROXY_FAILOVER_CAPS_NET/SERVER/RATE/QUOTA 按失败类型封顶换渠道次数（429/配额类瞬时失败时间预算形同虚设） - 预聚合：新增 usage_channel_hourlies（schema v10），记账事务内同步 upsert， 长窗口渠道统计走桶不再扫明细，spent_last_30d 改扫 usage_daily，v10 一次性回填幂等 - 守护：internal/arch_test.go 用 go/parser 强制分层依赖（fail-closed）
764 lines
28 KiB
Go
764 lines
28 KiB
Go
package proxy
|
||
|
||
import (
|
||
"context"
|
||
"encoding/json"
|
||
"errors"
|
||
"fmt"
|
||
"io"
|
||
"net/http"
|
||
"net/http/httptest"
|
||
"os"
|
||
"strings"
|
||
"sync/atomic"
|
||
"testing"
|
||
"time"
|
||
|
||
"github.com/gin-gonic/gin"
|
||
"github.com/openteam/server/internal/channel"
|
||
"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"
|
||
"github.com/shopspring/decimal"
|
||
"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: decimal.NewFromInt(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: decimal.NewFromFloat(0.15), OutputPrice: decimal.NewFromFloat(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())
|
||
}
|
||
}
|
||
|
||
// 429/配额类失败瞬时返回，时间预算对它们形同虚设；按失败类型封顶换渠道次数。
|
||
// rate 上限 2 → 首渠道 + 至多 2 次换渠道 = 3 次尝试即止；上限 0 = 不限（4 条全试）。
|
||
func TestFailoverSwitchCaps(t *testing.T) {
|
||
var hits atomic.Int64
|
||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||
hits.Add(1)
|
||
w.Header().Set("Retry-After", "0")
|
||
w.WriteHeader(http.StatusTooManyRequests)
|
||
_, _ = io.WriteString(w, `{"error":{"message":"rate limited, try again"}}`)
|
||
}))
|
||
defer srv.Close()
|
||
|
||
newMultiChannel := func(t *testing.T, n int) *harness {
|
||
t.Helper()
|
||
h := newHarness(t, srv.URL, 0, true, nil)
|
||
var m store.Model
|
||
h.db.Where("name = ?", "m1").First(&m)
|
||
encKey, _ := crypto.NewEncryptor("test-master-key-1234567890").Encrypt("upstream-key")
|
||
for i := 2; i <= n; i++ {
|
||
ch := store.Channel{Name: fmt.Sprintf("c%d", i), Provider: store.ChannelProviderOpenAI,
|
||
Formats: []string{store.FormatChat}, BaseURL: srv.URL, APIKeyEnc: encKey,
|
||
Weight: 1, TimeoutMS: 0, MaxConcurrency: 4,
|
||
HealthStatus: store.ChannelHealthHealthy, Enabled: true}
|
||
if err := h.db.Create(&ch).Error; err != nil {
|
||
t.Fatal(err)
|
||
}
|
||
if err := h.db.Create(&store.ChannelModelBinding{ChannelID: ch.ID, ModelID: m.ID, UpstreamModel: "m1"}).Error; err != nil {
|
||
t.Fatal(err)
|
||
}
|
||
}
|
||
return h
|
||
}
|
||
body := `{"model":"m1","messages":[{"role":"user","content":"hi"}]}`
|
||
|
||
// rate 上限 2：打满即停，第 4 条渠道不应被碰
|
||
h := newMultiChannel(t, 4)
|
||
h.gw.failoverCaps = map[channel.FailKind]int{channel.FailRate: 2}
|
||
w := h.run(t, nil, body)
|
||
if hits.Load() != 3 {
|
||
t.Fatalf("rate 上限 2 应恰好尝试 3 条渠道（1+2），实际 %d", hits.Load())
|
||
}
|
||
if w.Code != http.StatusTooManyRequests {
|
||
t.Fatalf("应透传最后的 429，实际 %d: %s", w.Code, w.Body.String())
|
||
}
|
||
if !strings.Contains(w.Body.String(), "failover switch cap reached (rate_limited)") {
|
||
t.Fatalf("错误体应说明按类型封顶：%s", w.Body.String())
|
||
}
|
||
|
||
// 上限 0 = 不限：4 条渠道全部试完
|
||
hits.Store(0)
|
||
h.gw.failoverCaps = map[channel.FailKind]int{channel.FailRate: 0}
|
||
if w := h.run(t, nil, body); w.Code != http.StatusTooManyRequests {
|
||
t.Fatalf("应 429，实际 %d", w.Code)
|
||
}
|
||
if hits.Load() != 4 {
|
||
t.Fatalf("不限次时应尝试全部 4 条渠道，实际 %d", hits.Load())
|
||
}
|
||
|
||
// quota 上限 1：402 打满即停（402 归为 FailQuota，第一次失败即长冷却，但换渠道仍受上限约束）
|
||
hits.Store(0)
|
||
h2 := newMultiChannel(t, 4)
|
||
h2.gw.failoverCaps = map[channel.FailKind]int{channel.FailQuota: 1}
|
||
quotaSrv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||
hits.Add(1)
|
||
w.WriteHeader(http.StatusPaymentRequired)
|
||
_, _ = io.WriteString(w, `{"error":{"message":"insufficient balance"}}`)
|
||
}))
|
||
defer quotaSrv.Close()
|
||
if err := h2.db.Model(&store.Channel{}).Where("1=1").Update("base_url", quotaSrv.URL).Error; err != nil {
|
||
t.Fatal(err)
|
||
}
|
||
w2 := h2.run(t, nil, body)
|
||
if hits.Load() != 2 {
|
||
t.Fatalf("quota 上限 1 应恰好尝试 2 条渠道，实际 %d", hits.Load())
|
||
}
|
||
if w2.Code != http.StatusPaymentRequired {
|
||
t.Fatalf("应透传 402，实际 %d", w2.Code)
|
||
}
|
||
}
|
||
|
||
// 流式空闲时补 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) decimal.Decimal {
|
||
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.LessThan(decimal.NewFromFloat(0.18)) || fallback.GreaterThan(decimal.NewFromFloat(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.LessThan(decimal.NewFromFloat(1.99)) || explicit.GreaterThan(decimal.NewFromFloat(2.02)) {
|
||
t.Fatalf("按 cache_write_price=2.0 计费应为 ~2.0，实际 %v", explicit)
|
||
}
|
||
if !explicit.GreaterThan(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.IsZero() {
|
||
t.Fatalf("失败不应产生扣费：%+v", l)
|
||
}
|
||
}
|
||
|
||
// Anthropic 原生上游必须收到 x-api-key（官方 API 不接受 API key 走 Bearer）。
|
||
func TestAnthropicUpstreamUsesXAPIKey(t *testing.T) {
|
||
var gotXAPIKey, gotAuth, gotVersion string
|
||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||
gotXAPIKey = r.Header.Get("x-api-key")
|
||
gotAuth = r.Header.Get("Authorization")
|
||
gotVersion = r.Header.Get("anthropic-version")
|
||
_, _ = io.WriteString(w, `{"id":"msg_1","type":"message","role":"assistant","model":"m1","content":[{"type":"text","text":"hi"}],"stop_reason":"end_turn","usage":{"input_tokens":1,"output_tokens":1}}`)
|
||
}))
|
||
defer srv.Close()
|
||
|
||
h := newHarness(t, srv.URL, 3000, true, nil)
|
||
// 把渠道改造成 Anthropic 原生（messages 协议）
|
||
if err := h.db.Model(&store.Channel{}).Where("id = ?", 1).Updates(map[string]any{
|
||
"provider": store.ChannelProviderAnthropic,
|
||
"formats": []string{store.FormatMessages},
|
||
}).Error; err != nil {
|
||
t.Fatal(err)
|
||
}
|
||
|
||
// Anthropic 客户端走 x-api-key 头 + /v1/messages
|
||
w := httptest.NewRecorder()
|
||
c, _ := gin.CreateTestContext(w)
|
||
req := httptest.NewRequest(http.MethodPost, "/v1/messages", strings.NewReader(`{"model":"m1","messages":[{"role":"user","content":"hi"}]}`))
|
||
req.Header.Set("x-api-key", h.key)
|
||
req.Header.Set("anthropic-version", "2023-06-01")
|
||
req.Header.Set("anthropic-beta", "prompt-caching-2024-07-31")
|
||
c.Request = req
|
||
h.gw.Auth(c)
|
||
if !c.IsAborted() {
|
||
h.gw.Handle(c)
|
||
}
|
||
if w.Code != http.StatusOK {
|
||
t.Fatalf("messages 请求应成功，实际 %d: %s", w.Code, w.Body.String())
|
||
}
|
||
if gotXAPIKey != "upstream-key" {
|
||
t.Fatalf("上游应收到 x-api-key=upstream-key，实际 %q", gotXAPIKey)
|
||
}
|
||
if gotAuth != "Bearer upstream-key" {
|
||
t.Fatalf("兼容网关仍需 Bearer 头，实际 %q", gotAuth)
|
||
}
|
||
if gotVersion == "" {
|
||
t.Fatal("上游应收到 anthropic-version")
|
||
}
|
||
}
|
||
|
||
// 禁用/不存在的模型必须 404：不能只是从 /v1/models 列表消失、请求照样计费。
|
||
func TestDisabledOrUnknownModelIsRejected(t *testing.T) {
|
||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||
t.Fatal("禁用模型不应触达上游")
|
||
}))
|
||
defer srv.Close()
|
||
|
||
h := newHarness(t, srv.URL, 3000, true, nil)
|
||
if err := h.db.Model(&store.Model{}).Where("id = ?", 1).Update("enabled", false).Error; err != nil {
|
||
t.Fatal(err)
|
||
}
|
||
|
||
w := h.run(t, nil, `{"model":"m1","messages":[{"role":"user","content":"hi"}]}`)
|
||
if w.Code != http.StatusNotFound {
|
||
t.Fatalf("禁用模型应 404，实际 %d: %s", w.Code, w.Body.String())
|
||
}
|
||
if !strings.Contains(w.Body.String(), "model_not_found") {
|
||
t.Fatalf("错误码应为 model_not_found: %s", w.Body.String())
|
||
}
|
||
|
||
// 空模型名直接 400
|
||
w = h.run(t, nil, `{"messages":[{"role":"user","content":"hi"}]}`)
|
||
if w.Code != http.StatusBadRequest {
|
||
t.Fatalf("空模型应 400，实际 %d: %s", w.Code, w.Body.String())
|
||
}
|
||
}
|