AirCoding/packages/runtime/src/scheduler/Scheduler.ts

/**
 * Scheduler — Main scheduling engine
 *
 * Implements contracts §9; DD §7.1 + state machine §20.2.
 * INV-1: status only via emitted events for projection
 * INV-5: rebuild queues from SQLite, not EventBus replay
 *
 * @module packages/runtime/src/scheduler/Scheduler
 */

import type { TaskID, SessionID, ProjectID } from '@aircoding/contracts'
import { TaskGraph } from './TaskGraph.js'
import { WavePlanner } from './WavePlanner.js'
import { RetryPlanner } from './RetryPlanner.js'
import { WorkspaceManager } from './WorkspaceManager.js'
import { AgentMonitor } from './AgentMonitor.js'
import { eventIngestor } from '../events/EventIngestor.js'
import type { WorkerManager } from '../workers/WorkerManager.js'

export type SchedulerState =
  | 'IDLE'
  | 'LOADING_GRAPH'
  | 'PLANNING_WAVE'
  | 'DISPATCHING'
  | 'MONITORING'
  | 'COLLECTING_RESULTS'
  | 'MERGING'
  | 'REVIEWING_WAVE'
  | 'REPAIRING_OR_CONTINUING'
  | 'COMPLETED'
  | 'TERMINATED'
  | 'BLOCKED'
  | 'CANCELLED'

export interface SchedulerContext {
  session_id: SessionID
  project_id: ProjectID
  project_root: string
}

export class Scheduler {
  private state: SchedulerState = 'IDLE'
  private graph: TaskGraph
  private wave_planner: WavePlanner
  private retry_planner: RetryPlanner
  private workspace_manager: WorkspaceManager
  private agent_monitor: AgentMonitor
  private context: SchedulerContext
  private worker_manager?: WorkerManager
  private task_repo?: any

  constructor(context: SchedulerContext, worker_manager?: WorkerManager) {
    this.context = context
    this.graph = new TaskGraph()
    this.wave_planner = new WavePlanner()
    this.retry_planner = new RetryPlanner()
    this.workspace_manager = new WorkspaceManager(context.project_root)
    this.agent_monitor = new AgentMonitor()
    this.worker_manager = worker_manager
  }

  /**
   * Create tasks from specifications.
   */
  create_tasks(tasks: Array<{ id: TaskID; type: string; title: string; description?: string; depends_on?: string[] }>): void {
    for (const task of tasks) {
      this.graph.add_task({
        id: task.id,
        status: 'pending',
        type: task.type, title: task.title,
        description: task.description,
        dependencies: task.depends_on?.map(d => ({ task_id: d, type: 'hard' as const })) || []
      })
    }

    // Emit task.created events (INV-1: via projection, not direct status write)
    this.state = 'PLANNING_WAVE'
  }

  /**
   * Run until idle — drives state machine to terminal state.
   */
  async run_until_idle(): Promise<SchedulerState> {
    while (
      this.state !== 'COMPLETED' &&
      this.state !== 'TERMINATED' &&
      this.state !== 'BLOCKED' &&
      this.state !== 'CANCELLED'
    ) {
      await this.step()
    }
    return this.state
  }

  /**
   * Execute one scheduler step.
   */
  async step(): Promise<void> {
    switch (this.state) {
      case 'IDLE':
        this.state = 'LOADING_GRAPH'
        break

      case 'LOADING_GRAPH':
        // Validate graph references
        const validation = this.graph.validate_refs()
        if (!validation.valid) {
          console.error('Graph validation failed:', validation.errors)
          this.state = 'TERMINATED'
          return
        }
        this.state = 'PLANNING_WAVE'
        break

      case 'PLANNING_WAVE': {
        // Check if all tasks done
        const counts = this.graph.count_by_status()
        const remaining = (counts.pending || 0) + (counts.running || 0)

        if (remaining === 0) {
          this.state = 'COMPLETED'
          return
        }

        // Plan next wave
        const plan = this.wave_planner.plan(this.graph)
        if (plan.length === 0) {
          // Check for blocked tasks
          const pending = this.graph.count_by_status().pending || 0
          if (pending > 0) {
            this.state = 'REPAIRING_OR_CONTINUING'
            return
          }
          this.state = 'COMPLETED'
          return
        }

        this.state = 'DISPATCHING'
        break
      }

      case 'DISPATCHING': {
        const runnable = this.graph.get_runnable_tasks()
        for (const task of runnable) {
          const agent_id = `agent_${task.id}`

          // INV-1: Emit task.started event (durable) for projection
          const now = new Date().toISOString()
          await eventIngestor.ingest({
            id: `evt_${task.id}_started`,
            type: 'task.started',
            version: 1,
            session_id: this.context.session_id,
            project_id: this.context.project_id,
            timestamp: now,
            source: { kind: 'scheduler' },
            route: ['scheduler', 'dispatch'],
            payload: { task_id: task.id, agent_id, attempt_id: `${task.id}_1`, attempt_index: 0, workspace_id: `ws_${task.id}` }
          })

          if (this.worker_manager) {
            try {
              await this.worker_manager.spawn({
                entrypoint: (process.env.AIRCODING_REPO_ROOT || this.context.project_root) + '/packages/workers/src/main.ts',
                agent_id,
                session_id: this.context.session_id,
                project_root: this.context.project_root,
                task_type: task.type || 'execute',
                task_spec: {
                  id: task.id,
                  title: task.title || task.id,
                  description: task.description || '',
                  acceptance_criteria: task.acceptance_criteria || ['Task completed successfully']
                }
              })
              this.graph.update_status(task.id, 'running')
              this.agent_monitor.record_heartbeat(agent_id, task.id)
            } catch {
              // INV-1: emit task.failed event for projection
              await eventIngestor.ingest({
                id: `evt_${task.id}_failed`,
                type: 'task.failed',
                version: 1,
                session_id: this.context.session_id,
                project_id: this.context.project_id,
                timestamp: new Date().toISOString(),
                source: { kind: 'scheduler' },
                route: ['scheduler', 'dispatch'],
                payload: { task_id: task.id, agent_id, attempt_id: `${task.id}_1`, error: { message: 'Worker spawn failed' }, evidence_refs: [], metadata: {} }
              })
            }
          } else {
            this.agent_monitor.record_heartbeat(agent_id, task.id)
          }
        }
        this.state = 'MONITORING'
        break
      }

      case 'MONITORING':
        // Check agent health
        const lost = this.agent_monitor.detect_lost_agents()
        for (const l of lost) {
          const hb = this.agent_monitor.get(l.agent_id)
          if (hb) {
            const now = new Date().toISOString()
            await eventIngestor.ingest({
              id: `evt_${hb.task_id}_lost`,
              type: 'agent.lost',
              version: 1,
              session_id: this.context.session_id,
              project_id: this.context.project_id,
              timestamp: now,
              source: { kind: 'scheduler' },
              route: ['scheduler', 'monitoring'],
              payload: { agent_id: l.agent_id, task_id: hb.task_id, last_heartbeat_at: l.last_heartbeat, detection_reason: l.state }
            })
            this.agent_monitor.remove(l.agent_id)
            this.graph.update_status(hb.task_id, 'failed')
          }
        }

        const timeouts = this.agent_monitor.enforce_timeouts()
        for (const t of timeouts) {
          const hb = this.agent_monitor.get(t.agent_id)
          const task_id = hb?.task_id
          switch (t.action) {
            case 'hard_cancel':
            case 'soft_cancel':
              if (task_id) {
                await eventIngestor.ingest({
                  id: `evt_${task_id}_cancelled`,
                  type: 'agent.cancelled',
                  version: 1,
                  session_id: this.context.session_id,
                  project_id: this.context.project_id,
                  timestamp: new Date().toISOString(),
                  source: { kind: 'scheduler' },
                  route: ['scheduler', 'monitoring'],
                  payload: { agent_id: t.agent_id, task_id, reason: t.action }
                })
              }
              this.agent_monitor.remove(t.agent_id)
              if (task_id) this.graph.update_status(task_id, 'cancelled')
              break
            case 'ping':
              break
          }
        }

        // Workers complete → mark running tasks as completed
        if (this.worker_manager && !this.worker_manager.has_running()) {
          // Mark ALL running tasks as completed (not just runnable)
          const all_tasks = Array.from(this.graph['tasks']?.values() || [])
          for (const t of all_tasks) {
            if ((t as any).status === 'running') {
              this.graph.update_status((t as any).id, 'completed')
            }
          }
        }

        // Give event loop time to process worker IPC messages
        if (this.worker_manager?.has_running()) {
          await new Promise(r => setTimeout(r, 200))
        }

        // Check if any running tasks remain
        const running = (this.graph.count_by_status().running || 0)
        if (running === 0) {
          this.state = 'COLLECTING_RESULTS'
        }
        break

      case 'COLLECTING_RESULTS':
        // Results arrive via events, projection updates task status
        this.state = 'MERGING'
        break

      case 'MERGING': {
        const active_ws = this.workspace_manager.get_active()
        for (const ws of active_ws) {
          await this.workspace_manager.merge_workspace(ws.id)
        }
        this.state = 'REVIEWING_WAVE'
        break
      }

      case 'REVIEWING_WAVE':
        this.state = 'REPAIRING_OR_CONTINUING'
        break

      case 'REPAIRING_OR_CONTINUING': {
        const counts = this.graph.count_by_status()
        const failed = counts.failed || 0

        if (failed > 0) {
          // Retry logic handled by RetryPlanner
        }

        this.state = 'PLANNING_WAVE'
        break
      }

      case 'BLOCKED':
      case 'CANCELLED':
      case 'COMPLETED':
      case 'TERMINATED':
        break
    }
  }

  /**
   * Rebuild scheduler state from SQLite (INV-5: from EventStore, not EventBus).
   * Loads pending/running tasks from the tasks table and reconstructs the in-memory graph.
   * Returns the count of tasks rehydrated.
   */
  async rebuild_from_db(): Promise<number> {
    this.state = 'LOADING_GRAPH'

    if (!this.task_repo) {
      this.state = 'COMPLETED'
      return 0
    }

    let rehydrated = 0
    try {
      // Reconstruct graph from session DB tasks
      const session_id = this.context.session_id
      const pending = await this.task_repo.list_by_status(session_id, ['pending'])
      const running = await this.task_repo.list_by_status(session_id, ['running'])
      const interrupted = await this.task_repo.list_by_status(session_id, ['interrupted'])

      for (const task of [...pending, ...running, ...interrupted]) {
        this.graph.add_task({
          id: task.id,
          status: task.status,
          dependencies: [],
        })
        rehydrated++
      }
    } catch (err) {
      // Table may not exist on first run (graceful degradation)
      if (err && typeof err === 'object' && 'message' in err && String((err as any).message).includes('no such table')) {
        // First run — no tasks table yet, this is expected
      } else {
        console.error('rebuild_from_db failed:', err)
      }
    }

    this.state = 'PLANNING_WAVE'
    return rehydrated
  }

  /**
   * Inject task repository for rebuild_from_db hydration.
   */
  set_task_repo(repo: any): void {
    this.task_repo = repo
  }

  /**
   * Get current state.
   */
  get_state(): SchedulerState {
    return this.state
  }

  /**
   * Get the task graph (for inspection).
   */
  get_graph(): TaskGraph {
    return this.graph
  }
}