Blame for source/study/pi-runner/runner.mjs

import { fork } from 'node:child_process';
import { appendFileSync } from 'node:fs';
import { mkdir, readFile, writeFile } from 'node:fs/promises';
import { resolve, join, dirname } from 'node:path';
import { fileURLToPath, pathToFileURL } from 'node:url';
import { ContainerBackend } from './container.mjs';

const HERE = dirname(fileURLToPath(import.meta.url));
const textResult = value => ({ content: [{ type: 'text', text: typeof value === 'string' ? value : JSON.stringify(value) }] });
const errorText = error => error instanceof Error ? error.message : String(error);

export function validateConfig(config) {
  for (const key of ['condition', 'run_id', 'output_dir', 'proxy_url', 'task_prompt', 'container_name']) {
    if (typeof config[key] !== 'string' || !config[key]) throw new Error(`Missing config ${key}`);
  }
  if (!['single', 'peers', 'async'].includes(config.condition)) throw new Error('condition must be single, peers, or async');
  if (config.max_active_agents != null && config.max_active_agents !== 5) throw new Error('Pilot maximum active agents must be five');
  const url = new URL(config.proxy_url);
  if (!['http:', 'https:'].includes(url.protocol) || url.username || url.password) throw new Error('Invalid proxy_url');
  if (config.wall_seconds != null && !(config.wall_seconds > 0)) throw new Error('Invalid wall_seconds');
  return { ...config, output_dir: resolve(config.output_dir), wall_seconds: config.wall_seconds ?? 7200, max_output_tokens: config.max_output_tokens ?? 16000 };
}

export function agentPrompt(config, id, task, parentId) {
  const common = `You are ${id} in a program reconstruction study. Work autonomously until you have completed the task or exhausted your budget.\n` +
    `Use only the bash tool to inspect files, edit code, compile and run tests inside the assigned container. The tool starts each command in your own checkout. Reference executable and supplied documentation are under /reference. No hidden tests or original source are available.\n` +
    `Your checkout is ${(config.work_root ?? '/workspace')}/.pi-study/agents/${id}. The origin remote is shared local Git. Only origin's main branch is graded. Commit your changes, integrate origin/main, and push your accepted submission with git push origin HEAD:main. A rejected non-fast-forward push means fetch and integrate before retrying. Do not force-push. Your final answer alone is not a submission. Create compile.sh producing ./executable as required by the task.\n` +
    `Lifetime context budget: ${config.agent_context_token_cap ?? (config.condition === 'single' ? 1000000 : 200000)} estimated tokens; shared episode budget: ${config.context_token_cap ?? 1000000} estimated context tokens. The proxy enforces model resource limits. Prioritize a working submission early.\n`;
  const mode = config.condition === 'single' ? '' : config.condition === 'peers'
    ? `You are one of five equally capable peers: agent-0, agent-1, agent-2, agent-3, agent-4. There is no lead. All receive the same task. Coordinate using send_message, inspect agent_status, and share work through Git. A final answer means you are idle; later messages can resume you.\n`
    : parentId
      ? `You were delegated this task by ${parentId}. Work on the instructions below; you have not received the lead's original task. You can message any known peer. Your final answer will be sent to the lead and you will remain available for resumption.\n`
      : `You are the lead. You may spawn persistent children for bounded tasks and continue immediately. At most five agents including you may be active at once. Children receive only your delegated instructions and the shared operating rules, not your original task. Use send_message or resume_agent for follow-up work. Integrate child contributions into main.\n`;
  return common + mode + '\nTask:\n' + task;
}

export class Supervisor {
  constructor(config, { backend, workerFactory, eventSink } = {}) {
    this.config = validateConfig(config);
    this.agents = new Map(); this.nextAgent = 0; this.sequence = 0;
    this.backend = backend ?? new ContainerBackend(this.config);
    this.workerFactory = workerFactory ?? createProcessWorker;
    this.eventSink = eventSink; this.snapshots = []; this.snapshotQueue = Promise.resolve();
    this.started = false; this.stopping = false; this.timers = [];
    this.finished = new Promise(resolve => { this.resolveFinished = resolve; });
  }
  event(type, data = {}) {
    const event = { ...data, type, sequence: ++this.sequence, run_id: this.config.run_id, utc: new Date().toISOString(), monotonic_ms: performance.now(), elapsed_ms: this.startTime == null ? 0 : performance.now() - this.startTime };
    if (this.eventSink) this.eventSink(event);
    else appendFileSync(join(this.config.output_dir, 'events.jsonl'), JSON.stringify(event) + '\n', { mode: 0o600 });
    return event;
  }
  activeCount() { return [...this.agents.values()].filter(a => ['queued', 'working'].includes(a.state)).length; }
  status() { return [...this.agents.values()].map(({ id, parentId, state, stopReason }) => ({ id, parent_id: parentId, state, stop_reason: stopReason })); }
  reserveAgent(task, parentId = null) {
    if (this.stopping) throw new Error('Episode is stopping');
    if (this.activeCount() >= 5) throw new Error('Five active agents already; wait for a result or resume an idle agent later');
    const id = `agent-${this.nextAgent++}`;
    const agent = { id, parentId, task, state: 'queued', worker: null, pending: [], waiter: null, controller: new AbortController() };
    this.agents.set(id, agent);
    this.event('agent_created', { agent_id: id, parent_id: parentId });
    agent.startPromise = this.launch(agent);
    return id;
  }
  async launch(agent) {
    try {
      await this.backend.createAgent(agent.id);
      if (this.stopping) return;
      const directory = join(this.config.output_dir, 'agents', agent.id);
      await mkdir(directory, { recursive: true });
      agent.worker = await this.workerFactory({ config: this.config, id: agent.id, parentId: agent.parentId, directory }, event => this.onWorkerEvent(agent, event));
      if (this.stopping || agent.state === 'terminated') { await agent.worker.dispose(); return; }
      agent.state = 'working'; this.event('agent_started', { agent_id: agent.id });
      agent.worker.prompt(agentPrompt(this.config, agent.id, agent.task, agent.parentId));
      for (const message of agent.pending.splice(0)) this.deliver(agent, message);
    } catch (error) { this.failAgent(agent, errorText(error)); }
  }
  failAgent(agent, message) {
    if (this.stopping || agent.state === 'terminated') return;
    const budgetReason = message.match(/\b(?:(?:agent|episode|global|category)_[a-z_]*budget[a-z_]*|daily_quota_headroom)\b/)?.[0];
    agent.state = 'terminated'; agent.stopReason = budgetReason ?? (/budget|token cap|resource limit|allocation exhausted/i.test(message) ? 'budget_exhausted' : 'agent_error');
    agent.controller.abort();
    void agent.worker?.dispose();
    this.event('agent_failed', { agent_id: agent.id, stop_reason: agent.stopReason, error: message });
    if (budgetReason?.startsWith('episode_') || budgetReason?.startsWith('global_') || budgetReason?.startsWith('category_') || budgetReason === 'daily_quota_headroom') { void this.stop(budgetReason); return; }
    if (this.config.condition === 'async' && agent.id === 'agent-0') { void this.stop(agent.stopReason); return; }
    if (agent.parentId) this.send(agent.id, agent.parentId, `Agent failed: ${message}`, 'child_result');
    this.maybeFinish();
  }
  onWorkerEvent(agent, event) {
    if (event.type === 'tool_request') {
      void this.tool(agent, event.name, event.args).then(result => agent.worker?.reply(event.request_id, result), error => agent.worker?.reply(event.request_id, textResult(errorText(error)), true));
      return;
    }
    if (event.type === 'turn_done') {
      if (this.stopping || agent.state === 'terminated') return;
      if (event.error) { this.failAgent(agent, event.error); return; }
      agent.state = 'idle'; this.event('agent_idle', { agent_id: agent.id, answer: event.answer });
      if (agent.parentId) this.send(agent.id, agent.parentId, event.answer || '(Agent completed without text.)', 'child_result');
      this.maybeFinish(); return;
    }
    if (event.type === 'worker_error') { this.failAgent(agent, event.error); return; }
    this.event(event.type, { ...event, agent_id: agent.id });
  }
  send(senderId, targetId, content, kind = 'peer_message') {
    const target = this.agents.get(targetId);
    if (!target || target.state === 'terminated') throw new Error(`Agent unavailable: ${targetId}`);
    if (target.state === 'idle' && this.activeCount() >= 5) throw new Error('Cannot resume recipient while five agents are active');
    const sent = this.event('message_sent', { agent_id: senderId, recipient_id: targetId, content, kind });
    const message = { sender_id: senderId, content, kind, causal_event_id: sent.sequence };
    if (!target.worker) target.pending.push(message);
    else this.deliver(target, message);
    return { delivered_to: targetId, message_id: sent.sequence };
  }
  deliver(target, message) {
    if (this.stopping) return;
    this.event(target.waiter ? 'message_delivered' : 'message_queued', { agent_id: target.id, ...message, delivery: target.waiter ? 'waiting_tool' : target.state === 'idle' ? 'new_turn' : 'steer_queue' });
    if (target.waiter) { const waiter = target.waiter; target.waiter = null; clearTimeout(waiter.timer); waiter.resolve(textResult(message)); }
    else { target.state = 'working'; target.worker.deliver(`Message from ${message.sender_id} (${message.kind}; event ${message.causal_event_id}):\n${message.content}`); }
  }
  async tool(agent, name, args) {
    if (this.stopping || agent.state === 'terminated') throw new Error('Agent is stopped');
    if (name === 'bash') {
      const start = this.event('tool_start', { agent_id: agent.id, tool: name, command: args.command });
      try {
        const result = await this.backend.bash(agent.id, args.command, { signal: agent.controller.signal, timeout: args.timeout });
        this.event('tool_end', { agent_id: agent.id, tool: name, causal_event_id: start.sequence, duration_ms: performance.now() - start.monotonic_ms, exit_code: result.exitCode, output: result.output });
        void this.snapshot('main_update');
        return { ...textResult(result.output || '(no output)'), isError: result.exitCode !== 0 };
      } catch (error) { this.event('tool_end', { agent_id: agent.id, tool: name, causal_event_id: start.sequence, duration_ms: performance.now() - start.monotonic_ms, error: errorText(error) }); throw error; }
    }
    if (name === 'agent_status') return textResult(this.status());
    if (name === 'send_message') return textResult(this.send(agent.id, args.agent_id, args.message));
    if (name === 'wait_for_message') {
      if (agent.waiter) throw new Error('Agent already has a pending wait_for_message tool');
      return new Promise(resolve => {
        const timer = setTimeout(() => { agent.waiter = null; resolve(textResult('No message arrived before timeout.')); }, Math.min(Math.max(args.timeout_seconds ?? 30, 1), 120) * 1000);
        agent.waiter = { resolve, timer };
      });
    }
    if (this.config.condition !== 'async' || agent.id !== 'agent-0') throw new Error('Only the async lead may manage agents');
    if (name === 'spawn_agent') return textResult({ agent_id: this.reserveAgent(args.task, agent.id), state: 'queued' });
    if (name === 'resume_agent') return textResult(this.send(agent.id, args.agent_id, args.message, 'resume'));
    if (name === 'terminate_agent') {
      const target = this.agents.get(args.agent_id);
      if (!target || target.id === agent.id) throw new Error('Invalid target');
      target.state = 'terminated'; target.stopReason = 'lead_terminated'; target.controller.abort();
      if (target.waiter) { clearTimeout(target.waiter.timer); target.waiter.resolve(textResult('Agent terminated')); target.waiter = null; }
      await target.worker?.abort(); await target.worker?.dispose(); this.event('agent_terminated', { agent_id: target.id, sender_id: agent.id });
      return textResult({ terminated: target.id });
    }
    throw new Error(`Unknown tool: ${name}`);
  }
  snapshot(reason) {
    this.snapshotQueue = this.snapshotQueue.then(async () => {
      if (!this.prepared) return;
      try {
        const revision = this.backend.revision ? await this.backend.revision() : null;
        if (reason === 'main_update' && revision && revision === this.snapshots.at(-1)?.revision) return;
        const path = join(this.config.output_dir, 'snapshots', `${String(this.snapshots.length).padStart(4, '0')}-${reason}.tar.gz`);
        const saved = await this.backend.snapshot(path, revision);
        const event = this.event('snapshot', { ...saved, reason });
        this.snapshots.push({ ...saved, reason, event_sequence: event.sequence, elapsed_ms: event.elapsed_ms });
      } catch (error) { this.event('snapshot_error', { reason, error: errorText(error) }); this.snapshotError = errorText(error); }
    });
    return this.snapshotQueue;
  }
  maybeFinish() {
    if (!this.started || this.stopping || this.activeCount()) return;
    if (!this.finishTimer) this.finishTimer = setTimeout(() => {
      this.finishTimer = null;
      if (!this.activeCount()) void this.stop([...this.agents.values()].some(a => a.stopReason === 'agent_error') ? 'agent_error' : [...this.agents.values()].some(a => a.stopReason?.includes('budget')) ? 'budget_exhausted' : 'completed');
    }, 50);
  }
  async run() {
    await mkdir(this.config.output_dir, { recursive: true });
    this.startTime = performance.now(); this.event('episode_start', { condition: this.config.condition });
    this.timers.push(setTimeout(() => { void this.stop('wall_time_limit'); }, this.config.wall_seconds * 1000));
    try {
      await this.backend.prepare(); this.prepared = true;
      if (this.stopping) return this.finished;
      await this.snapshot('initial');
      const count = this.config.condition === 'peers' ? 5 : 1;
      for (let i = 0; i < count; i++) this.reserveAgent(this.config.task_prompt);
      this.started = true;
      for (const seconds of this.config.snapshot_seconds ?? [30, 60, 120, 300, 600, 1200, 2400, 4800, 7200]) {
        if (seconds < this.config.wall_seconds) this.timers.push(setTimeout(() => { void this.snapshot('scheduled'); }, seconds * 1000));
      }
      this.maybeFinish();
    } catch (error) { this.event('episode_error', { error: errorText(error) }); void this.stop('setup_error'); }
    return this.finished;
  }
  async stop(reason) {
    if (this.stopping) return this.finished;
    this.stopping = true;
    for (const timer of this.timers) clearTimeout(timer);
    clearTimeout(this.finishTimer);
    this.event('episode_stopping', { stop_reason: reason });
    for (const agent of this.agents.values()) {
      agent.controller.abort();
      if (['queued', 'working'].includes(agent.state)) { agent.state = 'terminated'; agent.stopReason ??= reason; }
      if (agent.waiter) { clearTimeout(agent.waiter.timer); agent.waiter.resolve(textResult('Episode stopped')); agent.waiter = null; }
    }
    await Promise.allSettled([...this.agents.values()].map(a => a.worker?.abort()));
    await this.backend.stop();
    await Promise.allSettled([...this.agents.values()].map(a => a.startPromise));
    await Promise.allSettled([...this.agents.values()].map(a => a.worker?.dispose()));
    await this.snapshot('final');
    const result = { run_id: this.config.run_id, condition: this.config.condition, stop_reason: reason, elapsed_seconds: (performance.now() - this.startTime) / 1000, agents: this.status(), snapshots: this.snapshots, snapshot_error: this.snapshotError ?? null };
    this.event('episode_end', result);
    await writeFile(join(this.config.output_dir, 'result.json'), JSON.stringify(result, null, 2) + '\n', { mode: 0o600 });
    this.resolveFinished(result);
    return result;
  }
}

export async function createProcessWorker(meta, onEvent) {
  const safeEnv = { PATH: process.env.PATH, LANG: 'C.UTF-8', HOME: meta.directory, PI_STUDY_API_KEY: process.env.PI_STUDY_API_KEY ?? 'study-proxy' };
  const child = fork(join(HERE, 'worker.mjs'), [], { env: safeEnv, stdio: ['ignore', 'ignore', 'pipe', 'ipc'] });
  let expectedExit = false, initialized = false, readyResolve, readyReject;
  const ready = new Promise((r, j) => { readyResolve = r; readyReject = j; });
  const timer = setTimeout(() => { child.kill('SIGKILL'); readyReject(new Error('Pi worker startup timed out')); }, 60000);
  child.on('message', message => {
    if (message.type === 'ready') { initialized = true; clearTimeout(timer); readyResolve(); }
    else if (message.type === 'worker_error' && !initialized) { clearTimeout(timer); expectedExit = true; child.kill('SIGKILL'); readyReject(new Error(message.error)); }
    else onEvent(message);
  });
  child.stderr.on('data', chunk => onEvent({ type: 'worker_stderr', text: chunk.toString().slice(0, 10000) }));
  child.on('error', readyReject);
  child.on('exit', (code, signal) => {
    clearTimeout(timer); readyReject(new Error(`Pi worker exited ${code ?? signal}`));
    if (!expectedExit) onEvent({ type: 'worker_error', error: `Pi worker exited ${code ?? signal}` });
  });
  const send = message => { if (child.connected) child.send(message); };
  send({ type: 'init', ...meta });
  await ready;
  return {
    prompt: message => send({ type: 'prompt', message }),
    deliver: message => send({ type: 'prompt', message }),
    reply: (request_id, result, isError = false) => send({ type: 'tool_result', request_id, result: { ...result, isError: isError || result.isError } }),
    abort: async () => send({ type: 'abort' }),
    dispose: async () => {
      expectedExit = true;
      if (child.exitCode !== null || child.signalCode) return;
      await new Promise(resolve => {
        const timeout = setTimeout(() => child.kill('SIGKILL'), 5000);
        child.once('exit', () => { clearTimeout(timeout); resolve(); });
        send({ type: 'dispose' });
      });
    },
  };
}

if (process.argv[1] && import.meta.url === pathToFileURL(resolve(process.argv[1])).href) {
  let supervisor;
  try {
    const config = JSON.parse(await readFile(process.argv[2], 'utf8'));
    supervisor = new Supervisor(config);
    for (const signal of ['SIGINT', 'SIGTERM']) process.on(signal, () => { void supervisor.stop('interrupted'); });
    const result = await supervisor.run();
    process.stdout.write(JSON.stringify(result) + '\n');
    if (['agent_error', 'setup_error'].includes(result.stop_reason) || result.snapshot_error) process.exitCode = 1;
  } catch (error) { process.stderr.write(errorText(error) + '\n'); process.exitCode = 1; }
}