core/server/scheduler/scheduler.go

// Package scheduler runs the task polling loop. It checks for due tasks
// every 30 seconds, creates service channels, and dispatches execution.
//
// v0.27.1: Foundation — scheduler loop + service channel creation.
// v0.27.2: Adds global config checks (enabled, max_concurrent), full cron
//          parsing via robfig/cron/v3, and completion invocation via executor.
// v0.28.0: Adopt queued runs (webhook triggers), action tasks, C3/C4 audit fixes.
// v0.32.0: SKIP LOCKED atomic claim — every replica polls, PG serializes.
//          Replaces ListDue with ClaimDueTask, CreateRunExclusive for
//          belt-and-suspenders uniqueness. Startup jitter staggers replicas.
package scheduler

import (
	"context"
	"database/sql"
	"errors"
	"log"
	"math/rand"
	"time"

	"switchboard-core/models"
	"switchboard-core/store"
	"switchboard-core/taskutil"
)

// Scheduler polls for due tasks and dispatches execution.
type Scheduler struct {
	stores   store.Stores
	executor *Executor
	interval time.Duration
	stop     chan struct{}
	running  bool
}

// New creates a task scheduler. Call Run() in a goroutine to start.
// The executor is optional — if nil, tasks create channels and persist
// prompts but do not invoke completions (v0.27.1 behavior).
func New(stores store.Stores, executor *Executor) *Scheduler {
	return &Scheduler{
		stores:   stores,
		executor: executor,
		interval: 30 * time.Second,
		stop:     make(chan struct{}),
	}
}

// Run starts the scheduler loop. Blocks until Stop() is called.
func (s *Scheduler) Run() {
	if s.stores.Tasks == nil {
		log.Println("[scheduler] TaskStore not available — scheduler disabled")
		return
	}
	s.running = true

	// v0.32.0: Startup jitter — stagger replica polling to reduce lock contention.
	// Not strictly necessary with SKIP LOCKED but reduces unnecessary work.
	jitter := time.Duration(rand.Intn(15000)) * time.Millisecond
	time.Sleep(jitter)
	log.Printf("[scheduler] Started (jitter=%s, interval=%s)", jitter, s.interval)

	ticker := time.NewTicker(s.interval)
	defer ticker.Stop()

	for {
		select {
		case <-ticker.C:
			s.poll()
		case <-s.stop:
			log.Println("[scheduler] Stopped")
			return
		}
	}
}

// Stop signals the scheduler to exit.
func (s *Scheduler) Stop() {
	if s.running {
		close(s.stop)
		s.running = false
	}
}

// poll claims due tasks one at a time and dispatches them.
// v0.32.0: Each call to ClaimDueTask atomically locks and claims one task
// via FOR UPDATE SKIP LOCKED (PG). Multiple replicas poll concurrently;
// PG ensures each task is handed to exactly one replica.
func (s *Scheduler) poll() {
	ctx, cancel := context.WithTimeout(context.Background(), 20*time.Second)
	defer cancel()

	// v0.32.0: Piggyback cleanup of expired tickets and stale rate limit
	// counters on the scheduler tick. Cheap no-ops when tables are empty.
	if s.stores.Tickets != nil {
		if n, err := s.stores.Tickets.Reap(ctx); err == nil && n > 0 {
			log.Printf("[scheduler] Reaped %d expired WS tickets", n)
		}
	}
	if s.stores.RateLimits != nil {
		_ = s.stores.RateLimits.Cleanup(ctx, 5*time.Minute)
	}

	// Check global config — tasks may be disabled at runtime.
	cfg := taskutil.LoadTaskConfig(ctx, s.stores.GlobalConfig)
	if !cfg.Enabled {
		return
	}

	// Claim tasks one at a time until none remain or max_concurrent reached.
	claimed := 0
	for claimed < cfg.MaxConcurrent {
		task, err := s.stores.Tasks.ClaimDueTask(ctx)
		if err != nil {
			// sql.ErrNoRows = nothing due; any other error = log and stop.
			if !errors.Is(err, sql.ErrNoRows) {
				log.Printf("[scheduler] ClaimDueTask error: %v", err)
			}
			break
		}
		claimed++
		go s.execute(task)
	}
}

// execute runs a single task. The task has already been claimed
// (next_run_at set to NULL by ClaimDueTask).
func (s *Scheduler) execute(task *models.Task) {
	ctx, cancel := context.WithTimeout(context.Background(), time.Duration(task.MaxWallClock)*time.Second)
	defer cancel()

	// v0.28.0: Check for queued run (from webhook trigger) — adopt it instead
	// of creating a new one so the trigger_payload is preserved.
	run, _ := s.stores.Tasks.GetQueuedRun(ctx, task.ID)
	if run != nil {
		// Adopt: transition queued → running
		_ = s.stores.Tasks.TransitionRunStatus(ctx, run.ID, "running")
		run.Status = "running"
		log.Printf("[scheduler] Adopting queued run %s for task %s (%s)", run.ID, task.ID, task.Name)
	} else {
		// v0.32.0: Conditional insert — prevents double-execution if another
		// replica somehow also processes this task (belt-and-suspenders).
		var err error
		run, err = s.stores.Tasks.CreateRunExclusive(ctx, task.ID)
		if err != nil {
			if errors.Is(err, sql.ErrNoRows) {
				log.Printf("[scheduler] Skipping task %s (%s) — run already exists", task.ID, task.Name)
			} else {
				log.Printf("[scheduler] Failed to create run for task %s: %v", task.ID, err)
			}
			s.advanceNextRun(ctx, *task)
			return
		}
	}

	log.Printf("[scheduler] Executing task %s (%s) type=%s", task.ID, task.Name, task.TaskType)

	// Mark execution start
	_ = s.stores.Tasks.SetLastRun(ctx, task.ID)

	// Create or reuse service channel
	channelID, err := s.ensureServiceChannel(ctx, *task)
	if err != nil {
		log.Printf("[scheduler] Failed to create service channel for task %s: %v", task.ID, err)
		_ = s.stores.Tasks.UpdateRun(ctx, run.ID, "failed", 0, 0, 0, "channel creation failed: "+err.Error())
		s.advanceNextRun(ctx, *task)
		return
	}

	// Persist the user prompt as a message (prompt tasks only)
	if task.TaskType == "prompt" && task.UserPrompt != "" && s.stores.Messages != nil {
		content := task.UserPrompt
		// v0.28.0: Prepend trigger payload as context if present
		if run.TriggerPayload != "" {
			content = "[Webhook trigger data]\n```json\n" + run.TriggerPayload + "\n```\n\n[Task instructions]\n" + task.UserPrompt
		}
		_ = s.stores.Messages.Create(ctx, &models.Message{
			ChannelID: channelID,
			Role:      "user",
			Content:   content,
		})
	}

	// Invoke completion via executor
	if s.executor != nil {
		s.executor.Execute(ctx, *task, run, channelID)
	} else {
		// No executor — mark completed (channel + prompt persisted)
		_ = s.stores.Tasks.UpdateRun(ctx, run.ID, "completed", 0, 0, 0, "")
		_ = s.stores.Tasks.IncrementRunCount(ctx, task.ID)
	}

	log.Printf("[scheduler] Task %s finished (channel %s)", task.ID, channelID)

	s.advanceNextRun(ctx, *task)
}

// ensureServiceChannel creates a new service channel or reuses an existing one.
func (s *Scheduler) ensureServiceChannel(ctx context.Context, task models.Task) (string, error) {
	// If output_channel_id is set and valid, reuse it
	if task.OutputChannelID != nil && *task.OutputChannelID != "" {
		return *task.OutputChannelID, nil
	}

	// Create a new service channel
	ch := &models.Channel{
		UserID:      task.OwnerID,
		Title:       task.Name,
		Description: "Task output: " + task.Description,
		Type:        "service",
		TeamID:      task.TeamID,
	}
	if err := s.stores.Channels.Create(ctx, ch); err != nil {
		return "", err
	}

	// C3 fix: Persist output_channel_id on the task so future runs reuse this channel.
	channelID := ch.ID
	_ = s.stores.Tasks.Update(ctx, task.ID, models.TaskPatch{
		OutputChannelID: &channelID,
	})

	return channelID, nil
}

// advanceNextRun computes the next run time and updates the task.
func (s *Scheduler) advanceNextRun(ctx context.Context, task models.Task) {
	if task.Schedule == "once" {
		// One-shot task — deactivate after execution
		isActive := false
		_ = s.stores.Tasks.Update(ctx, task.ID, models.TaskPatch{IsActive: &isActive})
		_ = s.stores.Tasks.SetNextRun(ctx, task.ID, nil)
		return
	}

	if task.Schedule == "webhook" {
		// Webhook tasks have no cron schedule — clear next_run_at.
		// They only fire when triggered externally.
		_ = s.stores.Tasks.SetNextRun(ctx, task.ID, nil)
		return
	}

	// Full cron parsing via robfig/cron/v3.
	next := taskutil.NextRunFromSchedule(task.Schedule, task.Timezone)
	if next == nil {
		log.Printf("[scheduler] Failed to compute next run for task %s (schedule: %q) — deactivating", task.ID, task.Schedule)
		isActive := false
		_ = s.stores.Tasks.Update(ctx, task.ID, models.TaskPatch{IsActive: &isActive})
		_ = s.stores.Tasks.SetNextRun(ctx, task.ID, nil)
		return
	}
	_ = s.stores.Tasks.SetNextRun(ctx, task.ID, next)
}