package service import ( "context" "fmt" "io" "net/http" "sort" "strings" "sync" "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/model" "github.com/QuantumNous/new-api/pkg/billingexpr" "github.com/QuantumNous/new-api/relay/channel/task/taskcommon" relaycommon "github.com/QuantumNous/new-api/relay/common" "github.com/QuantumNous/new-api/relaykit/dto" "github.com/bytedance/gopkg/util/gopool" "github.com/samber/lo" ) // TaskPollingAdaptor 定义轮询所需的最小适配器接口,避免 service -> relay 的循环依赖 type TaskPollingAdaptor interface { Init(info *relaycommon.RelayInfo) FetchTask(baseURL string, key string, body map[string]any, proxy string) (*http.Response, error) ParseTaskResult(body []byte) (*relaycommon.TaskInfo, error) // AdjustBillingOnComplete 在任务到达终态(成功/失败)时由轮询循环调用。 // 返回正数触发差额结算(补扣/退还),返回 0 保持预扣费金额不变。 AdjustBillingOnComplete(task *model.Task, taskResult *relaycommon.TaskInfo) int } type BatchTaskPollingAdaptor interface { TaskPollingAdaptor FetchMode() string FetchBatchTasks(baseURL, key string, taskIDs []string, proxy string) (*http.Response, error) ParseBatchResult(body []byte) (map[string]*BatchTaskResult, error) } type BatchTaskResult struct { TaskInfo relaycommon.TaskInfo Action string SubmitTime int64 StartTime int64 FinishTime int64 Data any } // GetTaskAdaptorFunc 由 main 包注入,用于获取指定平台的任务适配器。 // 打破 service -> relay -> relay/channel -> service 的循环依赖。 var GetTaskAdaptorFunc func(platform constant.TaskPlatform) TaskPollingAdaptor // sweepTimedOutTasks 在主轮询之前独立清理超时任务。 // 每次最多处理 100 条,剩余的下个周期继续处理。 // 使用 per-task CAS (UpdateWithStatus) 防止覆盖被正常轮询已推进的任务。 func sweepTimedOutTasks(ctx context.Context) { if constant.TaskTimeoutMinutes <= 0 { return } cutoff := time.Now().Unix() - int64(constant.TaskTimeoutMinutes)*60 tasks := model.GetTimedOutUnfinishedTasks(cutoff, 100) if len(tasks) == 0 { return } reason := fmt.Sprintf("任务超时(%d分钟)", constant.TaskTimeoutMinutes) legacyReason := "任务超时(旧系统遗留任务,不进行退款,请联系管理员)" now := time.Now().Unix() timedOutCount := 0 for _, task := range tasks { isLegacy := task.SubmitTime > 0 && task.SubmitTime < model.TaskRefundLegacyCutoff oldStatus := task.Status task.Status = model.TaskStatusFailure task.Progress = "100%" task.FinishTime = now if isLegacy { task.FailReason = legacyReason // 旧系统任务明确不退款,随终态 CAS 一并清掉 quota, // 避免留下可再次退款的计费状态。 task.Quota = 0 } else { task.FailReason = reason } won, err := task.UpdateWithStatus(oldStatus) if err != nil { logger.LogError(ctx, fmt.Sprintf("sweepTimedOutTasks CAS update error for task %s: %v", task.TaskID, err)) continue } if !won { logger.LogInfo(ctx, fmt.Sprintf("sweepTimedOutTasks: task %s already transitioned, skip", task.TaskID)) continue } timedOutCount++ if !isLegacy && task.Quota != 0 { RefundTaskQuota(ctx, task, reason) } } if timedOutCount > 0 { logger.LogInfo(ctx, fmt.Sprintf("sweepTimedOutTasks: timed out %d tasks", timedOutCount)) } } // TaskPollSummary is the result recorded on an async_task_poll system task row, // summarizing one polling pass. type TaskPollSummary struct { UnfinishedTasks int `json:"unfinished_tasks"` PlatformsScanned int `json:"platforms_scanned"` NullTasksFailed int `json:"null_tasks_failed"` } // RunTaskPollingOnce performs one async-task (Suno/video) polling pass // synchronously. It honors ctx cancellation (the system-task runner cancels it // when the lease is lost) and, when report is non-nil, reports progress as // (processedPlatforms, totalPlatforms). It returns immediately if the task // adaptor factory has not been wired yet, to avoid a nil call during startup. func RunTaskPollingOnce(ctx context.Context, report func(processed, total int)) TaskPollSummary { summary := TaskPollSummary{} if GetTaskAdaptorFunc == nil { return summary } if ctx == nil { ctx = context.Background() } common.SysLog("任务进度轮询开始") sweepTimedOutTasks(ctx) allTasks := model.GetAllUnFinishSyncTasks(constant.TaskQueryLimit) summary.UnfinishedTasks = len(allTasks) platformTask := make(map[constant.TaskPlatform][]*model.Task) for _, t := range allTasks { platformTask[t.Platform] = append(platformTask[t.Platform], t) } totalPlatforms := len(platformTask) processedPlatforms := 0 for platform, tasks := range platformTask { if ctx.Err() != nil { break } if report != nil { report(processedPlatforms, totalPlatforms) } processedPlatforms++ if len(tasks) == 0 { continue } summary.PlatformsScanned++ taskChannelM := make(map[int][]string) taskM := make(map[string]*model.Task) nullTaskIds := make([]int64, 0) for _, task := range tasks { upstreamID := task.GetUpstreamTaskID() if upstreamID == "" { // 统计失败的未完成任务 nullTaskIds = append(nullTaskIds, task.ID) continue } taskM[upstreamID] = task taskChannelM[task.ChannelId] = append(taskChannelM[task.ChannelId], upstreamID) } if len(nullTaskIds) > 0 { summary.NullTasksFailed += len(nullTaskIds) err := model.TaskBulkUpdateByID(nullTaskIds, map[string]any{ "status": "FAILURE", "progress": "100%", }) if err != nil { logger.LogError(ctx, fmt.Sprintf("Fix null task_id task error: %v", err)) } else { logger.LogInfo(ctx, fmt.Sprintf("Fix null task_id task success: %v", nullTaskIds)) } } if len(taskChannelM) == 0 { continue } DispatchPlatformUpdate(ctx, platform, taskChannelM, taskM) } if report != nil && ctx.Err() == nil { report(totalPlatforms, totalPlatforms) } common.SysLog("任务进度轮询完成") return summary } // DispatchPlatformUpdate 按平台分发轮询更新 func DispatchPlatformUpdate(ctx context.Context, platform constant.TaskPlatform, taskChannelM map[int][]string, taskM map[string]*model.Task) { if ctx == nil { ctx = context.Background() } if platform == constant.TaskPlatformMidjourney { // MJ 轮询由其自身处理,这里预留入口 return } adaptor := GetTaskAdaptorFunc(platform) if batchAdaptor, ok := adaptor.(BatchTaskPollingAdaptor); ok && batchAdaptor.FetchMode() == "batch" { if err := UpdateBatchTasks(ctx, batchAdaptor, taskChannelM, taskM); err != nil { common.SysLog(fmt.Sprintf("UpdateBatchTasks fail: %s", err)) } return } if err := UpdateVideoTasks(ctx, platform, taskChannelM, taskM); err != nil { common.SysLog(fmt.Sprintf("UpdateVideoTasks fail: %s", err)) } } func UpdateBatchTasks(ctx context.Context, adaptor BatchTaskPollingAdaptor, taskChannelM map[int][]string, taskM map[string]*model.Task) error { for channelId, taskIds := range taskChannelM { if ctx.Err() != nil { return ctx.Err() } err := updateBatchTasks(ctx, adaptor, channelId, taskIds, taskM) if err != nil { logger.LogError(ctx, fmt.Sprintf("渠道 #%d 更新异步任务失败: %s", channelId, err.Error())) } } return nil } func updateBatchTasks(ctx context.Context, adaptor BatchTaskPollingAdaptor, channelId int, taskIds []string, taskM map[string]*model.Task) error { logger.LogInfo(ctx, fmt.Sprintf("渠道 #%d 未完成的任务有: %d", channelId, len(taskIds))) if ctx.Err() != nil { return ctx.Err() } if len(taskIds) == 0 { return nil } ch, err := model.CacheGetChannel(channelId) if err != nil { common.SysLog(fmt.Sprintf("CacheGetChannel: %v", err)) // Collect DB primary key IDs for bulk update (taskIds are upstream IDs, not task_id column values) var failedIDs []int64 for _, upstreamID := range taskIds { if t, ok := taskM[upstreamID]; ok { failedIDs = append(failedIDs, t.ID) } } err = model.TaskBulkUpdateByID(failedIDs, map[string]any{ "fail_reason": fmt.Sprintf("获取渠道信息失败,请联系管理员,渠道ID:%d", channelId), "status": "FAILURE", "progress": "100%", }) if err != nil { common.SysLog(fmt.Sprintf("UpdateSunoTask error: %v", err)) } return err } proxy := ch.GetSetting().Proxy baseURL := ch.GetBaseURL() if baseURL == "" { baseURL = constant.GetChannelBaseURL(ch.Type) } resp, err := adaptor.FetchBatchTasks(baseURL, ch.Key, taskIds, proxy) if err != nil { common.SysLog(fmt.Sprintf("Get Task Do req error: %v", err)) return err } if resp.StatusCode != http.StatusOK { logger.LogError(ctx, fmt.Sprintf("Get Task status code: %d", resp.StatusCode)) return fmt.Errorf("Get Task status code: %d", resp.StatusCode) } defer resp.Body.Close() responseBody, err := io.ReadAll(resp.Body) if err != nil { common.SysLog(fmt.Sprintf("Get Suno Task parse body error: %v", err)) return err } responseItems, err := adaptor.ParseBatchResult(responseBody) if err != nil { return fmt.Errorf("parse batch result: %w", err) } for upstreamID, responseItem := range responseItems { if ctx.Err() != nil { return ctx.Err() } task := taskM[upstreamID] if task == nil { logger.LogWarn(ctx, fmt.Sprintf("Batch task response ignored: unknown task_id=%s", upstreamID)) continue } snap := task.Snapshot() task.Status = lo.If(model.TaskStatus(responseItem.TaskInfo.Status) != "", model.TaskStatus(responseItem.TaskInfo.Status)).Else(task.Status) task.FailReason = lo.If(responseItem.TaskInfo.Reason != "", responseItem.TaskInfo.Reason).Else(task.FailReason) task.SubmitTime = lo.If(responseItem.SubmitTime != 0, responseItem.SubmitTime).Else(task.SubmitTime) task.StartTime = lo.If(responseItem.StartTime != 0, responseItem.StartTime).Else(task.StartTime) task.FinishTime = lo.If(responseItem.FinishTime != 0, responseItem.FinishTime).Else(task.FinishTime) if responseItem.TaskInfo.Progress != "" { task.Progress = responseItem.TaskInfo.Progress } if responseItem.TaskInfo.Reason != "" || task.Status == model.TaskStatusFailure { logger.LogInfo(ctx, task.TaskID+" 构建失败,"+task.FailReason) task.Status = model.TaskStatusFailure task.Progress = "100%" } if responseItem.TaskInfo.Status == model.TaskStatusSuccess { task.Progress = "100%" } if responseItem.Data != nil { task.SetData(responseItem.Data) } else if task.Status == model.TaskStatusSuccess || task.Status == model.TaskStatusFailure { logger.LogWarn(ctx, fmt.Sprintf( "Batch task %s reached terminal status without data; preserving existing task data", task.TaskID, )) } if responseItem.TaskInfo.Url != "" { task.PrivateData.ResultURL = responseItem.TaskInfo.Url } isDone := task.Status == model.TaskStatusSuccess || task.Status == model.TaskStatusFailure terminalTransition := isDone && snap.Status != task.Status won, updateErr := task.UpdateWithStatus(snap.Status) if updateErr != nil { common.SysLog("UpdateSunoTask task error: " + updateErr.Error()) continue } if !won { logger.LogWarn(ctx, fmt.Sprintf("Batch task %s already transitioned by another process, skip billing", task.TaskID)) continue } if terminalTransition { billingSettled := settleTaskBillingOnComplete(ctx, adaptor, task, &responseItem.TaskInfo) if task.Status == model.TaskStatusFailure && !billingSettled && task.Quota != 0 { RefundTaskQuota(ctx, task, task.FailReason) } } } return nil } // UpdateVideoTasks 按渠道更新所有视频任务 func UpdateVideoTasks(ctx context.Context, platform constant.TaskPlatform, taskChannelM map[int][]string, taskM map[string]*model.Task) error { channelIDs := make([]int, 0, len(taskChannelM)) for channelID := range taskChannelM { channelIDs = append(channelIDs, channelID) } sort.Ints(channelIDs) var wg sync.WaitGroup for _, channelId := range channelIDs { taskIds := taskChannelM[channelId] if len(taskIds) == 0 { continue } taskIds = append([]string(nil), taskIds...) wg.Add(1) gopool.Go(func() { defer wg.Done() if err := updateVideoTasks(ctx, platform, channelId, taskIds, taskM); err != nil { logger.LogError(ctx, fmt.Sprintf("Channel #%d failed to update video async tasks: %s", channelId, err.Error())) } }) } wg.Wait() if ctx.Err() != nil { return ctx.Err() } return nil } func updateVideoTasks(ctx context.Context, platform constant.TaskPlatform, channelId int, taskIds []string, taskM map[string]*model.Task) error { logger.LogInfo(ctx, fmt.Sprintf("Channel #%d pending video tasks: %d", channelId, len(taskIds))) if ctx.Err() != nil { return ctx.Err() } if len(taskIds) == 0 { return nil } cacheGetChannel, err := model.CacheGetChannel(channelId) if err != nil { // Collect DB primary key IDs for bulk update (taskIds are upstream IDs, not task_id column values) var failedIDs []int64 for _, upstreamID := range taskIds { if t, ok := taskM[upstreamID]; ok { failedIDs = append(failedIDs, t.ID) } } errUpdate := model.TaskBulkUpdateByID(failedIDs, map[string]any{ "fail_reason": fmt.Sprintf("Failed to get channel info, channel ID: %d", channelId), "status": "FAILURE", "progress": "100%", }) if errUpdate != nil { common.SysLog(fmt.Sprintf("UpdateVideoTask error: %v", errUpdate)) } return fmt.Errorf("CacheGetChannel failed: %w", err) } adaptor := GetTaskAdaptorFunc(platform) if adaptor == nil { return fmt.Errorf("video adaptor not found") } info := &relaycommon.RelayInfo{} info.ChannelMeta = &relaycommon.ChannelMeta{ ChannelBaseUrl: cacheGetChannel.GetBaseURL(), } info.ApiKey = cacheGetChannel.Key adaptor.Init(info) disablePollingSleep := cacheGetChannel.GetOtherSettings().DisableTaskPollingSleep for i, taskId := range taskIds { if ctx.Err() != nil { return ctx.Err() } if err := updateVideoSingleTask(ctx, adaptor, cacheGetChannel, taskId, taskM); err != nil { logger.LogError(ctx, fmt.Sprintf("Failed to update video task %s: %s", taskId, err.Error())) } if disablePollingSleep || i == len(taskIds)-1 { continue } // sleep 1 second between tasks for this channel only. select { case <-ctx.Done(): return ctx.Err() case <-time.After(1 * time.Second): } } return nil } func updateVideoSingleTask(ctx context.Context, adaptor TaskPollingAdaptor, ch *model.Channel, taskId string, taskM map[string]*model.Task) error { if ctx.Err() != nil { return ctx.Err() } baseURL := constant.GetChannelBaseURL(ch.Type) if ch.GetBaseURL() != "" { baseURL = ch.GetBaseURL() } proxy := ch.GetSetting().Proxy task := taskM[taskId] if task == nil { logger.LogError(ctx, fmt.Sprintf("Task %s not found in taskM", taskId)) return fmt.Errorf("task %s not found", taskId) } key := ch.Key privateData := task.PrivateData if privateData.Key != "" { key = privateData.Key } resp, err := adaptor.FetchTask(baseURL, key, map[string]any{ "task_id": task.GetUpstreamTaskID(), "action": constant.NormalizeTaskAction(task.Action), }, proxy) if err != nil { return fmt.Errorf("fetchTask failed for task %s: %w", taskId, err) } defer resp.Body.Close() responseBody, err := io.ReadAll(resp.Body) if err != nil { return fmt.Errorf("readAll failed for task %s: %w", taskId, err) } logger.LogDebug(ctx, "updateVideoSingleTask response: %s", responseBody) snap := task.Snapshot() taskResult := &relaycommon.TaskInfo{} // try parse as New API response format var responseItems taskdto.TaskResponse[model.Task] if err = common.Unmarshal(responseBody, &responseItems); err == nil && responseItems.IsSuccess() { logger.LogDebug(ctx, "updateVideoSingleTask parsed as new api response format: %+v", responseItems) t := responseItems.Data taskResult.TaskID = t.TaskID taskResult.Status = string(t.Status) taskResult.Url = t.GetResultURL() taskResult.Progress = t.Progress taskResult.Reason = t.FailReason task.Data = t.Data } else if taskResult, err = adaptor.ParseTaskResult(responseBody); err != nil { return fmt.Errorf("parseTaskResult failed for task %s: %w", taskId, err) } task.Data = redactVideoResponseBody(responseBody) logger.LogDebug(ctx, "updateVideoSingleTask taskResult: %+v", taskResult) now := time.Now().Unix() if taskResult.Status == "" { //taskResult = relaycommon.FailTaskInfo("upstream returned empty status") errorResult := &dto.GeneralErrorResponse{} if err = common.Unmarshal(responseBody, &errorResult); err == nil { openaiError := errorResult.TryToOpenAIError() if openaiError != nil { // 返回规范的 OpenAI 错误格式,提取错误信息,判断错误是否为任务失败 if openaiError.Code == "429" { // 429 错误通常表示请求过多或速率限制,暂时不认为是任务失败,保持原状态等待下一轮轮询 return nil } // 其他错误认为是任务失败,记录错误信息并更新任务状态 taskResult = relaycommon.FailTaskInfo("upstream returned error") } else { // unknown error format, log original response logger.LogError(ctx, fmt.Sprintf("Task %s returned empty status with unrecognized error format, response: %s", taskId, string(responseBody))) taskResult = relaycommon.FailTaskInfo("upstream returned unrecognized message") } } } shouldFinalizeBilling := false task.Status = model.TaskStatus(taskResult.Status) switch taskResult.Status { case model.TaskStatusSubmitted: task.Progress = taskcommon.ProgressSubmitted case model.TaskStatusQueued: task.Progress = taskcommon.ProgressQueued case model.TaskStatusInProgress: task.Progress = taskcommon.ProgressInProgress if task.StartTime == 0 { task.StartTime = now } case model.TaskStatusSuccess: task.Progress = taskcommon.ProgressComplete if task.FinishTime == 0 { task.FinishTime = now } if strings.HasPrefix(taskResult.Url, "data:") { // data: URI (e.g. Vertex base64 encoded video) — keep in Data, not in ResultURL task.PrivateData.ResultURL = taskcommon.BuildProxyURL(task.TaskID) } else if taskResult.Url != "" { // Direct upstream URL (e.g. Kling, Ali, Doubao, etc.) task.PrivateData.ResultURL = taskResult.Url } else { // No URL from adaptor — construct proxy URL using public task ID task.PrivateData.ResultURL = taskcommon.BuildProxyURL(task.TaskID) } shouldFinalizeBilling = true case model.TaskStatusFailure: logger.LogJson(ctx, fmt.Sprintf("Task %s failed", taskId), task) task.Status = model.TaskStatusFailure task.Progress = taskcommon.ProgressComplete if task.FinishTime == 0 { task.FinishTime = now } task.FailReason = taskResult.Reason logger.LogInfo(ctx, fmt.Sprintf("Task %s failed: %s", task.TaskID, task.FailReason)) taskResult.Progress = taskcommon.ProgressComplete shouldFinalizeBilling = true default: return fmt.Errorf("unknown task status %s for task %s", taskResult.Status, task.TaskID) } if taskResult.Progress != "" { task.Progress = taskResult.Progress } isDone := task.Status == model.TaskStatusSuccess || task.Status == model.TaskStatusFailure if isDone && snap.Status != task.Status { won, err := task.UpdateWithStatus(snap.Status) if err != nil { logger.LogError(ctx, fmt.Sprintf("UpdateWithStatus failed for task %s: %s", task.TaskID, err.Error())) shouldFinalizeBilling = false } else if !won { logger.LogWarn(ctx, fmt.Sprintf("Task %s CAS lost or no-op update, skip billing", task.TaskID)) shouldFinalizeBilling = false } } else if !snap.Equal(task.Snapshot()) { if _, err := task.UpdateWithStatus(snap.Status); err != nil { logger.LogError(ctx, fmt.Sprintf("Failed to update task %s: %s", task.TaskID, err.Error())) } } else { // No changes, skip update logger.LogDebug(ctx, "No update needed for task %s", task.TaskID) } if shouldFinalizeBilling { billingSettled := settleTaskBillingOnComplete(ctx, adaptor, task, taskResult) if task.Status == model.TaskStatusFailure && !billingSettled && task.Quota != 0 { RefundTaskQuota(ctx, task, task.FailReason) } } return nil } func redactVideoResponseBody(body []byte) []byte { var m map[string]any if err := common.Unmarshal(body, &m); err != nil { return body } resp, _ := m["response"].(map[string]any) if resp != nil { delete(resp, "bytesBase64Encoded") if v, ok := resp["video"].(string); ok { resp["video"] = truncateBase64(v) } if vs, ok := resp["videos"].([]any); ok { for i := range vs { if vm, ok := vs[i].(map[string]any); ok { delete(vm, "bytesBase64Encoded") } } } } b, err := common.Marshal(m) if err != nil { return body } return b } func truncateBase64(s string) string { const maxKeep = 256 if len(s) <= maxKeep { return s } return s[:maxKeep] + "..." } // settleTaskBillingOnComplete 任务完成时的统一计费调整。 // 返回 true 表示用量结算路径已接管最终计费;失败任务仅在返回 false 时补做全额退款。 // 优先级:1. tiered snapshot → 2. adaptor 调整 → 3. token 重算。 // // 表达式求值失败会保留预扣额度,因此也视为已接管,避免错误全退。 func settleTaskBillingOnComplete(ctx context.Context, adaptor TaskPollingAdaptor, task *model.Task, taskResult *relaycommon.TaskInfo) bool { if bc := task.PrivateData.BillingContext; bc != nil && bc.TieredSnapshot != nil { // 用量表达式结算只适用于成功任务;失败任务由调用方全额退款。 if task.Status == model.TaskStatusFailure { return false } usageFacts := make(map[string]any, len(bc.TieredSnapshot.UsageFacts)+len(taskResult.UsageFacts)) for key, value := range bc.TieredSnapshot.UsageFacts { usageFacts[key] = value } for key, value := range taskResult.UsageFacts { usageFacts[key] = value } result, err := billingexpr.ComputeTieredQuotaWithRequest(bc.TieredSnapshot, billingexpr.TokenParams{}, billingexpr.RequestInput{Usage: usageFacts}) if err != nil { logger.LogWarn(ctx, fmt.Sprintf("任务 %s 表达式结算失败,保留预扣额度: %v", task.TaskID, err)) return true } if result.Clamp != nil { logger.LogWarn(ctx, fmt.Sprintf("任务 %s 表达式结算额度发生饱和: %+v", task.TaskID, result.Clamp)) } bc.TieredSnapshot.UsageFacts = usageFacts bc.TieredSnapshot.EstimatedTier = result.MatchedTier RecalculateTaskQuota(ctx, task, result.ActualQuotaAfterGroup, "任务用量表达式结算", result.Clamp) return true } // 按次计费的成功任务保持预扣;失败任务由调用方全额退款。 if bc := task.PrivateData.BillingContext; bc != nil && bc.PerCallBilling { logger.LogInfo(ctx, fmt.Sprintf("任务 %s 按次计费,跳过差额结算", task.TaskID)) return false } // 优先让 adaptor 决定最终额度。 if actualQuota := adaptor.AdjustBillingOnComplete(task, taskResult); actualQuota > 0 { RecalculateTaskQuota(ctx, task, actualQuota, "adaptor计费调整") return true } // 回退到 token 重算。 tokens := taskResult.TotalTokens if tokens == 0 && taskResult.CompletionTokens > 0 { tokens = taskResult.CompletionTokens } if tokens > 0 { return RecalculateTaskQuotaByTokens(ctx, task, tokens) } return false }