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:
1 parent
d911881499
commit
e865f69a41
10 files changed
+659
-69
No files matched your search
@@ -27,6 +27,8 @@ OT_MASTER_KEY=change-me-master-key
|
|||||||
OT_PROXY_UPSTREAM_KEY=
|
OT_PROXY_UPSTREAM_KEY=
|
||||||
OT_PROXY_UPSTREAM_BASE_URL=https://api.openai.com
|
OT_PROXY_UPSTREAM_BASE_URL=https://api.openai.com
|
||||||
OT_PROXY_DEFAULT_MODEL=gpt-4o-mini
|
OT_PROXY_DEFAULT_MODEL=gpt-4o-mini
|
||||||
|
# 上游超时兜底值(渠道 timeout_ms 优先):
|
||||||
|
# 非流式 = 整段请求上限;流式 = 只约束首包,正文不限时长
|
||||||
OT_PROXY_TIMEOUT=300s
|
OT_PROXY_TIMEOUT=300s
|
||||||
|
|
||||||
# 渠道健康检查
|
# 渠道健康检查
|
||||||
|
|||||||
@@ -13,7 +13,8 @@
|
|||||||
- `GET /v1/models` — 可用模型列表
|
- `GET /v1/models` — 可用模型列表
|
||||||
- **三协议互转**:客户端协议 × 渠道协议不匹配时自动转换(如 Chat 调用 Claude、Messages 调用 OpenAI、Responses 调用 Claude),流式逐事件转换;协议匹配时直通
|
- **三协议互转**:客户端协议 × 渠道协议不匹配时自动转换(如 Chat 调用 Claude、Messages 调用 OpenAI、Responses 调用 Claude),流式逐事件转换;协议匹配时直通
|
||||||
- 错误按客户端协议返回(OpenAI 格式 / Anthropic 格式)
|
- 错误按客户端协议返回(OpenAI 格式 / Anthropic 格式)
|
||||||
- **渠道体系**:按模型绑定选渠道 + 加权负载均衡;每渠道并发信号量(满载溢出);后台健康检查(连续失败进 cooldown、恢复放回);可安全重试的失败自动故障转移(网络错误/429/5xx/超时且未写出响应头)
|
- **渠道体系**:按模型绑定选渠道 + 加权负载均衡;模型已有绑定时只用绑定渠道(都不可用则 503,不会把请求发给不认识它的渠道);每渠道并发信号量(满载溢出);后台健康检查(探测 `/models`,连续失败进 cooldown、恢复放回)+ 运行期健康(连续 3 次网络错误/超时/429/5xx 进 30s 冷却,成功即清零);可安全重试的失败自动故障转移(网络错误/429/5xx/超时且未写出响应头)
|
||||||
|
- **上游超时**:渠道 `timeout_ms` 生效——非流式为整段请求上限,流式仅约束首包(正文不限时长);上游请求绑定客户端 context,断连即取消
|
||||||
- **限流/配额**(内存计数):密钥级每日请求数 / Token 数配额、用户级每秒速率(`OT_RATELIMIT_USER_RPS`);超限返回 429
|
- **限流/配额**(内存计数):密钥级每日请求数 / Token 数配额、用户级每秒速率(`OT_RATELIMIT_USER_RPS`);超限返回 429
|
||||||
- **前端**:Dashboard / 管理总览骨架屏加载态
|
- **前端**:Dashboard / 管理总览骨架屏加载态
|
||||||
- **用户体系**:注册(开放/邀请码可切换,管理后台可改)、登录(JWT access + HttpOnly refresh cookie)、argon2id 密码
|
- **用户体系**:注册(开放/邀请码可切换,管理后台可改)、登录(JWT access + HttpOnly refresh cookie)、argon2id 密码
|
||||||
|
|||||||
+54
-25
@@ -13,11 +13,11 @@ import (
|
|||||||
)
|
)
|
||||||
|
|
||||||
type createKeyReq struct {
|
type createKeyReq struct {
|
||||||
Name *string `json:"name" binding:"required,min=1,max=64"`
|
Name *string `json:"name" binding:"required,min=1,max=64"`
|
||||||
QuotaTokensPerDay *int64 `json:"quota_tokens_per_day"`
|
QuotaTokensPerDay *int64 `json:"quota_tokens_per_day"`
|
||||||
QuotaRequestsPerDay *int `json:"quota_requests_per_day"`
|
QuotaRequestsPerDay *int `json:"quota_requests_per_day"`
|
||||||
AllowedModels []string `json:"allowed_models"`
|
AllowedModels []string `json:"allowed_models"`
|
||||||
ExpiresAt *string `json:"expires_at"` // RFC3339
|
ExpiresAt *string `json:"expires_at"` // RFC3339
|
||||||
}
|
}
|
||||||
|
|
||||||
// CreateKey POST /api/v1/keys — 创建密钥,明文仅此一次返回。
|
// CreateKey POST /api/v1/keys — 创建密钥,明文仅此一次返回。
|
||||||
@@ -47,6 +47,13 @@ func (h *Handler) CreateKey(c *gin.Context) {
|
|||||||
AllowedModels: req.AllowedModels,
|
AllowedModels: req.AllowedModels,
|
||||||
Status: store.KeyStatusActive,
|
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 {
|
if req.ExpiresAt != nil {
|
||||||
t, err := time.Parse(time.RFC3339, *req.ExpiresAt)
|
t, err := time.Parse(time.RFC3339, *req.ExpiresAt)
|
||||||
if err != nil {
|
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
|
// ListKeys GET /api/v1/keys
|
||||||
func (h *Handler) ListKeys(c *gin.Context) {
|
func (h *Handler) ListKeys(c *gin.Context) {
|
||||||
u, ok := userFromContext(c)
|
u, ok := userFromContext(c)
|
||||||
@@ -82,18 +105,7 @@ func (h *Handler) ListKeys(c *gin.Context) {
|
|||||||
}
|
}
|
||||||
out := make([]gin.H, 0, len(keys))
|
out := make([]gin.H, 0, len(keys))
|
||||||
for _, k := range keys {
|
for _, k := range keys {
|
||||||
out = append(out, gin.H{
|
out = append(out, keyJSON(k))
|
||||||
"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,
|
|
||||||
})
|
|
||||||
}
|
}
|
||||||
resp.OK(c, gin.H{"items": out})
|
resp.OK(c, gin.H{"items": out})
|
||||||
}
|
}
|
||||||
@@ -111,11 +123,11 @@ func (h *Handler) PatchKey(c *gin.Context) {
|
|||||||
return
|
return
|
||||||
}
|
}
|
||||||
var req struct {
|
var req struct {
|
||||||
Name *string `json:"name"`
|
Name *string `json:"name"`
|
||||||
QuotaTokensPerDay *int64 `json:"quota_tokens_per_day"`
|
QuotaTokensPerDay *int64 `json:"quota_tokens_per_day"`
|
||||||
QuotaRequestsPerDay *int `json:"quota_requests_per_day"`
|
QuotaRequestsPerDay *int `json:"quota_requests_per_day"`
|
||||||
AllowedModels *[]string `json:"allowed_models"`
|
AllowedModels *[]string `json:"allowed_models"`
|
||||||
Status *string `json:"status"`
|
Status *string `json:"status"`
|
||||||
}
|
}
|
||||||
if err := c.ShouldBindJSON(&req); err != nil {
|
if err := c.ShouldBindJSON(&req); err != nil {
|
||||||
resp.Fail(c, http.StatusBadRequest, "invalid input")
|
resp.Fail(c, http.StatusBadRequest, "invalid input")
|
||||||
@@ -130,11 +142,21 @@ func (h *Handler) PatchKey(c *gin.Context) {
|
|||||||
if req.Name != nil {
|
if req.Name != nil {
|
||||||
updates["name"] = *req.Name
|
updates["name"] = *req.Name
|
||||||
}
|
}
|
||||||
|
// 限额:>0 生效,0 或负数视为"不限制"(写 NULL)。
|
||||||
|
// 不能存 0:token 配额判定是 `已用 >= 上限`,0 会让密钥第一次请求就被 429 打死。
|
||||||
if req.QuotaTokensPerDay != nil {
|
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 {
|
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 != nil {
|
||||||
if *req.Status != store.KeyStatusActive && *req.Status != store.KeyStatusRevoked {
|
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)
|
// allowed_models 是 jsonb:手动序列化为 JSON 字符串写入(跨 SQLite/Postgres)
|
||||||
|
// 显式传 [] 表示"清空白名单"(前端编辑框预填现值,字段内容即最终意图)
|
||||||
if req.AllowedModels != nil {
|
if req.AllowedModels != nil {
|
||||||
raw, _ := json.Marshal(*req.AllowedModels)
|
raw, _ := json.Marshal(*req.AllowedModels)
|
||||||
if err := h.a.DB.Model(&k).Update("allowed_models", string(raw)).Error; err != nil {
|
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
|
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 — 删除(硬删,立即失效不可恢复)。
|
// DeleteKey DELETE /api/v1/keys/:id — 删除(硬删,立即失效不可恢复)。
|
||||||
|
|||||||
@@ -5,8 +5,10 @@ package channel
|
|||||||
import (
|
import (
|
||||||
"crypto/rand"
|
"crypto/rand"
|
||||||
"errors"
|
"errors"
|
||||||
|
"log"
|
||||||
"math/big"
|
"math/big"
|
||||||
"sync"
|
"sync"
|
||||||
|
"time"
|
||||||
|
|
||||||
"github.com/openteam/server/internal/pkg/crypto"
|
"github.com/openteam/server/internal/pkg/crypto"
|
||||||
"github.com/openteam/server/internal/store"
|
"github.com/openteam/server/internal/store"
|
||||||
@@ -21,10 +23,29 @@ type Service struct {
|
|||||||
|
|
||||||
mu sync.Mutex
|
mu sync.Mutex
|
||||||
sems map[uint64]chan struct{}
|
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 {
|
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 一个候选渠道 + 该模型的映射关系。
|
// Candidate 一个候选渠道 + 该模型的映射关系。
|
||||||
@@ -33,18 +54,20 @@ type Candidate struct {
|
|||||||
UpstreamModel string // 全局模型在此渠道的映射名(无绑定则为空,用客户端模型名)
|
UpstreamModel string // 全局模型在此渠道的映射名(无绑定则为空,用客户端模型名)
|
||||||
}
|
}
|
||||||
|
|
||||||
// Candidates 返回可用渠道候选:健康 + 启用,按优先级、权重降序、id 升序排列。
|
// Candidates 返回可用渠道候选:健康 + 启用 + 不在运行期冷却,按优先级、权重降序、id 升序排列。
|
||||||
// model 非空时优先取绑定该模型的渠道(携带 upstream_model 映射);无绑定则退回全局。
|
// model 非空且该模型有渠道绑定时,只返回绑定渠道(携带 upstream_model 映射);
|
||||||
|
// 绑定渠道全部不可用则返回空(上层报 no_channel),不会退化到不认识该模型的全局渠道。
|
||||||
|
// 仅当模型完全没有任何绑定时才退回全局渠道。
|
||||||
func (s *Service) Candidates(model string) []Candidate {
|
func (s *Service) Candidates(model string) []Candidate {
|
||||||
if model != "" {
|
if model != "" {
|
||||||
var b []store.ChannelModelBinding
|
|
||||||
var modelIDs []uint64
|
var modelIDs []uint64
|
||||||
|
var b []store.ChannelModelBinding
|
||||||
s.db.Model(&store.Model{}).Where("name = ? AND enabled = ?", model, true).Pluck("id", &modelIDs)
|
s.db.Model(&store.Model{}).Where("name = ? AND enabled = ?", model, true).Pluck("id", &modelIDs)
|
||||||
if len(modelIDs) > 0 {
|
if len(modelIDs) > 0 {
|
||||||
s.db.Where("model_id IN ?", modelIDs).Find(&b)
|
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
|
var chs []store.Channel
|
||||||
@@ -52,12 +75,15 @@ func (s *Service) Candidates(model string) []Candidate {
|
|||||||
Order("priority ASC, weight DESC, id ASC").Find(&chs)
|
Order("priority ASC, weight DESC, id ASC").Find(&chs)
|
||||||
out := make([]Candidate, 0, len(chs))
|
out := make([]Candidate, 0, len(chs))
|
||||||
for i := range chs {
|
for i := range chs {
|
||||||
|
if s.InCooldown(chs[i].ID) {
|
||||||
|
continue
|
||||||
|
}
|
||||||
out = append(out, Candidate{Channel: &chs[i]})
|
out = append(out, Candidate{Channel: &chs[i]})
|
||||||
}
|
}
|
||||||
return out
|
return out
|
||||||
}
|
}
|
||||||
|
|
||||||
// loadBound 按绑定顺序加载渠道候选,过滤健康/启用,携带 upstream_model 映射。
|
// loadBound 按绑定顺序加载渠道候选,过滤健康/启用/运行期冷却,携带 upstream_model 映射。
|
||||||
func (s *Service) loadBound(bindings []store.ChannelModelBinding) []Candidate {
|
func (s *Service) loadBound(bindings []store.ChannelModelBinding) []Candidate {
|
||||||
if len(bindings) == 0 {
|
if len(bindings) == 0 {
|
||||||
return nil
|
return nil
|
||||||
@@ -80,7 +106,7 @@ func (s *Service) loadBound(bindings []store.ChannelModelBinding) []Candidate {
|
|||||||
}
|
}
|
||||||
out := make([]Candidate, 0, len(ids))
|
out := make([]Candidate, 0, len(ids))
|
||||||
for _, id := range 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})
|
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) {
|
func (s *Service) UpstreamKey(ch *store.Channel) (string, error) {
|
||||||
return s.enc.Decrypt(ch.APIKeyEnc)
|
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
|
||||||
|
}
|
||||||
@@ -129,3 +129,79 @@ func TestHealthStartStops(t *testing.T) {
|
|||||||
m.Start(ctx)
|
m.Start(ctx)
|
||||||
cancel() // 不应 panic
|
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("非连续失败不应触发冷却")
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -6,6 +6,7 @@ import (
|
|||||||
"encoding/json"
|
"encoding/json"
|
||||||
"errors"
|
"errors"
|
||||||
"fmt"
|
"fmt"
|
||||||
|
"net"
|
||||||
"net/http"
|
"net/http"
|
||||||
"strings"
|
"strings"
|
||||||
"sync"
|
"sync"
|
||||||
@@ -23,9 +24,10 @@ import (
|
|||||||
)
|
)
|
||||||
|
|
||||||
const (
|
const (
|
||||||
CtxUserID = "proxy_user_id"
|
CtxUserID = "proxy_user_id"
|
||||||
CtxKeyID = "proxy_key_id"
|
CtxKeyID = "proxy_key_id"
|
||||||
CtxTrace = "proxy_trace_id"
|
CtxKeyModel = "proxy_key_models"
|
||||||
|
CtxTrace = "proxy_trace_id"
|
||||||
)
|
)
|
||||||
|
|
||||||
type Gateway struct {
|
type Gateway struct {
|
||||||
@@ -36,7 +38,11 @@ type Gateway struct {
|
|||||||
lim *ratelimit.Limiter
|
lim *ratelimit.Limiter
|
||||||
userRPS int
|
userRPS int
|
||||||
logRaw bool
|
logRaw bool
|
||||||
hc *http.Client
|
|
||||||
|
// 上游超时兜底值(渠道 timeout_ms 非法时使用);按超时值缓存 Transport 复用连接池。
|
||||||
|
gwTimeout time.Duration
|
||||||
|
tpMu sync.Mutex
|
||||||
|
tps map[time.Duration]*http.Transport
|
||||||
|
|
||||||
policyMu sync.Mutex
|
policyMu sync.Mutex
|
||||||
policy modelPolicy
|
policy modelPolicy
|
||||||
@@ -58,7 +64,7 @@ func (g *Gateway) ResetModelPolicy() {
|
|||||||
g.policyMu.Unlock()
|
g.policyMu.Unlock()
|
||||||
}
|
}
|
||||||
|
|
||||||
// globalModelRestrictions 读取全局模型允许/禁止列表(缓存 30s)。
|
// globalModelRestrictions 读取全局模型允许/禁止列表(缓存 modelPolicyTTL)。
|
||||||
func (g *Gateway) globalModelRestrictions() (allow, deny []string) {
|
func (g *Gateway) globalModelRestrictions() (allow, deny []string) {
|
||||||
g.policyMu.Lock()
|
g.policyMu.Lock()
|
||||||
defer g.policyMu.Unlock()
|
defer g.policyMu.Unlock()
|
||||||
@@ -75,18 +81,24 @@ func (g *Gateway) globalModelRestrictions() (allow, deny []string) {
|
|||||||
return
|
return
|
||||||
}
|
}
|
||||||
|
|
||||||
// checkModelAllowed 模型访问控制:用户级 > 全局。
|
// checkModelAllowed 模型访问控制:用户级 > 密钥级 > 全局。
|
||||||
// 1. 用户禁止列表命中 → 拒绝
|
// 1. 用户禁止列表命中 → 拒绝
|
||||||
// 2. 用户允许列表非空 → 仅列表内可访问(不再看全局)
|
// 2. 密钥允许列表非空且不含该模型 → 拒绝(密钥级只收窄,不放宽)
|
||||||
// 3. 全局禁止命中 → 拒绝
|
// 3. 用户允许列表非空 → 仅列表内可访问,不再看全局(管理员显式授权该用户,短路全局)
|
||||||
// 4. 全局允许列表非空 → 仅列表内可访问
|
// 4. 全局禁止命中 → 拒绝
|
||||||
func (g *Gateway) checkModelAllowed(u *store.User, model string) bool {
|
// 5. 全局允许列表非空 → 仅列表内可访问
|
||||||
|
func (g *Gateway) checkModelAllowed(c *gin.Context, u *store.User, model string) bool {
|
||||||
if model == "" {
|
if model == "" {
|
||||||
return true
|
return true
|
||||||
}
|
}
|
||||||
if contains(u.DeniedModels, model) {
|
if contains(u.DeniedModels, model) {
|
||||||
return false
|
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 {
|
if len(u.AllowedModels) > 0 {
|
||||||
return contains(u.AllowedModels, model)
|
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 {
|
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{
|
return &Gateway{
|
||||||
db: db,
|
db: db,
|
||||||
ch: channel.NewService(db, enc),
|
ch: channel.NewService(db, enc),
|
||||||
rec: rec,
|
rec: rec,
|
||||||
enc: enc,
|
enc: enc,
|
||||||
lim: lim,
|
lim: lim,
|
||||||
userRPS: userRPS,
|
userRPS: userRPS,
|
||||||
logRaw: logRaw,
|
logRaw: logRaw,
|
||||||
hc: &http.Client{Timeout: gw_timeout * time.Second},
|
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 → 哈希查表 → 校验状态/过期/模型白名单。
|
// Auth 代理鉴权中间件:Bearer sk-xxx → 哈希查表 → 校验状态/过期/模型白名单。
|
||||||
func (g *Gateway) Auth(c *gin.Context) {
|
func (g *Gateway) Auth(c *gin.Context) {
|
||||||
// 先按路径确定客户端协议,保证 Auth 阶段错误也按协议格式返回
|
// 先按路径确定客户端协议,保证 Auth 阶段错误也按协议格式返回
|
||||||
@@ -186,6 +238,7 @@ func (g *Gateway) Auth(c *gin.Context) {
|
|||||||
|
|
||||||
c.Set(CtxUserID, u.ID)
|
c.Set(CtxUserID, u.ID)
|
||||||
c.Set(CtxKeyID, k.ID)
|
c.Set(CtxKeyID, k.ID)
|
||||||
|
c.Set(CtxKeyModel, k.AllowedModels)
|
||||||
c.Set(CtxTrace, newTraceID())
|
c.Set(CtxTrace, newTraceID())
|
||||||
g.db.Model(&store.APIKey{}).Where("id = ?", k.ID).Update("last_used_at", time.Now())
|
g.db.Model(&store.APIKey{}).Where("id = ?", k.ID).Update("last_used_at", time.Now())
|
||||||
c.Next()
|
c.Next()
|
||||||
|
|||||||
@@ -25,7 +25,7 @@ func (g *Gateway) chatCompletions(c *gin.Context) {
|
|||||||
c.Set("protocol", convert.ProtoChat)
|
c.Set("protocol", convert.ProtoChat)
|
||||||
c.Set("model_name", br.Model)
|
c.Set("model_name", br.Model)
|
||||||
g.recordRawRequest(c, u, body)
|
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", "模型未对你开放,请联系管理员")
|
apiError(c, http.StatusForbidden, "model_not_allowed", "模型未对你开放,请联系管理员")
|
||||||
return
|
return
|
||||||
}
|
}
|
||||||
@@ -58,7 +58,7 @@ func (g *Gateway) responses(c *gin.Context) {
|
|||||||
c.Set("protocol", convert.ProtoResponses)
|
c.Set("protocol", convert.ProtoResponses)
|
||||||
c.Set("model_name", br.Model)
|
c.Set("model_name", br.Model)
|
||||||
g.recordRawRequest(c, u, body)
|
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", "模型未对你开放,请联系管理员")
|
apiError(c, http.StatusForbidden, "model_not_allowed", "模型未对你开放,请联系管理员")
|
||||||
return
|
return
|
||||||
}
|
}
|
||||||
@@ -91,7 +91,7 @@ func (g *Gateway) messages(c *gin.Context) {
|
|||||||
c.Set("protocol", convert.ProtoMessages)
|
c.Set("protocol", convert.ProtoMessages)
|
||||||
c.Set("model_name", br.Model)
|
c.Set("model_name", br.Model)
|
||||||
g.recordRawRequest(c, u, body)
|
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", "模型未对你开放,请联系管理员")
|
apiError(c, http.StatusForbidden, "model_not_allowed", "模型未对你开放,请联系管理员")
|
||||||
return
|
return
|
||||||
}
|
}
|
||||||
@@ -125,7 +125,7 @@ func apiError(c *gin.Context, status int, code, message string) {
|
|||||||
if p, _ := c.Get("protocol"); p == convert.ProtoMessages {
|
if p, _ := c.Get("protocol"); p == convert.ProtoMessages {
|
||||||
// Anthropic 格式
|
// Anthropic 格式
|
||||||
c.AbortWithStatusJSON(status, gin.H{
|
c.AbortWithStatusJSON(status, gin.H{
|
||||||
"type": "error",
|
"type": "error",
|
||||||
"error": gin.H{"type": errorTypeFor(status), "message": message},
|
"error": gin.H{"type": errorTypeFor(status), "message": message},
|
||||||
})
|
})
|
||||||
return
|
return
|
||||||
|
|||||||
@@ -3,9 +3,11 @@ package proxy
|
|||||||
import (
|
import (
|
||||||
"bufio"
|
"bufio"
|
||||||
"bytes"
|
"bytes"
|
||||||
|
"context"
|
||||||
"crypto/rand"
|
"crypto/rand"
|
||||||
"encoding/hex"
|
"encoding/hex"
|
||||||
"encoding/json"
|
"encoding/json"
|
||||||
|
"errors"
|
||||||
"fmt"
|
"fmt"
|
||||||
"io"
|
"io"
|
||||||
"net/http"
|
"net/http"
|
||||||
@@ -187,8 +189,17 @@ func (g *Gateway) doProxy(c *gin.Context, cands []channel.Candidate, clientProto
|
|||||||
idx := (startIdx + offset) % len(cands)
|
idx := (startIdx + offset) % len(cands)
|
||||||
cand := cands[idx]
|
cand := cands[idx]
|
||||||
ch := cand.Channel
|
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)
|
plan, err := prepareUpstream(ch, clientProto, body, cand.UpstreamModel)
|
||||||
if err != nil {
|
if err != nil {
|
||||||
|
// 请求体转换失败与渠道健康无关,不换渠道计数
|
||||||
lastStatus, lastBody = http.StatusInternalServerError, []byte("conversion error: "+err.Error())
|
lastStatus, lastBody = http.StatusInternalServerError, []byte("conversion error: "+err.Error())
|
||||||
continue
|
continue
|
||||||
}
|
}
|
||||||
@@ -196,9 +207,9 @@ func (g *Gateway) doProxy(c *gin.Context, cands []channel.Candidate, clientProto
|
|||||||
if !ok {
|
if !ok {
|
||||||
continue // 渠道满载,溢出到下一个
|
continue // 渠道满载,溢出到下一个
|
||||||
}
|
}
|
||||||
written, retry, st, b := g.proxyOne(c, ch, plan, stream, sink)
|
done, retry, st, b := g.proxyOne(c, ch, plan, stream, sink)
|
||||||
release()
|
release()
|
||||||
if written {
|
if done {
|
||||||
return
|
return
|
||||||
}
|
}
|
||||||
if !retry {
|
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))
|
g.recordError(c, ch, nil, now(), "upstream_http_"+strconv.Itoa(st))
|
||||||
return
|
return
|
||||||
}
|
}
|
||||||
|
g.ch.ReportFailure(ch) // 网络错误/超时/429/5xx:累计运行期失败
|
||||||
lastStatus, lastBody = st, b
|
lastStatus, lastBody = st, b
|
||||||
}
|
}
|
||||||
// 全部候选重试性失败
|
// 全部候选重试性失败
|
||||||
@@ -216,8 +228,9 @@ func (g *Gateway) doProxy(c *gin.Context, cands []channel.Candidate, clientProto
|
|||||||
}
|
}
|
||||||
|
|
||||||
// proxyOne 对单个渠道执行一次代理。
|
// proxyOne 对单个渠道执行一次代理。
|
||||||
// 返回:written=是否已写客户端响应;retry=是否可安全换渠道重试;status+respBody=失败信息。
|
// 返回:done=本次请求已终结(客户端响应已写出或已记账,停止故障转移);
|
||||||
func (g *Gateway) proxyOne(c *gin.Context, ch *store.Channel, plan *upstreamPlan, stream bool, sink *usageSink) (written bool, retry bool, status int, respBody []byte) {
|
// 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)
|
upKey, err := g.ch.UpstreamKey(ch)
|
||||||
if err != nil {
|
if err != nil {
|
||||||
return false, true, http.StatusInternalServerError, []byte("failed to decrypt channel key")
|
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 {
|
if err != nil {
|
||||||
return false, false, http.StatusInternalServerError, []byte("failed to build upstream request")
|
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()
|
start := time.Now()
|
||||||
resp, err := g.hc.Do(req)
|
resp, err := g.clientFor(limit).Do(req)
|
||||||
if err != nil {
|
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())
|
return false, true, http.StatusBadGateway, []byte("upstream request failed: " + err.Error())
|
||||||
}
|
}
|
||||||
defer resp.Body.Close()
|
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
|
return false, false, resp.StatusCode, errBody
|
||||||
}
|
}
|
||||||
|
g.ch.ReportSuccess(ch)
|
||||||
|
|
||||||
c.Header("Content-Type", resp.Header.Get("Content-Type"))
|
c.Header("Content-Type", resp.Header.Get("Content-Type"))
|
||||||
c.Status(http.StatusOK)
|
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) {
|
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)
|
data, err := io.ReadAll(r)
|
||||||
if err != nil {
|
if err != nil {
|
||||||
apiError(c, http.StatusBadGateway, "upstream_error", "failed reading upstream response")
|
// 走到这里响应头已按 200 写出,无法再回错误体,只按中断原因记账
|
||||||
g.recordError(c, ch, nil, start, "read_error")
|
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
|
return
|
||||||
}
|
}
|
||||||
if usageRaw := extractUsage(data); usageRaw != nil && sink != nil {
|
if usageRaw := extractUsage(data); usageRaw != nil && sink != nil {
|
||||||
|
|||||||
@@ -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)
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -107,7 +107,8 @@ async function saveEdit() {
|
|||||||
const m = parseModels(editForm.allowed_models)
|
const m = parseModels(editForm.allowed_models)
|
||||||
if (t !== undefined) payload.quota_tokens_per_day = t
|
if (t !== undefined) payload.quota_tokens_per_day = t
|
||||||
if (r !== undefined) payload.quota_requests_per_day = r
|
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)
|
await http.patch(`/keys/${editing.value.id}`, payload)
|
||||||
toast.ok('密钥已更新')
|
toast.ok('密钥已更新')
|
||||||
editOpen.value = false
|
editOpen.value = false
|
||||||
@@ -280,10 +281,15 @@ onMounted(load)
|
|||||||
<div class="space-y-4">
|
<div class="space-y-4">
|
||||||
<Input v-model="editForm.name" label="名称" />
|
<Input v-model="editForm.name" label="名称" />
|
||||||
<div class="grid grid-cols-1 sm:grid-cols-2 gap-3">
|
<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_tokens_per_day" label="每日 Token 上限" placeholder="留空不修改,0 = 不限制" />
|
||||||
<Input v-model="editForm.quota_requests_per_day" label="每日请求上限" placeholder="留空则不修改" />
|
<Input v-model="editForm.quota_requests_per_day" label="每日请求上限" placeholder="留空不修改,0 = 不限制" />
|
||||||
</div>
|
</div>
|
||||||
<Input v-model="editForm.allowed_models" label="模型白名单" placeholder="逗号分隔" hint="留空则不修改" />
|
<Input
|
||||||
|
v-model="editForm.allowed_models"
|
||||||
|
label="模型白名单"
|
||||||
|
placeholder="逗号分隔"
|
||||||
|
hint="留空 = 不限制模型(清空后保存即取消白名单)"
|
||||||
|
/>
|
||||||
<label class="block">
|
<label class="block">
|
||||||
<span class="mb-1.5 block text-xs font-medium text-muted">状态</span>
|
<span class="mb-1.5 block text-xs font-medium text-muted">状态</span>
|
||||||
<select
|
<select
|
||||||
|
|||||||
Reference in new issue
Block a user