new-api-custom/service/detection_schedule.go

package service

import (
	"context"
	"sort"
	"time"

	"github.com/QuantumNous/new-api/common"
	"github.com/QuantumNous/new-api/model"
)

type automaticDetectionResult struct {
	channelID int
	id        string
	at        int64
	outcome   string
}

func FingerprintScheduleOutcome(status string) string {
	switch status {
	case "match":
		return "normal"
	case "mismatch":
		return "degraded"
	default:
		return "unknown"
	}
}

// ReconcileAutomaticDetectionSchedules seeds pre-upgrade history once and
// recovers results persisted before a worker lost its lease or process.
func ReconcileAutomaticDetectionSchedules(ctx context.Context, channels []model.Channel, kind string) (map[int]model.DetectionSchedule, map[int]bool, error) {
	states := map[int]model.DetectionSchedule{}
	active := map[int]bool{}
	if len(channels) == 0 {
		return states, active, nil
	}
	ids := make([]int, 0, len(channels))
	for _, c := range channels {
		ids = append(ids, c.Id)
	}
	var rows []model.DetectionSchedule
	if err := model.DB.WithContext(ctx).Where("channel_id IN ? AND kind = ?", ids, kind).Find(&rows).Error; err != nil {
		return nil, nil, err
	}
	cutoff := time.Now().Unix()
	for _, row := range rows {
		states[row.ChannelID] = row
		cutoff = min(cutoff, row.LastFinishedAt)
	}
	if len(rows) != len(ids) {
		cutoff = 0
	}
	var results []automaticDetectionResult
	if kind == "modeltrace" {
		var tasks []model.SystemTask
		if err := model.DB.WithContext(ctx).Select("status", "payload", "state", "updated_at").Where("type = ? AND updated_at >= ?", AstraSamplingTaskType, cutoff).Find(&tasks).Error; err != nil {
			return nil, nil, err
		}
		for _, task := range tasks {
			var payload struct {
				Force bool `json:"force"`
			}
			if task.DecodePayload(&payload) != nil || payload.Force {
				continue
			}
			var state AstraSamplingState
			if task.DecodeState(&state) != nil {
				continue
			}
			terminal := task.Status == model.SystemTaskStatusSucceeded || task.Status == model.SystemTaskStatusFailed
			for _, round := range state.Rounds {
				if round.FinishedAt == 0 && !terminal {
					active[round.ChannelID] = true
					continue
				}
				at, outcome := round.FinishedAt, FingerprintScheduleOutcome(AstraSamplingResultStatus(round.Result))
				if at == 0 {
					at, outcome = task.UpdatedAt, "unknown"
				}
				results = append(results, automaticDetectionResult{round.ChannelID, round.ID, at, outcome})
			}
		}
	} else {
		var runs []model.DetectionRun
		if err := model.DB.WithContext(ctx).Select("id", "channel_id", "status", "quality", "settled", "finished_at", "kind", "error", "result").Where("channel_id IN ? AND user_id = 0 AND request_key NOT LIKE 'chain:%' AND kind = ? AND (finished_at >= ? OR settled = ?)", ids, kind, cutoff, false).Find(&runs).Error; err != nil {
			return nil, nil, err
		}
		for _, run := range runs {
			run = NormalizeVisionRun(run)
			if !run.Settled {
				active[run.ChannelID] = true
				continue
			}
			outcome := "unknown"
			if run.Status == "succeeded" && (run.Quality == "normal" || run.Quality == "degraded" || run.Quality == "incomplete") {
				outcome = run.Quality
			}
			results = append(results, automaticDetectionResult{run.ChannelID, run.ID, run.FinishedAt, outcome})
		}
	}
	sort.SliceStable(results, func(i, j int) bool {
		if results[i].at == results[j].at {
			return results[i].id < results[j].id
		}
		return results[i].at < results[j].at
	})
	for _, channel := range channels {
		if _, exists := states[channel.Id]; exists {
			continue
		}
		seed := model.DetectionSchedule{ChannelID: channel.Id, Kind: kind, IntervalMinutes: 30}
		for _, result := range results {
			if result.channelID != channel.Id {
				continue
			}
			seed.LastFinishedAt, seed.LastResultID = result.at, result.id
			seed.IntervalMinutes = 15
			if result.outcome == "normal" || result.outcome == "degraded" || result.outcome == "incomplete" {
				seed.LastValid = result.outcome
				seed.IntervalMinutes = 30
			}
			seed.NextDueAt = result.at + int64(seed.IntervalMinutes)*60
		}
		if kind == "ciallo" {
			seed.IntervalMinutes = 45
			if seed.LastFinishedAt > 0 {
				seed.NextDueAt = seed.LastFinishedAt + 45*60
			}
		}
		if err := model.EnsureDetectionSchedule(ctx, seed); err != nil {
			return nil, nil, err
		}
	}
	for _, result := range results {
		// Newly seeded clocks already represent the latest historical result.
		if state, existed := states[result.channelID]; !existed || result.at <= state.LastFinishedAt {
			continue
		}
		if err := model.ApplyDetectionScheduleResult(ctx, result.channelID, kind, result.id, result.at, result.outcome); err != nil {
			return nil, nil, err
		}
	}
	rows = nil
	if err := model.DB.WithContext(ctx).Where("channel_id IN ? AND kind = ?", ids, kind).Find(&rows).Error; err != nil {
		return nil, nil, err
	}
	for _, row := range rows {
		states[row.ChannelID] = row
	}
	return states, active, nil
}

func AutomaticDetectionChannelsDue(ctx context.Context, kind string, now time.Time) ([]model.Channel, map[int]model.DetectionSchedule, error) {
	channels, err := AstraSamplingChannels(ctx)
	if err != nil {
		return nil, nil, err
	}
	states, active, err := ReconcileAutomaticDetectionSchedules(ctx, channels, kind)
	if err != nil {
		return nil, nil, err
	}
	var due []model.Channel
	for _, channel := range channels {
		if !active[channel.Id] && states[channel.Id].NextDueAt <= now.Unix() {
			due = append(due, channel)
		}
	}
	return due, states, nil
}

// Public group summaries expose ranges, never channel identities.
func DetectionScheduleRanges(ctx context.Context, channels []model.Channel, kind string) (map[string][2]int, error) {
	var rows []model.DetectionSchedule
	if err := model.DB.WithContext(ctx).Where("kind = ?", kind).Find(&rows).Error; err != nil {
		return nil, err
	}
	intervals := map[int]int{}
	for _, row := range rows {
		intervals[row.ChannelID] = row.IntervalMinutes
	}
	ranges := map[string][2]int{}
	for _, channel := range channels {
		if channel.Status != common.ChannelStatusEnabled || !ChannelSupportsAstra(&channel) {
			continue
		}
		interval := intervals[channel.Id]
		if interval == 0 {
			interval = 30
		}
		for _, group := range channel.GetGroups() {
			r := ranges[group]
			if r[0] == 0 {
				r = [2]int{interval, interval}
			} else {
				r = [2]int{min(r[0], interval), max(r[1], interval)}
			}
			ranges[group] = r
		}
	}
	return ranges, nil
}