package controller import ( "bytes" "context" "fmt" "io" "net/http" "time" "github.com/QuantumNous/new-api/common" "github.com/QuantumNous/new-api/dto" "github.com/QuantumNous/new-api/logger" "github.com/QuantumNous/new-api/model" "github.com/QuantumNous/new-api/service" "github.com/QuantumNous/new-api/setting" "github.com/QuantumNous/new-api/setting/system_setting" "github.com/gin-gonic/gin" ) // midjourneyPollSummary is the result recorded on a midjourney_poll system task // row, summarizing one polling pass. type midjourneyPollSummary struct { UnfinishedTasks int `json:"unfinished_tasks"` ChannelsScanned int `json:"channels_scanned"` NullTasksFailed int `json:"null_tasks_failed"` } // runMidjourneyTaskUpdateOnce performs one Midjourney 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 (processedChannels, // totalChannels) so the system task surfaces a percentage. func runMidjourneyTaskUpdateOnce(ctx context.Context, report func(processed, total int)) midjourneyPollSummary { summary := midjourneyPollSummary{} if ctx == nil { ctx = context.Background() } tasks := model.GetAllUnFinishTasks() if len(tasks) == 0 { return summary } summary.UnfinishedTasks = len(tasks) logger.LogInfo(ctx, fmt.Sprintf("检测到未完成的任务数有: %v", len(tasks))) taskChannelM := make(map[int][]string) taskM := make(map[string]*model.Midjourney) nullTaskIds := make([]int, 0) for _, task := range tasks { if task.MjId == "" { // 统计失败的未完成任务 nullTaskIds = append(nullTaskIds, task.Id) continue } taskM[task.MjId] = task taskChannelM[task.ChannelId] = append(taskChannelM[task.ChannelId], task.MjId) } if len(nullTaskIds) > 0 { summary.NullTasksFailed = len(nullTaskIds) err := model.MjBulkUpdateByTaskIds(nullTaskIds, map[string]any{ "status": "FAILURE", "progress": "100%", }) if err != nil { logger.LogError(ctx, fmt.Sprintf("Fix null mj_id task error: %v", err)) } else { logger.LogInfo(ctx, fmt.Sprintf("Fix null mj_id task success: %v", nullTaskIds)) } } if len(taskChannelM) == 0 { return summary } totalChannels := len(taskChannelM) processedChannels := 0 for channelId, taskIds := range taskChannelM { if ctx != nil && ctx.Err() != nil { break } if report != nil { report(processedChannels, totalChannels) } processedChannels++ summary.ChannelsScanned++ logger.LogInfo(ctx, fmt.Sprintf("渠道 #%d 未完成的任务有: %d", channelId, len(taskIds))) if len(taskIds) == 0 { continue } midjourneyChannel, err := model.CacheGetChannel(channelId) if err != nil { logger.LogError(ctx, fmt.Sprintf("CacheGetChannel: %v", err)) err := model.MjBulkUpdate(taskIds, map[string]any{ "fail_reason": fmt.Sprintf("获取渠道信息失败,请联系管理员,渠道ID:%d", channelId), "status": "FAILURE", "progress": "100%", }) if err != nil { logger.LogInfo(ctx, fmt.Sprintf("UpdateMidjourneyTask error: %v", err)) } continue } requestUrl := fmt.Sprintf("%s/mj/task/list-by-condition", *midjourneyChannel.BaseURL) body, err := common.Marshal(map[string]any{ "ids": taskIds, }) if err != nil { logger.LogError(ctx, fmt.Sprintf("Get Task marshal body error: %v", err)) continue } timeout := time.Second * 15 requestCtx, cancel := context.WithTimeout(ctx, timeout) req, err := http.NewRequestWithContext(requestCtx, "POST", requestUrl, bytes.NewBuffer(body)) if err != nil { cancel() logger.LogError(ctx, fmt.Sprintf("Get Task error: %v", err)) continue } req.Header.Set("Content-Type", "application/json") req.Header.Set("mj-api-secret", midjourneyChannel.Key) resp, err := service.GetHttpClient().Do(req) if err != nil { logger.LogError(ctx, fmt.Sprintf("Get Task Do req error: %v", err)) cancel() continue } if resp.StatusCode != http.StatusOK { logger.LogError(ctx, fmt.Sprintf("Get Task status code: %d", resp.StatusCode)) resp.Body.Close() cancel() continue } responseBody, err := io.ReadAll(resp.Body) if err != nil { logger.LogError(ctx, fmt.Sprintf("Get Mjp Task parse body error: %v", err)) resp.Body.Close() cancel() continue } var responseItems []dto.MidjourneyDto err = common.Unmarshal(responseBody, &responseItems) if err != nil { logger.LogError(ctx, fmt.Sprintf("Get Mjp Task parse body error2: %v, body: %s", err, string(responseBody))) resp.Body.Close() cancel() continue } resp.Body.Close() req.Body.Close() cancel() for _, responseItem := range responseItems { task := taskM[responseItem.MjId] if task == nil { logger.LogWarn(ctx, fmt.Sprintf("Midjourney task response ignored: unknown mj_id=%s", responseItem.MjId)) continue } useTime := (time.Now().UnixNano() / int64(time.Millisecond)) - task.SubmitTime // 如果时间超过一小时,且进度不是100%,则认为任务失败 if useTime > 3600000 && task.Progress != "100%" { responseItem.FailReason = "上游任务超时(超过1小时)" responseItem.Status = "FAILURE" } if !checkMjTaskNeedUpdate(task, responseItem) { continue } preStatus := task.Status task.Code = 1 task.Progress = responseItem.Progress task.PromptEn = responseItem.PromptEn task.State = responseItem.State task.SubmitTime = responseItem.SubmitTime task.StartTime = responseItem.StartTime task.FinishTime = responseItem.FinishTime task.ImageUrl = responseItem.ImageUrl task.Status = responseItem.Status task.FailReason = responseItem.FailReason if responseItem.Properties != nil { propertiesStr, _ := common.Marshal(responseItem.Properties) task.Properties = string(propertiesStr) } if responseItem.Buttons != nil { buttonStr, _ := common.Marshal(responseItem.Buttons) task.Buttons = string(buttonStr) } // 映射 VideoUrl task.VideoUrl = responseItem.VideoUrl // 映射 VideoUrls - 将数组序列化为 JSON 字符串 if responseItem.VideoUrls != nil && len(responseItem.VideoUrls) > 0 { videoUrlsStr, err := common.Marshal(responseItem.VideoUrls) if err != nil { logger.LogError(ctx, fmt.Sprintf("序列化 VideoUrls 失败: %v", err)) task.VideoUrls = "[]" // 失败时设置为空数组 } else { task.VideoUrls = string(videoUrlsStr) } } else { task.VideoUrls = "" // 空值时清空字段 } shouldReturnQuota := false if (task.Progress != "100%" && responseItem.FailReason != "") || (task.Progress == "100%" && task.Status == "FAILURE") { logger.LogInfo(ctx, task.MjId+" 构建失败,"+task.FailReason) task.Progress = "100%" if task.Quota != 0 { shouldReturnQuota = true } } won, err := task.UpdateWithStatus(preStatus) if err != nil { logger.LogError(ctx, "UpdateMidjourneyTask task error: "+err.Error()) } else if won && shouldReturnQuota { err = model.IncreaseUserQuota(task.UserId, task.Quota, false) if err != nil { logger.LogError(ctx, "fail to increase user quota: "+err.Error()) } model.RecordTaskBillingLog(model.RecordTaskBillingLogParams{ UserId: task.UserId, LogType: model.LogTypeRefund, Content: "", ChannelId: task.ChannelId, ModelName: service.CovertMjpActionToModelName(task.Action), Quota: task.Quota, Other: map[string]interface{}{ "task_id": task.MjId, "reason": "构图失败", }, }) } } } if report != nil && (ctx == nil || ctx.Err() == nil) { report(totalChannels, totalChannels) } return summary } func checkMjTaskNeedUpdate(oldTask *model.Midjourney, newTask dto.MidjourneyDto) bool { if oldTask.Code != 1 { return true } if oldTask.Progress != newTask.Progress { return true } if oldTask.PromptEn != newTask.PromptEn { return true } if oldTask.State != newTask.State { return true } if oldTask.SubmitTime != newTask.SubmitTime { return true } if oldTask.StartTime != newTask.StartTime { return true } if oldTask.FinishTime != newTask.FinishTime { return true } if oldTask.ImageUrl != newTask.ImageUrl { return true } if oldTask.Status != newTask.Status { return true } if oldTask.FailReason != newTask.FailReason { return true } if oldTask.FinishTime != newTask.FinishTime { return true } if oldTask.Progress != "100%" && newTask.FailReason != "" { return true } // 检查 VideoUrl 是否需要更新 if oldTask.VideoUrl != newTask.VideoUrl { return true } // 检查 VideoUrls 是否需要更新 if newTask.VideoUrls != nil && len(newTask.VideoUrls) > 0 { newVideoUrlsStr, _ := common.Marshal(newTask.VideoUrls) if oldTask.VideoUrls != string(newVideoUrlsStr) { return true } } else if oldTask.VideoUrls != "" { // 如果新数据没有 VideoUrls 但旧数据有,需要更新(清空) return true } return false } func GetAllMidjourney(c *gin.Context) { pageInfo := common.GetPageQuery(c) // 解析其他查询参数 queryParams := model.TaskQueryParams{ ChannelID: c.Query("channel_id"), MjID: c.Query("mj_id"), StartTimestamp: c.Query("start_timestamp"), EndTimestamp: c.Query("end_timestamp"), } items := model.GetAllTasks(pageInfo.GetStartIdx(), pageInfo.GetPageSize(), queryParams) total := model.CountAllTasks(queryParams) if setting.MjForwardUrlEnabled { for i, midjourney := range items { midjourney.ImageUrl = system_setting.ServerAddress + "/mj/image/" + midjourney.MjId items[i] = midjourney } } pageInfo.SetTotal(int(total)) pageInfo.SetItems(items) common.ApiSuccess(c, pageInfo) } func GetUserMidjourney(c *gin.Context) { pageInfo := common.GetPageQuery(c) userId := c.GetInt("id") queryParams := model.TaskQueryParams{ MjID: c.Query("mj_id"), StartTimestamp: c.Query("start_timestamp"), EndTimestamp: c.Query("end_timestamp"), } items := model.GetAllUserTask(userId, pageInfo.GetStartIdx(), pageInfo.GetPageSize(), queryParams) total := model.CountAllUserTask(userId, queryParams) if setting.MjForwardUrlEnabled { for i, midjourney := range items { midjourney.ImageUrl = system_setting.ServerAddress + "/mj/image/" + midjourney.MjId items[i] = midjourney } } pageInfo.SetTotal(int(total)) pageInfo.SetItems(items) common.ApiSuccess(c, pageInfo) }