fix: 修掉模型调用链路上的四处失效逻辑
- 超时：NewGateway 里 gw_timeout*time.Second 把已是 Duration 的值二次换算， int64 溢出后无论 OT_PROXY_TIMEOUT 填多少都等于没有超时，渠道 timeout_ms 也从没被读过。现在非流式用 context deadline 约束整段时长，流式只约束首包 （Transport.ResponseHeaderTimeout），正文不限时长，长回答不会被砍。 - 断连：上游请求改用 NewRequestWithContext 绑定客户端 context，客户端断开即 取消上游并记 canceled，不再空转生成继续付费；PATCH 已写 200 头后不再往 响应里追加错误体。 - 白名单：APIKey.AllowedModels 以前只存不用，现在 Auth 注入 context 并参与 判定（用户黑名单 > 密钥白名单 > 用户白名单 > 全局）。 - 选路：模型已有渠道绑定时只用绑定渠道，都不可用直接 503 no_channel，不再 把请求发给不认识它的渠道；新增运行期健康（连续 3 次网络错误/超时/429/5xx 进 30s 冷却，成功清零），与 60s 一次的 /models 探测互补。 - 密钥 API：PATCH 回显更新后的完整对象；限额填 0 统一视为不限制（原来 token 配额存 0 会让密钥第一次请求就被 429 打死，而请求配额存 0 又等于 不限制，两个字段语义相反）；前端编辑框总是提交白名单，清空能真的取消。 测试：新增 proxy 链路集成测试（超时生效、流式不被截断、断连取消并记账、 密钥白名单 403、连续 429 后候选剔除）与 channel 用例（绑定模型不退全局、 冷却阈值与到期恢复、成功清零）。
This commit is contained in:
Sakurasan committed
1 parent
d911881499
commit
e865f69a41
10 files changed
+659
-69
No files matched your search
@@ -0,0 +1,307 @@
|
||||
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(), 0, false, 300*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, 0, false, 30*time.Second)
|
||||
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)
|
||||
}
|
||||
}
|
||||
Reference in new issue
Block a user