317 lines
10 KiB
Go
317 lines
10 KiB
Go
package handler
|
||
|
||
import (
|
||
"encoding/json"
|
||
"fmt"
|
||
"net/http"
|
||
"net/url"
|
||
"os"
|
||
"strings"
|
||
"time"
|
||
|
||
"github.com/gogf/gf/v2/frame/g"
|
||
"github.com/gorilla/websocket"
|
||
|
||
"github.com/qianfree/team-api/relay/channel"
|
||
"github.com/qianfree/team-api/relay/channel/openai"
|
||
"github.com/qianfree/team-api/relay/common"
|
||
"github.com/qianfree/team-api/relay/constant"
|
||
)
|
||
|
||
// websocketUpgrader HTTP → WebSocket 升级器
|
||
var websocketUpgrader = websocket.Upgrader{
|
||
Subprotocols: []string{"realtime"},
|
||
CheckOrigin: checkRealtimeOrigin,
|
||
}
|
||
|
||
// checkRealtimeOrigin 校验 WebSocket 握手来源，防跨站 WebSocket 劫持（CSWSH，B15）。
|
||
// 策略：
|
||
// - 无 Origin 头：放行。Realtime API 的主要客户端是原生 SDK/CLI/服务端，不发送 Origin；
|
||
// 浏览器才强制带 Origin，据此区分。
|
||
// - 同源（Origin.host == 请求 Host）：放行。
|
||
// - 命中 REALTIME_ALLOWED_ORIGINS（逗号分隔的可信来源，如租户控制台/playground 域名）：放行。
|
||
// - 其余跨源浏览器请求：拒绝。
|
||
func checkRealtimeOrigin(r *http.Request) bool {
|
||
origin := r.Header.Get("Origin")
|
||
if origin == "" {
|
||
return true // 非浏览器客户端
|
||
}
|
||
|
||
u, err := url.Parse(origin)
|
||
if err != nil || u.Host == "" {
|
||
return false
|
||
}
|
||
|
||
if strings.EqualFold(u.Host, r.Host) {
|
||
return true // 同源
|
||
}
|
||
|
||
for _, allowed := range strings.Split(os.Getenv("REALTIME_ALLOWED_ORIGINS"), ",") {
|
||
allowed = strings.TrimSpace(allowed)
|
||
if allowed == "" {
|
||
continue
|
||
}
|
||
if strings.EqualFold(origin, allowed) || strings.EqualFold(u.Host, allowed) {
|
||
return true
|
||
}
|
||
}
|
||
|
||
return false
|
||
}
|
||
|
||
// RealtimeContext Realtime 请求上下文
|
||
type RealtimeContext struct {
|
||
TenantID int64
|
||
UserID int64
|
||
ApiKeyID int64
|
||
ProjectID int64 // 通过 API Key 关联的项目 ID
|
||
RequestID string
|
||
ClientIP string
|
||
Scope string
|
||
KeyRateLimitQps int
|
||
KeyConcurrency int
|
||
KeyIpWhitelist string
|
||
KeyTotalQuota float64
|
||
KeyUsedQuota float64
|
||
}
|
||
|
||
// HandleRealtime 处理 /v1/realtime WebSocket 请求
|
||
func HandleRealtime(w http.ResponseWriter, r *http.Request, rc *RealtimeContext, provider common.DataProvider, billing common.BillingProvider) (*common.Usage, *BillingResult, error) {
|
||
ctx := r.Context()
|
||
|
||
// 1. WebSocket 升级
|
||
clientConn, err := websocketUpgrader.Upgrade(w, r, nil)
|
||
if err != nil {
|
||
g.Log().Errorf(ctx, "[HandleRealtime] WebSocket upgrade failed: %v", err)
|
||
return nil, nil, fmt.Errorf("websocket upgrade failed: %w", err)
|
||
}
|
||
defer clientConn.Close()
|
||
|
||
// 2. 读取第一条消息获取 model 名称
|
||
_, firstMsg, err := clientConn.ReadMessage()
|
||
if err != nil {
|
||
return nil, nil, fmt.Errorf("read first websocket message failed: %w", err)
|
||
}
|
||
|
||
var firstEvent map[string]json.RawMessage
|
||
if err := json.Unmarshal(firstMsg, &firstEvent); err != nil {
|
||
return nil, nil, fmt.Errorf("parse first websocket message failed: %w", err)
|
||
}
|
||
|
||
modelName := extractModelFromEvent(firstEvent)
|
||
if modelName == "" {
|
||
errMsg, _ := json.Marshal(map[string]any{
|
||
"type": "error",
|
||
"error": map[string]string{"message": "model is required in the first message"},
|
||
})
|
||
_ = clientConn.WriteMessage(websocket.TextMessage, errMsg)
|
||
return nil, nil, fmt.Errorf("model not found in first websocket message")
|
||
}
|
||
|
||
billingResult := &BillingResult{}
|
||
if billing != nil {
|
||
if !billing.CheckScope(rc.Scope, "realtime") {
|
||
return nil, billingResult, constant.NewAuthError("API key scope denied")
|
||
}
|
||
if !billing.CheckIPWhitelist(rc.KeyIpWhitelist, rc.ClientIP) {
|
||
return nil, billingResult, constant.NewAuthError("IP address is not allowed")
|
||
}
|
||
allowed, limitLevel, limit, remaining, resetAt := billing.CheckRateLimit(ctx, rc.TenantID, rc.UserID, rc.ApiKeyID, rc.KeyRateLimitQps)
|
||
if !allowed {
|
||
return nil, billingResult, &RelayErrorWithRateLimit{
|
||
StatusCode: 429,
|
||
Message: fmt.Sprintf("rate limit exceeded at %s level", limitLevel),
|
||
LimitLevel: limitLevel,
|
||
Remaining: remaining,
|
||
ResetAt: resetAt,
|
||
}
|
||
}
|
||
billingResult.RateLimitInfo = &common.RateLimitInfo{
|
||
Limit: limit,
|
||
Remaining: remaining,
|
||
ResetAt: resetAt,
|
||
}
|
||
}
|
||
|
||
// 3. 验证模型
|
||
_, _, err = provider.GetModelMapping(ctx, modelName)
|
||
if err != nil {
|
||
return nil, nil, fmt.Errorf("model not found: %s", modelName)
|
||
}
|
||
|
||
// 3.1 检查成员模型范围
|
||
if allowed, err := provider.CheckMemberModelAccess(ctx, rc.TenantID, rc.UserID, modelName); err != nil {
|
||
return nil, nil, err
|
||
} else if !allowed {
|
||
return nil, nil, constant.NewAuthError("model not allowed for this member")
|
||
}
|
||
if allowed, err := provider.CheckApiKeyModelAccess(ctx, rc.ApiKeyID, modelName); err != nil {
|
||
return nil, nil, err
|
||
} else if !allowed {
|
||
return nil, nil, constant.NewAuthError("model not allowed for this API key")
|
||
}
|
||
|
||
// 4. 渠道选择
|
||
var (
|
||
selection *common.ChannelSelection
|
||
lease *channelLease
|
||
excluded []int64
|
||
)
|
||
for {
|
||
selection, err = provider.GetChannelForModel(ctx, rc.TenantID, rc.UserID, rc.ApiKeyID, modelName, excluded)
|
||
if err != nil {
|
||
return nil, nil, constant.NewChannelError("no available channel for model: "+modelName, err)
|
||
}
|
||
var acquired bool
|
||
lease, acquired = acquireChannelLease(ctx, provider, selection, rc.RequestID)
|
||
if acquired {
|
||
break
|
||
}
|
||
excluded = append(excluded, selection.ChannelID)
|
||
}
|
||
defer lease.Release()
|
||
|
||
// 5. 构造 RelayInfo
|
||
info := &common.RelayInfo{
|
||
Context: ctx,
|
||
TenantID: rc.TenantID,
|
||
UserID: rc.UserID,
|
||
ApiKeyID: rc.ApiKeyID,
|
||
RequestID: rc.RequestID,
|
||
RelayMode: int(constant.RelayModeRealtime),
|
||
OriginModelName: modelName,
|
||
RequestURLPath: "/realtime",
|
||
RequestHeaders: r.Header,
|
||
StartTime: time.Now(),
|
||
InboundFormat: constant.RelayFormatOpenAI,
|
||
ClientFormat: constant.RelayFormatOpenAI,
|
||
ChannelMeta: &common.ChannelMeta{
|
||
ChannelID: selection.ChannelID,
|
||
ChannelType: selection.ChannelType,
|
||
ChannelName: selection.ChannelName,
|
||
BaseURL: selection.BaseURL,
|
||
ApiKey: selection.ApiKey,
|
||
UpstreamModelName: selection.UpstreamModelName,
|
||
IsModelMapped: selection.IsModelMapped,
|
||
Settings: selection.Settings,
|
||
},
|
||
ClientConn: clientConn,
|
||
}
|
||
|
||
// 6. 获取适配器
|
||
adaptor := channel.GetAdaptor(selection.ChannelType)
|
||
if adaptor == nil {
|
||
return nil, nil, fmt.Errorf("unsupported channel type for realtime: %d", selection.ChannelType)
|
||
}
|
||
adaptor.Init(info)
|
||
|
||
// 7. 预扣费用（Realtime 使用 0 input tokens 作为初始估算）
|
||
var preDeductAmount float64
|
||
if billing != nil {
|
||
if !billing.AcquireConcurrent(ctx, rc.TenantID, rc.UserID, rc.ApiKeyID, modelName) {
|
||
return nil, billingResult, constant.NewRateLimitError("concurrent request limit exceeded")
|
||
}
|
||
defer billing.ReleaseConcurrent(ctx, rc.TenantID, rc.UserID, rc.ApiKeyID, modelName)
|
||
|
||
if !billing.AcquireApiKeyConcurrent(ctx, rc.ApiKeyID, rc.KeyConcurrency) {
|
||
return nil, billingResult, constant.NewRateLimitError("API key concurrent request limit exceeded")
|
||
}
|
||
defer billing.ReleaseApiKeyConcurrent(ctx, rc.ApiKeyID)
|
||
|
||
if err := billing.CheckMemberQuota(ctx, rc.TenantID, rc.UserID, 0); err != nil {
|
||
return nil, billingResult, constant.NewQuotaError("member quota exceeded", err)
|
||
}
|
||
if err := billing.CheckApiKeyQuota(ctx, rc.ApiKeyID, 0); err != nil {
|
||
return nil, billingResult, constant.NewQuotaError("API key quota exceeded", err)
|
||
}
|
||
|
||
amt, billErr := billing.PreDeduct(ctx, rc.TenantID, modelName, 0, 0, false, rc.RequestID)
|
||
if billErr != nil {
|
||
return nil, nil, constant.NewQuotaError("insufficient balance", billErr)
|
||
}
|
||
preDeductAmount = amt
|
||
if err := billing.CheckApiKeyQuota(ctx, rc.ApiKeyID, amt); err != nil {
|
||
_ = billing.SettleFailed(ctx, rc.TenantID, rc.RequestID, amt)
|
||
return nil, billingResult, constant.NewQuotaError("API key quota exceeded", err)
|
||
}
|
||
}
|
||
|
||
billingResult.PreDeductAmount = preDeductAmount
|
||
|
||
// 8. 建立 Realtime 代理
|
||
proxy := openai.NewRealtimeProxy(info)
|
||
if err := proxy.DialUpstream(); err != nil {
|
||
provider.DeleteChannelAffinity(ctx, rc.TenantID, rc.UserID, rc.ApiKeyID, modelName)
|
||
if billing != nil && preDeductAmount > 0 {
|
||
_ = billing.SettleFailed(ctx, rc.TenantID, rc.RequestID, preDeductAmount)
|
||
}
|
||
return nil, billingResult, err
|
||
}
|
||
defer proxy.Close()
|
||
|
||
// 把第一条消息转发到上游
|
||
if err := proxy.GetTargetConn().WriteMessage(websocket.TextMessage, firstMsg); err != nil {
|
||
if billing != nil && preDeductAmount > 0 {
|
||
_ = billing.SettleFailed(ctx, rc.TenantID, rc.RequestID, preDeductAmount)
|
||
}
|
||
return nil, billingResult, err
|
||
}
|
||
|
||
// 9. 启动双向代理
|
||
usage, proxyErr := proxy.Proxy(ctx)
|
||
if proxyErr == nil && !selection.PreserveAffinity {
|
||
provider.SetChannelAffinity(ctx, rc.TenantID, rc.UserID, rc.ApiKeyID, modelName, selection.ChannelID)
|
||
}
|
||
|
||
// 10. 结算费用
|
||
if billing != nil && preDeductAmount > 0 {
|
||
if proxyErr != nil {
|
||
streamUsage := usage
|
||
if streamUsage == nil {
|
||
streamUsage = &common.Usage{}
|
||
}
|
||
settleResult, _ := billing.SettleStreamInterrupted(ctx, rc.TenantID, rc.UserID, rc.ApiKeyID, selection.ChannelID,
|
||
modelName, rc.RequestID, "realtime", streamUsage, preDeductAmount, rc.ProjectID)
|
||
if settleResult != nil && settleResult.ActualCost > 0 {
|
||
billing.IncrMemberQuotaUsed(ctx, rc.TenantID, rc.UserID, settleResult.ActualCost)
|
||
billing.IncrApiKeyQuotaUsed(ctx, rc.ApiKeyID, settleResult.ActualCost)
|
||
}
|
||
} else if usage != nil {
|
||
settleResult, _ := billing.Settle(ctx, rc.TenantID, rc.UserID, rc.ApiKeyID, selection.ChannelID,
|
||
modelName, rc.RequestID, "realtime", usage, preDeductAmount, rc.ProjectID)
|
||
if settleResult != nil && settleResult.ActualCost > 0 {
|
||
billing.IncrMemberQuotaUsed(ctx, rc.TenantID, rc.UserID, settleResult.ActualCost)
|
||
billing.IncrApiKeyQuotaUsed(ctx, rc.ApiKeyID, settleResult.ActualCost)
|
||
}
|
||
}
|
||
}
|
||
|
||
return usage, billingResult, proxyErr
|
||
}
|
||
|
||
// extractModelFromEvent 从 WebSocket 事件中提取模型名称
|
||
func extractModelFromEvent(event map[string]json.RawMessage) string {
|
||
// 顶层 model 字段
|
||
if m, ok := event["model"]; ok {
|
||
var s string
|
||
if json.Unmarshal(m, &s) == nil {
|
||
return strings.Trim(s, `"`)
|
||
}
|
||
}
|
||
|
||
// session.update 的 session.model
|
||
if sessionRaw, ok := event["session"]; ok {
|
||
var sess map[string]json.RawMessage
|
||
if json.Unmarshal(sessionRaw, &sess) == nil {
|
||
if m, ok := sess["model"]; ok {
|
||
var s string
|
||
if json.Unmarshal(m, &s) == nil {
|
||
return strings.Trim(s, `"`)
|
||
}
|
||
}
|
||
}
|
||
}
|
||
|
||
return ""
|
||
}
|