path: root/internal/scheduler/scheduler_test.go

package scheduler

import (
	"context"
	"encoding/json"
	"errors"
	"sync"
	"testing"
	"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"
)

// fakeStore is a minimal, in-memory implementation of Store for unit-testing
// the Scheduler without a real SQLite database.
type fakeStore struct {
	mu           sync.Mutex
	tasks        map[string]*task.Task
	execsByTask  map[string][]*storage.Execution // index 0 = most recent (matches ListExecutions' DESC order)
	roleConfigs  map[string]*storage.RoleConfigRow
	agentUpdates []task.AgentConfig
	stateUpdates []task.State
	events       []*event.Event

	// ask-user-timeout fakes (Phase 7c).
	questionUpdates    []string
	interactions       []task.Interaction
	needsReviewUpdates []bool
}

func newFakeStore() *fakeStore {
	return &fakeStore{
		tasks:       make(map[string]*task.Task),
		execsByTask: make(map[string][]*storage.Execution),
		roleConfigs: make(map[string]*storage.RoleConfigRow),
	}
}

func (f *fakeStore) ListTasks(filter storage.TaskFilter) ([]*task.Task, error) {
	f.mu.Lock()
	defer f.mu.Unlock()
	var out []*task.Task
	for _, t := range f.tasks {
		if filter.State != "" && t.State != filter.State {
			continue
		}
		out = append(out, t)
	}
	return out, nil
}

func (f *fakeStore) ListExecutions(taskID string) ([]*storage.Execution, error) {
	f.mu.Lock()
	defer f.mu.Unlock()
	return f.execsByTask[taskID], nil
}

func (f *fakeStore) GetActiveRoleConfig(roleName string) (*storage.RoleConfigRow, error) {
	f.mu.Lock()
	defer f.mu.Unlock()
	row, ok := f.roleConfigs[roleName]
	if !ok {
		return nil, errors.New("not found")
	}
	return row, nil
}

func (f *fakeStore) UpdateTaskAgent(id string, agent task.AgentConfig) error {
	f.mu.Lock()
	defer f.mu.Unlock()
	f.agentUpdates = append(f.agentUpdates, agent)
	if t, ok := f.tasks[id]; ok {
		t.Agent = agent
	}
	return nil
}

func (f *fakeStore) UpdateTaskState(id string, newState task.State) error {
	f.mu.Lock()
	defer f.mu.Unlock()
	f.stateUpdates = append(f.stateUpdates, newState)
	if t, ok := f.tasks[id]; ok {
		t.State = newState
	}
	return nil
}

func (f *fakeStore) CreateEvent(e *event.Event) error {
	f.mu.Lock()
	defer f.mu.Unlock()
	e.ID = uuid.NewString()
	f.events = append(f.events, e)
	return nil
}

func (f *fakeStore) UpdateTaskQuestion(taskID, questionJSON string) error {
	f.mu.Lock()
	defer f.mu.Unlock()
	f.questionUpdates = append(f.questionUpdates, questionJSON)
	if t, ok := f.tasks[taskID]; ok {
		t.QuestionJSON = questionJSON
	}
	return nil
}

func (f *fakeStore) AppendTaskInteraction(taskID string, interaction task.Interaction) error {
	f.mu.Lock()
	defer f.mu.Unlock()
	f.interactions = append(f.interactions, interaction)
	if t, ok := f.tasks[taskID]; ok {
		t.Interactions = append(t.Interactions, interaction)
	}
	return nil
}

func (f *fakeStore) UpdateTaskNeedsReview(id string, needsReview bool) error {
	f.mu.Lock()
	defer f.mu.Unlock()
	f.needsReviewUpdates = append(f.needsReviewUpdates, needsReview)
	if t, ok := f.tasks[id]; ok {
		t.NeedsReview = needsReview
	}
	return nil
}

func (f *fakeStore) eventsOfKind(k event.Kind) []*event.Event {
	f.mu.Lock()
	defer f.mu.Unlock()
	var out []*event.Event
	for _, e := range f.events {
		if e.Kind == k {
			out = append(out, e)
		}
	}
	return out
}

// fakePool records every task submitted to it.
type fakePool struct {
	mu        sync.Mutex
	submitted []*task.Task
	err       error

	// resumed records every SubmitResume call (Phase 7c's ask-user-timeout
	// path); resumeErr lets tests exercise the error-logging path.
	resumed   []resumeCall
	resumeErr error
}

// resumeCall captures one SubmitResume invocation's arguments.
type resumeCall struct {
	task *task.Task
	exec *storage.Execution
}

func (f *fakePool) Submit(_ context.Context, t *task.Task) error {
	f.mu.Lock()
	defer f.mu.Unlock()
	f.submitted = append(f.submitted, t)
	return f.err
}

func (f *fakePool) SubmitResume(_ context.Context, t *task.Task, exec *storage.Execution) error {
	f.mu.Lock()
	defer f.mu.Unlock()
	f.resumed = append(f.resumed, resumeCall{task: t, exec: exec})
	return f.resumeErr
}

func (f *fakePool) submitCount() int {
	f.mu.Lock()
	defer f.mu.Unlock()
	return len(f.submitted)
}

// fakeBudget is a configurable BudgetGate.
type fakeBudget struct {
	allow bool
	err   error
}

func (f *fakeBudget) Allow(_ string, _ float64) (bool, error) { return f.allow, f.err }

func twoTierLadder() role.RoleConfig {
	return role.RoleConfig{
		Role: "coder",
		EscalationLadder: []role.Tier{
			{
				Candidates:    []role.Rung{{Provider: "local", Model: "m0"}},
				SelectionMode: "single",
				MaxRetries:    2,
			},
			{
				Candidates: []role.Rung{{Provider: "anthropic", Model: "claude-sonnet-5"}},
				MaxRetries: 1,
			},
		},
	}
}

func seedRoleConfig(t *testing.T, f *fakeStore, rc role.RoleConfig) {
	t.Helper()
	b, err := json.Marshal(rc)
	if err != nil {
		t.Fatalf("marshal role config: %v", err)
	}
	f.roleConfigs[rc.Role] = &storage.RoleConfigRow{
		ID: uuid.NewString(), Role: rc.Role, Version: 1, Status: "active", ConfigJSON: string(b),
	}
}

func failedTask(id, roleName string, agentType string) *task.Task {
	return &task.Task{
		ID:    id,
		Name:  "test",
		Agent: task.AgentConfig{Type: agentType, Role: roleName, MaxBudgetUSD: 0.5},
		State: task.StateFailed,
	}
}

func failedExec(rung int) *storage.Execution {
	return &storage.Execution{
		ID:             uuid.NewString(),
		Status:         "FAILED",
		EscalationRung: rung,
		StartTime:      time.Now(),
	}
}

// blockedTask builds a role-typed task BLOCKED on an ask_user question,
// with UpdatedAt standing in for "outstanding since" (see
// tickAskUserTimeouts' doc comment on why task.UpdatedAt, not a new column).
func blockedTask(id, roleName, agentType, questionJSON string, updatedAt time.Time) *task.Task {
	return &task.Task{
		ID:           id,
		Name:         "test",
		Agent:        task.AgentConfig{Type: agentType, Role: roleName, MaxBudgetUSD: 0.5},
		State:        task.StateBlocked,
		QuestionJSON: questionJSON,
		UpdatedAt:    updatedAt,
	}
}

func blockedExec(rung int, sessionID string) *storage.Execution {
	return &storage.Execution{
		ID:             uuid.NewString(),
		Status:         "BLOCKED",
		EscalationRung: rung,
		SessionID:      sessionID,
		StartTime:      time.Now(),
	}
}

// TestScheduler_RetriesSameRung_WhileUnderMaxRetries proves that a task
// whose current rung has fewer attempts than tier.MaxRetries is resubmitted
// at the same rung, with no escalation event and no Agent.Type/Model change.
func TestScheduler_RetriesSameRung_WhileUnderMaxRetries(t *testing.T) {
	store := newFakeStore()
	seedRoleConfig(t, store, twoTierLadder())

	tk := failedTask("t1", "coder", "local")
	store.tasks[tk.ID] = tk
	// Only 1 attempt so far at rung 0; tier0.MaxRetries == 2, so 1 < 2 → retry.
	store.execsByTask[tk.ID] = []*storage.Execution{failedExec(0)}

	pool := &fakePool{}
	sch := &Scheduler{Store: store, Pool: pool}
	sch.Tick(context.Background())

	if pool.submitCount() != 1 {
		t.Fatalf("expected 1 submission, got %d", pool.submitCount())
	}
	resubmitted := pool.submitted[0]
	if resubmitted.Agent.Type != "local" {
		t.Errorf("retry should keep the same provider: got %q", resubmitted.Agent.Type)
	}
	if resubmitted.State != task.StateQueued {
		t.Errorf("resubmitted task should be QUEUED, got %v", resubmitted.State)
	}
	if len(store.eventsOfKind(event.KindEscalated)) != 0 {
		t.Errorf("a same-rung retry should not emit a KindEscalated event")
	}
	if len(store.agentUpdates) != 0 {
		t.Errorf("a same-rung retry should not change the task's Agent config, got %d UpdateTaskAgent calls", len(store.agentUpdates))
	}
}

// TestScheduler_EscalatesToNextRung_WhenBudgetAllows proves that once a
// tier's MaxRetries is exhausted, the scheduler escalates to the next tier
// (when the budget allows it), updates the task's Agent to the new rung's
// provider/model, resubmits, and emits a non-final KindEscalated event.
func TestScheduler_EscalatesToNextRung_WhenBudgetAllows(t *testing.T) {
	store := newFakeStore()
	seedRoleConfig(t, store, twoTierLadder())

	tk := failedTask("t1", "coder", "local")
	store.tasks[tk.ID] = tk
	// 2 consecutive failed attempts at rung 0; tier0.MaxRetries == 2, so
	// 2 < 2 is false → escalate.
	store.execsByTask[tk.ID] = []*storage.Execution{failedExec(0), failedExec(0)}

	pool := &fakePool{}
	sch := &Scheduler{Store: store, Pool: pool, Budget: &fakeBudget{allow: true}}
	sch.Tick(context.Background())

	if pool.submitCount() != 1 {
		t.Fatalf("expected 1 submission, got %d", pool.submitCount())
	}
	resubmitted := pool.submitted[0]
	if resubmitted.Agent.Type != "anthropic" || resubmitted.Agent.Model != "claude-sonnet-5" {
		t.Errorf("escalated task should carry tier 1's rung: got %q/%q", resubmitted.Agent.Type, resubmitted.Agent.Model)
	}
	if resubmitted.State != task.StateQueued {
		t.Errorf("resubmitted task should be QUEUED, got %v", resubmitted.State)
	}
	if len(store.agentUpdates) != 1 {
		t.Fatalf("expected 1 UpdateTaskAgent call, got %d", len(store.agentUpdates))
	}

	evs := store.eventsOfKind(event.KindEscalated)
	if len(evs) != 1 {
		t.Fatalf("expected 1 KindEscalated event, got %d", len(evs))
	}
	var payload struct {
		FromRung     int    `json:"from_rung"`
		ToRung       int    `json:"to_rung"`
		FromProvider string `json:"from_provider"`
		ToProvider   string `json:"to_provider"`
		Final        bool   `json:"final"`
	}
	if err := json.Unmarshal(evs[0].Payload, &payload); err != nil {
		t.Fatalf("unmarshal event payload: %v", err)
	}
	if payload.Final {
		t.Errorf("escalation event should have final=false")
	}
	if payload.FromRung != 0 || payload.ToRung != 1 {
		t.Errorf("expected from_rung=0 to_rung=1, got %d -> %d", payload.FromRung, payload.ToRung)
	}
	if payload.ToProvider != "anthropic" {
		t.Errorf("expected to_provider=anthropic, got %q", payload.ToProvider)
	}
}

// TestScheduler_DeclinesFinal_WhenBudgetDenies proves that when the budget
// denies an otherwise-due escalation, the task is left FAILED (no state
// change, no resubmission) and a final:true KindEscalated event is recorded.
func TestScheduler_DeclinesFinal_WhenBudgetDenies(t *testing.T) {
	store := newFakeStore()
	seedRoleConfig(t, store, twoTierLadder())

	tk := failedTask("t1", "coder", "local")
	store.tasks[tk.ID] = tk
	store.execsByTask[tk.ID] = []*storage.Execution{failedExec(0), failedExec(0)}

	pool := &fakePool{}
	sch := &Scheduler{Store: store, Pool: pool, Budget: &fakeBudget{allow: false}}
	sch.Tick(context.Background())

	if pool.submitCount() != 0 {
		t.Fatalf("budget-denied escalation should not resubmit, got %d submissions", pool.submitCount())
	}
	if len(store.stateUpdates) != 0 {
		t.Errorf("budget-denied escalation should not change task state, got %d state updates", len(store.stateUpdates))
	}

	evs := store.eventsOfKind(event.KindEscalated)
	if len(evs) != 1 {
		t.Fatalf("expected 1 KindEscalated event, got %d", len(evs))
	}
	var payload struct {
		Final bool `json:"final"`
	}
	if err := json.Unmarshal(evs[0].Payload, &payload); err != nil {
		t.Fatalf("unmarshal event payload: %v", err)
	}
	if !payload.Final {
		t.Errorf("budget-denied escalation event should have final=true")
	}
}

// TestScheduler_DeclinesFinal_WhenLadderExhausted proves the same
// final:true behavior when the task is already at the ladder's last tier.
func TestScheduler_DeclinesFinal_WhenLadderExhausted(t *testing.T) {
	store := newFakeStore()
	seedRoleConfig(t, store, twoTierLadder())

	tk := failedTask("t1", "coder", "anthropic")
	store.tasks[tk.ID] = tk
	// At rung 1 (the last tier), tier1.MaxRetries == 1, so 1 attempt already
	// made means 1 < 1 is false → would escalate, but there is no rung 2.
	store.execsByTask[tk.ID] = []*storage.Execution{failedExec(1)}

	pool := &fakePool{}
	sch := &Scheduler{Store: store, Pool: pool, Budget: &fakeBudget{allow: true}}
	sch.Tick(context.Background())

	if pool.submitCount() != 0 {
		t.Fatalf("exhausted ladder should not resubmit, got %d submissions", pool.submitCount())
	}
	evs := store.eventsOfKind(event.KindEscalated)
	if len(evs) != 1 {
		t.Fatalf("expected 1 KindEscalated event, got %d", len(evs))
	}
	var payload struct {
		Final bool `json:"final"`
	}
	if err := json.Unmarshal(evs[0].Payload, &payload); err != nil {
		t.Fatalf("unmarshal event payload: %v", err)
	}
	if !payload.Final {
		t.Errorf("exhausted-ladder decline event should have final=true")
	}
}

// TestScheduler_Convergence_DoesNotReprocessSameExecution proves the
// double-processing guard: ticking twice against the same unchanged FAILED
// execution only acts once (no duplicate submissions or events), which is
// what keeps a task stuck at the end of its ladder from generating a fresh
// "final" event (or worse, a fresh escalation) on every single poll forever.
func TestScheduler_Convergence_DoesNotReprocessSameExecution(t *testing.T) {
	store := newFakeStore()
	seedRoleConfig(t, store, twoTierLadder())

	tk := failedTask("t1", "coder", "anthropic")
	store.tasks[tk.ID] = tk
	store.execsByTask[tk.ID] = []*storage.Execution{failedExec(1)} // last tier, exhausted

	pool := &fakePool{}
	sch := &Scheduler{Store: store, Pool: pool, Budget: &fakeBudget{allow: true}}

	sch.Tick(context.Background())
	sch.Tick(context.Background())
	sch.Tick(context.Background())

	if got := len(store.eventsOfKind(event.KindEscalated)); got != 1 {
		t.Fatalf("expected exactly 1 KindEscalated event across 3 ticks, got %d", got)
	}
	if pool.submitCount() != 0 {
		t.Fatalf("expected 0 submissions, got %d", pool.submitCount())
	}
}

// TestScheduler_AskUserTimeout_EscalatesAndResumes proves the core
// ask-user-timeout flow: a BLOCKED role-typed task whose question has been
// outstanding longer than the configured timeout gets resumed at the next
// escalation tier with a system-authored answer, needs_review gets set, and
// a KindEscalated event with the ask-user-timeout trigger is emitted.
func TestScheduler_AskUserTimeout_EscalatesAndResumes(t *testing.T) {
	store := newFakeStore()
	seedRoleConfig(t, store, twoTierLadder())

	askedAt := time.Now().Add(-20 * time.Minute)
	tk := blockedTask("t1", "coder", "local", `{"text":"Which approach?","options":["a","b"]}`, askedAt)
	store.tasks[tk.ID] = tk
	store.execsByTask[tk.ID] = []*storage.Execution{blockedExec(0, "sess-1")}

	pool := &fakePool{}
	sch := &Scheduler{Store: store, Pool: pool, AskUserTimeout: 10 * time.Minute}
	sch.Tick(context.Background())

	if len(pool.resumed) != 1 {
		t.Fatalf("expected 1 resume submission, got %d", len(pool.resumed))
	}
	rc := pool.resumed[0]
	if rc.task.Agent.Type != "anthropic" || rc.task.Agent.Model != "claude-sonnet-5" {
		t.Errorf("resumed task should carry the escalated tier's provider/model: got %q/%q", rc.task.Agent.Type, rc.task.Agent.Model)
	}
	if rc.task.State != task.StateBlocked {
		t.Errorf("resumed task struct should retain BLOCKED state (matches api.answerTaskQuestion's convention for SubmitResume's resumable-state check), got %v", rc.task.State)
	}
	if rc.exec.ResumeSessionID != "sess-1" {
		t.Errorf("resume exec ResumeSessionID: got %q want sess-1", rc.exec.ResumeSessionID)
	}
	if rc.exec.ResumeAnswer != fallbackAnswer {
		t.Errorf("resume exec ResumeAnswer: got %q want the fallback answer marker", rc.exec.ResumeAnswer)
	}

	if len(store.needsReviewUpdates) != 1 || !store.needsReviewUpdates[0] {
		t.Errorf("expected needs_review to be set true exactly once, got %+v", store.needsReviewUpdates)
	}
	if len(store.interactions) != 1 || store.interactions[0].Answer != fallbackAnswer || store.interactions[0].QuestionText != "Which approach?" {
		t.Errorf("expected the fallback answer recorded as an interaction, got %+v", store.interactions)
	}
	if len(store.questionUpdates) != 1 || store.questionUpdates[0] != "" {
		t.Errorf("expected the question to be cleared, got %+v", store.questionUpdates)
	}
	if len(store.agentUpdates) != 1 {
		t.Fatalf("expected 1 UpdateTaskAgent call, got %d", len(store.agentUpdates))
	}

	evs := store.eventsOfKind(event.KindEscalated)
	if len(evs) != 1 {
		t.Fatalf("expected 1 KindEscalated event, got %d", len(evs))
	}
	var payload struct {
		FromRung     int    `json:"from_rung"`
		ToRung       int    `json:"to_rung"`
		FromProvider string `json:"from_provider"`
		ToProvider   string `json:"to_provider"`
		Final        bool   `json:"final"`
		Trigger      string `json:"trigger"`
	}
	if err := json.Unmarshal(evs[0].Payload, &payload); err != nil {
		t.Fatalf("unmarshal event payload: %v", err)
	}
	if payload.Trigger != "ask_user_timeout" {
		t.Errorf("expected trigger=ask_user_timeout, got %q", payload.Trigger)
	}
	if payload.Final {
		t.Errorf("escalating to a higher tier should not be final")
	}
	if payload.FromRung != 0 || payload.ToRung != 1 || payload.FromProvider != "local" || payload.ToProvider != "anthropic" {
		t.Errorf("expected escalation local(rung0) -> anthropic(rung1), got %+v", payload)
	}
}

// TestScheduler_AskUserTimeout_WithinWindow_LeftAlone proves that a BLOCKED
// task still within the configured timeout window is left completely alone:
// no resume, no needs_review, no event.
func TestScheduler_AskUserTimeout_WithinWindow_LeftAlone(t *testing.T) {
	store := newFakeStore()
	seedRoleConfig(t, store, twoTierLadder())

	tk := blockedTask("t1", "coder", "local", `{"text":"Which approach?"}`, time.Now())
	store.tasks[tk.ID] = tk
	store.execsByTask[tk.ID] = []*storage.Execution{blockedExec(0, "sess-1")}

	pool := &fakePool{}
	sch := &Scheduler{Store: store, Pool: pool, AskUserTimeout: 10 * time.Minute}
	sch.Tick(context.Background())

	if len(pool.resumed) != 0 {
		t.Fatalf("task within the timeout window should not be resumed, got %d resumes", len(pool.resumed))
	}
	if len(store.needsReviewUpdates) != 0 {
		t.Errorf("task within the timeout window should not be flagged needs_review, got %+v", store.needsReviewUpdates)
	}
	if len(store.eventsOfKind(event.KindEscalated)) != 0 {
		t.Errorf("task within the timeout window should not emit a KindEscalated event")
	}
	if len(store.questionUpdates) != 0 {
		t.Errorf("task within the timeout window should not have its question cleared")
	}
}

// TestScheduler_AskUserTimeout_IgnoresSubtaskBlocked proves that a BLOCKED
// task with no pending question (i.e. blocked waiting on subtasks, per
// task.go's state machine, not on ask_user) is never touched by
// tickAskUserTimeouts, however long it's been BLOCKED.
func TestScheduler_AskUserTimeout_IgnoresSubtaskBlocked(t *testing.T) {
	store := newFakeStore()
	seedRoleConfig(t, store, twoTierLadder())

	tk := blockedTask("t1", "coder", "local", "", time.Now().Add(-1*time.Hour))
	store.tasks[tk.ID] = tk
	store.execsByTask[tk.ID] = []*storage.Execution{blockedExec(0, "sess-1")}

	pool := &fakePool{}
	sch := &Scheduler{Store: store, Pool: pool, AskUserTimeout: 10 * time.Minute}
	sch.Tick(context.Background())

	if len(pool.resumed) != 0 {
		t.Fatalf("subtask-blocked task should not be resumed, got %d resumes", len(pool.resumed))
	}
}

// TestScheduler_AskUserTimeout_LadderExhausted_StillResumesAtSameTier proves
// the documented fallback for when no higher tier exists: the task is still
// unblocked (resumed) so it doesn't stay stuck forever, but at its current
// tier's provider/model (no UpdateTaskAgent call), and the escalation event
// is marked final:true to distinguish "resumed without escalating" from a
// genuine tier bump.
func TestScheduler_AskUserTimeout_LadderExhausted_StillResumesAtSameTier(t *testing.T) {
	store := newFakeStore()
	seedRoleConfig(t, store, twoTierLadder())

	askedAt := time.Now().Add(-20 * time.Minute)
	tk := blockedTask("t1", "coder", "anthropic", `{"text":"Which approach?"}`, askedAt)
	store.tasks[tk.ID] = tk
	store.execsByTask[tk.ID] = []*storage.Execution{blockedExec(1, "sess-1")} // already at the last tier

	pool := &fakePool{}
	sch := &Scheduler{Store: store, Pool: pool, AskUserTimeout: 10 * time.Minute}
	sch.Tick(context.Background())

	if len(pool.resumed) != 1 {
		t.Fatalf("expected the task to still be resumed even with no higher tier, got %d resumes", len(pool.resumed))
	}
	if len(store.agentUpdates) != 0 {
		t.Errorf("no higher tier exists, so Agent should be unchanged: got %d UpdateTaskAgent calls", len(store.agentUpdates))
	}
	rc := pool.resumed[0]
	if rc.task.Agent.Type != "anthropic" {
		t.Errorf("resumed task should keep its current provider, got %q", rc.task.Agent.Type)
	}

	evs := store.eventsOfKind(event.KindEscalated)
	if len(evs) != 1 {
		t.Fatalf("expected 1 KindEscalated event, got %d", len(evs))
	}
	var payload struct {
		Final   bool   `json:"final"`
		Trigger string `json:"trigger"`
	}
	if err := json.Unmarshal(evs[0].Payload, &payload); err != nil {
		t.Fatalf("unmarshal event payload: %v", err)
	}
	if !payload.Final {
		t.Errorf("resuming without an escalation should be marked final=true")
	}
	if payload.Trigger != "ask_user_timeout" {
		t.Errorf("expected trigger=ask_user_timeout, got %q", payload.Trigger)
	}
	if len(store.needsReviewUpdates) != 1 || !store.needsReviewUpdates[0] {
		t.Errorf("expected needs_review to still be set even without an escalation, got %+v", store.needsReviewUpdates)
	}
}