AirCoding/packages/runtime/src/workers/WorkerManager.ts

/**
 * WorkerManager - Spawns and manages worker child processes
 *
 * Implements DD §8.1. spawn (Bun child process + handshake), cancel.
 * INV-1: WorkerManager never writes agents.status directly.
 *
 * @module packages/runtime/src/workers/WorkerManager
 */

import { spawn } from 'child_process'
import { existsSync } from 'fs'
import type { ChildProcess } from 'child_process'
import { WorkerProcess, type WorkerExitCode } from './WorkerProcess.js'
import { WorkerProtocol, type WorkerMessage, type WorkerMessageType } from './WorkerProtocol.js'
import type { WorkerResult, WorkerStatus, AgentType } from '@aircoding/contracts'
import type { ToolRegistry } from '../tools/ToolRegistry.js'
import type { ProviderManager } from '@aircoding/llm'

export interface WorkerConfig {
  entrypoint: string       // Path to worker main.ts
  agent_id: string
  session_id: string
  project_root: string
  timeout_ms?: number
  env?: Record<string, string>
  task_type?: string       // DD §8.3: execute/review/debug/compact/mine_experience
  task_spec?: Record<string, unknown>  // DD §9: TaskSpec payload for worker
}

export interface WorkerHandle {
  worker_id: string
  process: WorkerProcess
  config: WorkerConfig
  state: 'starting' | 'ready' | 'running' | 'completed' | 'error' | 'cancelled'
  started_at: string
  completed_at?: string
  result?: WorkerResult<unknown>
}

export class WorkerManager {
  private protocol: WorkerProtocol
  private workers: Map<string, WorkerHandle> = new Map()
  private tool_registry?: ToolRegistry
  private provider_manager?: ProviderManager
  private execution_context?: { session_id: string; project_id: string; project_root: string }

  constructor(tool_registry?: ToolRegistry, provider_manager?: ProviderManager) {
    this.protocol = new WorkerProtocol()
    this.tool_registry = tool_registry
    this.provider_manager = provider_manager
  }

  /**
   * Set execution context (session_id, project_id, project_root).
   * Required for tool/LLM calls from workers.
   */
  set_context(ctx: { session_id: string; project_id: string; project_root: string }): void {
    this.execution_context = ctx
  }

  /**
   * Set tool registry (can be injected after construction).
   */
  set_tool_registry(registry: ToolRegistry): void {
    this.tool_registry = registry
  }

  /**
   * Set provider manager (can be injected after construction).
   */
  set_provider_manager(pm: ProviderManager): void {
    this.provider_manager = pm
  }

  /**
   * Spawn a worker child process and perform handshake.
   * INV-1: worker.ready handshake is a live signal, not a status write.
   */
  async spawn(config: WorkerConfig): Promise<WorkerHandle> {
    const proc = new WorkerProcess()
    const handle: WorkerHandle = {
      worker_id: config.agent_id,
      process: proc,
      config,
      state: 'starting',
      started_at: new Date().toISOString()
    }

    // Spawn worker process using Bun
    const bun_path = this.find_bun()
    const child = spawn(bun_path, ['run', config.entrypoint], {
      stdio: ['pipe', 'pipe', 'pipe'],
      env: {
        ...process.env,
        ...config.env,
        AIRCODING_AGENT_ID: config.agent_id,
        AIRCODING_SESSION_ID: config.session_id,
        AIRCODING_PROJECT_ROOT: config.project_root
      },
      cwd: config.project_root
    })

    proc.set_process(child)

    // Set up handlers for worker IPC messages
    this.setup_worker_handlers(proc, config.agent_id)

    // Wait for handshake: worker.ready
    await this.wait_for_handshake(proc, config)

    // Validate protocol version and dispatch task
    const ready_msg = this.send_and_wait(proc, 'agent.start', {
      protocol_version: this.protocol.get_version(),
      agent_id: config.agent_id,
      session_id: config.session_id,
      project_root: config.project_root,
      task_type: config.task_type || 'execute',
      task_spec: config.task_spec || { id: `${config.agent_id}_task`, title: 'Execute task', description: '' }
    })

    handle.state = 'ready'
    this.workers.set(config.agent_id, handle)

    // Set up timeout
    if (config.timeout_ms) {
      setTimeout(() => this.cancel(config.agent_id, 'timeout'), config.timeout_ms)
    }

    return handle
  }

  /**
   * Cancel a worker.
   */
  async cancel(agent_id: string, reason: string): Promise<void> {
    const handle = this.workers.get(agent_id)
    if (!handle) return

    const msg = this.protocol.create_message('agent.cancel', { reason }, 'parent_to_worker')
    handle.process.send(msg)
    handle.state = 'cancelled'

    // Wait briefly then force kill
    setTimeout(() => {
      if (handle.process.is_alive()) {
        handle.process.kill('SIGKILL')
      }
    }, 5000)
  }

  /**
   * Set up IPC message handlers for a worker process.
   * Handles tool.call and llm.request forwarded from worker to parent.
   */
  private setup_worker_handlers(proc: WorkerProcess, agent_id: string): void {
    // Handle tool.call from worker → execute via ToolRegistry
    proc.on_message('tool.call', async (msg) => {
      const call_id = msg.payload.call_id as string
      const tool_name = msg.payload.name as string
      const tool_args = (msg.payload.arguments || {}) as Record<string, unknown>

      try {
        if (!this.tool_registry) {
          this.send_to_worker(agent_id, 'tool.result', {
            call_id,
            type: 'error',
            content: { message: 'ToolRegistry not available' }
          })
          return
        }

        const result = await this.tool_registry.call(
          { call_id, name: tool_name, arguments: tool_args },
          {
            session_id: this.execution_context?.session_id || msg.session_id,
            project_id: this.execution_context?.project_id || '',
            agent_id,
            permission_template: 'executor',
            cwd: this.execution_context?.project_root || process.cwd()
          } as any
        )

        this.send_to_worker(agent_id, 'tool.result', {
          call_id,
          type: result.status === 'ok' ? 'text' : 'error',
          content: result.output || result.error || {}
        })
      } catch (e: any) {
        this.send_to_worker(agent_id, 'tool.result', {
          call_id,
          type: 'error',
          content: { message: e.message || 'Tool execution failed' }
        })
      }
    })

    // Handle llm.request from worker → execute via ProviderManager
    proc.on_message('llm.request', async (msg) => {
      const call_id = msg.payload.call_id as string
      try {
        if (!this.provider_manager) {
          this.send_to_worker(agent_id, 'llm.response', {
            call_id,
            content: '[No provider configured]',
            usage: undefined
          })
          return
        }

        const messages = (msg.payload.messages || []) as Array<{ role: string; content: unknown }>
        const response = await this.provider_manager.complete_text(messages, {
          model: msg.payload.model as string,
          max_tokens: msg.payload.max_tokens as number,
          temperature: msg.payload.temperature as number
        })

        this.send_to_worker(agent_id, 'llm.response', {
          call_id,
          content: response.content,
          usage: response.usage
        })
      } catch (e: any) {
        this.send_to_worker(agent_id, 'llm.response', {
          call_id,
          content: `[LLM Error: ${e.message}]`,
          usage: undefined
        })
      }
    })

    // Handle worker.result → update handle
    proc.on_message('worker.result', (msg) => {
      const handle = this.workers.get(agent_id)
      if (handle) {
        handle.state = 'completed'
        handle.result = this.wrap_worker_result(msg.payload, handle)
      }
    })

    // Handle worker.checkpoint
    proc.on_message('worker.checkpoint', (msg) => {
      const handle = this.workers.get(agent_id)
      if (handle) {
        handle.state = 'running'
      }
    })
  }

  /**
   * Send a message back to a specific worker.
   */
  private send_to_worker(agent_id: string, type: string, payload: Record<string, unknown>): void {
    const handle = this.workers.get(agent_id)
    if (!handle) return
    const msg = this.protocol.create_message(type as WorkerMessageType, payload, 'parent_to_worker', {
      session_id: this.execution_context?.session_id,
      agent_id
    })
    handle.process.send(msg)
  }

  /**
   * Send a message to a worker.
   */
  send(agent_id: string, type: WorkerMessageType, payload: Record<string, unknown>): void {
    const handle = this.workers.get(agent_id)
    if (!handle) throw new Error(`Worker not found: ${agent_id}`)

    const msg = this.protocol.create_message(type, payload, 'parent_to_worker')
    handle.process.send(msg)
  }

  /**
   * Get a worker handle.
   */
  get(agent_id: string): WorkerHandle | undefined {
    return this.workers.get(agent_id)
  }

  /**
   * List all workers.
   */
  list(): WorkerHandle[] {
    return Array.from(this.workers.values())
  }

  /**
   * List workers by state.
   */
  list_by_state(state: WorkerHandle['state']): WorkerHandle[] {
    return this.list().filter(w => w.state === state)
  }

  /**
   * Check if any workers are running.
   */
  has_running(): boolean {
    return this.list().some(w => w.state === 'running' || w.state === 'ready' || w.state === 'starting')
  }

  /**
   * Get the stored WorkerResult for an agent.
   */
  get_result(agent_id: string): WorkerResult<unknown> | undefined {
    const handle = this.workers.get(agent_id)
    if (!handle) return undefined
    return handle.result
  }

  // ============================================================================
  // Private
  // ============================================================================

  /**
   * Wrap a raw worker payload into a properly typed WorkerResult envelope.
   * Provides safe defaults for any missing fields.
   */
  private wrap_worker_result(payload: Record<string, unknown>, handle: WorkerHandle): WorkerResult<unknown> {
    return {
      task_id: (payload.task_id as string) || '' as any,
      agent_id: (payload.agent_id as string) || handle.config.agent_id as any,
      agent_type: (payload.agent_type as AgentType) || 'executor',
      status: (payload.status as WorkerStatus) || 'completed',
      summary: (payload.summary as string) || '',
      changed_files: (payload.changed_files as string[]) || [],
      diff_ref: (payload.diff_ref as string | undefined) || undefined,
      artifacts: (payload.artifacts as any[]) || [],
      verification: (payload.verification as any[]) || [],
      risks: (payload.risks as any[]) || [],
      follow_up_tasks: (payload.follow_up_tasks as any[]) || [],
      evidence_refs: (payload.evidence_refs as any[]) || [],
      result: (payload.result as unknown) || null,
    }
  }

  private async wait_for_handshake(proc: WorkerProcess, config: WorkerConfig): Promise<void> {
    return new Promise((resolve, reject) => {
      const timeout = setTimeout(() => {
        reject(new Error(`Worker handshake timeout: ${config.agent_id}`))
      }, 30000)

      proc.on_message('worker.ready', (msg) => {
        clearTimeout(timeout)
        const version = msg.payload.protocol_version as number
        const check = this.protocol.check_version(version)
        if (!check.compatible) {
          reject(new Error(check.error))
          return
        }
        resolve()
      })

      // Also handle worker.error
      proc.on_message('worker.error', (msg) => {
        clearTimeout(timeout)
        reject(new Error(`Worker error during handshake: ${msg.payload.message}`))
      })
    })
  }

  private async send_and_wait(
    proc: WorkerProcess,
    type: WorkerMessageType,
    payload: Record<string, unknown>
  ): Promise<WorkerMessage> {
    const msg = this.protocol.create_message(type, payload, 'parent_to_worker')

    return new Promise((resolve, reject) => {
      // Wait for worker.send callback (worker acknowledges agent.start)
      // For now just send and resolve after a short delay
      proc.send(msg)
      setTimeout(() => resolve(msg), 100)
    })
  }

  private find_bun(): string {
    // Try common paths first (no shell, no string interpolation)
    const candidates = [
      process.env.BUN_INSTALL ? `${process.env.BUN_INSTALL}/bin/bun` : null,
      `${process.env.HOME || '/root'}/.bun/bin/bun`,
      '/usr/local/bin/bun',
      '/usr/bin/bun',
    ].filter((p): p is string => Boolean(p))

    for (const candidate of candidates) {
      if (existsSync(candidate)) return candidate
    }
    return 'bun'  // PATH fallback
  }
}