mirror of
https://github.com/QuantumNous/new-api.git
synced 2026-03-30 04:22:58 +00:00
refactor(task): extract billing and polling logic from controller to service layer
Restructure the task relay system for better separation of concerns: - Extract task billing into service/task_billing.go with unified settlement flow - Move task polling loop from controller to service/task_polling.go (supports Suno + video platforms) - Split RelayTask into fetch/submit paths with dedicated retry logic (taskSubmitWithRetry) - Add TaskDto, TaskResponse generics, and FetchReq to dto/task.go - Add taskcommon/helpers.go for shared task adaptor utilities - Remove controller/task_video.go (logic consolidated into service layer) - Update all task adaptors (ali, doubao, gemini, hailuo, jimeng, kling, sora, suno, vertex, vidu) - Simplify frontend task logs to use new TaskDto response format
This commit is contained in:
@@ -451,17 +451,102 @@ func RelayNotFound(c *gin.Context) {
|
||||
}
|
||||
|
||||
func RelayTask(c *gin.Context) {
|
||||
retryTimes := common.RetryTimes
|
||||
channelId := c.GetInt("channel_id")
|
||||
c.Set("use_channel", []string{fmt.Sprintf("%d", channelId)})
|
||||
relayInfo, err := relaycommon.GenRelayInfo(c, types.RelayFormatTask, nil, nil)
|
||||
if err != nil {
|
||||
c.JSON(http.StatusInternalServerError, &dto.TaskError{
|
||||
Code: "gen_relay_info_failed",
|
||||
Message: err.Error(),
|
||||
StatusCode: http.StatusInternalServerError,
|
||||
})
|
||||
return
|
||||
}
|
||||
taskErr := taskRelayHandler(c, relayInfo)
|
||||
if taskErr == nil {
|
||||
retryTimes = 0
|
||||
|
||||
// Fetch 操作是纯 DB 查询(或 task 自带 channelId 的上游查询),不依赖上下文 channel,无需重试
|
||||
// TODO: 在video-route层面优化,避免无谓的 channel 选择和上下文设置,也没必要吧代码放到这里来写这么多屎山
|
||||
switch relayInfo.RelayMode {
|
||||
case relayconstant.RelayModeSunoFetch, relayconstant.RelayModeSunoFetchByID, relayconstant.RelayModeVideoFetchByID:
|
||||
if taskErr := relay.RelayTaskFetch(c, relayInfo.RelayMode); taskErr != nil {
|
||||
respondTaskError(c, taskErr)
|
||||
}
|
||||
return
|
||||
}
|
||||
|
||||
// ── Submit 路径 ─────────────────────────────────────────────────
|
||||
|
||||
// 1. 解析原始任务(remix / continuation),一次性,可能锁定渠道并禁止重试
|
||||
if taskErr := relay.ResolveOriginTask(c, relayInfo); taskErr != nil {
|
||||
respondTaskError(c, taskErr)
|
||||
return
|
||||
}
|
||||
|
||||
// 2. defer Refund(全部失败时回滚预扣费)
|
||||
var result *relay.TaskSubmitResult
|
||||
var taskErr *dto.TaskError
|
||||
defer func() {
|
||||
if taskErr != nil && relayInfo.Billing != nil {
|
||||
relayInfo.Billing.Refund(c)
|
||||
}
|
||||
}()
|
||||
|
||||
// 3. 执行 + 重试(RelayTaskSubmit 内部在首次调用时自动预扣费)
|
||||
taskErr = taskSubmitWithRetry(c, relayInfo, channelId, common.RetryTimes, func() *dto.TaskError {
|
||||
var te *dto.TaskError
|
||||
result, te = relay.RelayTaskSubmit(c, relayInfo)
|
||||
return te
|
||||
})
|
||||
|
||||
// 4. 成功:结算 + 日志 + 插入任务
|
||||
if taskErr == nil {
|
||||
if settleErr := service.SettleBilling(c, relayInfo, result.Quota); settleErr != nil {
|
||||
common.SysError("settle task billing error: " + settleErr.Error())
|
||||
}
|
||||
service.LogTaskConsumption(c, relayInfo, result.ModelName)
|
||||
|
||||
task := model.InitTask(result.Platform, relayInfo)
|
||||
task.PrivateData.UpstreamTaskID = result.UpstreamTaskID
|
||||
task.PrivateData.BillingSource = relayInfo.BillingSource
|
||||
task.PrivateData.SubscriptionId = relayInfo.SubscriptionId
|
||||
task.PrivateData.TokenId = relayInfo.TokenId
|
||||
task.Quota = result.Quota
|
||||
task.Data = result.TaskData
|
||||
task.Action = relayInfo.Action
|
||||
if insertErr := task.Insert(); insertErr != nil {
|
||||
//taskErr = service.TaskErrorWrapper(insertErr, "insert_task_failed", http.StatusInternalServerError)
|
||||
common.SysError("insert task error: " + insertErr.Error())
|
||||
}
|
||||
}
|
||||
|
||||
if taskErr != nil {
|
||||
respondTaskError(c, taskErr)
|
||||
}
|
||||
}
|
||||
|
||||
// respondTaskError 统一输出 Task 错误响应(含 429 限流提示改写)
|
||||
func respondTaskError(c *gin.Context, taskErr *dto.TaskError) {
|
||||
if taskErr.StatusCode == http.StatusTooManyRequests {
|
||||
taskErr.Message = "当前分组上游负载已饱和,请稍后再试"
|
||||
}
|
||||
c.JSON(taskErr.StatusCode, taskErr)
|
||||
}
|
||||
|
||||
// taskSubmitWithRetry 执行首次尝试并在失败时切换渠道重试,返回最终的 taskErr。
|
||||
// attempt 闭包负责实际的上游请求,不涉及计费。
|
||||
func taskSubmitWithRetry(c *gin.Context, relayInfo *relaycommon.RelayInfo,
|
||||
channelId int, retryTimes int, attempt func() *dto.TaskError) *dto.TaskError {
|
||||
|
||||
taskErr := attempt()
|
||||
if taskErr == nil {
|
||||
return nil
|
||||
}
|
||||
if !taskErr.LocalError {
|
||||
processChannelError(c,
|
||||
*types.NewChannelError(channelId, c.GetInt("channel_type"), c.GetString("channel_name"), common.GetContextKeyBool(c, constant.ContextKeyChannelIsMultiKey),
|
||||
common.GetContextKeyString(c, constant.ContextKeyChannelKey), common.GetContextKeyBool(c, constant.ContextKeyChannelAutoBan)),
|
||||
types.NewOpenAIError(taskErr.Error, types.ErrorCodeBadResponseStatusCode, taskErr.StatusCode))
|
||||
}
|
||||
|
||||
retryParam := &service.RetryParam{
|
||||
Ctx: c,
|
||||
TokenGroup: relayInfo.TokenGroup,
|
||||
@@ -480,7 +565,7 @@ func RelayTask(c *gin.Context) {
|
||||
useChannel = append(useChannel, fmt.Sprintf("%d", channelId))
|
||||
c.Set("use_channel", useChannel)
|
||||
logger.LogInfo(c, fmt.Sprintf("using channel #%d to retry (remain times %d)", channel.Id, retryParam.GetRetry()))
|
||||
//middleware.SetupContextForSelectedChannel(c, channel, originalModel)
|
||||
middleware.SetupContextForSelectedChannel(c, channel, c.GetString("original_model"))
|
||||
|
||||
bodyStorage, err := common.GetBodyStorage(c)
|
||||
if err != nil {
|
||||
@@ -492,30 +577,21 @@ func RelayTask(c *gin.Context) {
|
||||
break
|
||||
}
|
||||
c.Request.Body = io.NopCloser(bodyStorage)
|
||||
taskErr = taskRelayHandler(c, relayInfo)
|
||||
taskErr = attempt()
|
||||
if taskErr != nil && !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))
|
||||
}
|
||||
}
|
||||
|
||||
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 {
|
||||
if taskErr.StatusCode == http.StatusTooManyRequests {
|
||||
taskErr.Message = "当前分组上游负载已饱和,请稍后再试"
|
||||
}
|
||||
c.JSON(taskErr.StatusCode, taskErr)
|
||||
}
|
||||
}
|
||||
|
||||
func taskRelayHandler(c *gin.Context, relayInfo *relaycommon.RelayInfo) *dto.TaskError {
|
||||
var err *dto.TaskError
|
||||
switch relayInfo.RelayMode {
|
||||
case relayconstant.RelayModeSunoFetch, relayconstant.RelayModeSunoFetchByID, relayconstant.RelayModeVideoFetchByID:
|
||||
err = relay.RelayTaskFetch(c, relayInfo.RelayMode)
|
||||
default:
|
||||
err = relay.RelayTaskSubmit(c, relayInfo)
|
||||
}
|
||||
return err
|
||||
return taskErr
|
||||
}
|
||||
|
||||
func shouldRetryTaskRelay(c *gin.Context, channelId int, taskErr *dto.TaskError, retryTimes int) bool {
|
||||
|
||||
Reference in New Issue
Block a user