diff --git a/agents/adapters/supervisor.ts b/agents/adapters/supervisor.ts index 4f7c8672..e5c3d72e 100644 --- a/agents/adapters/supervisor.ts +++ b/agents/adapters/supervisor.ts @@ -14,6 +14,7 @@ export const OUTPUT_LIMITS = Object.freeze({ }); const DEFAULT_TIMEOUT_MS = 10 * 60_000; const CAPTURE_ABORT_GRACE_MS = 1_000; +const ACKNOWLEDGEMENT_RETRY_MS = 250; const DIAGNOSTIC_BYTES = 1024; const active = new Map(); const activeProfiles = new WeakSet(); @@ -363,6 +364,27 @@ export function readBoundedContainerFile(container: string, source: string, maxi }); } +/** + * Run as root inside the container: write the deferred output acknowledgement `collected-` that the wrapper + * waits for. It is idempotent, so a retry succeeds after an attempt that died once it had created the file: a + * complete file is kept, and a partial one is unlinked and written again. Only root can create entries in the + * control directory, and without capabilities root cannot reopen its own read-only file for writing. + */ +export const ACKNOWLEDGEMENT_SCRIPT = [ + "const fs=require('node:fs'),c=fs.constants,token=process.argv[1],directory='/run/codeboost-control';", + "let dirfd,fd,existing;try{dirfd=fs.openSync(directory,c.O_RDONLY|c.O_DIRECTORY|c.O_NOFOLLOW);", + "const name='/proc/self/fd/'+dirfd+'/collected-'+token,create=()=>fs.openSync(name,", + 'c.O_WRONLY|c.O_CREAT|c.O_EXCL|c.O_NOFOLLOW,0o444);', + "try{fd=create()}catch(error){if(error.code!=='EEXIST')throw error;", + 'existing=fs.openSync(name,c.O_RDONLY|c.O_NOFOLLOW|c.O_NONBLOCK);const stat=fs.fstatSync(existing,{bigint:true});', + "if(!stat.isFile()||stat.nlink!==1n)throw new Error('UNSAFE_ACK');", + "const complete=fs.readFileSync(existing,'utf8')===token;fs.closeSync(existing);existing=undefined;", + 'if(!complete){fs.unlinkSync(name);fd=create()}}', + "if(fd!==undefined){fs.writeFileSync(fd,token);fs.fsyncSync(fd);const stat=fs.fstatSync(fd,{bigint:true});", + "if(!stat.isFile()||stat.nlink!==1n)throw new Error('UNSAFE_ACK')}}finally{if(existing!==undefined)fs.closeSync(existing);", + 'if(fd!==undefined)fs.closeSync(fd);if(dirfd!==undefined)fs.closeSync(dirfd)}', +].join(''); + export function isInvocationActive(attemptId: string): boolean { return ownsAttempt(attemptId); } @@ -417,6 +439,7 @@ export function startProfileInvocation(profile: ContainerProfile, options: Super let decodedOutput: DecodedOutput | undefined, decodePromise: Promise | undefined; let decodeAbort: AbortController | undefined; let protocolToken: string | undefined, protocolStarted = false, protocolReady = false; + let readyStatus: number | undefined, acknowledgementFailures = 0; let protocolBuffer = Buffer.alloc(0); // Setup (container create and validation) runs inside the handle; a stop kills its in-flight Docker call. const setupAbort = new AbortController(); @@ -443,22 +466,28 @@ export function startProfileInvocation(profile: ContainerProfile, options: Super void operation.finally(() => controls.delete(operation)); return operation; }; - const acknowledgeDeferredOutput = (token: string) => { - const script = [ - "const fs=require('node:fs'),token=process.argv[1],directory='/run/codeboost-control';", - 'let dirfd,fd;try{dirfd=fs.openSync(directory,fs.constants.O_RDONLY|fs.constants.O_DIRECTORY|fs.constants.O_NOFOLLOW);', - "fd=fs.openSync('/proc/self/fd/'+dirfd+'/collected-'+token,", - 'fs.constants.O_WRONLY|fs.constants.O_CREAT|fs.constants.O_EXCL|fs.constants.O_NOFOLLOW,0o444);', - "fs.writeFileSync(fd,token);fs.fsyncSync(fd);const stat=fs.fstatSync(fd,{bigint:true});", - "if(!stat.isFile()||stat.nlink!==1n)throw new Error('UNSAFE_ACK')}finally{if(fd!==undefined)fs.closeSync(fd);", - 'if(dirfd!==undefined)fs.closeSync(dirfd)}', - ].join(''); - return runControl(['exec', '--user', '0', container, 'node', '-e', script, token]); - }; + const acknowledgeDeferredOutput = (token: string) => + runControl(['exec', '--user', '0', container, 'node', '-e', ACKNOWLEDGEMENT_SCRIPT, token]); const later = (callback: () => void, delay: number) => { const timer = setTimeout(() => { timers.delete(timer); callback(); }, delay); timer.unref(); timers.add(timer); return timer; }; + /** + * The wrapper exits as soon as it sees the acknowledgement, and that exit can kill the acknowledging `docker exec` + * before it reports success. The exec's status therefore decides nothing: a failed attempt is retried until one + * succeeds, the container exits, or a stop ends the invocation. The close handler accepts the output only if the + * container exited with the READY status, as the wrapper does after it has seen the acknowledgement; any other + * exit status, or a signal to the attached client, fails closed. A container killed from outside cannot be told + * apart when its kill status equals the READY status, but its output was already collected in full by then. + */ + const acknowledge = (token: string) => { + if (closed || stopReason) return; + void acknowledgeDeferredOutput(token).then(success => { + if (success || closed || stopReason) return; + acknowledgementFailures += 1; + later(() => acknowledge(token), ACKNOWLEDGEMENT_RETRY_MS); + }); + }; const terminate = () => { if (terminating || closed || !child) return; terminating = true; @@ -600,18 +629,9 @@ export function startProfileInvocation(profile: ContainerProfile, options: Super return true; } protocolReady = true; + readyStatus = Number(ready[2]); const readyToken = ready[1]!; - void decodeOutput().then(() => { - if (decodedOutput && !stopReason) { - void acknowledgeDeferredOutput(readyToken) - .then(success => { - if (!success) { - failureDetail ??= 'Deferred output acknowledgement failed.'; - stop('capture-failure'); - } - }); - } - }); + void decodeOutput().then(() => { if (decodedOutput) acknowledge(readyToken); }); return true; }; const captureStderr = (value: Buffer | string) => { @@ -733,6 +753,12 @@ export function startProfileInvocation(profile: ContainerProfile, options: Super if (!stopReason && profile.deferredOutput && !decodedOutput) { stopReason = 'capture-failure'; failureDetail ??= 'Deferred output protocol did not complete.'; } + if (stopReason === 'timeout' && acknowledgementFailures) + failureDetail ??= `Deferred output acknowledgement failed ${acknowledgementFailures} times.`; + if (!stopReason && profile.deferredOutput && (signal !== null || exitCode !== readyStatus)) { + stopReason = 'capture-failure'; + failureDetail ??= 'Container exited without taking the deferred output acknowledgement.'; + } if (decodedOutput && !stopReason) { finalStdout = Buffer.from(decodedOutput.text); if (decodedOutput.providerFailed) exitCode = exitCode === 0 ? 1 : exitCode; diff --git a/test/agent-deferred-acknowledgement.test.ts b/test/agent-deferred-acknowledgement.test.ts new file mode 100644 index 00000000..aa4361e7 --- /dev/null +++ b/test/agent-deferred-acknowledgement.test.ts @@ -0,0 +1,153 @@ +import { EventEmitter } from 'node:events'; +import { PassThrough } from 'node:stream'; +import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest'; +import type { InvocationResult } from '../agents/contract.ts'; +import type { ContainerProfile } from '../agents/container/profile.ts'; + +// Replays, without Docker, the interleavings of the deferred-output acknowledgement with the container's exit. The +// wrapper (probe.sh) exits as soon as it sees the acknowledgement, and that exit can kill the acknowledging +// `docker exec` before it reports success; CI run 36470233043 failed that way. +type FakeChild = EventEmitter & { stdout: PassThrough; stderr: PassThrough; kill: ReturnType }; +const state = vi.hoisted(() => ({ + attach: undefined as FakeChild | undefined, + acknowledgements: [] as { child: FakeChild; args: readonly string[] }[], +})); +vi.mock('node:child_process', () => ({ + execFile: vi.fn(), + spawn: vi.fn((_command: string, args: readonly string[]) => { + const child = Object.assign(new EventEmitter(), { stdout: new PassThrough(), stderr: new PassThrough(), kill: vi.fn() }); + if (args[0] === 'start') state.attach = child; + else if (args[0] === 'exec') state.acknowledgements.push({ child, args }); + else process.nextTick(() => child.emit('close', 0, null)); + return child; + }), +})); +vi.mock('../agents/policy.ts', () => ({ assertPhasePolicy: (policy: { invocation: unknown }) => policy.invocation })); +vi.mock('../agents/container/profile.ts', () => ({ + assertContainerProfileAuthenticity: () => {}, + isContainerProfileAuthentic: () => true, + disposeContainerProfile: async () => {}, + containerProfileResources: () => [], +})); +vi.mock('../agents/container/run.ts', () => ({ + createValidatedContainer: async () => {}, + validateContainer: async () => {}, + disposeValidatedContainer: async () => {}, + retireContainerProfile: () => {}, + isContainerProfileRetired: () => false, + agentContainerId: () => 'c'.repeat(64), + agentContainerResources: () => [], +})); +const { startProfileInvocation } = await import('../agents/adapters/supervisor.ts'); + +const token = '11111111-2222-3333-4444-555555555555'; +const fakeProfile = (attemptId: string) => ({ + name: `codeboost-agent-${attemptId}`, deferredOutput: true, resources: [], + policy: { invocation: { attemptId, deadline: Date.now() + 10 * 60_000, + context: { snapshotId: 's', planId: 'p', planRevision: 1, assignmentId: 'a', referencedCodeHash: 'c', stateVersion: 1 } } }, +}) as unknown as ContainerProfile; + +describe('deferred output acknowledgement', () => { + beforeEach(() => { + vi.useFakeTimers({ toFake: ['setTimeout', 'clearTimeout', 'performance'] }); + state.attach = undefined; + state.acknowledgements = []; + }); + afterEach(() => { vi.useRealTimers(); }); + + /** Start an invocation whose agent writes newline-free stderr, then READY with `status`. */ + const readyInvocation = async (attemptId: string, status = 0, timeoutMs?: number) => { + const box: { result?: InvocationResult } = {}; + void startProfileInvocation(fakeProfile(attemptId), { timeoutMs, decode: () => ({ text: 'captured' }) }) + .settled.then(result => { box.result = result; }); + await vi.advanceTimersByTimeAsync(0); + const stderr = state.attach!.stderr; + stderr.emit('data', Buffer.from(`\x1eCODEBOOST_START:${token}\x1e\n`)); + stderr.emit('data', Buffer.from('trailing-diagnostic')); + stderr.emit('data', Buffer.from(`\n\x1eCODEBOOST_READY:${token}:${status}\x1e\n`)); + await vi.advanceTimersByTimeAsync(0); + // The decoded output was collected, so the supervisor acknowledges it with the READY token. + expect(state.acknowledgements).toHaveLength(1); + expect(state.acknowledgements[0]!.args.at(-1)).toBe(token); + return box; + }; + const acknowledgementExits = (index: number, code: number | null, signal: NodeJS.Signals | null = null) => + state.acknowledgements[index]!.child.emit('close', code, signal); + const containerExits = (code: number | null, signal: NodeJS.Signals | null = null) => + state.attach!.emit('close', code, signal); + + it.each([ + ['before the container exit is observed', 'acknowledgement-first'], + ['after the container exit is observed', 'exit-first'], + ] as const)('keeps the result when the wrapper takes the acknowledgement and its exec then fails %s', + async (_label, order) => { + const box = await readyInvocation(`ack-killed-${order}`); + // The wrapper saw the acknowledgement and exited with the READY status; the teardown killed the exec. + if (order === 'acknowledgement-first') { acknowledgementExits(0, 137); containerExits(0); } + else { containerExits(0); acknowledgementExits(0, 137); } + await vi.advanceTimersByTimeAsync(0); + expect(box.result?.stopReason, box.result?.stderr).toBeUndefined(); + expect(box.result?.exitCode).toBe(0); + expect(box.result?.stdout).toBe('captured'); + expect(box.result?.stderr).toBe('trailing-diagnostic\n'); + // Nothing retries once the container has exited. + await vi.advanceTimersByTimeAsync(5_000); + expect(state.acknowledgements).toHaveLength(1); + }); + + it('retries an acknowledgement that failed while the container keeps waiting for it', async () => { + const box = await readyInvocation('ack-retried'); + acknowledgementExits(0, 1); + await vi.advanceTimersByTimeAsync(0); + expect(box.result).toBeUndefined(); + await vi.advanceTimersByTimeAsync(250); + expect(state.acknowledgements).toHaveLength(2); + expect(state.acknowledgements[1]!.args.at(-1)).toBe(token); + acknowledgementExits(1, 0); + containerExits(0); + await vi.advanceTimersByTimeAsync(0); + expect(box.result?.stopReason, box.result?.stderr).toBeUndefined(); + expect(box.result?.stderr).toBe('trailing-diagnostic\n'); + }); + + it('reports the READY status of an agent that failed', async () => { + const box = await readyInvocation('ack-nonzero', 3); + acknowledgementExits(0, 0); + containerExits(3); + await vi.advanceTimersByTimeAsync(0); + expect(box.result?.stopReason, box.result?.stderr).toBeUndefined(); + expect(box.result?.exitCode).toBe(3); + }); + + it.each([ + ['an exit status other than READY', 137, null], + ['a signal', null, 'SIGKILL'], + ] as const)('fails closed when the container exits with %s instead of taking the acknowledgement', + async (_label, code, signal) => { + const box = await readyInvocation(`ack-bypassed-${signal ?? code}`); + acknowledgementExits(0, 137); + containerExits(code, signal); + await vi.advanceTimersByTimeAsync(0); + expect(box.result?.stopReason).toBe('capture-failure'); + expect(box.result?.stderr).toContain('trailing-diagnostic\n'); + expect(box.result?.stderr).toContain('Container exited without taking the deferred output acknowledgement.'); + }); + + it('times out instead of succeeding when no acknowledgement ever takes effect', async () => { + const box = await readyInvocation('ack-never', 0, 5_000); + for (let index = 0; index < 4; index += 1) { + acknowledgementExits(index, 1); + await vi.advanceTimersByTimeAsync(250); + } + expect(state.acknowledgements).toHaveLength(5); + await vi.advanceTimersByTimeAsync(5_000); + expect(box.result).toBeUndefined(); + // The deadline stopped the container, which exits by signal and ends the last attempt without a retry. + acknowledgementExits(4, 137); + containerExits(null, 'SIGKILL'); + await vi.advanceTimersByTimeAsync(0); + expect(box.result?.stopReason).toBe('timeout'); + expect(box.result?.stderr).toContain('trailing-diagnostic\n'); + expect(box.result?.stderr).toContain('Deferred output acknowledgement failed 4 times.'); + }); +}); diff --git a/test/agent-supervisor.test.ts b/test/agent-supervisor.test.ts index 2bfa7f3c..68b97fd7 100644 --- a/test/agent-supervisor.test.ts +++ b/test/agent-supervisor.test.ts @@ -6,7 +6,7 @@ import { dirname, join } from 'node:path'; import { afterAll, beforeAll, describe, expect, it } from 'vitest'; import { startClaudeInvocation } from '../agents/adapters/claude.ts'; import { readCodexOutput, startCodexInvocation } from '../agents/adapters/codex.ts'; -import { isInvocationActive, readBoundedContainerFile, retainSetupCleanup, +import { ACKNOWLEDGEMENT_SCRIPT, isInvocationActive, readBoundedContainerFile, retainSetupCleanup, startProfileInvocation } from '../agents/adapters/supervisor.ts'; import { captureInvocation, type InvocationInput, type InvocationResult } from '../agents/contract.ts'; import { buildAgentImage } from '../agents/container/image.ts'; @@ -543,6 +543,38 @@ describe('container invocation supervisor', () => { expect(result.stderr).toContain('trailing-diagnostic'); }, 60_000); + it('writes the deferred output acknowledgement idempotently, repairing a partial one', () => { + const token = randomUUID(), file = `/run/codeboost-control/collected-${token}`; + const name = `codeboost-ack-script-${randomUUID()}`; + execFileSync('docker', ['run', '-d', '--rm', '--name', name, '--network', 'none', '--cap-drop=ALL', + '--security-opt=no-new-privileges', '--user', '10001:10001', '--entrypoint', 'sleep', + '--tmpfs', '/run/codeboost-control:rw,nosuid,nodev,noexec,size=65536,nr_inodes=16,uid=0,gid=0,mode=0711', + imageId, '120'], { stdio: 'ignore' }); + try { + const acknowledge = () => spawnSync('docker', ['exec', '--user', '0', name, 'node', '-e', ACKNOWLEDGEMENT_SCRIPT, + token], { encoding: 'utf8' }); + const content = () => execFileSync('docker', ['exec', name, 'cat', file], { encoding: 'utf8' }); + const root = (script: string) => execFileSync('docker', ['exec', '--user', '0', name, 'sh', '-c', script]); + expect(acknowledge().status).toBe(0); + expect(content()).toBe(token); + // A retry after an attempt that completed. + expect(acknowledge().status).toBe(0); + expect(content()).toBe(token); + // A retry after an attempt that died once it had created the file. + for (const partial of ['', token.slice(0, 8)]) { + root(`rm ${file} && printf '${partial}' > ${file} && chmod 444 ${file}`); + expect(acknowledge().status).toBe(0); + expect(content()).toBe(token); + } + expect(execFileSync('docker', ['exec', name, 'stat', '-c', '%a %u %h', file], { encoding: 'utf8' }).trim()) + .toBe('444 0 1'); + root(`rm ${file} && ln -s /etc/passwd ${file}`); + expect(acknowledge().status).not.toBe(0); + } finally { + spawnSync('docker', ['rm', '--force', name], { stdio: 'ignore' }); + } + }, 60_000); + if (process.env.CODEBOOST_RUN_AUTH_PROBES === '1') { const schemaPrompt = 'Read /run/codeboost-input/schema.json and reply only with the exact value of its probe field, ' + 'without quotes or Markdown formatting.';