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 2026-10-16 19:15:51 +08:00
1 parent d911881499
commit e865f69a41
10 files changed
+659 -69

No files matched your search

+2
View File
@@ -27,6 +27,8 @@ OT_MASTER_KEY=change-me-master-key
OT_PROXY_UPSTREAM_KEY=
OT_PROXY_UPSTREAM_BASE_URL=https://api.openai.com
OT_PROXY_DEFAULT_MODEL=gpt-4o-mini
# 上游超时兜底值(渠道 timeout_ms 优先):
# 非流式 = 整段请求上限;流式 = 只约束首包,正文不限时长
OT_PROXY_TIMEOUT=300s
# 渠道健康检查
+2 -1
View File
@@ -13,7 +13,8 @@
- `GET /v1/models` — 可用模型列表
- **三协议互转**:客户端协议 × 渠道协议不匹配时自动转换(如 Chat 调用 Claude、Messages 调用 OpenAI、Responses 调用 Claude),流式逐事件转换;协议匹配时直通
- 错误按客户端协议返回(OpenAI 格式 / Anthropic 格式)
- **渠道体系**:按模型绑定选渠道 + 加权负载均衡;每渠道并发信号量(满载溢出);后台健康检查(连续失败进 cooldown、恢复放回);可安全重试的失败自动故障转移(网络错误/429/5xx/超时且未写出响应头)
- **渠道体系**:按模型绑定选渠道 + 加权负载均衡;模型已有绑定时只用绑定渠道(都不可用则 503,不会把请求发给不认识它的渠道);每渠道并发信号量(满载溢出);后台健康检查(探测 `/models`,连续失败进 cooldown、恢复放回)+ 运行期健康(连续 3 次网络错误/超时/429/5xx 进 30s 冷却,成功即清零);可安全重试的失败自动故障转移(网络错误/429/5xx/超时且未写出响应头)
- **上游超时**:渠道 `timeout_ms` 生效——非流式为整段请求上限,流式仅约束首包(正文不限时长);上游请求绑定客户端 context,断连即取消
- **限流/配额**(内存计数):密钥级每日请求数 / Token 数配额、用户级每秒速率(`OT_RATELIMIT_USER_RPS`);超限返回 429
- **前端**:Dashboard / 管理总览骨架屏加载态
- **用户体系**:注册(开放/邀请码可切换,管理后台可改)、登录(JWT access + HttpOnly refresh cookie)、argon2id 密码
+54 -25
View File
@@ -13,11 +13,11 @@ import (
)
type createKeyReq struct {
Name *string `json:"name" binding:"required,min=1,max=64"`
QuotaTokensPerDay *int64 `json:"quota_tokens_per_day"`
QuotaRequestsPerDay *int `json:"quota_requests_per_day"`
AllowedModels []string `json:"allowed_models"`
ExpiresAt *string `json:"expires_at"` // RFC3339
Name *string `json:"name" binding:"required,min=1,max=64"`
QuotaTokensPerDay *int64 `json:"quota_tokens_per_day"`
QuotaRequestsPerDay *int `json:"quota_requests_per_day"`
AllowedModels []string `json:"allowed_models"`
ExpiresAt *string `json:"expires_at"` // RFC3339
}
// CreateKey POST /api/v1/keys — 创建密钥,明文仅此一次返回。
@@ -47,6 +47,13 @@ func (h *Handler) CreateKey(c *gin.Context) {
AllowedModels: req.AllowedModels,
Status: store.KeyStatusActive,
}
// 限额填 0 视为不限制(同 PatchKey),存成 0 会让密钥一发都发不出去
if k.QuotaTokensPerDay != nil && *k.QuotaTokensPerDay <= 0 {
k.QuotaTokensPerDay = nil
}
if k.QuotaRequestsPerDay != nil && *k.QuotaRequestsPerDay <= 0 {
k.QuotaRequestsPerDay = nil
}
if req.ExpiresAt != nil {
t, err := time.Parse(time.RFC3339, *req.ExpiresAt)
if err != nil {
@@ -68,6 +75,22 @@ func (h *Handler) CreateKey(c *gin.Context) {
})
}
// keyJSON 密钥的对外表示(列表与 PATCH 回显共用)。
func keyJSON(k store.APIKey) gin.H {
return gin.H{
"id": k.ID,
"name": k.Name,
"key_prefix": k.KeyPrefix,
"quota_tokens_per_day": k.QuotaTokensPerDay,
"quota_requests_per_day": k.QuotaRequestsPerDay,
"allowed_models": k.AllowedModels,
"expires_at": k.ExpiresAt,
"status": k.Status,
"last_used_at": k.LastUsedAt,
"created_at": k.CreatedAt,
}
}
// ListKeys GET /api/v1/keys
func (h *Handler) ListKeys(c *gin.Context) {
u, ok := userFromContext(c)
@@ -82,18 +105,7 @@ func (h *Handler) ListKeys(c *gin.Context) {
}
out := make([]gin.H, 0, len(keys))
for _, k := range keys {
out = append(out, gin.H{
"id": k.ID,
"name": k.Name,
"key_prefix": k.KeyPrefix,
"quota_tokens_per_day": k.QuotaTokensPerDay,
"quota_requests_per_day": k.QuotaRequestsPerDay,
"allowed_models": k.AllowedModels,
"expires_at": k.ExpiresAt,
"status": k.Status,
"last_used_at": k.LastUsedAt,
"created_at": k.CreatedAt,
})
out = append(out, keyJSON(k))
}
resp.OK(c, gin.H{"items": out})
}
@@ -111,11 +123,11 @@ func (h *Handler) PatchKey(c *gin.Context) {
return
}
var req struct {
Name *string `json:"name"`
QuotaTokensPerDay *int64 `json:"quota_tokens_per_day"`
QuotaRequestsPerDay *int `json:"quota_requests_per_day"`
AllowedModels *[]string `json:"allowed_models"`
Status *string `json:"status"`
Name *string `json:"name"`
QuotaTokensPerDay *int64 `json:"quota_tokens_per_day"`
QuotaRequestsPerDay *int `json:"quota_requests_per_day"`
AllowedModels *[]string `json:"allowed_models"`
Status *string `json:"status"`
}
if err := c.ShouldBindJSON(&req); err != nil {
resp.Fail(c, http.StatusBadRequest, "invalid input")
@@ -130,11 +142,21 @@ func (h *Handler) PatchKey(c *gin.Context) {
if req.Name != nil {
updates["name"] = *req.Name
}
// 限额:>0 生效,0 或负数视为"不限制"(写 NULL)。
// 不能存 0:token 配额判定是 `已用 >= 上限`,0 会让密钥第一次请求就被 429 打死。
if req.QuotaTokensPerDay != nil {
updates["quota_tokens_per_day"] = *req.QuotaTokensPerDay
if *req.QuotaTokensPerDay > 0 {
updates["quota_tokens_per_day"] = *req.QuotaTokensPerDay
} else {
updates["quota_tokens_per_day"] = nil
}
}
if req.QuotaRequestsPerDay != nil {
updates["quota_requests_per_day"] = *req.QuotaRequestsPerDay
if *req.QuotaRequestsPerDay > 0 {
updates["quota_requests_per_day"] = *req.QuotaRequestsPerDay
} else {
updates["quota_requests_per_day"] = nil
}
}
if req.Status != nil {
if *req.Status != store.KeyStatusActive && *req.Status != store.KeyStatusRevoked {
@@ -150,6 +172,7 @@ func (h *Handler) PatchKey(c *gin.Context) {
}
}
// allowed_models 是 jsonb:手动序列化为 JSON 字符串写入(跨 SQLite/Postgres)
// 显式传 [] 表示"清空白名单"(前端编辑框预填现值,字段内容即最终意图)
if req.AllowedModels != nil {
raw, _ := json.Marshal(*req.AllowedModels)
if err := h.a.DB.Model(&k).Update("allowed_models", string(raw)).Error; err != nil {
@@ -157,7 +180,13 @@ func (h *Handler) PatchKey(c *gin.Context) {
return
}
}
resp.OK(c, gin.H{"ok": true})
// 重新读出,回显更新后的完整密钥(含被清零的限额)
var updated store.APIKey
if err := h.a.DB.First(&updated, k.ID).Error; err != nil {
resp.Fail(c, http.StatusInternalServerError, "failed to reload key")
return
}
resp.OK(c, keyJSON(updated))
}
// DeleteKey DELETE /api/v1/keys/:id — 删除(硬删,立即失效不可恢复)。
+88 -9
View File
@@ -5,8 +5,10 @@ package channel
import (
"crypto/rand"
"errors"
"log"
"math/big"
"sync"
"time"
"github.com/openteam/server/internal/pkg/crypto"
"github.com/openteam/server/internal/store"
@@ -21,10 +23,29 @@ type Service struct {
mu sync.Mutex
sems map[uint64]chan struct{}
runt map[uint64]*runtimeHealth
// 运行期冷却参数;NewService 给默认值,测试可覆盖。
failThreshold int
cooldownFor time.Duration
}
// runtimeHealth 运行期健康:连续可重试失败(网络错误/超时/429/5xx)累计到阈值进入冷却。
// 与 HealthMonitor 的探测结果相互独立——探测只看 /models 是否 200,
// 覆盖不了"models 正常但对话接口一直报错"的渠道。
type runtimeHealth struct {
fails int
until time.Time
}
func NewService(db *gorm.DB, enc *crypto.Encryptor) *Service {
return &Service{db: db, enc: enc, sems: map[uint64]chan struct{}{}}
return &Service{
db: db, enc: enc,
sems: map[uint64]chan struct{}{},
runt: map[uint64]*runtimeHealth{},
failThreshold: 3,
cooldownFor: 30 * time.Second,
}
}
// Candidate 一个候选渠道 + 该模型的映射关系。
@@ -33,18 +54,20 @@ type Candidate struct {
UpstreamModel string // 全局模型在此渠道的映射名(无绑定则为空,用客户端模型名)
}
// Candidates 返回可用渠道候选:健康 + 启用,按优先级、权重降序、id 升序排列。
// model 非空时优先取绑定该模型的渠道(携带 upstream_model 映射);无绑定则退回全局。
// Candidates 返回可用渠道候选:健康 + 启用 + 不在运行期冷却,按优先级、权重降序、id 升序排列。
// model 非空且该模型有渠道绑定时,只返回绑定渠道(携带 upstream_model 映射);
// 绑定渠道全部不可用则返回空(上层报 no_channel),不会退化到不认识该模型的全局渠道。
// 仅当模型完全没有任何绑定时才退回全局渠道。
func (s *Service) Candidates(model string) []Candidate {
if model != "" {
var b []store.ChannelModelBinding
var modelIDs []uint64
var b []store.ChannelModelBinding
s.db.Model(&store.Model{}).Where("name = ? AND enabled = ?", model, true).Pluck("id", &modelIDs)
if len(modelIDs) > 0 {
s.db.Where("model_id IN ?", modelIDs).Find(&b)
if cands := s.loadBound(b); len(cands) > 0 {
return cands
}
}
if len(b) > 0 {
return s.loadBound(b)
}
}
var chs []store.Channel
@@ -52,12 +75,15 @@ func (s *Service) Candidates(model string) []Candidate {
Order("priority ASC, weight DESC, id ASC").Find(&chs)
out := make([]Candidate, 0, len(chs))
for i := range chs {
if s.InCooldown(chs[i].ID) {
continue
}
out = append(out, Candidate{Channel: &chs[i]})
}
return out
}
// loadBound 按绑定顺序加载渠道候选,过滤健康/启用,携带 upstream_model 映射。
// loadBound 按绑定顺序加载渠道候选,过滤健康/启用/运行期冷却,携带 upstream_model 映射。
func (s *Service) loadBound(bindings []store.ChannelModelBinding) []Candidate {
if len(bindings) == 0 {
return nil
@@ -80,7 +106,7 @@ func (s *Service) loadBound(bindings []store.ChannelModelBinding) []Candidate {
}
out := make([]Candidate, 0, len(ids))
for _, id := range ids {
if ch, ok := byID[id]; ok {
if ch, ok := byID[id]; ok && !s.InCooldown(id) {
out = append(out, Candidate{Channel: ch, UpstreamModel: byChannel[id].UpstreamModel})
}
}
@@ -154,3 +180,56 @@ func (s *Service) TryAcquire(ch *store.Channel) (func(), bool) {
func (s *Service) UpstreamKey(ch *store.Channel) (string, error) {
return s.enc.Decrypt(ch.APIKeyEnc)
}
// ---------------------------------------------------------------------------
// 运行期健康(进程内,与 /v1/models 探测独立)
// ReportFailure 记一次可重试的上游失败;连续达 failThreshold 次则该渠道进入冷却,
// 期间不出现在候选集里。冷却到期后需再攒满阈值才会被剔除。
func (s *Service) ReportFailure(ch *store.Channel) {
s.mu.Lock()
st := s.runt[ch.ID]
if st == nil {
st = &runtimeHealth{}
s.runt[ch.ID] = st
}
st.fails++
tripped := false
if st.fails >= s.failThreshold {
st.fails = 0
st.until = time.Now().Add(s.cooldownFor)
tripped = true
}
s.mu.Unlock()
if tripped {
log.Printf("channel: %q entered runtime cooldown (%s)", ch.Name, s.cooldownFor)
}
}
// ReportSuccess 上游正常返回响应头,清空该渠道的连续失败计数。
func (s *Service) ReportSuccess(ch *store.Channel) {
s.mu.Lock()
if st, ok := s.runt[ch.ID]; ok {
st.fails = 0
}
s.mu.Unlock()
}
// InCooldown 渠道是否处于运行期冷却;过期则清掉状态、恢复可用。
// until 为零值表示只在累计失败、尚未触发冷却,此时不能清状态(否则计数永远攒不满阈值)。
func (s *Service) InCooldown(id uint64) bool {
s.mu.Lock()
defer s.mu.Unlock()
st, ok := s.runt[id]
if !ok {
return false
}
if st.until.IsZero() {
return false
}
if time.Now().Before(st.until) {
return true
}
delete(s.runt, id)
return false
}
+76
View File
@@ -129,3 +129,79 @@ func TestHealthStartStops(t *testing.T) {
m.Start(ctx)
cancel() // 不应 panic
}
// 模型被渠道绑定后只允许走绑定渠道;绑定渠道不可用时返回空候选(上层 503),
// 不能退化到"不认识这个模型"的全局渠道。
func TestBoundModelDoesNotFallBackToGlobal(t *testing.T) {
db := newTestDB(t)
bound := seedChannel(t, db, "bound", 1, 0, 4)
seedChannel(t, db, "other", 1, 1, 4)
m := store.Model{Name: "gpt-m", Enabled: true}
if err := db.Create(&m).Error; err != nil {
t.Fatal(err)
}
if err := db.Create(&store.ChannelModelBinding{ChannelID: bound, ModelID: m.ID, UpstreamModel: "gpt-m"}).Error; err != nil {
t.Fatal(err)
}
s := NewService(db, crypto.NewEncryptor("test-master-key-1234567890"))
cands := s.Candidates("gpt-m")
if len(cands) != 1 || cands[0].Channel.Name != "bound" {
t.Fatalf("绑定模型应只命中 bound,实际 %+v", cands)
}
if cands[0].UpstreamModel != "gpt-m" {
t.Fatalf("映射名丢失:%s", cands[0].UpstreamModel)
}
db.Model(&store.Channel{}).Where("id = ?", bound).Update("enabled", false)
if got := s.Candidates("gpt-m"); len(got) != 0 {
t.Fatalf("绑定渠道停用后应为空候选,实际 %d", len(got))
}
// 完全未绑定的模型仍然走全局
if got := s.Candidates("unmanaged"); len(got) != 1 || got[0].Channel.Name != "other" {
t.Fatalf("未纳管模型应退回全局渠道,实际 %+v", got)
}
}
// 运行期冷却:连续失败才踢,成功即清零,冷却到期自动恢复。
func TestRuntimeCooldownThresholdAndExpiry(t *testing.T) {
db := newTestDB(t)
id := seedChannel(t, db, "r", 1, 0, 4)
var ch store.Channel
db.First(&ch, id)
s := NewService(db, crypto.NewEncryptor("test-master-key-1234567890"))
s.cooldownFor = 60 * time.Millisecond
// 默认阈值 3:前两次失败不该踢,第三次才进冷却
s.ReportFailure(&ch)
s.ReportFailure(&ch)
if s.InCooldown(id) {
t.Fatal("未达阈值(3 次连续)不应进冷却")
}
s.ReportFailure(&ch)
if !s.InCooldown(id) {
t.Fatal("连续达阈值应进冷却")
}
if got := s.Candidates(""); len(got) != 0 {
t.Fatalf("冷却渠道不应进候选,实际 %d", len(got))
}
if got := s.Candidates("whatever"); len(got) != 0 {
t.Fatalf("全局候选也应过滤冷却渠道,实际 %d", len(got))
}
time.Sleep(80 * time.Millisecond)
if s.InCooldown(id) {
t.Fatal("冷却到期应恢复可用")
}
if got := s.Candidates(""); len(got) != 1 {
t.Fatalf("恢复后应重新进候选,实际 %d", len(got))
}
s.ReportFailure(&ch)
s.ReportSuccess(&ch)
s.ReportFailure(&ch)
if s.InCooldown(id) {
t.Fatal("非连续失败不应触发冷却")
}
}
+71 -18
View File
@@ -6,6 +6,7 @@ import (
"encoding/json"
"errors"
"fmt"
"net"
"net/http"
"strings"
"sync"
@@ -23,9 +24,10 @@ import (
)
const (
CtxUserID = "proxy_user_id"
CtxKeyID = "proxy_key_id"
CtxTrace = "proxy_trace_id"
CtxUserID = "proxy_user_id"
CtxKeyID = "proxy_key_id"
CtxKeyModel = "proxy_key_models"
CtxTrace = "proxy_trace_id"
)
type Gateway struct {
@@ -36,7 +38,11 @@ type Gateway struct {
lim *ratelimit.Limiter
userRPS int
logRaw bool
hc *http.Client
// 上游超时兜底值(渠道 timeout_ms 非法时使用);按超时值缓存 Transport 复用连接池。
gwTimeout time.Duration
tpMu sync.Mutex
tps map[time.Duration]*http.Transport
policyMu sync.Mutex
policy modelPolicy
@@ -58,7 +64,7 @@ func (g *Gateway) ResetModelPolicy() {
g.policyMu.Unlock()
}
// globalModelRestrictions 读取全局模型允许/禁止列表(缓存 30s)。
// globalModelRestrictions 读取全局模型允许/禁止列表(缓存 modelPolicyTTL)。
func (g *Gateway) globalModelRestrictions() (allow, deny []string) {
g.policyMu.Lock()
defer g.policyMu.Unlock()
@@ -75,18 +81,24 @@ func (g *Gateway) globalModelRestrictions() (allow, deny []string) {
return
}
// checkModelAllowed 模型访问控制:用户级 > 全局。
// checkModelAllowed 模型访问控制:用户级 > 密钥级 > 全局。
// 1. 用户禁止列表命中 → 拒绝
// 2. 用户允许列表非空 → 仅列表内可访问(不再看全局)
// 3. 全局禁止命中 → 拒绝
// 4. 全局允许列表非空 → 仅列表内可访问
func (g *Gateway) checkModelAllowed(u *store.User, model string) bool {
// 2. 密钥允许列表非空且不含该模型 → 拒绝(密钥级只收窄,不放宽)
// 3. 用户允许列表非空 → 仅列表内可访问,不再看全局(管理员显式授权该用户,短路全局)
// 4. 全局禁止命中 → 拒绝
// 5. 全局允许列表非空 → 仅列表内可访问
func (g *Gateway) checkModelAllowed(c *gin.Context, u *store.User, model string) bool {
if model == "" {
return true
}
if contains(u.DeniedModels, model) {
return false
}
if keyAllow, ok := c.Get(CtxKeyModel); ok {
if list, _ := keyAllow.([]string); len(list) > 0 && !contains(list, model) {
return false
}
}
if len(u.AllowedModels) > 0 {
return contains(u.AllowedModels, model)
}
@@ -111,17 +123,57 @@ func contains(list []string, s string) bool {
func NewGateway(db *gorm.DB, enc *crypto.Encryptor, rec *usage.Recorder, lim *ratelimit.Limiter, userRPS int, logRaw bool, gw_timeout time.Duration) *Gateway {
return &Gateway{
db: db,
ch: channel.NewService(db, enc),
rec: rec,
enc: enc,
lim: lim,
userRPS: userRPS,
logRaw: logRaw,
hc: &http.Client{Timeout: gw_timeout * time.Second},
db: db,
ch: channel.NewService(db, enc),
rec: rec,
enc: enc,
lim: lim,
userRPS: userRPS,
logRaw: logRaw,
gwTimeout: gw_timeout,
tps: map[time.Duration]*http.Transport{},
}
}
// channelTimeout 该渠道的上游超时:timeout_ms 优先,未配置则用全局兜底值。
func (g *Gateway) channelTimeout(ch *store.Channel) time.Duration {
if ch.TimeoutMS > 0 {
return time.Duration(ch.TimeoutMS) * time.Millisecond
}
return g.gwTimeout
}
// clientFor 返回带该超时的上游 client。Transport 按超时值缓存以复用连接池,
// ResponseHeaderTimeout 只约束"连上之后多久要回首包"——流式正文不限时长,长回答不会被砍。
// 整段时长(含读 body)由请求 context 的 deadline 控制,见 proxyOne。
func (g *Gateway) clientFor(d time.Duration) *http.Client {
if d < 0 {
d = 0
}
g.tpMu.Lock()
defer g.tpMu.Unlock()
tp, ok := g.tps[d]
if !ok {
tp = &http.Transport{
Proxy: http.ProxyFromEnvironment,
ForceAttemptHTTP2: true,
DialContext: (&net.Dialer{
Timeout: 30 * time.Second,
KeepAlive: 30 * time.Second,
}).DialContext,
MaxIdleConns: 100,
MaxIdleConnsPerHost: 32,
IdleConnTimeout: 90 * time.Second,
TLSHandshakeTimeout: 10 * time.Second,
}
if d > 0 {
tp.ResponseHeaderTimeout = d
}
g.tps[d] = tp
}
return &http.Client{Transport: tp}
}
// Auth 代理鉴权中间件:Bearer sk-xxx → 哈希查表 → 校验状态/过期/模型白名单。
func (g *Gateway) Auth(c *gin.Context) {
// 先按路径确定客户端协议,保证 Auth 阶段错误也按协议格式返回
@@ -186,6 +238,7 @@ func (g *Gateway) Auth(c *gin.Context) {
c.Set(CtxUserID, u.ID)
c.Set(CtxKeyID, k.ID)
c.Set(CtxKeyModel, k.AllowedModels)
c.Set(CtxTrace, newTraceID())
g.db.Model(&store.APIKey{}).Where("id = ?", k.ID).Update("last_used_at", time.Now())
c.Next()
+4 -4
View File
@@ -25,7 +25,7 @@ func (g *Gateway) chatCompletions(c *gin.Context) {
c.Set("protocol", convert.ProtoChat)
c.Set("model_name", br.Model)
g.recordRawRequest(c, u, body)
if !g.checkModelAllowed(u, br.Model) {
if !g.checkModelAllowed(c, u, br.Model) {
apiError(c, http.StatusForbidden, "model_not_allowed", "模型未对你开放,请联系管理员")
return
}
@@ -58,7 +58,7 @@ func (g *Gateway) responses(c *gin.Context) {
c.Set("protocol", convert.ProtoResponses)
c.Set("model_name", br.Model)
g.recordRawRequest(c, u, body)
if !g.checkModelAllowed(u, br.Model) {
if !g.checkModelAllowed(c, u, br.Model) {
apiError(c, http.StatusForbidden, "model_not_allowed", "模型未对你开放,请联系管理员")
return
}
@@ -91,7 +91,7 @@ func (g *Gateway) messages(c *gin.Context) {
c.Set("protocol", convert.ProtoMessages)
c.Set("model_name", br.Model)
g.recordRawRequest(c, u, body)
if !g.checkModelAllowed(u, br.Model) {
if !g.checkModelAllowed(c, u, br.Model) {
apiError(c, http.StatusForbidden, "model_not_allowed", "模型未对你开放,请联系管理员")
return
}
@@ -125,7 +125,7 @@ func apiError(c *gin.Context, status int, code, message string) {
if p, _ := c.Get("protocol"); p == convert.ProtoMessages {
// Anthropic 格式
c.AbortWithStatusJSON(status, gin.H{
"type": "error",
"type": "error",
"error": gin.H{"type": errorTypeFor(status), "message": message},
})
return
+45 -8
View File
@@ -3,9 +3,11 @@ package proxy
import (
"bufio"
"bytes"
"context"
"crypto/rand"
"encoding/hex"
"encoding/json"
"errors"
"fmt"
"io"
"net/http"
@@ -187,8 +189,17 @@ func (g *Gateway) doProxy(c *gin.Context, cands []channel.Candidate, clientProto
idx := (startIdx + offset) % len(cands)
cand := cands[idx]
ch := cand.Channel
// 客户端已断开:上游请求会随 context 立刻失败,换渠道也没有意义
if c.Request.Context().Err() != nil {
g.finishUsage(c, ch, now(), store.UsageStatusCanceled, "client_disconnect")
return
}
if g.ch.InCooldown(ch.ID) {
continue // 本次请求期间已被并发请求判入运行期冷却
}
plan, err := prepareUpstream(ch, clientProto, body, cand.UpstreamModel)
if err != nil {
// 请求体转换失败与渠道健康无关,不换渠道计数
lastStatus, lastBody = http.StatusInternalServerError, []byte("conversion error: "+err.Error())
continue
}
@@ -196,9 +207,9 @@ func (g *Gateway) doProxy(c *gin.Context, cands []channel.Candidate, clientProto
if !ok {
continue // 渠道满载,溢出到下一个
}
written, retry, st, b := g.proxyOne(c, ch, plan, stream, sink)
done, retry, st, b := g.proxyOne(c, ch, plan, stream, sink)
release()
if written {
if done {
return
}
if !retry {
@@ -208,6 +219,7 @@ func (g *Gateway) doProxy(c *gin.Context, cands []channel.Candidate, clientProto
g.recordError(c, ch, nil, now(), "upstream_http_"+strconv.Itoa(st))
return
}
g.ch.ReportFailure(ch) // 网络错误/超时/429/5xx:累计运行期失败
lastStatus, lastBody = st, b
}
// 全部候选重试性失败
@@ -216,8 +228,9 @@ func (g *Gateway) doProxy(c *gin.Context, cands []channel.Candidate, clientProto
}
// proxyOne 对单个渠道执行一次代理。
// 返回:written=是否已写客户端响应;retry=是否可安全换渠道重试;status+respBody=失败信息。
func (g *Gateway) proxyOne(c *gin.Context, ch *store.Channel, plan *upstreamPlan, stream bool, sink *usageSink) (written bool, retry bool, status int, respBody []byte) {
// 返回:done=本次请求已终结(客户端响应已写出或已记账,停止故障转移);
// retry=是否可安全换渠道重试;status+respBody=失败信息。
func (g *Gateway) proxyOne(c *gin.Context, ch *store.Channel, plan *upstreamPlan, stream bool, sink *usageSink) (done bool, retry bool, status int, respBody []byte) {
upKey, err := g.ch.UpstreamKey(ch)
if err != nil {
return false, true, http.StatusInternalServerError, []byte("failed to decrypt channel key")
@@ -235,7 +248,18 @@ func (g *Gateway) proxyOne(c *gin.Context, ch *store.Channel, plan *upstreamPlan
}
}
req, err := http.NewRequest(http.MethodPost, upstreamURL(ch, plan.proto, plan.path), bytes.NewReader(upBody))
// 上游请求绑定客户端 context:客户端断开即取消上游,不再空转生成、继续付费。
// 非流式再叠加渠道超时作为整段硬上限;流式只靠 Transport.ResponseHeaderTimeout
// 约束首包,正文时长不设限(长回答不能被超时砍断)。
limit := g.channelTimeout(ch)
ctx := c.Request.Context()
if !stream && limit > 0 {
var cancel context.CancelFunc
ctx, cancel = context.WithTimeout(ctx, limit)
defer cancel()
}
req, err := http.NewRequestWithContext(ctx, http.MethodPost, upstreamURL(ch, plan.proto, plan.path), bytes.NewReader(upBody))
if err != nil {
return false, false, http.StatusInternalServerError, []byte("failed to build upstream request")
}
@@ -255,8 +279,13 @@ func (g *Gateway) proxyOne(c *gin.Context, ch *store.Channel, plan *upstreamPlan
}
start := time.Now()
resp, err := g.hc.Do(req)
resp, err := g.clientFor(limit).Do(req)
if err != nil {
// 客户端断开导致的取消:就地记账并终结本次请求(不计入渠道失败,也不换渠道)
if c.Request.Context().Err() != nil {
g.finishUsage(c, ch, start, store.UsageStatusCanceled, "client_disconnect")
return true, false, 0, nil
}
return false, true, http.StatusBadGateway, []byte("upstream request failed: " + err.Error())
}
defer resp.Body.Close()
@@ -269,6 +298,7 @@ func (g *Gateway) proxyOne(c *gin.Context, ch *store.Channel, plan *upstreamPlan
}
return false, false, resp.StatusCode, errBody
}
g.ch.ReportSuccess(ch)
c.Header("Content-Type", resp.Header.Get("Content-Type"))
c.Status(http.StatusOK)
@@ -284,8 +314,15 @@ func (g *Gateway) proxyOne(c *gin.Context, ch *store.Channel, plan *upstreamPlan
func (g *Gateway) copyAndCapture(c *gin.Context, ch *store.Channel, r io.Reader, start time.Time, bodyConv func([]byte) ([]byte, error), sink *usageSink) {
data, err := io.ReadAll(r)
if err != nil {
apiError(c, http.StatusBadGateway, "upstream_error", "failed reading upstream response")
g.recordError(c, ch, nil, start, "read_error")
// 走到这里响应头已按 200 写出,无法再回错误体,只按中断原因记账
switch {
case c.Request.Context().Err() != nil:
g.finishUsage(c, ch, start, store.UsageStatusCanceled, "client_disconnect")
case errors.Is(err, context.DeadlineExceeded):
g.finishUsage(c, ch, start, store.UsageStatusError, "upstream_timeout")
default:
g.finishUsage(c, ch, start, store.UsageStatusError, "read_error")
}
return
}
if usageRaw := extractUsage(data); usageRaw != nil && sink != nil {
+307
View File
@@ -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)
}
}
+10 -4
View File
@@ -107,7 +107,8 @@ async function saveEdit() {
const m = parseModels(editForm.allowed_models)
if (t !== undefined) payload.quota_tokens_per_day = t
if (r !== undefined) payload.quota_requests_per_day = r
if (m !== undefined) payload.allowed_models = m
// 白名单框预填的是当前值,内容即最终意图:清空要能取消限制,所以总是提交
payload.allowed_models = m ?? []
await http.patch(`/keys/${editing.value.id}`, payload)
toast.ok('密钥已更新')
editOpen.value = false
@@ -280,10 +281,15 @@ onMounted(load)
<div class="space-y-4">
<Input v-model="editForm.name" label="名称" />
<div class="grid grid-cols-1 sm:grid-cols-2 gap-3">
<Input v-model="editForm.quota_tokens_per_day" label="每日 Token 上限" placeholder="留空则不修改" />
<Input v-model="editForm.quota_requests_per_day" label="每日请求上限" placeholder="留空则不修改" />
<Input v-model="editForm.quota_tokens_per_day" label="每日 Token 上限" placeholder="留空不修改,0 = 不限制" />
<Input v-model="editForm.quota_requests_per_day" label="每日请求上限" placeholder="留空不修改,0 = 不限制" />
</div>
<Input v-model="editForm.allowed_models" label="模型白名单" placeholder="逗号分隔" hint="留空则不修改" />
<Input
v-model="editForm.allowed_models"
label="模型白名单"
placeholder="逗号分隔"
hint="留空 = 不限制模型(清空后保存即取消白名单)"
/>
<label class="block">
<span class="mb-1.5 block text-xs font-medium text-muted">状态</span>
<select