Spaces:
Running
Running
| import { spawn } from 'node:child_process'; | |
| import { createWriteStream } from 'node:fs'; | |
| import { mkdir } from 'node:fs/promises'; | |
| import { dirname } from 'node:path'; | |
| export const quote = value => "'" + String(value).replaceAll("'", "'\\''") + "'"; | |
| export function command(argv, { signal, outputPath, maxBytes = 100000 } = {}) { | |
| return new Promise((resolve, reject) => { | |
| const child = spawn(argv[0], argv.slice(1), { stdio: ['ignore', 'pipe', 'pipe'] }); | |
| let output = '', bytes = 0, truncated = false, killTimer; | |
| const file = outputPath ? createWriteStream(outputPath, { mode: 0o600 }) : null; | |
| if (file) child.stdout.pipe(file); | |
| const collect = chunk => { | |
| const remaining = Math.max(0, maxBytes - bytes); | |
| output += chunk.subarray(0, remaining).toString(); | |
| bytes += chunk.length; | |
| if (bytes > maxBytes) truncated = true; | |
| }; | |
| if (!file) child.stdout.on('data', collect); | |
| child.stderr.on('data', collect); | |
| const abort = () => { child.kill('SIGTERM'); killTimer = setTimeout(() => child.kill('SIGKILL'), 5000); }; | |
| signal?.addEventListener('abort', abort, { once: true }); | |
| if (signal?.aborted) abort(); | |
| child.on('error', reject); | |
| file?.on('error', reject); | |
| child.on('close', async (code, terminationSignal) => { | |
| clearTimeout(killTimer); | |
| signal?.removeEventListener('abort', abort); | |
| if (file && !file.writableFinished) await new Promise(r => file.once('finish', r)); | |
| resolve({ output: output + (truncated ? '\n[output truncated]' : ''), exitCode: code ?? 128, signal: terminationSignal }); | |
| }); | |
| }); | |
| } | |
| export class ContainerBackend { | |
| constructor(config, event = () => {}, execute = command) { | |
| this.config = config; this.event = event; this.execute = execute; | |
| this.root = `${config.work_root ?? '/workspace'}/.pi-study`; | |
| this.bare = `${this.root}/shared.git`; | |
| this.active = new Map(); this.nextTool = 0; | |
| } | |
| argv(script) { | |
| return [...(this.config.docker_prefix ?? ['docker']), 'exec', '--user', this.config.container_user ?? 'agent', this.config.container_name, 'bash', '-lc', script]; | |
| } | |
| async checked(script, options) { | |
| const result = await this.execute(this.argv(script), options); | |
| if (result.exitCode) throw new Error(`Container command failed (${result.exitCode}): ${result.output}`); | |
| return result; | |
| } | |
| agentPath(id) { | |
| if (!/^agent-\d+$/.test(id)) throw new Error('Invalid agent ID'); | |
| return `${this.root}/agents/${id}`; | |
| } | |
| async prepare() { | |
| const seed = this.config.seed_dir ?? `${this.config.work_root ?? '/workspace'}/solution`; | |
| await this.checked(`set -euo pipefail\ntest ! -e ${quote(this.root)}\nmkdir -p ${quote(this.root + '/agents')} ${quote(this.root + '/jobs')}\ngit clone --bare --no-hardlinks ${quote(seed)} ${quote(this.bare)}\ngit --git-dir=${quote(this.bare)} rev-parse --verify main\ngit --git-dir=${quote(this.bare)} symbolic-ref HEAD refs/heads/main`); | |
| } | |
| async createAgent(id) { | |
| const path = this.agentPath(id); | |
| await this.checked(`set -euo pipefail\ngit clone ${quote(this.bare)} ${quote(path)}\ngit -C ${quote(path)} config user.name ${quote(id)}\ngit -C ${quote(path)} config user.email ${quote(id + '@study.invalid')}\ngit -C ${quote(path)} checkout -b ${quote(id)} origin/main`); | |
| } | |
| async bash(id, script, { signal, timeout = 120 } = {}) { | |
| if (signal?.aborted) throw new Error('Tool cancelled'); | |
| const job = `${this.root}/jobs/${id}-${++this.nextTool}.pid`; | |
| const cancelled = `${job}.cancelled`; | |
| const seconds = Math.max(1, Math.min(Number(timeout) || 120, 600)); | |
| const wrapped = `set -uo pipefail\ncd ${quote(this.agentPath(id))}\ntest ! -e ${quote(cancelled)} || exit 130\nsetsid timeout --signal=TERM --kill-after=5s ${seconds}s bash -lc ${quote(script)} &\nchild=$!\nprintf '%s' "$child" > ${quote(job)}\nif test -e ${quote(cancelled)}; then kill -TERM -- "-$child" 2>/dev/null || true; sleep 1; kill -KILL -- "-$child" 2>/dev/null || true; fi\nwait "$child"\nstatus=$?\nrm -f ${quote(job)}\nexit "$status"`; | |
| const kill = async () => { | |
| await this.execute(this.argv(`touch ${quote(cancelled)}; if test -f ${quote(job)}; then pid=$(cat ${quote(job)}); case "$pid" in ''|*[!0-9]*) exit 1;; esac; kill -TERM -- "-$pid" 2>/dev/null || true; sleep 1; kill -KILL -- "-$pid" 2>/dev/null || true; rm -f ${quote(job)}; fi`)); | |
| }; | |
| this.active.set(job, kill); | |
| const onAbort = () => { void kill(); }; | |
| signal?.addEventListener('abort', onAbort, { once: true }); | |
| try { return await this.execute(this.argv(wrapped), { signal }); } | |
| finally { signal?.removeEventListener('abort', onAbort); this.active.delete(job); if (signal?.aborted) await kill(); } | |
| } | |
| async revision() { | |
| const revision = (await this.checked(`git --git-dir=${quote(this.bare)} rev-parse refs/heads/main`)).output.trim(); | |
| if (!/^[0-9a-f]{40,64}$/.test(revision)) throw new Error('Invalid snapshot revision'); | |
| return revision; | |
| } | |
| async snapshot(outputPath, revision = null) { | |
| await mkdir(dirname(outputPath), { recursive: true }); | |
| revision ??= await this.revision(); | |
| if (!/^[0-9a-f]{40,64}$/.test(revision)) throw new Error('Invalid snapshot revision'); | |
| await this.checked(`git --git-dir=${quote(this.bare)} archive --format=tar.gz ${quote(revision)}`, { outputPath }); | |
| return { path: outputPath, revision }; | |
| } | |
| async stop() { await Promise.allSettled([...this.active.values()].map(kill => kill())); } | |
| } | |