From e865f69a4167925f4e1cf6a0817dafec044e26cf Mon Sep 17 00:00:00 2001 From: Sakurasan <26715255+Sakurasan@users.noreply.github.com> Date: Fri, 16 Oct 2026 19:15:51 +0800 Subject: [PATCH] =?UTF-8?q?fix:=20=E4=BF=AE=E6=8E=89=E6=A8=A1=E5=9E=8B?= =?UTF-8?q?=E8=B0=83=E7=94=A8=E9=93=BE=E8=B7=AF=E4=B8=8A=E7=9A=84=E5=9B=9B?= =?UTF-8?q?=E5=A4=84=E5=A4=B1=E6=95=88=E9=80=BB=E8=BE=91?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - 超时：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 用例（绑定模型不退全局、 冷却阈值与到期恢复、成功清零）。 --- .env.example | 2 + README.md | 3 +- server/internal/api/keys.go | 79 ++++-- server/internal/channel/channel.go | 97 +++++++- server/internal/channel/channel_test.go | 76 ++++++ server/internal/proxy/gateway.go | 89 +++++-- server/internal/proxy/handlers.go | 8 +- server/internal/proxy/passthrough.go | 53 +++- server/internal/proxy/proxy_test.go | 307 ++++++++++++++++++++++++ web/src/views/console/KeysView.vue | 14 +- 10 files changed, 659 insertions(+), 69 deletions(-) create mode 100644 server/internal/proxy/proxy_test.go diff --git a/.env.example b/.env.example index c7c12eb..0033aca 100644 --- a/.env.example +++ b/.env.example @@ -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 # 渠道健康检查 diff --git a/README.md b/README.md index 48c9150..d22f60e 100644 --- a/README.md +++ b/README.md @@ -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 密码 diff --git a/server/internal/api/keys.go b/server/internal/api/keys.go index 9a1703c..e2db368 100644 --- a/server/internal/api/keys.go +++ b/server/internal/api/keys.go @@ -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 — 删除（硬删，立即失效不可恢复）。 diff --git a/server/internal/channel/channel.go b/server/internal/channel/channel.go index 422f9ef..693152e 100644 --- a/server/internal/channel/channel.go +++ b/server/internal/channel/channel.go @@ -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 +} diff --git a/server/internal/channel/channel_test.go b/server/internal/channel/channel_test.go index d72e9c9..4424c12 100644 --- a/server/internal/channel/channel_test.go +++ b/server/internal/channel/channel_test.go @@ -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("非连续失败不应触发冷却") + } +} diff --git a/server/internal/proxy/gateway.go b/server/internal/proxy/gateway.go index b1adc47..edd7713 100644 --- a/server/internal/proxy/gateway.go +++ b/server/internal/proxy/gateway.go @@ -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() diff --git a/server/internal/proxy/handlers.go b/server/internal/proxy/handlers.go index 8e6cd64..e4e6c80 100644 --- a/server/internal/proxy/handlers.go +++ b/server/internal/proxy/handlers.go @@ -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 diff --git a/server/internal/proxy/passthrough.go b/server/internal/proxy/passthrough.go index 68842aa..1938948 100644 --- a/server/internal/proxy/passthrough.go +++ b/server/internal/proxy/passthrough.go @@ -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 { diff --git a/server/internal/proxy/proxy_test.go b/server/internal/proxy/proxy_test.go new file mode 100644 index 0000000..d55fc5e --- /dev/null +++ b/server/internal/proxy/proxy_test.go @@ -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) + } +} diff --git a/web/src/views/console/KeysView.vue b/web/src/views/console/KeysView.vue index 57fef49..93083d2 100644 --- a/web/src/views/console/KeysView.vue +++ b/web/src/views/console/KeysView.vue @@ -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)
- - + +
- +