mirror of
https://gh-proxy.org/https://github.com/klinsmaya/new-api.git
synced 2026-11-09 18:07:04 +08:00
CaIon
d8cb177440
Move account security settings into a dedicated page and add token status, rotation, revocation, and access history. Store audit events with role snapshots and JSON metadata, add audit.read authorization and an independent audit page, and upgrade the ClickHouse driver to v2.46.0.
905 lines
31 KiB
Go
905 lines
31 KiB
Go
package controller
|
||
|
||
import (
|
||
"errors"
|
||
"fmt"
|
||
"io"
|
||
"log"
|
||
"net/http"
|
||
"strings"
|
||
"time"
|
||
|
||
"github.com/QuantumNous/new-api/common"
|
||
"github.com/QuantumNous/new-api/constant"
|
||
taskdto "github.com/QuantumNous/new-api/dto"
|
||
"github.com/QuantumNous/new-api/logger"
|
||
"github.com/QuantumNous/new-api/middleware"
|
||
"github.com/QuantumNous/new-api/model"
|
||
pluginruntime "github.com/QuantumNous/new-api/pkg/jsplugin"
|
||
perfmetrics "github.com/QuantumNous/new-api/pkg/perf_metrics"
|
||
"github.com/QuantumNous/new-api/relay"
|
||
"github.com/QuantumNous/new-api/relay/channel/task/taskcommon"
|
||
relaycommon "github.com/QuantumNous/new-api/relay/common"
|
||
relayconstant "github.com/QuantumNous/new-api/relay/constant"
|
||
"github.com/QuantumNous/new-api/relay/helper"
|
||
"github.com/QuantumNous/new-api/relaykit/dto"
|
||
"github.com/QuantumNous/new-api/relaykit/types"
|
||
"github.com/QuantumNous/new-api/service"
|
||
"github.com/QuantumNous/new-api/setting"
|
||
"github.com/QuantumNous/new-api/setting/operation_setting"
|
||
|
||
"github.com/bytedance/gopkg/util/gopool"
|
||
"github.com/samber/lo"
|
||
|
||
"github.com/gin-gonic/gin"
|
||
"github.com/gorilla/websocket"
|
||
)
|
||
|
||
func relayHandler(c *gin.Context, info *relaycommon.RelayInfo) *types.NewAPIError {
|
||
var err *types.NewAPIError
|
||
switch info.RelayMode {
|
||
case relayconstant.RelayModeImagesGenerations, relayconstant.RelayModeImagesEdits:
|
||
err = relay.ImageHelper(c, info)
|
||
case relayconstant.RelayModeAudioSpeech:
|
||
fallthrough
|
||
case relayconstant.RelayModeAudioTranslation:
|
||
fallthrough
|
||
case relayconstant.RelayModeAudioTranscription:
|
||
err = relay.AudioHelper(c, info)
|
||
case relayconstant.RelayModeRerank:
|
||
err = relay.RerankHelper(c, info)
|
||
case relayconstant.RelayModeEmbeddings:
|
||
err = relay.EmbeddingHelper(c, info)
|
||
case relayconstant.RelayModeResponses, relayconstant.RelayModeResponsesCompact:
|
||
err = relay.ResponsesHelper(c, info)
|
||
case relayconstant.RelayModeAlphaSearch:
|
||
err = relay.AlphaSearchHelper(c, info)
|
||
default:
|
||
err = relay.TextHelper(c, info)
|
||
}
|
||
return err
|
||
}
|
||
|
||
func geminiRelayHandler(c *gin.Context, info *relaycommon.RelayInfo) *types.NewAPIError {
|
||
var err *types.NewAPIError
|
||
if strings.Contains(c.Request.URL.Path, "embed") {
|
||
err = relay.GeminiEmbeddingHandler(c, info)
|
||
} else {
|
||
err = relay.GeminiHelper(c, info)
|
||
}
|
||
return err
|
||
}
|
||
|
||
func Relay(c *gin.Context, relayFormat types.RelayFormat) {
|
||
|
||
requestId := c.GetString(common.RequestIdKey)
|
||
//group := common.GetContextKeyString(c, constant.ContextKeyUsingGroup)
|
||
//originalModel := common.GetContextKeyString(c, constant.ContextKeyOriginalModel)
|
||
|
||
var (
|
||
newAPIError *types.NewAPIError
|
||
ws *websocket.Conn
|
||
)
|
||
|
||
if relayFormat == types.RelayFormatOpenAIRealtime {
|
||
var err error
|
||
ws, err = upgrader.Upgrade(c.Writer, c.Request, nil)
|
||
if err != nil {
|
||
helper.WssError(c, ws, types.NewError(err, types.ErrorCodeGetChannelFailed, types.ErrOptionWithSkipRetry()).ToOpenAIError())
|
||
return
|
||
}
|
||
defer ws.Close()
|
||
}
|
||
|
||
defer func() {
|
||
if newAPIError != nil {
|
||
logger.LogError(c, fmt.Sprintf("relay error: %s", common.LocalLogPreview(newAPIError.Error())))
|
||
newAPIError.SetMessage(common.MessageWithRequestId(newAPIError.Error(), requestId))
|
||
switch relayFormat {
|
||
case types.RelayFormatOpenAIRealtime:
|
||
helper.WssError(c, ws, newAPIError.ToOpenAIError())
|
||
case types.RelayFormatClaude:
|
||
c.JSON(newAPIError.StatusCode, gin.H{
|
||
"type": "error",
|
||
"error": newAPIError.ToClaudeError(),
|
||
})
|
||
default:
|
||
c.JSON(newAPIError.StatusCode, gin.H{
|
||
"error": newAPIError.ToOpenAIError(),
|
||
})
|
||
}
|
||
}
|
||
}()
|
||
|
||
request, err := helper.GetAndValidateRequest(c, relayFormat)
|
||
if err != nil {
|
||
// Map "request body too large" to 413 so clients can handle it correctly
|
||
if common.IsRequestBodyTooLargeError(err) || errors.Is(err, common.ErrRequestBodyTooLarge) {
|
||
newAPIError = types.NewErrorWithStatusCode(err, types.ErrorCodeReadRequestBodyFailed, http.StatusRequestEntityTooLarge, types.ErrOptionWithSkipRetry())
|
||
} else {
|
||
newAPIError = types.NewError(err, types.ErrorCodeInvalidRequest, types.ErrOptionWithStatusCode(http.StatusBadRequest), types.ErrOptionWithSkipRetry())
|
||
}
|
||
return
|
||
}
|
||
|
||
relayInfo, err := relaycommon.GenRelayInfo(c, relayFormat, request, ws)
|
||
if err != nil {
|
||
newAPIError = types.NewError(err, types.ErrorCodeGenRelayInfoFailed)
|
||
return
|
||
}
|
||
|
||
needSensitiveCheck := setting.ShouldCheckPromptSensitive()
|
||
needCountToken := constant.CountToken
|
||
// Avoid building huge CombineText (strings.Join) when token counting and sensitive check are both disabled.
|
||
var meta *types.TokenCountMeta
|
||
if needSensitiveCheck || needCountToken {
|
||
meta = request.GetTokenCountMeta()
|
||
} else {
|
||
meta = fastTokenCountMetaForPricing(request)
|
||
}
|
||
|
||
if needSensitiveCheck && meta != nil {
|
||
contains, words := service.CheckSensitiveText(meta.CombineText)
|
||
if contains {
|
||
logger.LogWarn(c, fmt.Sprintf("user sensitive words detected: %s", strings.Join(words, ", ")))
|
||
newAPIError = types.NewError(err, types.ErrorCodeSensitiveWordsDetected)
|
||
return
|
||
}
|
||
}
|
||
|
||
tokens, err := service.EstimateRequestToken(c, meta, relayInfo)
|
||
if err != nil {
|
||
newAPIError = types.NewError(err, types.ErrorCodeCountTokenFailed)
|
||
return
|
||
}
|
||
|
||
relayInfo.SetEstimatePromptTokens(tokens)
|
||
|
||
priceData, err := helper.ModelPriceHelper(c, relayInfo, tokens, meta)
|
||
if err != nil {
|
||
newAPIError = types.NewError(err, types.ErrorCodeModelPriceError, types.ErrOptionWithStatusCode(http.StatusBadRequest))
|
||
return
|
||
}
|
||
|
||
// common.SetContextKey(c, constant.ContextKeyTokenCountMeta, meta)
|
||
|
||
if priceData.FreeModel {
|
||
logger.LogInfo(c, fmt.Sprintf("模型 %s 免费，跳过预扣费", relayInfo.OriginModelName))
|
||
} else {
|
||
newAPIError = service.PreConsumeBilling(c, priceData.QuotaToPreConsume, relayInfo)
|
||
if newAPIError != nil {
|
||
return
|
||
}
|
||
}
|
||
|
||
defer func() {
|
||
// Only return quota if downstream failed and quota was actually pre-consumed
|
||
if newAPIError != nil {
|
||
newAPIError = service.NormalizeViolationFeeError(newAPIError)
|
||
if relayInfo.Billing != nil {
|
||
relayInfo.Billing.Refund(c)
|
||
}
|
||
service.ChargeViolationFeeIfNeeded(c, relayInfo, newAPIError)
|
||
}
|
||
}()
|
||
|
||
retryParam := &service.RetryParam{
|
||
Ctx: c,
|
||
TokenGroup: relayInfo.TokenGroup,
|
||
ModelName: relayInfo.OriginModelName,
|
||
RequestPath: c.Request.URL.Path,
|
||
Retry: common.GetPointer(0),
|
||
}
|
||
relayInfo.RetryIndex = 0
|
||
relayInfo.LastError = nil
|
||
|
||
for ; retryParam.GetRetry() <= common.RetryTimes; retryParam.IncreaseRetry() {
|
||
relayInfo.RetryIndex = retryParam.GetRetry()
|
||
channel, channelErr := getChannel(c, relayInfo, retryParam)
|
||
if channelErr != nil {
|
||
logger.LogError(c, channelErr.Error())
|
||
newAPIError = channelErr
|
||
break
|
||
}
|
||
addUsedChannel(c, channel.Id)
|
||
if billingErr := service.PrepareTieredBillingForSelectedGroup(c, relayInfo); billingErr != nil {
|
||
newAPIError = billingErr
|
||
break
|
||
}
|
||
|
||
bodyStorage, bodyErr := common.GetBodyStorage(c)
|
||
if bodyErr != nil {
|
||
// Ensure consistent 413 for oversized bodies even when error occurs later (e.g., retry path)
|
||
if common.IsRequestBodyTooLargeError(bodyErr) || errors.Is(bodyErr, common.ErrRequestBodyTooLarge) {
|
||
newAPIError = types.NewErrorWithStatusCode(bodyErr, types.ErrorCodeReadRequestBodyFailed, http.StatusRequestEntityTooLarge, types.ErrOptionWithSkipRetry())
|
||
} else {
|
||
newAPIError = types.NewErrorWithStatusCode(bodyErr, types.ErrorCodeReadRequestBodyFailed, http.StatusBadRequest, types.ErrOptionWithSkipRetry())
|
||
}
|
||
break
|
||
}
|
||
c.Request.Body = io.NopCloser(bodyStorage)
|
||
|
||
switch relayFormat {
|
||
case types.RelayFormatOpenAIRealtime:
|
||
newAPIError = relay.WssHelper(c, relayInfo)
|
||
case types.RelayFormatClaude:
|
||
newAPIError = relay.ClaudeHelper(c, relayInfo)
|
||
case types.RelayFormatGemini:
|
||
newAPIError = geminiRelayHandler(c, relayInfo)
|
||
default:
|
||
newAPIError = relayHandler(c, relayInfo)
|
||
}
|
||
|
||
if newAPIError == nil {
|
||
relayInfo.LastError = nil
|
||
return
|
||
}
|
||
|
||
newAPIError = service.NormalizeViolationFeeError(newAPIError)
|
||
relayInfo.LastError = newAPIError
|
||
|
||
processChannelError(c, *types.NewChannelError(channel.Id, channel.Type, channel.Name, channel.ChannelInfo.IsMultiKey, common.GetContextKeyString(c, constant.ContextKeyChannelKey), channel.GetAutoBan()), newAPIError, relayInfo)
|
||
|
||
if !shouldRetry(c, newAPIError, common.RetryTimes-retryParam.GetRetry()) {
|
||
break
|
||
}
|
||
}
|
||
|
||
useChannel := c.GetStringSlice("use_channel")
|
||
if len(useChannel) > 1 {
|
||
retryLogStr := fmt.Sprintf("重试：%s", strings.Trim(strings.Join(strings.Fields(fmt.Sprint(useChannel)), "->"), "[]"))
|
||
logger.LogInfo(c, retryLogStr)
|
||
}
|
||
if newAPIError != nil {
|
||
gopool.Go(func() {
|
||
perfmetrics.RecordRelaySample(relayInfo, false, 0)
|
||
})
|
||
}
|
||
}
|
||
|
||
// CountClaudeTokens implements Anthropic's token-counting utility endpoint.
|
||
// It deliberately skips upstream generation and billing; callers use this
|
||
// endpoint to size prompts before creating a Message.
|
||
func CountClaudeTokens(c *gin.Context) {
|
||
request, err := helper.GetAndValidateClaudeRequest(c)
|
||
if err != nil {
|
||
c.JSON(http.StatusBadRequest, gin.H{
|
||
"type": "error",
|
||
"error": gin.H{
|
||
"type": "invalid_request_error",
|
||
"message": common.MessageWithRequestId(err.Error(), c.GetString(common.RequestIdKey)),
|
||
},
|
||
})
|
||
return
|
||
}
|
||
|
||
info := relaycommon.GenRelayInfoClaude(c, request)
|
||
inputTokens, err := service.CountRequestToken(c, request.GetTokenCountMeta(), info)
|
||
if err != nil {
|
||
c.JSON(http.StatusInternalServerError, gin.H{
|
||
"type": "error",
|
||
"error": gin.H{
|
||
"type": "api_error",
|
||
"message": common.MessageWithRequestId(err.Error(), c.GetString(common.RequestIdKey)),
|
||
},
|
||
})
|
||
return
|
||
}
|
||
|
||
c.JSON(http.StatusOK, gin.H{"input_tokens": inputTokens})
|
||
}
|
||
|
||
var upgrader = websocket.Upgrader{
|
||
Subprotocols: []string{"realtime"}, // WS 握手支持的协议，如果有使用 Sec-WebSocket-Protocol，则必须在此声明对应的 Protocol TODO add other protocol
|
||
CheckOrigin: func(r *http.Request) bool {
|
||
return true // 允许跨域
|
||
},
|
||
}
|
||
|
||
func addUsedChannel(c *gin.Context, channelId int) {
|
||
useChannel := c.GetStringSlice("use_channel")
|
||
useChannel = append(useChannel, fmt.Sprintf("%d", channelId))
|
||
c.Set("use_channel", useChannel)
|
||
}
|
||
|
||
func fastTokenCountMetaForPricing(request dto.Request) *types.TokenCountMeta {
|
||
if request == nil {
|
||
return &types.TokenCountMeta{}
|
||
}
|
||
meta := &types.TokenCountMeta{
|
||
TokenType: types.TokenTypeTokenizer,
|
||
}
|
||
switch r := request.(type) {
|
||
case *dto.GeneralOpenAIRequest:
|
||
maxCompletionTokens := lo.FromPtrOr(r.MaxCompletionTokens, uint(0))
|
||
maxTokens := lo.FromPtrOr(r.MaxTokens, uint(0))
|
||
if maxCompletionTokens > maxTokens {
|
||
meta.MaxTokens = int(maxCompletionTokens)
|
||
} else {
|
||
meta.MaxTokens = int(maxTokens)
|
||
}
|
||
case *dto.OpenAIResponsesRequest:
|
||
meta.MaxTokens = int(lo.FromPtrOr(r.MaxOutputTokens, uint(0)))
|
||
case *dto.ClaudeRequest:
|
||
meta.MaxTokens = int(lo.FromPtr(r.MaxTokens))
|
||
case *dto.ImageRequest:
|
||
// Pricing for image requests depends on ImagePriceRatio; safe to compute even when CountToken is disabled.
|
||
return r.GetTokenCountMeta()
|
||
default:
|
||
// Best-effort: leave CombineText empty to avoid large allocations.
|
||
}
|
||
return meta
|
||
}
|
||
|
||
func getChannel(c *gin.Context, info *relaycommon.RelayInfo, retryParam *service.RetryParam) (*model.Channel, *types.NewAPIError) {
|
||
if info.ChannelMeta == nil {
|
||
autoBan := c.GetBool("auto_ban")
|
||
autoBanInt := 1
|
||
if !autoBan {
|
||
autoBanInt = 0
|
||
}
|
||
return &model.Channel{
|
||
Id: c.GetInt("channel_id"),
|
||
Type: c.GetInt("channel_type"),
|
||
Name: c.GetString("channel_name"),
|
||
AutoBan: &autoBanInt,
|
||
}, nil
|
||
}
|
||
channel, selectGroup, err := service.CacheGetRandomSatisfiedChannel(retryParam)
|
||
if err != nil {
|
||
return nil, types.NewError(fmt.Errorf("获取分组 %s 下模型 %s 的可用渠道失败（retry）: %s", selectGroup, info.OriginModelName, err.Error()), types.ErrorCodeGetChannelFailed, types.ErrOptionWithSkipRetry())
|
||
}
|
||
if channel == nil {
|
||
return nil, types.NewError(fmt.Errorf("分组 %s 下模型 %s 的可用渠道不存在（retry）", selectGroup, info.OriginModelName), types.ErrorCodeGetChannelFailed, types.ErrOptionWithSkipRetry())
|
||
}
|
||
|
||
info.PriceData.GroupRatioInfo = helper.HandleGroupRatio(c, info)
|
||
|
||
newAPIError := middleware.SetupContextForSelectedChannel(c, channel, info.OriginModelName)
|
||
if newAPIError != nil {
|
||
return nil, newAPIError
|
||
}
|
||
return channel, nil
|
||
}
|
||
|
||
func shouldRetry(c *gin.Context, openaiErr *types.NewAPIError, retryTimes int) bool {
|
||
if openaiErr == nil {
|
||
return false
|
||
}
|
||
if service.ShouldSkipRetryAfterChannelAffinityFailure(c) {
|
||
return false
|
||
}
|
||
if types.IsChannelError(openaiErr) {
|
||
return true
|
||
}
|
||
if types.IsSkipRetryError(openaiErr) {
|
||
return false
|
||
}
|
||
if retryTimes <= 0 {
|
||
return false
|
||
}
|
||
if service.GetChannelConstraints(c).SuppressesRetry() {
|
||
return false
|
||
}
|
||
code := openaiErr.StatusCode
|
||
if code >= 200 && code < 300 {
|
||
return false
|
||
}
|
||
if code < 100 || code > 599 {
|
||
return true
|
||
}
|
||
if operation_setting.IsAlwaysSkipRetryCode(openaiErr.GetErrorCode()) {
|
||
return false
|
||
}
|
||
return operation_setting.ShouldRetryByStatusCode(code)
|
||
}
|
||
|
||
func processChannelError(c *gin.Context, channelError types.ChannelError, err *types.NewAPIError, relayInfo *relaycommon.RelayInfo) {
|
||
logger.LogError(c, fmt.Sprintf("channel error (channel #%d, status code: %d): %s", channelError.ChannelId, err.StatusCode, common.LocalLogPreview(err.Error())))
|
||
// 不要使用context获取渠道信息，异步处理时可能会出现渠道信息不一致的情况
|
||
// do not use context to get channel info, there may be inconsistent channel info when processing asynchronously
|
||
if service.ShouldDisableChannel(err) && channelError.AutoBan {
|
||
gopool.Go(func() {
|
||
service.DisableChannel(channelError, err.ErrorWithStatusCode())
|
||
})
|
||
}
|
||
|
||
if constant.ErrorLogEnabled && types.IsRecordErrorLog(err) {
|
||
// 保存错误日志到mysql中
|
||
userId := c.GetInt("id")
|
||
tokenName := c.GetString("token_name")
|
||
modelName := c.GetString("original_model")
|
||
tokenId := c.GetInt("token_id")
|
||
userGroup := c.GetString("group")
|
||
other := model.NewLogOther()
|
||
if c.Request != nil && c.Request.URL != nil {
|
||
other.SetPublic("request_path", c.Request.URL.Path)
|
||
}
|
||
other.SetPublic("error_type", err.GetErrorType())
|
||
other.SetPublic("error_code", err.GetErrorCode())
|
||
other.SetPublic("status_code", err.StatusCode)
|
||
service.AppendRelayLogAdminInfo(c, relayInfo, other)
|
||
service.AppendTaskPluginContextAuditInfo(c, other)
|
||
startTime := common.GetContextKeyTime(c, constant.ContextKeyRequestStartTime)
|
||
if startTime.IsZero() {
|
||
startTime = time.Now()
|
||
}
|
||
useTimeSeconds := int(time.Since(startTime).Seconds())
|
||
model.RecordErrorLog(c, userId, channelError.ChannelId, modelName, tokenName, err.MaskSensitiveErrorWithStatusCode(), tokenId, useTimeSeconds, common.GetContextKeyBool(c, constant.ContextKeyIsStream), userGroup, other)
|
||
}
|
||
|
||
}
|
||
|
||
func RelayMidjourney(c *gin.Context) {
|
||
relayInfo, err := relaycommon.GenRelayInfo(c, types.RelayFormatMjProxy, nil, nil)
|
||
|
||
if err != nil {
|
||
c.JSON(http.StatusInternalServerError, gin.H{
|
||
"description": fmt.Sprintf("failed to generate relay info: %s", err.Error()),
|
||
"type": "upstream_error",
|
||
"code": 4,
|
||
})
|
||
return
|
||
}
|
||
|
||
var mjErr *taskdto.MidjourneyResponse
|
||
switch relayInfo.RelayMode {
|
||
case relayconstant.RelayModeMidjourneyNotify:
|
||
mjErr = relay.RelayMidjourneyNotify(c)
|
||
case relayconstant.RelayModeMidjourneyTaskFetch, relayconstant.RelayModeMidjourneyTaskFetchByCondition:
|
||
mjErr = relay.RelayMidjourneyTask(c, relayInfo.RelayMode)
|
||
case relayconstant.RelayModeMidjourneyTaskImageSeed:
|
||
mjErr = relay.RelayMidjourneyTaskImageSeed(c)
|
||
case relayconstant.RelayModeSwapFace:
|
||
mjErr = relay.RelaySwapFace(c, relayInfo)
|
||
default:
|
||
mjErr = relay.RelayMidjourneySubmit(c, relayInfo)
|
||
}
|
||
//err = relayMidjourneySubmit(c, relayMode)
|
||
log.Println(mjErr)
|
||
if mjErr != nil {
|
||
statusCode := http.StatusBadRequest
|
||
if mjErr.Code == 30 {
|
||
mjErr.Result = "当前分组负载已饱和，请稍后再试，或升级账户以提升服务质量。"
|
||
statusCode = http.StatusTooManyRequests
|
||
}
|
||
c.JSON(statusCode, gin.H{
|
||
"description": fmt.Sprintf("%s %s", mjErr.Description, mjErr.Result),
|
||
"type": "upstream_error",
|
||
"code": mjErr.Code,
|
||
})
|
||
channelId := c.GetInt("channel_id")
|
||
logger.LogError(c, fmt.Sprintf("relay error (channel #%d, status code %d): %s", channelId, statusCode, fmt.Sprintf("%s %s", mjErr.Description, mjErr.Result)))
|
||
}
|
||
}
|
||
|
||
func RelayNotImplemented(c *gin.Context) {
|
||
err := types.OpenAIError{
|
||
Message: "API not implemented",
|
||
Type: "new_api_error",
|
||
Param: "",
|
||
Code: "api_not_implemented",
|
||
}
|
||
c.JSON(http.StatusNotImplemented, gin.H{
|
||
"error": err,
|
||
})
|
||
}
|
||
|
||
func RelayNotFound(c *gin.Context) {
|
||
// The web fallback may already have applied static-asset cache headers.
|
||
// A missing API or asset can appear after an upgrade; never cache its 404.
|
||
c.Header("Cache-Control", "no-store, no-cache, must-revalidate, private, max-age=0")
|
||
c.Header("Pragma", "no-cache")
|
||
c.Header("Expires", "0")
|
||
err := types.OpenAIError{
|
||
Message: fmt.Sprintf("Invalid URL (%s %s)", c.Request.Method, c.Request.URL.Path),
|
||
Type: "invalid_request_error",
|
||
Param: "",
|
||
Code: "",
|
||
}
|
||
c.JSON(http.StatusNotFound, gin.H{
|
||
"error": err,
|
||
})
|
||
}
|
||
|
||
// RelayTaskPluginEndpoint keeps unclaimed shared-endpoint traffic on its
|
||
// existing handler while claimed requests enter the generation-pinned
|
||
// host-owned protocol bridge.
|
||
func RelayTaskPluginEndpoint(c *gin.Context, fallback gin.HandlerFunc) {
|
||
pinnedValue, exists := c.Get(pluginruntime.ContextKeyPinnedEndpoint)
|
||
if !exists {
|
||
fallback(c)
|
||
return
|
||
}
|
||
pinned, ok := pinnedValue.(pluginruntime.PinnedEndpoint)
|
||
if !ok || pinned.Plugin == nil || pinned.Generation == nil {
|
||
c.JSON(http.StatusInternalServerError, gin.H{
|
||
"error": gin.H{
|
||
"message": "Task protocol request failed",
|
||
"type": "new_api_error",
|
||
"code": "task_protocol_error",
|
||
},
|
||
})
|
||
return
|
||
}
|
||
if pinned.Protocol != "openai_responses" {
|
||
fallback(c)
|
||
return
|
||
}
|
||
serveTaskPluginProtocol(c, pinned, defaultPluginProtocolBridgeDeps())
|
||
}
|
||
|
||
func RelayTaskFetch(c *gin.Context) {
|
||
relayInfo, err := relaycommon.GenRelayInfo(c, types.RelayFormatTask, nil, nil)
|
||
if err != nil {
|
||
c.JSON(http.StatusInternalServerError, &taskdto.TaskError{
|
||
Code: "gen_relay_info_failed",
|
||
Message: err.Error(),
|
||
StatusCode: http.StatusInternalServerError,
|
||
})
|
||
return
|
||
}
|
||
if taskErr := relay.RelayTaskFetch(c, relayInfo.RelayMode); taskErr != nil {
|
||
respondTaskError(c, taskErr)
|
||
}
|
||
}
|
||
|
||
type taskSubmissionOutcome struct {
|
||
Result *relay.TaskSubmitResult
|
||
Task *model.Task
|
||
RelayInfo *relaycommon.RelayInfo
|
||
}
|
||
|
||
func RelayTask(c *gin.Context) {
|
||
relayInfo, err := relaycommon.GenRelayInfo(c, types.RelayFormatTask, nil, nil)
|
||
if err != nil {
|
||
respondTaskSubmissionError(c, &taskdto.TaskError{
|
||
Code: "gen_relay_info_failed",
|
||
Message: err.Error(),
|
||
StatusCode: http.StatusInternalServerError,
|
||
})
|
||
return
|
||
}
|
||
if action := c.GetString("task_action"); action != "" {
|
||
relayInfo.Action = action
|
||
}
|
||
|
||
if taskErr := relay.ResolveOriginTask(c, relayInfo); taskErr != nil {
|
||
respondTaskSubmissionError(c, taskErr)
|
||
return
|
||
}
|
||
if taskErr := relay.ApplyOriginTaskAffinity(c, relayInfo); taskErr != nil {
|
||
respondTaskSubmissionError(c, taskErr)
|
||
return
|
||
}
|
||
|
||
outcome, taskErr := executeTaskSubmission(c, relayInfo)
|
||
if taskErr != nil {
|
||
respondTaskSubmissionError(c, taskErr)
|
||
return
|
||
}
|
||
presentTaskSubmission(c, outcome)
|
||
}
|
||
|
||
// executeTaskSubmission owns the retry, billing, and persistence lifecycle.
|
||
// It deliberately performs no client response writes so JSON and protocol
|
||
// presenters share the same durable task barrier. Its cancellation semantics
|
||
// come from c.Request.Context: native task endpoints use the client context,
|
||
// while the Responses bridge supplies an independently bounded context.
|
||
func executeTaskSubmission(c *gin.Context, relayInfo *relaycommon.RelayInfo) (*taskSubmissionOutcome, *taskdto.TaskError) {
|
||
return executeTaskSubmissionWith(c, relayInfo, relay.RelayTaskSubmit)
|
||
}
|
||
|
||
type taskSubmitAttempt func(*gin.Context, *relaycommon.RelayInfo) (*relay.TaskSubmitResult, *taskdto.TaskError)
|
||
|
||
func executeTaskSubmissionWith(
|
||
c *gin.Context,
|
||
relayInfo *relaycommon.RelayInfo,
|
||
submit taskSubmitAttempt,
|
||
) (*taskSubmissionOutcome, *taskdto.TaskError) {
|
||
diagnostics := newTaskPluginSubmitDiagnostics(c)
|
||
diagnostics.start(relayInfo)
|
||
var result *relay.TaskSubmitResult
|
||
var taskErr *taskdto.TaskError
|
||
durable := false
|
||
stage := "start"
|
||
defer func() {
|
||
if !durable && relayInfo.Billing != nil {
|
||
diagnostics.refund(stage)
|
||
relayInfo.Billing.Refund(c)
|
||
}
|
||
}()
|
||
stage = "before_attempt"
|
||
if requestErr := c.Request.Context().Err(); requestErr != nil {
|
||
diagnostics.cancelled("before_attempt", 0)
|
||
return nil, service.TaskErrorWrapperLocal(requestErr, "request_cancelled", http.StatusRequestTimeout)
|
||
}
|
||
|
||
retryParam := &service.RetryParam{
|
||
Ctx: c,
|
||
TokenGroup: relayInfo.TokenGroup,
|
||
ModelName: relayInfo.OriginModelName,
|
||
RequestPath: c.Request.URL.Path,
|
||
Retry: common.GetPointer(0),
|
||
}
|
||
|
||
for ; retryParam.GetRetry() <= common.RetryTimes; retryParam.IncreaseRetry() {
|
||
stage = "select_channel"
|
||
if requestErr := c.Request.Context().Err(); requestErr != nil {
|
||
diagnostics.cancelled("before_attempt", retryParam.GetRetry()+1)
|
||
taskErr = service.TaskErrorWrapperLocal(requestErr, "request_cancelled", http.StatusRequestTimeout)
|
||
break
|
||
}
|
||
var channel *model.Channel
|
||
|
||
if lockedCh, ok := relayInfo.LockedChannel.(*model.Channel); ok && lockedCh != nil {
|
||
channel = lockedCh
|
||
if retryParam.GetRetry() > 0 {
|
||
if setupErr := middleware.SetupContextForSelectedChannel(c, channel, relayInfo.OriginModelName); setupErr != nil {
|
||
taskErr = service.TaskErrorWrapperLocal(setupErr.Err, "setup_locked_channel_failed", http.StatusInternalServerError)
|
||
break
|
||
}
|
||
}
|
||
} else {
|
||
var channelErr *types.NewAPIError
|
||
channel, channelErr = getChannel(c, relayInfo, retryParam)
|
||
if channelErr != nil {
|
||
logger.LogError(c, channelErr.Error())
|
||
taskErr = service.TaskErrorWrapperLocal(channelErr.Err, "get_channel_failed", http.StatusInternalServerError)
|
||
break
|
||
}
|
||
}
|
||
diagnostics.attempt(retryParam.GetRetry()+1, channel, relayInfo.LockedChannel != nil)
|
||
|
||
addUsedChannel(c, channel.Id)
|
||
bodyStorage, bodyErr := common.GetBodyStorage(c)
|
||
if bodyErr != nil {
|
||
stage = "read_body"
|
||
if common.IsRequestBodyTooLargeError(bodyErr) || errors.Is(bodyErr, common.ErrRequestBodyTooLarge) {
|
||
taskErr = service.TaskErrorWrapperLocal(bodyErr, "read_request_body_failed", http.StatusRequestEntityTooLarge)
|
||
} else {
|
||
taskErr = service.TaskErrorWrapperLocal(bodyErr, "read_request_body_failed", http.StatusBadRequest)
|
||
}
|
||
break
|
||
}
|
||
c.Request.Body = io.NopCloser(bodyStorage)
|
||
|
||
stage = "submit"
|
||
result, taskErr = submit(c, relayInfo)
|
||
if requestErr := c.Request.Context().Err(); requestErr != nil {
|
||
diagnostics.cancelled("after_submit", retryParam.GetRetry()+1)
|
||
taskErr = service.TaskErrorWrapperLocal(requestErr, "request_cancelled", http.StatusRequestTimeout)
|
||
break
|
||
}
|
||
if taskErr == nil {
|
||
diagnostics.attemptSucceeded(retryParam.GetRetry()+1, result)
|
||
break
|
||
}
|
||
|
||
if !taskErr.LocalError {
|
||
processChannelError(c,
|
||
*types.NewChannelError(channel.Id, channel.Type, channel.Name, channel.ChannelInfo.IsMultiKey,
|
||
common.GetContextKeyString(c, constant.ContextKeyChannelKey), channel.GetAutoBan()),
|
||
types.NewOpenAIError(taskErr.Error, types.ErrorCodeBadResponseStatusCode, taskErr.StatusCode),
|
||
relayInfo)
|
||
}
|
||
|
||
willRetry := shouldRetryTaskRelay(c, channel.Id, taskErr, common.RetryTimes-retryParam.GetRetry())
|
||
diagnostics.attemptFailed(retryParam.GetRetry()+1, channel, taskErr, willRetry)
|
||
if !willRetry {
|
||
break
|
||
}
|
||
}
|
||
|
||
useChannel := c.GetStringSlice("use_channel")
|
||
if len(useChannel) > 1 {
|
||
retryLogStr := fmt.Sprintf("重试：%s", strings.Trim(strings.Join(strings.Fields(fmt.Sprint(useChannel)), "->"), "[]"))
|
||
logger.LogInfo(c, retryLogStr)
|
||
}
|
||
|
||
if taskErr != nil {
|
||
diagnostics.failed(stage, "task_error", taskErr, false)
|
||
return nil, taskErr
|
||
}
|
||
if result == nil {
|
||
taskErr = service.TaskErrorWrapperLocal(errors.New("task submission returned no result"), "task_submit_failed", http.StatusInternalServerError)
|
||
diagnostics.failed("submit", "missing_result", taskErr, false)
|
||
return nil, taskErr
|
||
}
|
||
if requestErr := c.Request.Context().Err(); requestErr != nil {
|
||
diagnostics.cancelled("before_reserve", retryParam.GetRetry()+1)
|
||
return nil, service.TaskErrorWrapperLocal(requestErr, "request_cancelled", http.StatusRequestTimeout)
|
||
}
|
||
|
||
// Reserve any submit-time upward billing adjustment before persistence.
|
||
// This keeps insertion failures fully refundable while ensuring settlement
|
||
// after the barrier normally has a zero positive delta.
|
||
if relayInfo.Billing != nil {
|
||
stage = "reserve"
|
||
diagnostics.reserve("reserve_start", result.Quota)
|
||
if reserveErr := relayInfo.Billing.Reserve(result.Quota); reserveErr != nil {
|
||
common.SysError("reserve adjusted task billing error: " + reserveErr.Error())
|
||
taskErr = service.TaskErrorWrapperLocal(errors.New("insufficient quota for adjusted task cost"), string(types.ErrorCodeInsufficientUserQuota), http.StatusForbidden)
|
||
diagnostics.failed("reserve", "insufficient_quota", taskErr, false)
|
||
return nil, taskErr
|
||
}
|
||
diagnostics.reserve("reserve_complete", result.Quota)
|
||
}
|
||
if requestErr := c.Request.Context().Err(); requestErr != nil {
|
||
diagnostics.cancelled("before_insert", retryParam.GetRetry()+1)
|
||
return nil, service.TaskErrorWrapperLocal(requestErr, "request_cancelled", http.StatusRequestTimeout)
|
||
}
|
||
|
||
stage = "insert"
|
||
task := model.InitTask(result.Platform, relayInfo)
|
||
task.PrivateData.Execution = service.TaskExecutionSnapshotFromContext(c)
|
||
task.PrivateData.UpstreamTaskID = result.UpstreamTaskID
|
||
task.PrivateData.BillingSource = relayInfo.BillingSource
|
||
task.PrivateData.SubscriptionId = relayInfo.SubscriptionId
|
||
task.PrivateData.TokenId = relayInfo.TokenId
|
||
task.PrivateData.NodeName = common.NodeName
|
||
task.PrivateData.BillingContext = &model.TaskBillingContext{
|
||
ModelPrice: relayInfo.PriceData.ModelPrice,
|
||
GroupRatio: relayInfo.PriceData.GroupRatioInfo.GroupRatio,
|
||
ModelRatio: relayInfo.PriceData.ModelRatio,
|
||
OtherRatios: relayInfo.PriceData.OtherRatios(),
|
||
OriginModelName: relayInfo.OriginModelName,
|
||
PerCallBilling: common.StringsContains(constant.TaskPricePatches, relayInfo.OriginModelName) || relayInfo.PriceData.UsePrice,
|
||
TieredSnapshot: relayInfo.TieredBillingSnapshot,
|
||
}
|
||
task.Quota = result.Quota
|
||
task.Data = result.TaskData
|
||
if len(result.PluginState) > 0 {
|
||
task.PrivateData.PluginState = result.PluginState
|
||
}
|
||
task.Action = relayInfo.Action
|
||
if immediate := result.Immediate; immediate != nil {
|
||
task.Status = model.TaskStatus(immediate.Status)
|
||
task.Progress = immediate.Progress
|
||
if immediate.Status == model.TaskStatusSuccess || immediate.Status == model.TaskStatusFailure {
|
||
task.FinishTime = time.Now().Unix()
|
||
}
|
||
if immediate.Status == model.TaskStatusFailure {
|
||
task.FailReason = immediate.Reason
|
||
}
|
||
if immediate.Url != "" {
|
||
task.PrivateData.ResultURL = immediate.Url
|
||
} else if immediate.Status == model.TaskStatusSuccess {
|
||
task.PrivateData.ResultURL = taskcommon.BuildProxyURL(task.TaskID)
|
||
}
|
||
}
|
||
diagnostics.insertStart(task)
|
||
if insertErr := task.InsertWithContext(c.Request.Context()); insertErr != nil {
|
||
common.SysError("insert task error: " + insertErr.Error())
|
||
taskErr = service.TaskErrorWrapperLocal(errors.New("failed to persist task"), "task_insert_failed", http.StatusInternalServerError)
|
||
diagnostics.failed("insert", "database_error", taskErr, false)
|
||
return nil, taskErr
|
||
}
|
||
durable = true
|
||
stage = "settle"
|
||
diagnostics.durable(task)
|
||
diagnostics.settleStart(task, result.Quota)
|
||
|
||
if settleErr := service.SettleBilling(c, relayInfo, result.Quota); settleErr != nil {
|
||
common.SysError("settle task billing error: " + settleErr.Error())
|
||
taskErr = service.TaskErrorWrapperLocal(errors.New("failed to settle task billing"), "task_billing_settlement_failed", http.StatusInternalServerError)
|
||
diagnostics.failed("settle", "billing_error", taskErr, true)
|
||
return nil, taskErr
|
||
}
|
||
service.LogTaskConsumption(c, relayInfo, task)
|
||
diagnostics.complete(task, result.Quota)
|
||
|
||
return &taskSubmissionOutcome{Result: result, Task: task, RelayInfo: relayInfo}, nil
|
||
}
|
||
|
||
func presentTaskSubmission(c *gin.Context, outcome *taskSubmissionOutcome) {
|
||
diagnostics := newTaskPluginSubmitDiagnostics(c)
|
||
otherRatios := outcome.RelayInfo.PriceData.OtherRatios()
|
||
if otherRatios == nil {
|
||
otherRatios = map[string]float64{}
|
||
}
|
||
if ratiosJSON, err := common.Marshal(otherRatios); err == nil {
|
||
c.Header("X-New-Api-Other-Ratios", string(ratiosJSON))
|
||
}
|
||
if pinnedValue, exists := c.Get(pluginruntime.ContextKeyPinnedRoute); exists {
|
||
if pinned, ok := pinnedValue.(pluginruntime.PinnedRoute); ok && pinned.Plugin != nil && pinned.Route.Render != "" {
|
||
view, err := service.BuildTaskPluginView(outcome.Task)
|
||
requestValue, _ := c.Get(pluginruntime.ContextKeyRouteRequest)
|
||
requestContext, _ := requestValue.(pluginruntime.RouteRequestContext)
|
||
if err == nil {
|
||
viewValue, valueErr := taskPluginProtocolJSONValue(view)
|
||
if valueErr == nil {
|
||
if body, callErr := pinned.Plugin.Engine.CallPath(c.Request.Context(), "native", []string{pinned.Route.Render}, requestContext.JSValue(), viewValue); callErr == nil {
|
||
diagnostics.present(outcome.Task, "native_presenter")
|
||
c.JSON(http.StatusOK, body)
|
||
return
|
||
} else {
|
||
logger.LogError(c, "task plugin native submit presenter failed: "+callErr.Error())
|
||
}
|
||
} else {
|
||
logger.LogError(c, "encode task plugin native submit view failed: "+valueErr.Error())
|
||
}
|
||
} else {
|
||
logger.LogError(c, "build task plugin native submit view failed: "+err.Error())
|
||
}
|
||
}
|
||
}
|
||
if pinnedValue, exists := c.Get(pluginruntime.ContextKeyPinnedEndpoint); exists {
|
||
if pinned, ok := pinnedValue.(pluginruntime.PinnedEndpoint); ok && pinned.Protocol == "openai_video" && pinned.Operation.Name == "create" {
|
||
diagnostics.present(outcome.Task, "openai_video_create")
|
||
c.JSON(http.StatusOK, outcome.Task.ToOpenAIVideo())
|
||
return
|
||
}
|
||
}
|
||
createdAt := outcome.Task.CreatedAt
|
||
if createdAt == 0 {
|
||
createdAt = outcome.Task.SubmitTime
|
||
}
|
||
diagnostics.present(outcome.Task, "host_fallback")
|
||
c.JSON(http.StatusOK, map[string]any{
|
||
"id": outcome.Task.TaskID,
|
||
"task_id": outcome.Task.TaskID,
|
||
"status": "queued",
|
||
"model": outcome.RelayInfo.OriginModelName,
|
||
"created_at": createdAt,
|
||
})
|
||
}
|
||
|
||
func respondTaskSubmissionError(c *gin.Context, taskErr *taskdto.TaskError) {
|
||
newTaskPluginSubmitDiagnostics(c).presentError(taskErr)
|
||
if middleware.RespondTaskPluginError(c, taskErr) {
|
||
return
|
||
}
|
||
respondTaskError(c, taskErr)
|
||
}
|
||
|
||
// respondTaskError 统一输出 Task 错误响应（含 429 限流提示改写）
|
||
func respondTaskError(c *gin.Context, taskErr *taskdto.TaskError) {
|
||
if taskErr.StatusCode == http.StatusTooManyRequests {
|
||
taskErr.Message = "当前分组上游负载已饱和，请稍后再试"
|
||
}
|
||
c.JSON(taskErr.StatusCode, taskErr)
|
||
}
|
||
|
||
func shouldRetryTaskRelay(c *gin.Context, channelId int, taskErr *taskdto.TaskError, retryTimes int) bool {
|
||
if taskErr == nil {
|
||
return false
|
||
}
|
||
if service.ShouldSkipRetryAfterChannelAffinityFailure(c) {
|
||
return false
|
||
}
|
||
if retryTimes <= 0 {
|
||
return false
|
||
}
|
||
if service.GetChannelConstraints(c).SuppressesRetry() {
|
||
return false
|
||
}
|
||
if taskErr.StatusCode == http.StatusTooManyRequests {
|
||
return true
|
||
}
|
||
if taskErr.StatusCode == 307 {
|
||
return true
|
||
}
|
||
if taskErr.StatusCode/100 == 5 {
|
||
// 超时不重试
|
||
if operation_setting.IsAlwaysSkipRetryStatusCode(taskErr.StatusCode) {
|
||
return false
|
||
}
|
||
return true
|
||
}
|
||
if taskErr.StatusCode == http.StatusBadRequest {
|
||
return false
|
||
}
|
||
if taskErr.StatusCode == 408 {
|
||
// azure处理超时不重试
|
||
return false
|
||
}
|
||
if taskErr.LocalError {
|
||
return false
|
||
}
|
||
if taskErr.StatusCode/100 == 2 {
|
||
return false
|
||
}
|
||
return true
|
||
}
|