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()) } }