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:
CaIon
2026-02-10 20:40:33 +08:00
parent e0a6ee1cb8
commit 9e3954428d
34 changed files with 1465 additions and 1191 deletions

View File

@@ -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 {