367 lines
13 KiB
Go
367 lines
13 KiB
Go
package handler
|
||
|
||
import (
|
||
"context"
|
||
"encoding/json"
|
||
"errors"
|
||
"fmt"
|
||
"net/http"
|
||
"net/url"
|
||
"os"
|
||
"strings"
|
||
"time"
|
||
|
||
"github.com/gogf/gf/v2/frame/g"
|
||
"github.com/gorilla/websocket"
|
||
|
||
"github.com/qianfree/team-api/internal/dispatchadapter"
|
||
"github.com/qianfree/team-api/internal/logic/monitor"
|
||
"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"
|
||
"github.com/qianfree/team-api/relaykit/dispatch"
|
||
)
|
||
|
||
// 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()
|
||
|
||
// 0. 前置检查（scope / IP 白名单 / QPS 限流）：只依赖认证上下文，放在 WebSocket
|
||
// 升级之前——升级后连接已被劫持，无法再写标准 HTTP 状态码；前置后被拒请求能拿到
|
||
// 正常的 401/429 响应，也省掉无谓的升级握手和首消息读取。
|
||
// 注意：上层调用方（internal/handler/relay）对 realtime 的返回错误只记日志、不写响应
|
||
//（因升级后的错误无法写回），故升级前的拒绝必须在此处自行写回 HTTP 错误响应，
|
||
// 否则客户端会收到空 body 的 200。
|
||
billingResult := &BillingResult{}
|
||
if billing != nil {
|
||
if !billing.CheckScope(rc.Scope, "realtime") {
|
||
err := constant.NewAuthError("API key scope denied")
|
||
WriteRelayError(w, err)
|
||
return nil, billingResult, err
|
||
}
|
||
if !billing.CheckIPWhitelist(rc.KeyIpWhitelist, rc.ClientIP) {
|
||
err := constant.NewAuthError("IP address is not allowed")
|
||
WriteRelayError(w, err)
|
||
return nil, billingResult, err
|
||
}
|
||
allowed, limitLevel, limit, remaining, resetAt := billing.CheckRateLimit(ctx, rc.TenantID, rc.UserID, rc.ApiKeyID, rc.KeyRateLimitQps)
|
||
if !allowed {
|
||
err := &RelayErrorWithRateLimit{
|
||
StatusCode: 429,
|
||
Message: fmt.Sprintf("rate limit exceeded at %s level", limitLevel),
|
||
LimitLevel: limitLevel,
|
||
Remaining: remaining,
|
||
ResetAt: resetAt,
|
||
}
|
||
WriteRelayError(w, err)
|
||
return nil, billingResult, err
|
||
}
|
||
billingResult.RateLimitInfo = &common.RateLimitInfo{
|
||
Limit: limit,
|
||
Remaining: remaining,
|
||
ResetAt: resetAt,
|
||
}
|
||
}
|
||
|
||
// 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")
|
||
}
|
||
|
||
// 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. 渠道选择（新调度引擎；websocket 长连接由续期器保活租约）
|
||
enabled, channelScope, err := provider.CheckTenantModelAccess(ctx, rc.TenantID, modelName)
|
||
if err != nil {
|
||
return nil, nil, constant.NewChannelError("no available channel for model: "+modelName, err)
|
||
}
|
||
if !enabled {
|
||
return nil, nil, constant.NewChannelError("no available channel for model: "+modelName, common.ErrTenantModelNotEnabled)
|
||
}
|
||
sess := dispatchadapter.Coordinator(ctx).Route(ctx, dispatch.RequestProfile{
|
||
RequestID: rc.RequestID,
|
||
TenantID: rc.TenantID,
|
||
UserID: rc.UserID,
|
||
APIKeyID: rc.ApiKeyID,
|
||
Model: modelName,
|
||
Scope: channelScope,
|
||
Replay: dispatch.ReplayCostly,
|
||
Signals: dispatch.SessionSignals{HeaderSessionID: r.Header.Get("X-Session-Id")},
|
||
Policy: dispatchadapter.TenantRoutingPolicy(ctx, rc.TenantID),
|
||
})
|
||
defer sess.Finish(context.WithoutCancel(ctx), false, 0)
|
||
|
||
var selection *common.ChannelSelection
|
||
for {
|
||
d := sess.Next(ctx)
|
||
if d == nil {
|
||
monitor.TrackDispatchNoCandidate()
|
||
return nil, nil, constant.NewChannelError("no available channel for model: "+modelName, common.ErrChannelUnavailable)
|
||
}
|
||
monitor.TrackDispatchSelection(string(d.Reason), string(d.Channel.Tier), string(d.SessionKey.Source))
|
||
sel, mErr := provider.MaterializeSelection(ctx, d.Channel.ID, d.KeyID, modelName)
|
||
if mErr != nil {
|
||
if decision := reportMaterializeFailure(ctx, sess, mErr); decision == dispatch.DecisionAbort {
|
||
return nil, nil, constant.NewChannelError("no available channel for model: "+modelName, mErr)
|
||
}
|
||
continue
|
||
}
|
||
sel.SelectionReason = string(d.Reason)
|
||
selection = sel
|
||
break
|
||
}
|
||
leaseRefresh := startDispatchLeaseRefresher(selection.ChannelID, rc.RequestID)
|
||
defer leaseRefresh.Stop()
|
||
realtimeStart := time.Now()
|
||
|
||
// 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)
|
||
|
||
amt, billErr := billing.PreDeduct(ctx, rc.TenantID, modelName, 0, 0, false, rc.RequestID)
|
||
if billErr != nil {
|
||
// 未配价模型 fail-closed：返回明确的请求错误，而非误导性的"余额不足"
|
||
if errors.Is(billErr, common.ErrModelPricingNotConfigured) {
|
||
return nil, nil, constant.NewRequestError("model pricing not configured: "+modelName, billErr)
|
||
}
|
||
return nil, nil, constant.NewQuotaError("insufficient balance", billErr)
|
||
}
|
||
preDeductAmount = amt
|
||
// 成员/Key 额度检查只做这一次（带实际预扣额），与 RelayHandler 一致：
|
||
// 预扣前的 Check(0) 快速闸门对全量请求多付 Redis 往返，收益不成比例，已移除。
|
||
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)
|
||
}
|
||
if err := billing.CheckMemberQuota(ctx, rc.TenantID, rc.UserID, amt); err != nil {
|
||
_ = billing.SettleFailed(ctx, rc.TenantID, rc.RequestID, amt)
|
||
return nil, billingResult, constant.NewQuotaError("member quota exceeded", err)
|
||
}
|
||
}
|
||
|
||
billingResult.PreDeductAmount = preDeductAmount
|
||
|
||
// 8. 建立 Realtime 代理
|
||
proxy := openai.NewRealtimeProxy(info)
|
||
if err := proxy.DialUpstream(); err != nil {
|
||
// 拨号失败上报健康（绑定不删除，由守卫自然失效）；realtime 不做跨渠道重试，保持原语义
|
||
_, _ = sess.Report(ctx, 0, err, deliveryStateOfRequestErr(err), float64(time.Since(realtimeStart).Milliseconds()), 0)
|
||
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 {
|
||
// 成功：绑定续期 + 健康上报由 Finish 完成
|
||
sess.Finish(context.WithoutCancel(ctx), true, float64(time.Since(realtimeStart).Milliseconds()))
|
||
} else {
|
||
// 会话中断：响应已开始，不可重试，仅上报健康
|
||
_, _ = sess.Report(context.WithoutCancel(ctx), 0, proxyErr, dispatch.DeliveryResponseStarted,
|
||
float64(time.Since(realtimeStart).Milliseconds()), 0)
|
||
}
|
||
|
||
// 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, info.StartTime)
|
||
if settleResult != nil && settleResult.ActualCost > 0 && !settleResult.DuplicateSkip {
|
||
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, info.StartTime)
|
||
if settleResult != nil && settleResult.ActualCost > 0 && !settleResult.DuplicateSkip {
|
||
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 ""
|
||
}
|