Repository Public
1/** Test kit of the engine: throw-away database, workspaces, in-process worker, polling helpers. */2import { execFileSync, spawn, type ChildProcess } from 'node:child_process';3import { randomBytes } from 'node:crypto';4import fs from 'node:fs';5import os from 'node:os';6import path from 'node:path';7import { fileURLToPath } from 'node:url';8import { createKernel, createMission, migrate, requestCancel, runWorkerLoop, type Kernel, type WorkerLoopOptions } from '@aos/kernel';9import type { ProviderAdapter } from '@aos/providers';10import type { DoneCriterion, ResultEnvelope } from '@aos/shared';11import { agentLogic, type EngineDeps } from '../src/index.js';1213export const HERE = path.dirname(fileURLToPath(import.meta.url));14const PKG = path.join(HERE, '..');1516export interface TestDb {17 name: string;18 url: string;19 k: Kernel;20 ownerId: string;21 projectId: string;22 /** Real path of the temp root; workspaces live under `<tmp>/root`. */23 tmp: string;24 root: string;25 cleanup(): Promise<void>;26}2728function dropDb(name: string): void {29 try {30 execFileSync('dropdb', ['--force', '--if-exists', name], { stdio: 'pipe' });31 } catch {32 /* best effort */33 }34}3536export async function createTestDb(): Promise<TestDb> {37 const name = `aos_test_${randomBytes(6).toString('hex')}`;38 execFileSync('createdb', [name], { stdio: 'pipe' });39 const url = `postgres:///${name}?host=/tmp`;40 const tmp = fs.mkdtempSync(path.join(fs.realpathSync(os.tmpdir()), 'aos-engine-test-'));41 const root = path.join(tmp, 'root');42 fs.mkdirSync(root);43 try {44 await migrate(url);45 const k = createKernel({ databaseUrl: url, dataDir: path.join(tmp, 'data'), poolMax: 12 });46 const u = await k.pool.query(`INSERT INTO users (username, password_hash) VALUES ('owner', 'x') RETURNING id`);47 const ownerId = (u.rows[0] as { id: string }).id;48 const p = await k.pool.query(`INSERT INTO projects (owner_id, name, rel_path) VALUES ($1, 'test project', 'p0') RETURNING id`, [ownerId]);49 const projectId = (p.rows[0] as { id: string }).id;50 const runIds = async (): Promise<string[]> => ((await k.pool.query('SELECT id FROM runs')).rows as { id: string }[]).map((r) => r.id);51 return {52 name, url, k, ownerId, projectId, tmp, root,53 cleanup: async () => {54 // per-run sandbox temp dirs of runs that never reached the engine's own cleanup (waiting, cancelled by the kernel)55 for (const id of await runIds().catch(() => [] as string[])) fs.rmSync(path.join(fs.realpathSync(os.tmpdir()), `aos-run-${id}`), { recursive: true, force: true });56 await k.close().catch(() => undefined);57 dropDb(name);58 fs.rmSync(tmp, { recursive: true, force: true });59 },60 };61 } catch (err) {62 dropDb(name);63 fs.rmSync(tmp, { recursive: true, force: true });64 throw err;65 }66}6768export function git(cwd: string, ...args: string[]): string {69 return execFileSync('git', args, { cwd, encoding: 'utf8', env: { PATH: process.env.PATH ?? '', HOME: os.homedir(), GIT_TERMINAL_PROMPT: '0', LANG: 'C' } });70}7172let wsCount = 0;73export function makeWorkspace(db: TestDb, opts: { git?: boolean; files?: Record<string, string>; name?: string } = {}): string {74 const ws = path.join(db.root, opts.name ?? `ws${++wsCount}-${randomBytes(3).toString('hex')}`);75 fs.mkdirSync(ws, { recursive: true });76 for (const [rel, content] of Object.entries(opts.files ?? {})) {77 fs.mkdirSync(path.dirname(path.join(ws, rel)), { recursive: true });78 fs.writeFileSync(path.join(ws, rel), content);79 }80 if (opts.git) {81 git(ws, 'init', '-q', '-b', 'main', '.');82 git(ws, 'config', 'user.name', 'Test');83 git(ws, 'config', 'user.email', 'test@example.invalid');84 git(ws, 'add', '-A');85 git(ws, 'commit', '-q', '--allow-empty', '-m', 'init');86 }87 return ws;88}8990export interface MissionRef {91 missionId: string;92 runId: string;93 ws: string;94}95export async function newMission(db: TestDb, ws: string, over: Record<string, unknown> = {}): Promise<MissionRef> {96 const { missionId, rootRunId } = await createMission(97 {98 ownerId: db.ownerId, projectId: db.projectId, objective: 'Test objective', doneCriteria: ['It works'], provider: 'anthropic', model: 'claude-haiku-5-5',99 budgetUsd: 5, workspaceAbs: ws, ...over,100 } as Parameters<typeof createMission>[0],101 db.k,102 );103 return { missionId, runId: rootRunId, ws };104}105106export interface TestWorker {107 stop(): Promise<void>;108}109/** In-process worker running the real engine with fast timings. */110export function startEngine(db: TestDb, provider: ProviderAdapter, deps: Partial<EngineDeps> = {}, over: Partial<WorkerLoopOptions> = {}): TestWorker {111 const ac = new AbortController();112 const logic = agentLogic({ kernel: db.k, provider: () => provider, workspaceRoot: db.root, denyOwnerFolders: false, extraDenyRead: [], ...deps });113 const done = runWorkerLoop({114 kernel: db.k, logic, signal: ac.signal, concurrency: 4, leaseMs: 3000, heartbeatMs: 300, pollMs: 30, inlineSleepMs: 100,115 janitorMs: false, dailyBudgetUsd: null, stopGraceMs: 5000, workerId: `w-${randomBytes(3).toString('hex')}`, ...over,116 });117 let stopped = false;118 return {119 stop: async () => {120 if (stopped) return;121 stopped = true;122 ac.abort();123 await done;124 },125 };126}127128export const sleep = (ms: number): Promise<void> => new Promise((r) => setTimeout(r, ms));129export async function waitFor<T>(fn: () => Promise<T | null | undefined | false>, timeoutMs = 30_000, what = 'condition'): Promise<T> {130 const end = Date.now() + timeoutMs;131 for (;;) {132 const v = await fn();133 if (v) return v;134 if (Date.now() > end) throw new Error(`timeout waiting for ${what}`);135 await sleep(20);136 }137}138139export interface RunView {140 status: string;141 result: ResultEnvelope | null;142 usage: Record<string, number>;143 prompt_hash: string | null;144}145export async function rows<T = Record<string, unknown>>(db: TestDb, sql: string, params: unknown[] = []): Promise<T[]> {146 return (await db.k.pool.query(sql, params)).rows as T[];147}148export async function getRun(db: TestDb, runId: string): Promise<RunView> {149 return (await rows<RunView>(db, 'SELECT status, result, usage, prompt_hash FROM runs WHERE id = $1', [runId]))[0] as RunView;150}151export const TERMINAL = ['done', 'partial', 'blocked', 'failed', 'cancelled', 'budget_exhausted'];152export async function waitStatus(db: TestDb, runId: string, statuses: string[] = TERMINAL, timeoutMs = 30_000): Promise<RunView> {153 return waitFor(async () => {154 const r = await getRun(db, runId);155 return statuses.includes(r.status) ? r : null;156 }, timeoutMs, `run ${runId} to be ${statuses.join('|')}`);157}158export interface StepView {159 seq: number;160 kind: string;161 name: string;162 effect: string;163 status: string;164 request: Record<string, unknown>;165 result: Record<string, unknown> | null;166 error: { name: string; message: string } | null;167}168export async function steps(db: TestDb, runId: string): Promise<StepView[]> {169 return rows(db, 'SELECT seq, kind, name, effect, status, request, result, error FROM steps WHERE run_id = $1 ORDER BY seq', [runId]);170}171export async function toolSteps(db: TestDb, runId: string): Promise<StepView[]> {172 return (await steps(db, runId)).filter((s) => s.kind === 'tool_call');173}174export interface WaitView {175 id: string;176 kind: string;177 payload: Record<string, unknown>;178 step_seq: number;179}180export async function pendingWait(db: TestDb, runId: string, timeoutMs = 30_000): Promise<WaitView> {181 return waitFor(async () => (await rows<WaitView>(db, `SELECT id, kind, payload, step_seq FROM waits WHERE run_id = $1 AND status = 'pending' ORDER BY created_at LIMIT 1`, [runId]))[0], timeoutMs, 'a pending wait');182}183export async function events(db: TestDb, missionId: string, kind: string): Promise<Record<string, unknown>[]> {184 return (await rows<{ payload: Record<string, unknown> }>(db, 'SELECT payload FROM events WHERE mission_id = $1 AND kind = $2 ORDER BY id', [missionId, kind])).map((r) => r.payload);185}186export async function criteriaOf(db: TestDb, missionId: string): Promise<DoneCriterion[]> {187 return (await rows<{ done_criteria: DoneCriterion[] }>(db, 'SELECT done_criteria FROM missions WHERE id = $1', [missionId]))[0]?.done_criteria ?? [];188}189export async function cancelQuietly(db: TestDb, missionId: string): Promise<void> {190 await requestCancel(missionId, db.k).catch(() => undefined);191}192export const read = (ws: string, rel: string): string => fs.readFileSync(path.join(ws, rel), 'utf8');193export const exists = (ws: string, rel: string): boolean => fs.existsSync(path.join(ws, rel));194195// ---- child processes (crash tests) -----------------------------------------------------------------------196const children = new Set<ChildProcess>();197export function spawnFixture(file: string, env: Record<string, string>): ChildProcess {198 // `node --import tsx` keeps the fixture in ONE process (the tsx CLI would fork a child that survives a SIGKILL).199 const child = spawn(process.execPath, ['--import', 'tsx', path.join(HERE, 'fixtures', file)], {200 cwd: PKG,201 env: { ...process.env, AOS_LOG_LEVEL: 'silent', ...env },202 stdio: ['ignore', 'ignore', 'inherit'],203 });204 children.add(child);205 child.once('exit', () => children.delete(child));206 return child;207}208export function exited(child: ChildProcess): Promise<{ code: number | null; signal: NodeJS.Signals | null }> {209 if (child.exitCode !== null || child.signalCode !== null) return Promise.resolve({ code: child.exitCode, signal: child.signalCode });210 return new Promise((resolve) => child.once('exit', (code, signal) => resolve({ code, signal })));211}212export async function killAllChildren(): Promise<void> {213 const all = [...children];214 for (const c of all) c.kill('SIGKILL');215 await Promise.all(all.map(exited));216}217