package controller import ( "context" "fmt" "time" "github.com/QuantumNous/new-api/common" "github.com/QuantumNous/new-api/constant" "github.com/QuantumNous/new-api/model" "github.com/QuantumNous/new-api/service" "github.com/QuantumNous/new-api/setting/operation_setting" ) // RegisterScheduledSystemTasks wires the periodic channel test, upstream model // update, and async task polling (Midjourney / Suno / video) jobs into the // system task framework so a DB lease dedups execution across multiple master // instances and each run is recorded as one task row. Call this before // service.StartSystemTaskRunner. func RegisterScheduledSystemTasks() { service.RegisterSystemTaskHandler(channelTestHandler{}) service.RegisterSystemTaskHandler(modelUpdateHandler{}) service.RegisterSystemTaskHandler(midjourneyPollHandler{}) service.RegisterSystemTaskHandler(asyncTaskPollHandler{}) } // channelTestHandler runs the scheduled "test all channels" job. Enablement and // cadence still come from the monitor settings; only the execution path moved // into the system task runner. type channelTestHandler struct{} func (channelTestHandler) Type() string { return model.SystemTaskTypeChannelTest } func (channelTestHandler) Enabled() bool { return operation_setting.GetMonitorSetting().AutoTestChannelEnabled } func (channelTestHandler) Interval() time.Duration { minutes := operation_setting.GetMonitorSetting().AutoTestChannelMinutes if minutes <= 0 { minutes = 10 } return time.Duration(minutes * float64(time.Minute)) } func (channelTestHandler) NewPayload() any { return nil } // channelTestTaskPayload controls one channel_test run. A nil/empty payload is a // scheduled run, which uses the configured monitor ChannelTestMode and does not // notify. A manual "test all channels" trigger sets Mode=scheduled_all and // Notify=true to reproduce the legacy manual behavior (test every channel and // notify root on completion). type channelTestTaskPayload struct { Mode string `json:"mode,omitempty"` Notify bool `json:"notify,omitempty"` } func (channelTestHandler) Run(ctx context.Context, task *model.SystemTask, runnerID string) { payload := channelTestTaskPayload{} if err := task.DecodePayload(&payload); err != nil { finishSystemTaskHandler(task, runnerID, model.SystemTaskStatusFailed, nil, err) return } summary, err := runChannelTestTask(ctx, payload.Mode, payload.Notify, service.NewSystemTaskProgressReporter(task, runnerID)) if err != nil { finishSystemTaskHandler(task, runnerID, model.SystemTaskStatusFailed, nil, err) return } finishSystemTaskHandler(task, runnerID, model.SystemTaskStatusSucceeded, summary, nil) } // modelUpdateHandler runs the scheduled upstream model update detection job. type modelUpdateHandler struct{} func (modelUpdateHandler) Type() string { return model.SystemTaskTypeModelUpdate } func (modelUpdateHandler) Enabled() bool { return common.GetEnvOrDefaultBool("CHANNEL_UPSTREAM_MODEL_UPDATE_TASK_ENABLED", true) } func (modelUpdateHandler) Interval() time.Duration { intervalMinutes := common.GetEnvOrDefault( "CHANNEL_UPSTREAM_MODEL_UPDATE_TASK_INTERVAL_MINUTES", channelUpstreamModelUpdateTaskDefaultIntervalMinutes, ) if intervalMinutes < 1 { intervalMinutes = channelUpstreamModelUpdateTaskDefaultIntervalMinutes } return time.Duration(intervalMinutes) * time.Minute } func (modelUpdateHandler) NewPayload() any { return nil } // modelUpdateTaskPayload controls one model_update run. A scheduled run // (Manual=false) respects the per-channel minimum check interval and may // auto-apply detected models when a channel has auto-sync enabled. A manual // "detect all" trigger sets Manual=true to reproduce the legacy detect-all // semantics: force a re-check regardless of the interval and never auto-apply, // so the admin reviews and applies changes explicitly. type modelUpdateTaskPayload struct { Manual bool `json:"manual,omitempty"` } func (modelUpdateHandler) Run(ctx context.Context, task *model.SystemTask, runnerID string) { payload := modelUpdateTaskPayload{} if err := task.DecodePayload(&payload); err != nil { finishSystemTaskHandler(task, runnerID, model.SystemTaskStatusFailed, nil, err) return } summary := runChannelUpstreamModelUpdateTaskOnce(ctx, payload.Manual, !payload.Manual, service.NewSystemTaskProgressReporter(task, runnerID)) finishSystemTaskHandler(task, runnerID, model.SystemTaskStatusSucceeded, summary, nil) } // midjourneyPollHandler runs one Midjourney polling pass per scheduled run. // Enabled() folds the "are there unfinished tasks?" check into enablement so the // scheduler creates no row when the system is idle; only when at least one // Midjourney task is in progress does a row get scheduled. type midjourneyPollHandler struct{} func (midjourneyPollHandler) Type() string { return model.SystemTaskTypeMidjourneyPoll } func (midjourneyPollHandler) Enabled() bool { return constant.UpdateTask && model.HasUnfinishedMidjourneyTasks() } func (midjourneyPollHandler) Interval() time.Duration { return 15 * time.Second } func (midjourneyPollHandler) NewPayload() any { return nil } func (midjourneyPollHandler) Run(ctx context.Context, task *model.SystemTask, runnerID string) { summary := runMidjourneyTaskUpdateOnce(ctx, service.NewSystemTaskProgressReporter(task, runnerID)) finishSystemTaskHandler(task, runnerID, model.SystemTaskStatusSucceeded, summary, nil) } // asyncTaskPollHandler runs one async-task (Suno/video) polling pass per // scheduled run. Like midjourneyPollHandler, Enabled() folds in the unfinished // task existence check so an idle system schedules no rows. type asyncTaskPollHandler struct{} func (asyncTaskPollHandler) Type() string { return model.SystemTaskTypeAsyncTaskPoll } func (asyncTaskPollHandler) Enabled() bool { return constant.UpdateTask && model.HasUnfinishedSyncTasks() } func (asyncTaskPollHandler) Interval() time.Duration { return 15 * time.Second } func (asyncTaskPollHandler) NewPayload() any { return nil } func (asyncTaskPollHandler) Run(ctx context.Context, task *model.SystemTask, runnerID string) { summary := service.RunTaskPollingOnce(ctx, service.NewSystemTaskProgressReporter(task, runnerID)) finishSystemTaskHandler(task, runnerID, model.SystemTaskStatusSucceeded, summary, nil) } func finishSystemTaskHandler(task *model.SystemTask, runnerID string, status model.SystemTaskStatus, result any, runErr error) { errorMessage := "" if runErr != nil { errorMessage = runErr.Error() } if err := model.FinishSystemTask(task.TaskID, runnerID, status, result, errorMessage); err != nil { common.SysLog(fmt.Sprintf("system task %s failed to persist result: %v", task.TaskID, err)) } }