path: root/internal/scheduler/scheduler.go

// Package scheduler implements a first-pass retry-then-escalate loop for
// role-typed tasks (task.AgentConfig.Role != ""). On a poll interval it looks
// for tasks whose most recent execution ended FAILED, and either resubmits
// them at the same escalation-ladder tier (if under that tier's MaxRetries)
// or escalates them to the next tier (if the budget allows) — recording an
// event.KindEscalated event either way. If the ladder is exhausted or the
// budget denies the escalation, the task is left FAILED for human attention.
//
// Phase 7c extended this same "watch for stuck role-typed tasks and take
// escalation action" responsibility to a second trigger: a role-typed task
// BLOCKED on an ask_user question nobody has answered within
// SchedulerConfig.AskUserTimeoutSeconds (see tickAskUserTimeouts below). This
// is a natural extension of what Scheduler already does, not a new
// component — it owns retry/escalation for role-typed tasks generally, and a
// stuck question is just another way a role-typed task gets stuck.
//
// Explicit non-goal still remaining: no DAG/cascade-fail logic here (that's
// executor.Pool.cascadeFail's job). Handling for TIMED_OUT/CANCELLED/
// BUDGET_EXCEEDED tasks follows the same shape as FAILED but isn't
// implemented yet — only FAILED is polled for the failure-retry path.
package scheduler

import (
	"context"
	"encoding/json"
	"log/slog"
	"sync"
	"time"

	"github.com/google/uuid"
	"github.com/thepeterstone/claudomator/internal/event"
	"github.com/thepeterstone/claudomator/internal/role"
	"github.com/thepeterstone/claudomator/internal/storage"
	"github.com/thepeterstone/claudomator/internal/task"
)

// Store is the subset of storage.DB methods the Scheduler needs.
type Store interface {
	ListTasks(filter storage.TaskFilter) ([]*task.Task, error)
	ListExecutions(taskID string) ([]*storage.Execution, error)
	GetActiveRoleConfig(role string) (*storage.RoleConfigRow, error)
	UpdateTaskAgent(id string, agent task.AgentConfig) error
	UpdateTaskState(id string, newState task.State) error
	CreateEvent(e *event.Event) error
	// UpdateTaskQuestion, AppendTaskInteraction, and UpdateTaskNeedsReview
	// back tickAskUserTimeouts' resume-with-fallback-answer flow (Phase 7c):
	// clearing the stale question, recording the system-authored answer as an
	// interaction (mirroring api.answerTaskQuestion's audit trail for a real
	// human answer), and flagging the task for later human review.
	UpdateTaskQuestion(taskID, questionJSON string) error
	AppendTaskInteraction(taskID string, interaction task.Interaction) error
	UpdateTaskNeedsReview(id string, needsReview bool) error
}

// Pool is the subset of *executor.Pool the Scheduler needs. Satisfied by
// *executor.Pool directly (see internal/cli/serve.go); declared as an
// interface here purely so tests can supply a fake without dragging in the
// executor package's runner/sandbox machinery.
type Pool interface {
	Submit(ctx context.Context, t *task.Task) error
	// SubmitResume re-queues a BLOCKED (or otherwise interrupted) task using
	// a resume execution carrying ResumeSessionID/ResumeAnswer. Used by
	// tickAskUserTimeouts to resume a task with a system-authored fallback
	// answer the same way api.answerTaskQuestion resumes it with a real one.
	SubmitResume(ctx context.Context, t *task.Task, exec *storage.Execution) error
}

// BudgetGate reports whether an escalation to provider estimated at estCost
// is allowed. Satisfied by *budget.Accountant.
type BudgetGate interface {
	Allow(provider string, estCost float64) (bool, error)
}

// Scheduler polls for role-typed FAILED tasks and retries or escalates them
// per their active role_configs escalation ladder. It also polls for
// role-typed tasks BLOCKED on a stale ask_user question (see
// tickAskUserTimeouts).
type Scheduler struct {
	Store  Store
	Pool   Pool
	Budget BudgetGate // nil means "no budget gating" (always allow)
	Logger *slog.Logger

	// AskUserTimeout is how long a role-typed task may sit BLOCKED on an
	// unanswered ask_user question before tickAskUserTimeouts resumes it with
	// a system-authored fallback answer. <= 0 means DefaultAskUserTimeout
	// (see askUserTimeout()); set from config.SchedulerConfig.AskUserTimeout()
	// in production (internal/cli/serve.go).
	AskUserTimeout time.Duration

	// handled dedupes processing within a single running process: once a
	// decision (retry/escalate/decline) has been made for a given
	// execution ID, it is never reconsidered again by this Scheduler
	// instance. This is what keeps Run's poll loop convergent — a task left
	// FAILED after its ladder is exhausted (or an escalation is budget-
	// denied) has the same "latest execution" on every subsequent tick, so
	// without this it would emit a fresh "final" KindEscalated event, and
	// re-run the same decision, every single poll forever. A restart clears
	// this map, so a task can be reconsidered once more after a restart —
	// intentional: it's an idempotent bookkeeping decision, not orchestration
	// state, so re-deriving it once is harmless.
	//
	// tickAskUserTimeouts does NOT need an equivalent guard: successfully
	// resuming a task moves it out of BLOCKED (to QUEUED), which structurally
	// removes it from the next tick's BLOCKED query — the same
	// "idempotency via a state the next poll won't rediscover" pattern
	// StoryOrchestrator's ensureEvaluators/ensureArbitration use, rather than
	// an in-memory marker.
	mu      sync.Mutex
	handled map[string]bool
}

// DefaultPollInterval is used by Run when pollInterval <= 0.
const DefaultPollInterval = 30 * time.Second

// DefaultAskUserTimeout is used when Scheduler.AskUserTimeout <= 0. Mirrors
// config.SchedulerConfig.AskUserTimeout's default and reasoning (10 minutes:
// long enough for an actively-working human to notice and answer a
// clarification request, short enough that a role-typed task doesn't stall
// for hours on input that may never come — the common case for a
// self-hosted, typically single-operator deployment).
const DefaultAskUserTimeout = 10 * time.Minute

// askUserTimeout returns the effective ask-user timeout, defaulting to
// DefaultAskUserTimeout when unset.
func (s *Scheduler) askUserTimeout() time.Duration {
	if s.AskUserTimeout <= 0 {
		return DefaultAskUserTimeout
	}
	return s.AskUserTimeout
}

// fallbackAnswer is the system-authored answer injected into a role-typed
// task's ask_user question once it has gone unanswered for longer than
// askUserTimeout(). Clearly marked as a system fallback, not a real human
// answer, so anyone reading the task's interaction history or event stream
// later understands why the agent proceeded without a real decision.
const fallbackAnswer = "[auto-escalated: no human response within timeout; proceeding with best judgment]"

// Run polls for role-typed FAILED tasks every pollInterval until ctx is
// cancelled.
func (s *Scheduler) Run(ctx context.Context, pollInterval time.Duration) {
	if pollInterval <= 0 {
		pollInterval = DefaultPollInterval
	}
	ticker := time.NewTicker(pollInterval)
	defer ticker.Stop()
	for {
		select {
		case <-ctx.Done():
			return
		case <-ticker.C:
			s.Tick(ctx)
		}
	}
}

// Tick runs a single poll pass: FAILED-task retry/escalation, then
// ask_user-timeout escalation. Exported so tests can drive it directly
// without waiting on a ticker.
func (s *Scheduler) Tick(ctx context.Context) {
	tasks, err := s.Store.ListTasks(storage.TaskFilter{State: task.StateFailed})
	if err != nil {
		s.logf("scheduler: list failed tasks", "error", err)
	} else {
		for _, t := range tasks {
			if t.Agent.Role == "" {
				continue
			}
			s.processTask(ctx, t)
		}
	}
	s.tickAskUserTimeouts(ctx)
}

func (s *Scheduler) logf(msg string, args ...any) {
	if s.Logger != nil {
		s.Logger.Warn(msg, args...)
	}
}

func (s *Scheduler) markHandled(execID string) bool {
	s.mu.Lock()
	defer s.mu.Unlock()
	if s.handled == nil {
		s.handled = make(map[string]bool)
	}
	if s.handled[execID] {
		return false
	}
	s.handled[execID] = true
	return true
}

func (s *Scheduler) processTask(ctx context.Context, t *task.Task) {
	execs, err := s.Store.ListExecutions(t.ID)
	if err != nil || len(execs) == 0 {
		return
	}
	latest := execs[0] // ListExecutions orders DESC by start_time.
	if latest.Status != "FAILED" {
		return
	}
	if !s.markHandled(latest.ID) {
		return // already decided for this execution — converged, nothing to do.
	}

	row, err := s.Store.GetActiveRoleConfig(t.Agent.Role)
	if err != nil {
		s.logf("scheduler: no active role config for role", "role", t.Agent.Role, "taskID", t.ID, "error", err)
		return
	}
	var rc role.RoleConfig
	if err := json.Unmarshal([]byte(row.ConfigJSON), &rc); err != nil {
		s.logf("scheduler: decode role config", "role", t.Agent.Role, "taskID", t.ID, "error", err)
		return
	}
	if len(rc.EscalationLadder) == 0 {
		return
	}

	currentRung := latest.EscalationRung
	if currentRung < 0 {
		currentRung = 0
	}
	if currentRung >= len(rc.EscalationLadder) {
		// Ladder already exhausted (e.g. the ladder was shortened after this
		// task started climbing it) — nothing more to do.
		return
	}
	tier := rc.EscalationLadder[currentRung]
	attempts := attemptsAtRung(execs, currentRung)

	if attempts < tier.MaxRetries {
		s.retrySameRung(ctx, t, currentRung)
		return
	}

	nextRung := currentRung + 1
	if nextRung >= len(rc.EscalationLadder) || len(rc.EscalationLadder[nextRung].Candidates) == 0 {
		s.decline(ctx, t, currentRung, "", "escalation ladder exhausted")
		return
	}
	nextTier := rc.EscalationLadder[nextRung]
	target := nextTier.Candidates[0]

	estCost := rc.DefaultBudgetUSD
	if estCost <= 0 {
		estCost = t.Agent.MaxBudgetUSD
	}
	allowed := true
	if s.Budget != nil {
		var berr error
		allowed, berr = s.Budget.Allow(target.Provider, estCost)
		if berr != nil {
			s.logf("scheduler: budget check failed; declining escalation", "taskID", t.ID, "error", berr)
			allowed = false
		}
	}
	if !allowed {
		s.decline(ctx, t, currentRung, target.Provider, "budget denied")
		return
	}
	s.escalate(ctx, t, currentRung, nextRung, target)
}

// attemptsAtRung counts how many executions, starting from the most recent
// (execs[0]) and moving backward, ran at rung consecutively — i.e. how many
// attempts have already been made at the task's current tier since it last
// moved to (or started at) that tier.
func attemptsAtRung(execs []*storage.Execution, rung int) int {
	n := 0
	for _, e := range execs {
		if e.EscalationRung != rung {
			break
		}
		n++
	}
	return n
}

func (s *Scheduler) retrySameRung(ctx context.Context, t *task.Task, rung int) {
	if err := s.Store.UpdateTaskState(t.ID, task.StateQueued); err != nil {
		s.logf("scheduler: retry: update task state", "taskID", t.ID, "error", err)
		return
	}
	resubmit := *t
	resubmit.State = task.StateQueued
	if err := s.Pool.Submit(ctx, &resubmit); err != nil {
		s.logf("scheduler: retry: submit", "taskID", t.ID, "error", err)
	}
}

func (s *Scheduler) escalate(ctx context.Context, t *task.Task, fromRung, toRung int, target role.Rung) {
	fromProvider := t.Agent.Type
	newAgent := t.Agent
	newAgent.Type = target.Provider
	newAgent.Model = target.Model
	if err := s.Store.UpdateTaskAgent(t.ID, newAgent); err != nil {
		s.logf("scheduler: escalate: update task agent", "taskID", t.ID, "error", err)
		return
	}
	if err := s.Store.UpdateTaskState(t.ID, task.StateQueued); err != nil {
		s.logf("scheduler: escalate: update task state", "taskID", t.ID, "error", err)
		return
	}
	s.emitEscalated(t.ID, fromRung, toRung, fromProvider, target.Provider, false, "", "failure")

	resubmit := *t
	resubmit.Agent = newAgent
	resubmit.State = task.StateQueued
	if err := s.Pool.Submit(ctx, &resubmit); err != nil {
		s.logf("scheduler: escalate: submit", "taskID", t.ID, "error", err)
	}
}

// decline records that no further escalation will happen for t right now
// (ladder exhausted or budget denied) and leaves it FAILED for human
// attention.
func (s *Scheduler) decline(ctx context.Context, t *task.Task, atRung int, consideredProvider, reason string) {
	s.emitEscalated(t.ID, atRung, atRung, t.Agent.Type, consideredProvider, true, reason, "failure")
}

// emitEscalated records an event.KindEscalated event. trigger distinguishes
// what caused the scheduler to reconsider this task's tier: "failure" for
// the FAILED-task retry/escalation path above, "ask_user_timeout" for
// escalateAskUserTimeout below — so someone reading the event stream later
// can tell a stuck-question escalation from a stuck-failure one.
func (s *Scheduler) emitEscalated(taskID string, fromRung, toRung int, fromProvider, toProvider string, final bool, reason, trigger string) {
	payload, _ := json.Marshal(struct {
		FromRung     int    `json:"from_rung"`
		ToRung       int    `json:"to_rung"`
		FromProvider string `json:"from_provider"`
		ToProvider   string `json:"to_provider,omitempty"`
		Final        bool   `json:"final"`
		Reason       string `json:"reason,omitempty"`
		Trigger      string `json:"trigger,omitempty"`
	}{FromRung: fromRung, ToRung: toRung, FromProvider: fromProvider, ToProvider: toProvider, Final: final, Reason: reason, Trigger: trigger})
	if err := s.Store.CreateEvent(&event.Event{
		TaskID:  taskID,
		Kind:    event.KindEscalated,
		Actor:   event.ActorSystem,
		Payload: payload,
	}); err != nil {
		s.logf("scheduler: emit escalated event", "taskID", taskID, "error", err)
	}
}

// tickAskUserTimeouts finds role-typed tasks BLOCKED on an ask_user question
// that has been outstanding longer than askUserTimeout() and resumes each
// with a system-authored fallback answer, escalating to the next tier of the
// role's ladder where possible so a different (hopefully more capable)
// provider/model picks up where the original agent stalled.
//
// The "outstanding since" timestamp is task.UpdatedAt, not a new column:
// storage.DB.UpdateTaskQuestion (the last write made to a task's row on its
// way into BLOCKED — see executor.Pool.handleRunResult's BlockedError
// branch, which calls UpdateTaskState(BLOCKED) then
// UpdateTaskQuestion(questionJSON) in that order) stamps updated_at at the
// exact moment the question was recorded, and nothing else touches the row
// while it sits BLOCKED awaiting an answer. Reusing it avoids an entirely
// redundant "asked_at" column carrying the same information a second time.
func (s *Scheduler) tickAskUserTimeouts(ctx context.Context) {
	tasks, err := s.Store.ListTasks(storage.TaskFilter{State: task.StateBlocked})
	if err != nil {
		s.logf("scheduler: list blocked tasks", "error", err)
		return
	}
	timeout := s.askUserTimeout()
	for _, t := range tasks {
		if t.Agent.Role == "" || t.QuestionJSON == "" {
			// Not role-typed, or BLOCKED on pending subtasks rather than a
			// question (see task.go's state machine: BLOCKED covers both).
			continue
		}
		if time.Since(t.UpdatedAt) < timeout {
			continue // still within the grace period
		}
		s.escalateAskUserTimeout(ctx, t)
	}
}

// escalateAskUserTimeout resumes a single BLOCKED, role-typed task whose
// question has timed out. It resolves the role's escalation ladder from the
// tier the task was dispatched at (latest execution's EscalationRung) and
// picks the next tier up, mirroring processTask's escalate() above but
// applied to "stuck on a question" rather than "stuck on a failure" — if no
// higher tier exists (ladder exhausted, no active role config, etc.) it
// still resumes the task (unblocking it is the priority) at its current
// tier, just without a provider/model change, and marks the escalation
// event final:true so that distinction is visible in the event stream.
func (s *Scheduler) escalateAskUserTimeout(ctx context.Context, t *task.Task) {
	execs, err := s.Store.ListExecutions(t.ID)
	if err != nil || len(execs) == 0 {
		s.logf("scheduler: ask-user-timeout: no executions", "taskID", t.ID)
		return
	}
	latest := execs[0] // ListExecutions orders DESC by start_time.
	if latest.SessionID == "" {
		s.logf("scheduler: ask-user-timeout: no resumable session", "taskID", t.ID)
		return
	}

	currentRung := latest.EscalationRung
	if currentRung < 0 {
		currentRung = 0
	}

	// Captured before any Store mutation below: fakeStore-style test doubles
	// (and, in principle, a caching Store) may hand back the same *task.Task
	// pointer from ListTasks that UpdateTaskAgent then mutates in place, so
	// reading t.Agent.Type *after* that call would silently pick up the new
	// value instead of the original one — the same "read fromProvider before
	// mutating" care processTask's escalate() takes above.
	fromProvider := t.Agent.Type

	newAgent := t.Agent
	toRung := currentRung
	toProvider := fromProvider
	reason := ""
	final := false

	row, rcErr := s.Store.GetActiveRoleConfig(t.Agent.Role)
	if rcErr != nil {
		reason = "no active role config; resuming at same tier"
		final = true
	} else {
		var rc role.RoleConfig
		if jsonErr := json.Unmarshal([]byte(row.ConfigJSON), &rc); jsonErr != nil {
			reason = "failed to decode role config; resuming at same tier"
			final = true
		} else if len(rc.EscalationLadder) == 0 {
			reason = "empty escalation ladder; resuming at same tier"
			final = true
		} else {
			nextRung := currentRung + 1
			if nextRung >= len(rc.EscalationLadder) || len(rc.EscalationLadder[nextRung].Candidates) == 0 {
				reason = "escalation ladder exhausted; resuming at same tier"
				final = true
			} else {
				target := rc.EscalationLadder[nextRung].Candidates[0]
				newAgent.Type = target.Provider
				newAgent.Model = target.Model
				toRung = nextRung
				toProvider = target.Provider
			}
		}
	}

	// Record the system-authored fallback answer as an interaction before
	// clearing the question, mirroring the audit trail api.answerTaskQuestion
	// leaves for a real human answer.
	if t.QuestionJSON != "" {
		var qData struct {
			Text    string   `json:"text"`
			Options []string `json:"options"`
		}
		if json.Unmarshal([]byte(t.QuestionJSON), &qData) == nil {
			if err := s.Store.AppendTaskInteraction(t.ID, task.Interaction{
				QuestionText: qData.Text,
				Options:      qData.Options,
				Answer:       fallbackAnswer,
				AskedAt:      t.UpdatedAt,
			}); err != nil {
				s.logf("scheduler: ask-user-timeout: append interaction", "taskID", t.ID, "error", err)
			}
		}
	}
	if err := s.Store.UpdateTaskQuestion(t.ID, ""); err != nil {
		s.logf("scheduler: ask-user-timeout: clear question", "taskID", t.ID, "error", err)
		return
	}
	if newAgent.Type != fromProvider || newAgent.Model != t.Agent.Model {
		if err := s.Store.UpdateTaskAgent(t.ID, newAgent); err != nil {
			s.logf("scheduler: ask-user-timeout: update task agent", "taskID", t.ID, "error", err)
			return
		}
	}
	if err := s.Store.UpdateTaskState(t.ID, task.StateQueued); err != nil {
		s.logf("scheduler: ask-user-timeout: update task state", "taskID", t.ID, "error", err)
		return
	}
	if err := s.Store.UpdateTaskNeedsReview(t.ID, true); err != nil {
		s.logf("scheduler: ask-user-timeout: mark needs_review", "taskID", t.ID, "error", err)
	}

	s.emitEscalated(t.ID, currentRung, toRung, fromProvider, toProvider, final, reason, "ask_user_timeout")

	// SubmitResume requires the task passed in to carry a resumable State
	// (see executor.resumablePoolStates, which includes BLOCKED) — mirroring
	// api.answerTaskQuestion, which passes the pre-transition BLOCKED task
	// struct even though the DB row has already moved to QUEUED above.
	resume := *t
	resume.Agent = newAgent
	resume.State = task.StateBlocked
	resumeExec := &storage.Execution{
		ID:              uuid.NewString(),
		TaskID:          t.ID,
		ResumeSessionID: latest.SessionID,
		ResumeAnswer:    fallbackAnswer,
		SandboxDir:      latest.SandboxDir,
	}
	if err := s.Pool.SubmitResume(ctx, &resume, resumeExec); err != nil {
		s.logf("scheduler: ask-user-timeout: submit resume", "taskID", t.ID, "error", err)
	}
}