openteam/server/internal/proxy/proxy_test.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)
	}
}

// 预扣（reserve/capture）：在途请求的冻结额让「余额 - 冻结 ≤ 0」触发 402，
// 释放后恢复放行。并发的长流式请求不再互相透支。
func TestReserveFreezesBalance(t *testing.T) {
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		_, _ = io.WriteString(w, `{"id":"x","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)
	h.gw.reserveMaxTokens = 4096
	// 余额压到 0.001：4096 out × 0.6/1e6 = 0.00246 的冻结上界必然覆盖它
	if err := h.db.Model(&store.User{}).Where("username = ?", "u1").
		Update("balance", 0.001).Error; err != nil {
		t.Fatal(err)
	}
	var u store.User
	h.db.Where("username = ?", "u1").First(&u)
	var model store.Model
	h.db.Where("name = ?", "m1").First(&model)

	// 模拟一条在途请求的冻结（不释放）
	fr, _ := gin.CreateTestContext(httptest.NewRecorder())
	fr.Request = httptest.NewRequest(http.MethodPost, "/v1/chat/completions", nil)
	h.gw.freezeFor(fr, u.ID, model, &bodyReq{MaxTokens: nil}, make([]byte, 64))
	if w := h.run(t, nil, `{"model":"m1","messages":[{"role":"user","content":"hi"}]}`); w.Code != http.StatusPaymentRequired {
		t.Fatalf("冻结额覆盖余额时应 402，实际 %d: %s", w.Code, w.Body.String())
	}

	// capture（释放）后恢复放行
	h.gw.unfreeze(fr, u.ID)
	if a := h.gw.frozenOf(u.ID); a != nil && a.Load() != 0 {
		t.Fatalf("释放后冻结额应为 0，实际 %d", a.Load())
	}
	w := h.run(t, nil, `{"model":"m1","messages":[{"role":"user","content":"hi"}]}`)
	if w.Code != http.StatusOK {
		t.Fatalf("释放后应正常代理，实际 %d: %s", w.Code, w.Body.String())
	}

	// 幂等：重复 unfreeze 不产生负冻结
	h.gw.unfreeze(fr, u.ID)
	if a := h.gw.frozenOf(u.ID); a != nil && a.Load() < 0 {
		t.Fatalf("重复释放不应产生负冻结，实际 %d", a.Load())
	}
}

// 冷却兜底宽限：唯一渠道进运行期冷却后，请求 503；配置 grace 且上游已恢复时，
// 带上冷却中的渠道兜底重试一次即成功。
func TestCooldownGraceFallback(t *testing.T) {
	var healthy atomic.Bool
	healthy.Store(false)
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if !healthy.Load() {
			w.WriteHeader(http.StatusInternalServerError)
			_, _ = io.WriteString(w, `{"error":{"message":"boom"}}`)
			return
		}
		_, _ = io.WriteString(w, `{"id":"x","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"}]}`

	// 连续 3 次 5xx：触发运行期冷却（failThreshold=3）
	for i := 0; i < 3; i++ {
		h.run(t, nil, body)
	}
	// 默认 grace=0（关闭）：冷却中 → 无候选 → 503
	if w := h.run(t, nil, body); w.Code != http.StatusServiceUnavailable {
		t.Fatalf("冷却中且未配置宽限时应 503，实际 %d: %s", w.Code, w.Body.String())
	}

	// 上游恢复 + 配置宽限：兜底重试冷却中的渠道 → 200
	healthy.Store(true)
	h.gw.cooldownGrace = time.Minute
	if w := h.run(t, nil, body); w.Code != http.StatusOK {
		t.Fatalf("宽限期内应兜底重试恢复的渠道，实际 %d: %s", w.Code, w.Body.String())
	}
	// 冷却截止时间未被成功清除（由 InCooldown 到期惰性恢复，语义不变），
	// 但宽限期内每次请求都能经兜底打到这条渠道
	if w := h.run(t, nil, body); w.Code != http.StatusOK {
		t.Fatalf("宽限期内后续请求也应走兜底，实际 %d: %s", w.Code, w.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) 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())
	}
}