package controller import ( "context" "sync" "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(detectionHandler{}) service.RegisterSystemTaskHandler(astraSamplingHandler{}) service.RegisterSystemTaskHandler(channelAvailabilityHandler{}) 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{} type channelAvailabilityHandler struct{} func (channelAvailabilityHandler) Type() string { return model.SystemTaskTypeChannelAvailability } func (channelAvailabilityHandler) Enabled() bool { return service.AstraAvailabilityEnabled() } func (channelAvailabilityHandler) Interval() time.Duration { return time.Duration(service.AstraAvailabilityIntervalMinutes()) * time.Minute } func (channelAvailabilityHandler) ScheduleFromCreation() bool { return true } func (channelAvailabilityHandler) NewPayload() any { return nil } func (channelAvailabilityHandler) Run(ctx context.Context, task *model.SystemTask, runnerID string) { testUserID, err := resolveChannelTestUserID(nil) if err != nil { finishSystemTaskHandler(task, runnerID, model.SystemTaskStatusFailed, nil, err) return } channels, err := model.GetAllChannels(0, 0, true, false) if err != nil { finishSystemTaskHandler(task, runnerID, model.SystemTaskStatusFailed, nil, err) return } selected := make([]*model.Channel, 0, len(channels)) for _, channel := range channels { if channel.Status == common.ChannelStatusEnabled { selected = append(selected, channel) } } settings := operation_setting.GetMonitorSetting() state := service.ChannelAvailabilityState{Rounds: []service.ChannelAvailabilityRound{}} var mu sync.Mutex var persistErr error ctx, cancel := context.WithCancel(ctx) defer cancel() disableThreshold := int64(common.ChannelDisableThreshold * 1000) if disableThreshold == 0 { disableThreshold = 10000000 } summary := runChannelTestWorkers(ctx, selected, settings.ChannelTestConcurrency, func(ctx context.Context, channel *model.Channel) channelTestSummary { result := testChannelForHealthCheck(ctx, channel, testUserID, false, disableThreshold) if result.Tested == 0 { return result } mu.Lock() defer mu.Unlock() if persistErr == nil { state.Rounds = append(state.Rounds, service.ChannelAvailabilityRound{ChannelID: channel.Id, At: time.Now().Unix(), Success: result.Succeeded == 1}) persistErr = model.UpdateSystemTaskState(task.TaskID, runnerID, state) if persistErr != nil { cancel() } } return result }, nil) if persistErr != nil { finishSystemTaskHandler(task, runnerID, model.SystemTaskStatusFailed, summary, persistErr) return } if ctx.Err() != nil { finishSystemTaskHandler(task, runnerID, model.SystemTaskStatusFailed, summary, ctx.Err()) return } finishSystemTaskHandler(task, runnerID, model.SystemTaskStatusSucceeded, summary, nil) } 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(common.LogText("system task %s failed to persist result: %v", task.TaskID, err)) } }