kilo-loop/test/orchestrator-guard.test.mjs

import { test } from 'node:test'
import assert from 'node:assert/strict'
import fs from 'node:fs'
import os from 'node:os'
import path from 'node:path'

import { Orchestrator } from '../src/orchestrator.mjs'

const sink = { write: () => true }

function makeProject() {
  const project = fs.mkdtempSync(path.join(os.tmpdir(), 'kilo-guard-orch-'))
  fs.mkdirSync(path.join(project, '.kilocode-loop'), { recursive: true })
  fs.writeFileSync(path.join(project, 'goal.json'), JSON.stringify({ title: 'Guard goal', stages: [{ id: 's1', title: 'Do work' }] }))
  return project
}

function makeConfig(project, overrides = {}) {
  return {
    project,
    goal: 'goal.json',
    goalText: '',
    iterations: 3,
    agent: 'code-design',
    model: '',
    variant: '',
    auto: true,
    sharedContext: false,
    hitl: 'off',
    confirm: false,
    port: 7997,
    host: '127.0.0.1',
    dryRun: true,
    keepOpen: false,
    quiet: true,
    color: false,
    promptExtra: '',
    maxIterationMinutes: 0,
    contextWarnTokens: 150000,
    verify: '',
    verifyTimeoutMinutes: 30,
    maxCost: 0,
    maxTokens: 0,
    maxStaleIterations: 0,
    checkpoint: false,
    guard: true,
    guardStop: false,
    maxGuardBlocks: 3,
    redact: true,
    review: false,
    reviewAgent: '',
    autoFreshTokens: 0,
    notify: '',
    syncGoal: true,
    ...overrides,
  }
}

/**
 * Runner that emits a destructive bash tool_use only on the requested iteration,
 * so the loop's runtime guard aborts that session but (by default) keeps going.
 */
function guardOnceRunner(guardIterations = [1]) {
  return async ({ prompt, onEvent, signal }) => {
    const iteration = Number(prompt.match(/iteration (\d+)\//)?.[1] || 1)
    if (guardIterations.includes(iteration)) {
      onEvent({ type: 'tool_use', part: { tool: 'bash', state: { input: { command: 'cd /repo && rm -rf src' } } } })
    }
    return {
      sessionID: `ses_iter_${iteration}`,
      exitCode: signal.aborted ? 130 : 0,
      error: null,
      texts: [],
      reasoning: [],
      toolCalls: [],
      stderr: '',
      durationMs: 5,
      interrupted: signal.aborted,
    }
  }
}

test('a guard hit aborts the session but the run continues to the next iteration', async () => {
  const project = makeProject()
  const final = await new Orchestrator(makeConfig(project, { runner: guardOnceRunner([1]), iterations: 2, runId: 'guard-continue' }), { out: sink }).run()

  assert.notEqual(final.status, 'stopped-guard')
  assert.equal(final.stopReason, null)
  assert.equal(final.iterations.length, 2)
  assert.equal(final.iterations[0].guardViolation.id, 'rm-recursive-force')
  // The safety feedback must reach the next iteration's prompt.
  const prompt2 = fs.readFileSync(path.join(project, '.kilocode-loop', 'runs', 'guard-continue', 'iterations', '02-prompt.md'), 'utf8')
  assert.match(prompt2, /Loop safety feedback/)
  assert.match(prompt2, /rm-recursive-force/)
})

test('--guard-stop restores the hard stop on the first hit', async () => {
  const project = makeProject()
  const final = await new Orchestrator(makeConfig(project, { runner: guardOnceRunner(), guardStop: true, runId: 'guard-stop' }), { out: sink }).run()

  assert.equal(final.status, 'stopped-guard')
  assert.equal(final.stopReason, 'guard')
  assert.equal(final.iterations.length, 1)
})

test('a persistently blocked agent stops after maxGuardBlocks iterations', async () => {
  const project = makeProject()
  const final = await new Orchestrator(
    makeConfig(project, { runner: guardOnceRunner([1, 2, 3]), maxGuardBlocks: 2, iterations: 5, runId: 'guard-limit' }),
    { out: sink },
  ).run()

  assert.equal(final.status, 'stopped-guard')
  assert.equal(final.stopReason, 'guard')
  assert.equal(final.iterations.length, 2)
})