Sakurasan
0e10f0bbde
- 预扣：每请求按 body 估算成本上界计入冻结额（1e-8 USD 原子计数），402 判定 改为「余额 - 在途冻结 ≤ 0」，并发的长流式请求不再互相透支；OT_PROXY_RESERVE_MAX_TOKENS 控制估算基数（默认 4096，0=关闭），handler defer 释放，失败/中断/panic 均覆盖 - 调度：全部候选都在运行期冷却时，剩余冷却 ≤ OT_PROXY_COOLDOWN_GRACE 的最早恢复渠道 兜底重试一次而不是 503（429 Retry-After 语义）；DB 层不可用的渠道不参与兜底 - 指标：openteam_proxy_requests_total / request_duration / channel_failures / usage_queue_depth / frozen_balance，/metrics 默认关闭（OT_METRICS_ENABLED 开启）
719 lines
26 KiB
Go
719 lines
26 KiB
Go
// Package proxy API 网关核心：代理 /v1/chat/completions、/v1/responses、/v1/messages、/v1/models。
|
||
// M1 直通 OpenAI 渠道；M4 起按客户端协议 × 渠道协议自动转换（见 convert）。
|
||
package proxy
|
||
|
||
import (
|
||
"encoding/json"
|
||
"errors"
|
||
"fmt"
|
||
"net"
|
||
"net/http"
|
||
"strings"
|
||
"sync"
|
||
"time"
|
||
|
||
"github.com/gin-gonic/gin"
|
||
"github.com/openteam/server/internal/channel"
|
||
"github.com/openteam/server/internal/pkg/apikey"
|
||
"github.com/openteam/server/internal/pkg/crypto"
|
||
"github.com/openteam/server/internal/pkg/metrics"
|
||
"github.com/openteam/server/internal/pkg/ratelimit"
|
||
"github.com/openteam/server/internal/proxy/convert"
|
||
"github.com/openteam/server/internal/store"
|
||
"github.com/openteam/server/internal/usage"
|
||
"github.com/shopspring/decimal"
|
||
"gorm.io/gorm"
|
||
"sync/atomic"
|
||
)
|
||
|
||
const (
|
||
CtxUserID = "proxy_user_id"
|
||
CtxUser = "proxy_user" // Auth 阶段加载的完整用户（含余额），resolveUser 直接复用
|
||
CtxKeyID = "proxy_key_id"
|
||
CtxKeyModel = "proxy_key_models"
|
||
CtxTrace = "proxy_trace_id"
|
||
|
||
// CtxReserve 本请求的预扣（冻结）金额，单位 1e-8 USD（与 numeric(20,8) 对齐）。
|
||
CtxReserve = "proxy_reserve_1e8"
|
||
// CtxForceCooling 候选集来自「全冷却时最早恢复」的兜底，选路时放行冷却中的渠道。
|
||
CtxForceCooling = "proxy_force_cooling"
|
||
)
|
||
|
||
type Gateway struct {
|
||
db *gorm.DB
|
||
ch *channel.Service
|
||
rec *usage.Recorder
|
||
enc *crypto.Encryptor
|
||
lim *ratelimit.Limiter
|
||
userRPS int
|
||
logRaw bool
|
||
|
||
// 上游超时兜底值（渠道 timeout_ms 非法时使用）；按超时值缓存 Transport 复用连接池。
|
||
gwTimeout time.Duration
|
||
maxBodyMB int // /v1 请求体上限（MB），仅用于超限报错文案
|
||
// 故障转移给"第二个及以后的渠道"留的时间预算；超了就直接返回最后一个错误，
|
||
// 不限制首个渠道（长回答本来就可能跑很久）。
|
||
failoverBudget time.Duration
|
||
// 按失败类型限制换渠道次数（0=不限）：429/配额类失败是瞬时返回，
|
||
// 时间预算对它们形同虚设，按次数封顶才有意义。
|
||
failoverCaps map[channel.FailKind]int
|
||
// 冷却兜底宽限：全部候选都在运行期冷却时，最早恢复的渠道剩余冷却 ≤ 此值
|
||
// 就带上它兜底重试一次（429 的 Retry-After 语义本就是"到点再试"）；0=关闭。
|
||
cooldownGrace time.Duration
|
||
// 预扣（reserve/capture）：每请求按 body 估算成本上界计入冻结额，
|
||
// 402 判定改为「余额 - 在途冻结 ≤ 0」，并发的长流式请求不再能把余额拖到深度负值。
|
||
// reserveMaxTokens 是输出 token 估算基数（0=关闭预扣）；冻结额按用户分桶存
|
||
// 1e-8 USD 整数（每用户一个 atomic.Int64），冻结/释放 O(1) 无锁竞争。
|
||
reserveMaxTokens int
|
||
frozen sync.Map // uint64(用户ID) → *atomic.Int64（1e-8 USD）
|
||
frozenTotal atomic.Int64 // 全局冻结合计（指标用）
|
||
// 流式空闲多久补一行 SSE 注释心跳（0=关闭）。反向代理的空闲超时会掐掉长思考的流。
|
||
streamHeartbeat time.Duration
|
||
// 流式放行前等多久内容（0=不限制）。超时未出内容取消本次尝试换下一个渠道。
|
||
firstTokenWait time.Duration
|
||
// 会话亲和 TTL（0=关闭）：同一会话粘住上次应答的渠道，保住上游提示词缓存。
|
||
affinityTTL time.Duration
|
||
// 渠道并发满载时每个候选最多排队多久（0=立刻跳过）。
|
||
channelAcquireWait time.Duration
|
||
// 鉴权缓存（密钥 + 用户行），0=关闭。
|
||
auth *authCache
|
||
aff *affinityTable
|
||
tpMu sync.Mutex
|
||
tps map[time.Duration]*http.Transport
|
||
|
||
policyMu sync.Mutex
|
||
policy modelPolicy
|
||
|
||
// 模型目录缓存：每请求原本要查「是否存在」与「计价」两次同名查询，
|
||
// 这里按名称缓存 catalogTTL；管理端写操作立即失效。
|
||
modelMu sync.Mutex
|
||
modelCat map[string]modelEntry
|
||
modelAt time.Time
|
||
}
|
||
|
||
// catalogTTL 目录缓存有效期。运行期状态（冷却、健康）不受它影响，
|
||
// 它们每次读取时重新判定；这里只缓存「配置类」数据（模型、渠道、绑定）。
|
||
const catalogTTL = 10 * time.Second
|
||
|
||
type modelEntry struct {
|
||
m store.Model
|
||
ok bool
|
||
}
|
||
|
||
// modelByName 取模型（带进程内缓存）；第二个返回值表示模型存在且启用。
|
||
func (g *Gateway) modelByName(name string) (store.Model, bool) {
|
||
g.modelMu.Lock()
|
||
if g.modelCat != nil && time.Since(g.modelAt) < catalogTTL {
|
||
if e, ok := g.modelCat[name]; ok {
|
||
g.modelMu.Unlock()
|
||
return e.m, e.ok
|
||
}
|
||
}
|
||
g.modelMu.Unlock()
|
||
|
||
var m store.Model
|
||
err := g.db.Where("name = ? AND enabled = ?", name, true).First(&m).Error
|
||
e := modelEntry{m: m, ok: err == nil}
|
||
|
||
g.modelMu.Lock()
|
||
if g.modelCat == nil || time.Since(g.modelAt) >= catalogTTL {
|
||
g.modelCat = map[string]modelEntry{}
|
||
g.modelAt = time.Now()
|
||
}
|
||
g.modelCat[name] = e
|
||
g.modelMu.Unlock()
|
||
return e.m, e.ok
|
||
}
|
||
|
||
// InvalidateAuth 清空鉴权缓存（密钥/用户/密钥白名单等写操作后调用）。
|
||
func (g *Gateway) InvalidateAuth() { g.auth.clear() }
|
||
|
||
// InvalidateCatalog 清空模型与渠道目录缓存（管理端改动、健康状态变化后调用）。
|
||
func (g *Gateway) InvalidateCatalog() {
|
||
g.modelMu.Lock()
|
||
g.modelCat = nil
|
||
g.modelAt = time.Time{}
|
||
g.modelMu.Unlock()
|
||
g.ch.InvalidateCatalog()
|
||
}
|
||
|
||
// modelPolicy 全局模型限制策略（来自系统配置，短时缓存）。
|
||
type modelPolicy struct {
|
||
allow []string
|
||
deny []string
|
||
at time.Time
|
||
}
|
||
|
||
const modelPolicyTTL = 5 * time.Second
|
||
|
||
// ResetModelPolicy 清空全局模型限制缓存（系统配置保存后调用）。
|
||
func (g *Gateway) ResetModelPolicy() {
|
||
g.policyMu.Lock()
|
||
g.policy = modelPolicy{}
|
||
g.policyMu.Unlock()
|
||
}
|
||
|
||
// ChannelRuntime 渠道运行期健康快照：连续失败次数 + 冷却剩余时间 + 最近失败原因（管理后台展示）。
|
||
func (g *Gateway) ChannelRuntime(id uint64) (fails int, remaining time.Duration, reason string) {
|
||
return g.ch.RuntimeState(id)
|
||
}
|
||
|
||
// RecoverChannel 立即清除渠道的运行期冷却与失败计数（管理后台手动恢复）。
|
||
func (g *Gateway) RecoverChannel(id uint64) {
|
||
g.ch.ClearRuntime(id)
|
||
}
|
||
|
||
// globalModelRestrictions 读取全局模型允许/禁止列表（缓存 modelPolicyTTL）。
|
||
func (g *Gateway) globalModelRestrictions() (allow, deny []string) {
|
||
g.policyMu.Lock()
|
||
defer g.policyMu.Unlock()
|
||
if time.Since(g.policy.at) < modelPolicyTTL {
|
||
return g.policy.allow, g.policy.deny
|
||
}
|
||
var raw string
|
||
g.db.Model(&store.SystemConfig{}).Where("key = ?", "model_allowlist").Pluck("value", &raw)
|
||
_ = json.Unmarshal([]byte(raw), &allow)
|
||
raw = ""
|
||
g.db.Model(&store.SystemConfig{}).Where("key = ?", "model_denylist").Pluck("value", &raw)
|
||
_ = json.Unmarshal([]byte(raw), &deny)
|
||
g.policy = modelPolicy{allow: allow, deny: deny, at: time.Now()}
|
||
return
|
||
}
|
||
|
||
// checkModelAllowed 模型访问控制：用户级 > 密钥级 > 全局。
|
||
// 1. 用户禁止列表命中 → 拒绝
|
||
// 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)
|
||
}
|
||
allow, deny := g.globalModelRestrictions()
|
||
if contains(deny, model) {
|
||
return false
|
||
}
|
||
if len(allow) > 0 {
|
||
return contains(allow, model)
|
||
}
|
||
return true
|
||
}
|
||
|
||
// checkModelExists 拦截不存在或已禁用的模型。/v1/models 只列启用的模型，请求侧口径必须一致：
|
||
// 禁用模型不能只是从列表消失——否则照样能调用、照样计费，管理员的下架操作形同虚设。
|
||
func (g *Gateway) checkModelExists(c *gin.Context, model string) bool {
|
||
if model == "" {
|
||
apiError(c, http.StatusBadRequest, "invalid_request_error", "model is required")
|
||
return false
|
||
}
|
||
if _, ok := g.modelByName(model); !ok {
|
||
apiError(c, http.StatusNotFound, "model_not_found",
|
||
fmt.Sprintf("The model %q does not exist or is not enabled", model))
|
||
return false
|
||
}
|
||
return true
|
||
}
|
||
|
||
func contains(list []string, s string) bool {
|
||
for _, v := range list {
|
||
if v == s {
|
||
return true
|
||
}
|
||
}
|
||
return false
|
||
}
|
||
|
||
// Options 网关配置项。用结构体而不是十个位置参数：里面三个 time.Duration、
|
||
// 两个 int，传错的编译期发现不了。
|
||
type Options struct {
|
||
UserRPS int
|
||
LogRaw bool
|
||
GWTimeout time.Duration // 上游超时兜底（渠道 timeout_ms 优先）
|
||
MaxBodyMB int // /v1 请求体上限（MB）
|
||
FailoverBudget time.Duration // 故障转移重试的时间预算
|
||
FailoverSwitchCaps map[string]int // 按失败类型（net|server|rate|quota）限制换渠道次数，0=不限
|
||
CooldownGrace time.Duration // 全冷却时最早恢复渠道的兜底宽限（0=关闭）
|
||
ReserveMaxTokens int // 预扣估算的最大输出 token（0=关闭预扣）
|
||
StreamHeartbeat time.Duration // 流式空闲心跳间隔
|
||
FirstTokenWait time.Duration // 流式放行前等内容的最长时间（0=不限）
|
||
AffinityTTL time.Duration // 会话亲和 TTL（0=关闭）
|
||
ChannelAcquireWait time.Duration // 渠道满载排队时长（0=立刻跳过）
|
||
AuthCacheTTL time.Duration // 鉴权缓存 TTL（0=关闭）
|
||
}
|
||
|
||
// failoverCapsFrom 把字符串键的失败类型上限配置转成 FailKind 键；未知键忽略（配置容错）。
|
||
func failoverCapsFrom(m map[string]int) map[channel.FailKind]int {
|
||
if len(m) == 0 {
|
||
return nil
|
||
}
|
||
out := make(map[channel.FailKind]int, len(m))
|
||
for k, v := range m {
|
||
switch k {
|
||
case "net":
|
||
out[channel.FailNet] = v
|
||
case "server":
|
||
out[channel.FailServer] = v
|
||
case "rate":
|
||
out[channel.FailRate] = v
|
||
case "quota":
|
||
out[channel.FailQuota] = v
|
||
}
|
||
}
|
||
return out
|
||
}
|
||
|
||
func NewGateway(db *gorm.DB, enc *crypto.Encryptor, rec *usage.Recorder, lim *ratelimit.Limiter, o Options) *Gateway {
|
||
return &Gateway{
|
||
db: db,
|
||
ch: channel.NewService(db, enc),
|
||
rec: rec,
|
||
enc: enc,
|
||
lim: lim,
|
||
userRPS: o.UserRPS,
|
||
logRaw: o.LogRaw,
|
||
gwTimeout: o.GWTimeout,
|
||
maxBodyMB: o.MaxBodyMB,
|
||
failoverBudget: o.FailoverBudget,
|
||
failoverCaps: failoverCapsFrom(o.FailoverSwitchCaps),
|
||
cooldownGrace: o.CooldownGrace,
|
||
reserveMaxTokens: o.ReserveMaxTokens,
|
||
streamHeartbeat: o.StreamHeartbeat,
|
||
firstTokenWait: o.FirstTokenWait,
|
||
affinityTTL: o.AffinityTTL,
|
||
channelAcquireWait: o.ChannelAcquireWait,
|
||
auth: newAuthCache(o.AuthCacheTTL),
|
||
aff: newAffinityTable(),
|
||
modelCat: map[string]modelEntry{},
|
||
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 阶段错误也按协议格式返回
|
||
switch c.Request.URL.Path {
|
||
case "/v1/messages":
|
||
c.Set("protocol", convert.ProtoMessages)
|
||
case "/v1/responses":
|
||
c.Set("protocol", convert.ProtoResponses)
|
||
case "/v1/models":
|
||
// 模型目录两条协议共用同一路径，只能靠头区分：Anthropic 客户端带
|
||
// x-api-key / anthropic-version，OpenAI 客户端从不发这两个头。
|
||
if c.GetHeader("x-api-key") != "" || c.GetHeader("anthropic-version") != "" {
|
||
c.Set("protocol", convert.ProtoMessages)
|
||
} else {
|
||
c.Set("protocol", convert.ProtoChat)
|
||
}
|
||
default:
|
||
c.Set("protocol", convert.ProtoChat)
|
||
}
|
||
|
||
auth := c.GetHeader("Authorization")
|
||
key := strings.TrimPrefix(auth, "Bearer ")
|
||
key = strings.TrimSpace(key)
|
||
if key == "" {
|
||
// Anthropic 客户端（Claude Code / SDK）用 x-api-key 头而不是 Authorization
|
||
key = strings.TrimSpace(c.GetHeader("x-api-key"))
|
||
}
|
||
if !apikey.Valid(key) {
|
||
apiError(c, http.StatusUnauthorized, "invalid_api_key", "Invalid API key format. Expected: Authorization: Bearer sk-ot-... or x-api-key: sk-ot-...")
|
||
c.Abort()
|
||
return
|
||
}
|
||
hash := apikey.Hash(key)
|
||
// 命中缓存时省掉两条随机读；状态/过期仍在下面按当前时间判定
|
||
k, u, ok := g.auth.get(hash)
|
||
if !ok {
|
||
if err := g.db.Where("key_hash = ? AND status = ?", hash, store.KeyStatusActive).First(&k).Error; err != nil {
|
||
apiError(c, http.StatusUnauthorized, "invalid_api_key", "Invalid API key")
|
||
c.Abort()
|
||
return
|
||
}
|
||
if err := g.db.First(&u, k.UserID).Error; err != nil {
|
||
apiError(c, http.StatusForbidden, "user_disabled", "User account is disabled")
|
||
c.Abort()
|
||
return
|
||
}
|
||
g.auth.put(hash, k, u)
|
||
}
|
||
// 缓存值可能是 TTL 内的旧状态：每次请求都按当前时间重判
|
||
if k.Status != store.KeyStatusActive {
|
||
apiError(c, http.StatusUnauthorized, "invalid_api_key", "Invalid API key")
|
||
c.Abort()
|
||
return
|
||
}
|
||
if u.Status != store.UserStatusActive {
|
||
apiError(c, http.StatusForbidden, "user_disabled", "User account is disabled")
|
||
c.Abort()
|
||
return
|
||
}
|
||
if k.ExpiresAt != nil && time.Now().After(*k.ExpiresAt) {
|
||
apiError(c, http.StatusUnauthorized, "key_expired", "API key has expired")
|
||
c.Abort()
|
||
return
|
||
}
|
||
|
||
// 限流与配额（内存计数）
|
||
if g.lim != nil {
|
||
if k.QuotaRequestsPerDay != nil && !g.lim.AllowRequestDaily(k.ID, *k.QuotaRequestsPerDay) {
|
||
apiError(c, http.StatusTooManyRequests, "rate_limit_exceeded", "Daily request quota exceeded for this API key")
|
||
c.Abort()
|
||
return
|
||
}
|
||
if k.QuotaTokensPerDay != nil && g.lim.TokensUsed(k.ID) >= *k.QuotaTokensPerDay {
|
||
apiError(c, http.StatusTooManyRequests, "rate_limit_exceeded", "Daily token quota exceeded for this API key")
|
||
c.Abort()
|
||
return
|
||
}
|
||
if !g.lim.AllowUserRate(u.ID, g.userRPS) {
|
||
apiError(c, http.StatusTooManyRequests, "rate_limit_exceeded", "Too many requests. Please slow down.")
|
||
c.Abort()
|
||
return
|
||
}
|
||
}
|
||
|
||
c.Set(CtxUser, &u)
|
||
c.Set(CtxUserID, u.ID)
|
||
c.Set(CtxKeyID, k.ID)
|
||
c.Set(CtxKeyModel, k.AllowedModels)
|
||
c.Set(CtxTrace, newTraceID())
|
||
// last_used_at 按分钟节流：纯展示字段，不值得每请求一次 DB 写（高 QPS 下是可观开销）。
|
||
// 密钥行可能来自缓存（时间戳不会自己变），所以节流按缓存里的 touch 记录判断。
|
||
if g.auth.shouldTouch(k.ID, time.Minute) {
|
||
g.db.Model(&store.APIKey{}).Where("id = ?", k.ID).Update("last_used_at", time.Now())
|
||
}
|
||
c.Next()
|
||
}
|
||
|
||
// Handle 路由到对应协议处理器。
|
||
func (g *Gateway) Handle(c *gin.Context) {
|
||
switch {
|
||
case c.Request.URL.Path == "/v1/chat/completions":
|
||
g.chatCompletions(c)
|
||
case c.Request.URL.Path == "/v1/responses":
|
||
g.responses(c)
|
||
case c.Request.URL.Path == "/v1/messages":
|
||
g.messages(c)
|
||
case c.Request.URL.Path == "/v1/models" && c.Request.Method == http.MethodGet:
|
||
g.models(c)
|
||
default:
|
||
apiError(c, http.StatusNotFound, "not_found", "Unknown endpoint: "+c.Request.URL.Path)
|
||
}
|
||
}
|
||
|
||
// models GET /v1/models：返回调用方真正能用的模型。
|
||
// 口径与选路一致（启用 + 至少绑定一个启用且健康的渠道），再叠加该用户/密钥的模型白名单——
|
||
// 否则被限制的模型名照样泄漏在列表里，客户端试了才知道 403。
|
||
// 按客户端协议返回对应形状，并带上目录元数据（上下文长度、能力位）；价格不外泄。
|
||
func (g *Gateway) models(c *gin.Context) {
|
||
u, ok := g.resolveUser(c)
|
||
if !ok {
|
||
return
|
||
}
|
||
ids := g.ch.AvailableModelIDs()
|
||
if len(ids) == 0 {
|
||
g.writeModelList(c, nil)
|
||
return
|
||
}
|
||
var ms []store.Model
|
||
if err := g.db.Where("id IN ? AND enabled = ?", ids, true).Order("sort ASC, id ASC").Find(&ms).Error; err != nil {
|
||
apiError(c, http.StatusInternalServerError, "internal_error", "failed to load models")
|
||
return
|
||
}
|
||
out := make([]store.Model, 0, len(ms))
|
||
for _, m := range ms {
|
||
if g.checkModelAllowed(c, u, m.Name) {
|
||
out = append(out, m)
|
||
}
|
||
}
|
||
g.writeModelList(c, out)
|
||
}
|
||
|
||
func (g *Gateway) writeModelList(c *gin.Context, ms []store.Model) {
|
||
if p, _ := c.Get("protocol"); p == convert.ProtoMessages {
|
||
data := make([]gin.H, 0, len(ms))
|
||
for _, m := range ms {
|
||
data = append(data, gin.H{
|
||
"type": "model",
|
||
"id": m.Name,
|
||
"display_name": m.Name,
|
||
"created_at": m.CreatedAt.UTC().Format(time.RFC3339),
|
||
"context_length": m.ContextLength,
|
||
"capabilities": gin.H{
|
||
"tools": m.SupportsTools, "vision": m.SupportsVision, "reasoning": m.SupportsReasoning,
|
||
},
|
||
})
|
||
}
|
||
c.JSON(http.StatusOK, gin.H{"data": data, "has_more": false})
|
||
return
|
||
}
|
||
data := make([]gin.H, 0, len(ms))
|
||
for _, m := range ms {
|
||
item := gin.H{
|
||
"id": m.Name, "object": "model",
|
||
"created": m.CreatedAt.Unix(), "owned_by": "openteam",
|
||
}
|
||
if m.ContextLength > 0 {
|
||
item["context_length"] = m.ContextLength
|
||
}
|
||
if m.SupportsTools || m.SupportsVision || m.SupportsReasoning {
|
||
item["capabilities"] = gin.H{
|
||
"tools": m.SupportsTools, "vision": m.SupportsVision, "reasoning": m.SupportsReasoning,
|
||
}
|
||
}
|
||
data = append(data, item)
|
||
}
|
||
c.JSON(http.StatusOK, gin.H{"object": "list", "data": data})
|
||
}
|
||
|
||
// candidateChannels 返回可用渠道候选（按模型绑定优先，退化全局；携带模型映射）。
|
||
func (g *Gateway) candidateChannels(model string) []channel.Candidate {
|
||
return g.ch.Candidates(model)
|
||
}
|
||
|
||
// candidatesFor 候选获取的统一入口：健康候选为空且配置了冷却宽限时，
|
||
// 找冷却最早结束的候选兜底重试一次（429 的 Retry-After 语义本就是"到点再试"，
|
||
// 与其立刻 503 不如试试最先恢复的那条）。命中兜底时在 context 打标，
|
||
// doProxy 选路会放行这个冷却中的渠道。
|
||
func (g *Gateway) candidatesFor(c *gin.Context, model string) []channel.Candidate {
|
||
cands := g.candidateChannels(model)
|
||
if len(cands) > 0 || g.cooldownGrace <= 0 {
|
||
return cands
|
||
}
|
||
cand, remaining, ok := g.ch.EarliestRecovery(model)
|
||
if !ok || remaining > g.cooldownGrace {
|
||
return nil
|
||
}
|
||
c.Set(CtxForceCooling, true)
|
||
return []channel.Candidate{cand}
|
||
}
|
||
|
||
// resolveUser 取当前用户（含余额）。Auth 中间件已查过一次并放入上下文，
|
||
// handler 里直接复用，省掉每请求一次重复查询。
|
||
func (g *Gateway) resolveUser(c *gin.Context) (*store.User, bool) {
|
||
if u, ok := c.Get(CtxUser); ok {
|
||
if u2, ok := u.(*store.User); ok {
|
||
return u2, true
|
||
}
|
||
}
|
||
uid, _ := c.Get(CtxUserID)
|
||
var u store.User
|
||
if err := g.db.First(&u, uid).Error; err != nil {
|
||
apiError(c, http.StatusUnauthorized, "invalid_api_key", "Invalid API key")
|
||
return nil, false
|
||
}
|
||
return &u, true
|
||
}
|
||
|
||
// checkBalance 余额不足返回 402（PLANNING §4.4.3）。
|
||
// 可用余额 = 余额 - 在途预扣（冻结）额：并发流式请求各持一份冻结，
|
||
// 不会互相透支把余额拖到深度负值（冻结语义见 freezeFor）。
|
||
func (g *Gateway) checkBalance(c *gin.Context, u *store.User) bool {
|
||
avail := u.Balance
|
||
if a := g.frozenOf(u.ID); a != nil && a.Load() > 0 {
|
||
avail = avail.Sub(decimal.NewFromInt(a.Load()).Shift(-8))
|
||
}
|
||
if !avail.IsPositive() {
|
||
apiError(c, http.StatusPaymentRequired, "insufficient_balance", "Insufficient balance. Please recharge or contact admin.")
|
||
return false
|
||
}
|
||
return true
|
||
}
|
||
|
||
// frozenOf 取用户的在途冻结额桶；无冻结返回 nil（零分配快路径）。
|
||
func (g *Gateway) frozenOf(uid uint64) *atomic.Int64 {
|
||
if v, ok := g.frozen.Load(uid); ok {
|
||
return v.(*atomic.Int64)
|
||
}
|
||
return nil
|
||
}
|
||
|
||
// freezeFor 预扣（reserve）：按请求体估算本请求的成本上界并计入冻结额，随 gin context
|
||
// 携带，请求终结（含失败/中断/panic，由 handler 的 defer 释放）时按同一数值 capture。
|
||
// 估算只求上界不求准：输入按 ~4 字节/token 粗估并用缓存写价（1.25×/2× 输入价）封顶，
|
||
// 输出取 max_tokens（缺省用 reserveMaxTokens 基数）。
|
||
// 冻结只在本进程内存里，单实例语义；多实例部署时各实例独立冻结（余额本身允许为负兜底）。
|
||
// 检查（checkBalance）与冻结之间存在乐观竞窗：最坏超额 = 同时穿过竞窗的请求数，
|
||
// 有界且被「余额允许为负」兜底；预扣把无界的流式透支压成了这个常数。
|
||
func (g *Gateway) freezeFor(c *gin.Context, uid uint64, model store.Model, br *bodyReq, body []byte) {
|
||
if g.reserveMaxTokens <= 0 || model.ID == 0 {
|
||
return
|
||
}
|
||
outTok := int64(g.reserveMaxTokens)
|
||
if br != nil {
|
||
if br.MaxTokens != nil && *br.MaxTokens > 0 && int64(*br.MaxTokens) < outTok {
|
||
outTok = int64(*br.MaxTokens)
|
||
}
|
||
if br.MaxOutputTokens != nil && *br.MaxOutputTokens > 0 && int64(*br.MaxOutputTokens) < outTok {
|
||
outTok = int64(*br.MaxOutputTokens)
|
||
}
|
||
}
|
||
inPrice := model.InputPrice
|
||
if cwp := model.CacheWritePriceOr(); cwp.GreaterThan(inPrice) {
|
||
inPrice = cwp
|
||
}
|
||
million := decimal.NewFromInt(1e6)
|
||
est := decimal.NewFromInt(int64(len(body) / 4)).Mul(inPrice).
|
||
Add(decimal.NewFromInt(outTok).Mul(model.OutputPrice)).
|
||
Div(million).Round(8)
|
||
n := est.Shift(8).Ceil().BigInt().Int64() // 1e-8 USD 整数，向上取整保上界
|
||
if n <= 0 {
|
||
return
|
||
}
|
||
v, _ := g.frozen.LoadOrStore(uid, &atomic.Int64{})
|
||
v.(*atomic.Int64).Add(n)
|
||
g.frozenTotal.Add(n)
|
||
metrics.FrozenBalance.Set(float64(g.frozenTotal.Load()) / 1e8)
|
||
c.Set(CtxReserve, n)
|
||
}
|
||
|
||
// unfreeze 释放冻结额（capture）。幂等：CtxReserve 置零后重复调用为 no-op，
|
||
// handler 的 defer 与任何终结路径的双释放只生效一次。
|
||
func (g *Gateway) unfreeze(c *gin.Context, uid uint64) {
|
||
v, ok := c.Get(CtxReserve)
|
||
if !ok {
|
||
return
|
||
}
|
||
n, _ := v.(int64)
|
||
c.Set(CtxReserve, int64(0))
|
||
if n == 0 {
|
||
return
|
||
}
|
||
if a := g.frozenOf(uid); a != nil {
|
||
a.Add(-n)
|
||
}
|
||
total := g.frozenTotal.Add(-n)
|
||
metrics.FrozenBalance.Set(float64(total) / 1e8)
|
||
}
|
||
|
||
// ---------------------------------------------------------------------------
|
||
// 协议分派
|
||
|
||
func upstreamPath(proto string) string {
|
||
switch proto {
|
||
case convert.ProtoMessages:
|
||
return "/messages"
|
||
case convert.ProtoResponses:
|
||
return "/responses"
|
||
default:
|
||
return "/chat/completions"
|
||
}
|
||
}
|
||
|
||
// upstreamPlan 描述一次代理请求的上游访问方式。
|
||
type upstreamPlan struct {
|
||
proto string // 上游协议（用于分协议 base_url）
|
||
path string // 上游路径
|
||
body []byte // 已转换的请求体
|
||
lineConv func([]byte) []byte // 流式逐行转换（nil=直通）
|
||
bodyConv func([]byte) ([]byte, error) // 非流式响应体转换（nil=直通）
|
||
}
|
||
|
||
// conversionTarget 决定客户端协议在渠道上的处理方式：
|
||
// 渠道声明支持该协议则直通（返回原协议）；否则转为其首选支持协议（chat > messages > responses）。
|
||
func conversionTarget(formats []string, clientProto string) string {
|
||
for _, f := range formats {
|
||
if f == clientProto {
|
||
return clientProto
|
||
}
|
||
}
|
||
for _, p := range []string{convert.ProtoChat, convert.ProtoMessages, convert.ProtoResponses} {
|
||
for _, f := range formats {
|
||
if f == p {
|
||
return p
|
||
}
|
||
}
|
||
}
|
||
return ""
|
||
}
|
||
|
||
// prepareUpstream 计算上游访问计划：渠道声明支持客户端协议则直通，否则转换；
|
||
// 应用模型名称映射（upstream_model）。
|
||
func prepareUpstream(ch *store.Channel, clientProto string, body []byte, upstreamModel string) (*upstreamPlan, error) {
|
||
target := conversionTarget(ch.FormatsEffective(), clientProto)
|
||
if target == "" {
|
||
return nil, fmt.Errorf("channel %q declares no supported protocol format", ch.Name)
|
||
}
|
||
plan := &upstreamPlan{proto: target, path: upstreamPath(target), body: body}
|
||
if target != clientProto {
|
||
converted, err := convert.ConvertRequest(body, clientProto, target)
|
||
if err != nil {
|
||
return nil, err
|
||
}
|
||
plan.body = converted
|
||
plan.lineConv = convert.NewStreamTransformer(target, clientProto)
|
||
plan.bodyConv = func(b []byte) ([]byte, error) { return convert.ConvertResponse(b, target, clientProto) }
|
||
}
|
||
// 模型名称映射：把请求体 model 字段改写为渠道侧的 upstream_model
|
||
if upstreamModel != "" {
|
||
if out, err := rewriteModel(plan.body, upstreamModel); err == nil {
|
||
plan.body = out
|
||
}
|
||
}
|
||
return plan, nil
|
||
}
|
||
|
||
// rewriteModel 改写请求体顶层的 model 字段（模型映射）。
|
||
// 定点替换而不是整包 unmarshal/marshal：后者在 100KB 请求体上实测 ~585µs、
|
||
// 470KB 分配，而这里只有一次扫描 + 拼接。
|
||
func rewriteModel(body []byte, model string) ([]byte, error) {
|
||
out, ok := rewriteModelField(body, model)
|
||
if !ok {
|
||
return body, nil // 没有顶层 model（非法或空体）：原样转发，交给上游判
|
||
}
|
||
return out, nil
|
||
}
|
||
|
||
var errNoChannel = errors.New("no available channel")
|