source/study/pi-runner/runner.test.mjs

import test from 'node:test';
import assert from 'node:assert/strict';
import { mkdtemp, rm } from 'node:fs/promises';
import { tmpdir } from 'node:os';
import { join } from 'node:path';
import { createServer } from 'node:http';
import { Supervisor, agentPrompt } from './runner.mjs';
import { toolSpecifications, createSession } from './worker.mjs';

async function until(predicate) {
  for (let i = 0; i < 200; i++) { if (predicate()) return; await new Promise(resolve => setTimeout(resolve, 5)); }
  throw new Error('Timed out waiting for test state');
}
async function fixture(t, condition) {
  const output = await mkdtemp(join(tmpdir(), 'pi-study-test-'));
  const events = [], workers = new Map();
  const backend = { prepared: false, stopped: false, paths: [], async prepare() { this.prepared = true; }, async createAgent(id) { this.paths.push(id); }, async bash() { return { output: 'ok', exitCode: 0 }; }, async snapshot(path) { return { path, revision: 'a'.repeat(40) }; }, async stop() { this.stopped = true; } };
  const config = { condition, run_id: 'test', output_dir: output, proxy_url: 'http://127.0.0.1:9999', task_prompt: 'ORIGINAL TASK SENTINEL', container_name: 'unused', wall_seconds: 10, snapshot_seconds: [] };
  const supervisor = new Supervisor(config, { backend, eventSink: event => events.push(event), workerFactory: async (meta, onEvent) => {
    const worker = { meta, prompts: [], deliveries: [], replies: [], aborted: false, disposed: false,
      prompt(message) { this.prompts.push(message); }, deliver(message) { this.deliveries.push(message); }, reply(...args) { this.replies.push(args); },
      async abort() { this.aborted = true; }, async dispose() { this.disposed = true; }, done(answer = 'done', error) { onEvent({ type: 'turn_done', answer, error }); },
    };
    workers.set(meta.id, worker); return worker;
  } });
  t.after(async () => { if (!supervisor.stopping && supervisor.startTime != null) await supervisor.stop('test_cleanup'); await rm(output, { recursive: true, force: true }); });
  return { supervisor, events, workers, backend, config };
}

test('single exposes only container bash and preserves its persistent session until finish', async t => {
  const f = await fixture(t, 'single');
  const result = f.supervisor.run();
  await until(() => f.workers.get('agent-0')?.prompts.length === 1);
  assert.deepEqual(toolSpecifications('single', 'agent-0').map(t => t.name), ['bash']);
  assert.match(f.workers.get('agent-0').prompts[0], /ORIGINAL TASK SENTINEL/);
  f.workers.get('agent-0').done();
  assert.equal((await result).stop_reason, 'completed');
  assert.equal(f.workers.size, 1);
  assert.equal(f.workers.get('agent-0').disposed, true);
});

test('five peers receive full task; idle messages reuse existing session; all must finish', async t => {
  const f = await fixture(t, 'peers');
  const result = f.supervisor.run();
  await until(() => f.workers.size === 5 && [...f.workers.values()].every(w => w.prompts.length));
  for (const worker of f.workers.values()) assert.match(worker.prompts[0], /ORIGINAL TASK SENTINEL/);
  const original = f.workers.get('agent-1'); original.done('first pass');
  assert.equal(f.supervisor.agents.get('agent-1').state, 'idle');
  await f.supervisor.tool(f.supervisor.agents.get('agent-0'), 'send_message', { agent_id: 'agent-1', message: 'please continue' });
  assert.equal(f.workers.get('agent-1'), original);
  assert.match(original.deliveries[0], /please continue/);
  assert.equal(f.supervisor.agents.get('agent-1').state, 'working');
  await assert.rejects(f.supervisor.tool(f.supervisor.agents.get('agent-0'), 'spawn_agent', { task: 'no' }), /Only the async lead/);
  for (const worker of f.workers.values()) worker.done();
  assert.equal((await result).stop_reason, 'completed');
  assert.equal(f.backend.paths.length, 5);
});

test('async spawn returns before setup; task-only child idles and resumes with same identity', async t => {
  const f = await fixture(t, 'async');
  const run = f.supervisor.run();
  await until(() => f.workers.get('agent-0')?.prompts.length);
  let release;
  f.backend.createAgent = id => id === 'agent-1' ? new Promise(r => { release = r; }) : Promise.resolve();
  const lead = f.supervisor.agents.get('agent-0');
  const result = await f.supervisor.tool(lead, 'spawn_agent', { task: 'DELEGATED ONLY' });
  assert.equal(JSON.parse(result.content[0].text).agent_id, 'agent-1');
  assert.equal(f.workers.has('agent-1'), false);
  assert.equal(lead.state, 'working'); release();
  await until(() => f.workers.get('agent-1')?.prompts.length);
  const child = f.workers.get('agent-1');
  assert.match(child.prompts[0], /DELEGATED ONLY/); assert.doesNotMatch(child.prompts[0], /ORIGINAL TASK SENTINEL/);
  child.done('child finding');
  assert.equal(f.supervisor.agents.get('agent-1').state, 'idle');
  assert.match(f.workers.get('agent-0').deliveries[0], /child finding/);
  await f.supervisor.tool(lead, 'resume_agent', { agent_id: 'agent-1', message: 'second assignment' });
  assert.equal(f.workers.get('agent-1'), child);
  assert.match(child.deliveries[0], /second assignment/);
  child.done('again'); f.workers.get('agent-0').done('final');
  assert.equal((await run).stop_reason, 'completed');
});

test('async admission caps active agents at five including lead', async t => {
  const f = await fixture(t, 'async'); f.supervisor.run();
  await until(() => f.workers.get('agent-0')?.prompts.length);
  const lead = f.supervisor.agents.get('agent-0');
  for (let i = 0; i < 4; i++) await f.supervisor.tool(lead, 'spawn_agent', { task: 'child' });
  await assert.rejects(f.supervisor.tool(lead, 'spawn_agent', { task: 'sixth' }), /Five active agents/);
  assert.equal(f.supervisor.activeCount(), 5);
});

test('message wakes waiting tool without duplicate steering injection', async t => {
  const f = await fixture(t, 'peers'); f.supervisor.run();
  await until(() => f.workers.size === 5 && [...f.workers.values()].every(w => w.prompts.length));
  const waiting = f.supervisor.tool(f.supervisor.agents.get('agent-1'), 'wait_for_message', { timeout_seconds: 5 });
  f.supervisor.send('agent-0', 'agent-1', 'wake up');
  assert.match((await waiting).content[0].text, /wake up/);
  assert.equal(f.workers.get('agent-1').deliveries.length, 0);
  assert.equal(f.events.filter(e => e.type === 'message_delivered').length, 1);
});

test('episode budget denial stops all agents and preserves reason and final snapshot', async t => {
  const f = await fixture(t, 'peers'); const run = f.supervisor.run();
  await until(() => f.workers.size === 5 && [...f.workers.values()].every(w => w.prompts.length));
  f.workers.get('agent-1').done('', '402 study_budget_exhausted: episode_context_budget_estimated');
  const result = await run;
  assert.equal(result.stop_reason, 'episode_context_budget_estimated');
  assert.equal(result.snapshots.at(-1).reason, 'final');
  assert.ok([...f.workers.values()].every(w => w.aborted && w.disposed));
  assert.equal(f.backend.stopped, true);
});

test('daily quota headroom denial stops the entire team', async t => {
  const f = await fixture(t, 'peers'); const run = f.supervisor.run();
  await until(() => f.workers.size === 5 && [...f.workers.values()].every(w => w.prompts.length));
  f.workers.get('agent-2').done('', '402 study_budget_exhausted: daily_quota_headroom');
  const result = await run;
  assert.equal(result.stop_reason, 'daily_quota_headroom');
  assert.ok([...f.workers.values()].every(w => w.aborted && w.disposed));
});

test('cancellation aborts in-flight bash and resolves waiting tools', async t => {
  const f = await fixture(t, 'peers'); const run = f.supervisor.run();
  await until(() => f.workers.size === 5 && [...f.workers.values()].every(w => w.prompts.length));
  let cancelled = false;
  f.backend.bash = (_id, _command, { signal }) => new Promise(resolve => signal.addEventListener('abort', () => { cancelled = true; resolve({ output: 'cancelled', exitCode: 143 }); }, { once: true }));
  const bash = f.supervisor.tool(f.supervisor.agents.get('agent-0'), 'bash', { command: 'sleep 100' });
  const wait = f.supervisor.tool(f.supervisor.agents.get('agent-1'), 'wait_for_message', { timeout_seconds: 120 });
  await f.supervisor.stop('interrupted');
  await Promise.all([bash, wait, run]); assert.equal(cancelled, true);
});

test('real Pi SDK session has no native host filesystem tools and no discovered resources', async t => {
  const directory = await mkdtemp(join(tmpdir(), 'pi-sdk-isolation-'));
  t.after(() => rm(directory, { recursive: true, force: true }));
  const sdk = await import('@earendil-works/pi-coding-agent');
  const session = await createSession({ id: 'agent-0', directory, config: { condition: 'async', run_id: 'test', proxy_url: 'http://127.0.0.1:1' } }, sdk);
  t.after(() => session.dispose());
  assert.deepEqual(session.getActiveToolNames().sort(), toolSpecifications('async', 'agent-0').map(t => t.name).sort());
  assert.match(session.agent.state.tools.find(t => t.name === 'bash').description, /assigned isolated container/);
  assert.ok(session.sessionFile.startsWith(directory));
  assert.equal(session.model.baseUrl, 'http://127.0.0.1:1/runs/test/agents/agent-0/v1');
});

test('real SDK forced compaction uses instrumented stream, proxy route and identical Beam settings', async t => {
  const directory = await mkdtemp(join(tmpdir(), 'pi-compaction-test-'));
  t.after(() => rm(directory, { recursive: true, force: true }));
  const requests = [], events = [];
  const server = createServer(async (request, response) => {
    const chunks = []; for await (const chunk of request) chunks.push(chunk);
    const body = JSON.parse(Buffer.concat(chunks).toString());
    requests.push({ path: request.url, headers: request.headers, body });
    response.writeHead(200, { 'content-type': 'text/event-stream' });
    const base = { id: `test-${requests.length}`, object: 'chat.completion.chunk', model: 'Beam-501B-A23B', created: 1 };
    for (const payload of [
      { ...base, choices: [{ index: 0, delta: { role: 'assistant', content: request.headers['x-study-purpose'] === 'compaction' ? 'Summary of previous work.' : 'Done.' }, finish_reason: null }] },
      { ...base, choices: [{ index: 0, delta: {}, finish_reason: 'stop' }] },
      { ...base, choices: [], usage: { prompt_tokens: 30000, completion_tokens: 10, total_tokens: 30010, completion_tokens_details: { reasoning_tokens: 4 } } },
    ]) response.write(`data: ${JSON.stringify(payload)}\n\n`);
    response.end('data: [DONE]\n\n');
  });
  await new Promise(resolve => server.listen(0, '127.0.0.1', resolve));
  t.after(() => new Promise(resolve => server.close(resolve)));
  const sdk = await import('@earendil-works/pi-coding-agent');
  const session = await createSession({ id: 'agent-0', directory, config: { condition: 'single', run_id: 'compaction', proxy_url: `http://127.0.0.1:${server.address().port}`, max_output_tokens: 16000 } }, sdk, event => events.push(event));
  t.after(() => session.dispose());
  await session.prompt('First conversation segment. ' + 'a'.repeat(60000));
  await session.prompt('Second conversation segment. ' + 'b'.repeat(60000));
  const compacted = await session.compact();
  assert.match(compacted.summary, /Summary/);
  await session.prompt('Message from agent-1 (peer_message; event 123):\nContinue after compaction.');
  assert.ok(requests.some(request => request.headers['x-study-purpose'] === 'compaction'));
  assert.equal(requests.at(-1).headers['x-study-context-window'], 'agent-0-1');
  for (const request of requests) {
    assert.equal(request.path, '/runs/compaction/agents/agent-0/v1/chat/completions');
    assert.equal(request.body.temperature, 0.7); assert.equal(request.body.top_p, 0.9);
    assert.equal(request.body.max_tokens, 16000); assert.equal(request.body.reasoning_effort, 'medium');
    assert.ok(request.headers['x-study-request-id']);
  }
  const compactionStart = events.find(event => event.type === 'model_request_start' && event.purpose === 'compaction');
  assert.ok(compactionStart);
  assert.equal(events.find(event => event.type === 'model_request_end' && event.request_id === compactionStart.request_id).usage.output, 10);
  assert.equal(events.filter(event => event.type === 'model_request_end').length, requests.length);
  assert.ok(events.some(event => event.type === 'message_delivered' && event.causal_event_id === 123 && event.delivery === 'context_insertion'));
});

test('actual worker process completes an SDK tool loop through supervisor IPC with fake model', async t => {
  const f = await fixture(t, 'single');
  const requests = [];
  const server = createServer(async (request, response) => {
    const chunks = []; for await (const chunk of request) chunks.push(chunk);
    requests.push(JSON.parse(Buffer.concat(chunks).toString()));
    const useTool = requests.length === 1;
    response.writeHead(200, { 'content-type': 'text/event-stream' });
    const base = { id: `ipc-${requests.length}`, object: 'chat.completion.chunk', model: 'Beam-501B-A23B', created: 1 };
    const delta = useTool
      ? { role: 'assistant', tool_calls: [{ index: 0, id: 'call-one', type: 'function', function: { name: 'bash', arguments: '{"command":"printf container-only-result"}' } }] }
      : { role: 'assistant', content: 'Completed after container tool.' };
    for (const payload of [
      { ...base, choices: [{ index: 0, delta, finish_reason: null }] },
      { ...base, choices: [{ index: 0, delta: {}, finish_reason: useTool ? 'tool_calls' : 'stop' }] },
      { ...base, choices: [], usage: { prompt_tokens: 100, completion_tokens: 10, total_tokens: 110 } },
    ]) response.write(`data: ${JSON.stringify(payload)}\n\n`);
    response.end('data: [DONE]\n\n');
  });
  await new Promise(resolve => server.listen(0, '127.0.0.1', resolve));
  t.after(() => new Promise(resolve => server.close(resolve)));
  f.backend.bash = async (_id, command) => { assert.equal(command, 'printf container-only-result'); return { output: 'container-only-result', exitCode: 0 }; };
  const supervisor = new Supervisor({ ...f.config, proxy_url: `http://127.0.0.1:${server.address().port}` }, { backend: f.backend, eventSink: event => f.events.push(event) });
  const result = await supervisor.run();
  assert.equal(result.stop_reason, 'completed');
  assert.equal(requests.length, 2);
  assert.ok(requests[1].messages.some(message => message.role === 'tool' && message.content.includes('container-only-result')));
  assert.ok(f.events.some(event => event.type === 'tool_end' && event.exit_code === 0));
  assert.equal(f.events.filter(event => event.type === 'session_created').length, 1);
});