feat: add Pi subagent scheduler and child runtime
This commit is contained in:
241
tests/unit/pi-subagent.test.ts
Normal file
241
tests/unit/pi-subagent.test.ts
Normal file
@@ -0,0 +1,241 @@
|
||||
// @vitest-environment node
|
||||
|
||||
import { describe, expect, it } from 'vitest';
|
||||
import {
|
||||
PiSubagentChildError,
|
||||
PiSubagentScheduler,
|
||||
parsePiSubagentDispatchRequest,
|
||||
type PiSubagentChild,
|
||||
type PiSubagentChildOpenInput,
|
||||
} from '../../electron/coding-runtime/pi/subagent';
|
||||
import { PiProcessBudget } from '../../electron/coding-runtime/pi/worker-pool';
|
||||
|
||||
function parent(runId = 'run-a') {
|
||||
return {
|
||||
conversationId: 'conversation-a',
|
||||
workerGeneration: 1,
|
||||
runId,
|
||||
projectId: 'project-a',
|
||||
};
|
||||
}
|
||||
|
||||
function task(agentId: string, toolProfile: 'read-only' | 'coding' = 'read-only') {
|
||||
return { agentId, task: `Inspect ${agentId}`, toolProfile };
|
||||
}
|
||||
|
||||
function deferred() {
|
||||
let resolve!: () => void;
|
||||
const promise = new Promise<void>((done) => { resolve = done; });
|
||||
return { promise, resolve };
|
||||
}
|
||||
|
||||
describe('Pi subagent scheduler', () => {
|
||||
it('validates one bounded dispatch mode and rejects a ninth task', () => {
|
||||
expect(parsePiSubagentDispatchRequest({ mode: 'single', tasks: [task('one')] })).toEqual({
|
||||
mode: 'single', tasks: [task('one')],
|
||||
});
|
||||
expect(parsePiSubagentDispatchRequest({
|
||||
mode: 'parallel', tasks: Array.from({ length: 8 }, (_, index) => task(`agent-${index}`)),
|
||||
}).tasks).toHaveLength(8);
|
||||
expect(() => parsePiSubagentDispatchRequest({
|
||||
mode: 'parallel', tasks: Array.from({ length: 9 }, (_, index) => task(`agent-${index}`)),
|
||||
})).toThrowError('Subagent dispatch accepts at most 8 tasks');
|
||||
expect(() => parsePiSubagentDispatchRequest({
|
||||
mode: 'single', tasks: [task('one'), task('two')],
|
||||
})).toThrowError('Single subagent dispatch requires exactly one task');
|
||||
expect(() => parsePiSubagentDispatchRequest({ mode: 'unknown', tasks: [task('one')] }))
|
||||
.toThrowError('Subagent dispatch mode is invalid');
|
||||
});
|
||||
|
||||
it('shares one FIFO four-child cap across two parents and releases the process budget', async () => {
|
||||
const gate = deferred();
|
||||
const processBudget = new PiProcessBudget(8);
|
||||
let running = 0;
|
||||
let maxRunning = 0;
|
||||
const opened: PiSubagentChildOpenInput[] = [];
|
||||
const scheduler = new PiSubagentScheduler({
|
||||
processBudget,
|
||||
createId: (() => {
|
||||
let id = 0;
|
||||
return (kind) => `${kind}-${++id}`;
|
||||
})(),
|
||||
openChild: async (input) => {
|
||||
opened.push(structuredClone(input));
|
||||
return {
|
||||
id: input.taskId,
|
||||
async run() {
|
||||
running += 1;
|
||||
maxRunning = Math.max(maxRunning, running);
|
||||
await gate.promise;
|
||||
running -= 1;
|
||||
return {
|
||||
summary: `done ${input.agentId}`,
|
||||
usage: { inputTokens: 1, outputTokens: 2 },
|
||||
};
|
||||
},
|
||||
async stop() {},
|
||||
} satisfies PiSubagentChild;
|
||||
},
|
||||
});
|
||||
const request = {
|
||||
mode: 'parallel' as const,
|
||||
tasks: Array.from({ length: 4 }, (_, index) => task(`agent-${index}`)),
|
||||
};
|
||||
const left = scheduler.dispatch({ ...parent('run-left'), request });
|
||||
const right = scheduler.dispatch({ ...parent('run-right'), request });
|
||||
await expect.poll(() => running).toBe(4);
|
||||
expect(maxRunning).toBe(4);
|
||||
expect(processBudget.activeCount).toBe(4);
|
||||
gate.resolve();
|
||||
const [leftResult, rightResult] = await Promise.all([left, right]);
|
||||
expect(leftResult.details.tasks).toHaveLength(4);
|
||||
expect(rightResult.details.tasks).toHaveLength(4);
|
||||
expect(opened).toHaveLength(8);
|
||||
expect(maxRunning).toBe(4);
|
||||
expect(processBudget.activeCount).toBe(0);
|
||||
await scheduler.close();
|
||||
});
|
||||
|
||||
it('preserves parallel sibling results and stops a chain at its first failure', async () => {
|
||||
const scheduler = new PiSubagentScheduler({
|
||||
processBudget: new PiProcessBudget(8),
|
||||
createId: (() => {
|
||||
let id = 0;
|
||||
return (kind) => `${kind}-${++id}`;
|
||||
})(),
|
||||
openChild: async (input) => ({
|
||||
id: input.taskId,
|
||||
async run(prompt) {
|
||||
if (input.agentId === 'broken') {
|
||||
throw new PiSubagentChildError('SUBAGENT_CHILD_FAILED');
|
||||
}
|
||||
return {
|
||||
summary: `${input.agentId}:${prompt}`,
|
||||
usage: { inputTokens: 3, outputTokens: 5 },
|
||||
};
|
||||
},
|
||||
async stop() {},
|
||||
}),
|
||||
});
|
||||
const parallel = await scheduler.dispatch({
|
||||
...parent('parallel'),
|
||||
request: {
|
||||
mode: 'parallel',
|
||||
tasks: [task('left'), task('broken'), task('right')],
|
||||
},
|
||||
});
|
||||
expect(parallel.details.tasks).toEqual(expect.arrayContaining([
|
||||
expect.objectContaining({ agentId: 'left', status: 'complete', summary: 'left:Inspect left' }),
|
||||
expect.objectContaining({ agentId: 'broken', status: 'error', errorCode: 'SUBAGENT_CHILD_FAILED' }),
|
||||
expect.objectContaining({ agentId: 'right', status: 'complete', usage: { inputTokens: 3, outputTokens: 5 } }),
|
||||
]));
|
||||
|
||||
const chain = await scheduler.dispatch({
|
||||
...parent('chain'),
|
||||
request: {
|
||||
mode: 'chain',
|
||||
tasks: [
|
||||
{ agentId: 'first', task: 'first', toolProfile: 'read-only' },
|
||||
{ agentId: 'broken', task: 'review {previous}', toolProfile: 'read-only' },
|
||||
{ agentId: 'never', task: 'never', toolProfile: 'coding' },
|
||||
],
|
||||
},
|
||||
});
|
||||
expect(chain.details.tasks.map(({ status }) => status)).toEqual(['complete', 'error', 'skipped']);
|
||||
expect(chain.details.tasks[1]).toMatchObject({ errorCode: 'SUBAGENT_CHILD_FAILED' });
|
||||
await scheduler.close();
|
||||
});
|
||||
|
||||
it('aborts every unfinished child with the parent and leaves no permit or process lease', async () => {
|
||||
const processBudget = new PiProcessBudget(8);
|
||||
const stopped: string[] = [];
|
||||
const scheduler = new PiSubagentScheduler({
|
||||
processBudget,
|
||||
openChild: async (input) => ({
|
||||
id: input.taskId,
|
||||
async run(_prompt, signal) {
|
||||
await new Promise<void>((_resolve, reject) => {
|
||||
const abort = () => reject(new PiSubagentChildError('SUBAGENT_ABORTED'));
|
||||
if (signal.aborted) abort();
|
||||
else signal.addEventListener('abort', abort, { once: true });
|
||||
});
|
||||
return { summary: 'unreachable' };
|
||||
},
|
||||
async stop() { stopped.push(input.taskId); },
|
||||
}),
|
||||
});
|
||||
const controller = new AbortController();
|
||||
const flight = scheduler.dispatch({
|
||||
...parent('abort'),
|
||||
request: { mode: 'parallel', tasks: [task('left'), task('right', 'coding')] },
|
||||
}, { signal: controller.signal });
|
||||
await expect.poll(() => processBudget.activeCount).toBe(2);
|
||||
controller.abort();
|
||||
const result = await flight;
|
||||
expect(result.details.tasks.map(({ status }) => status)).toEqual(['aborted', 'aborted']);
|
||||
expect(stopped).toHaveLength(2);
|
||||
expect(processBudget.activeCount).toBe(0);
|
||||
await scheduler.close();
|
||||
});
|
||||
|
||||
it('marks the rest of an in-flight chain aborted when its parent aborts', async () => {
|
||||
const processBudget = new PiProcessBudget(8);
|
||||
const scheduler = new PiSubagentScheduler({
|
||||
processBudget,
|
||||
openChild: async (input) => ({
|
||||
id: input.taskId,
|
||||
async run(_prompt, signal) {
|
||||
await new Promise<void>((_resolve, reject) => {
|
||||
const abort = () => reject(new PiSubagentChildError('SUBAGENT_ABORTED'));
|
||||
if (signal.aborted) abort();
|
||||
else signal.addEventListener('abort', abort, { once: true });
|
||||
});
|
||||
return { summary: 'unreachable' };
|
||||
},
|
||||
async stop() {},
|
||||
}),
|
||||
});
|
||||
const controller = new AbortController();
|
||||
const flight = scheduler.dispatch({
|
||||
...parent('abort-chain'),
|
||||
request: { mode: 'chain', tasks: [task('one'), task('two'), task('three')] },
|
||||
}, { signal: controller.signal });
|
||||
await expect.poll(() => processBudget.activeCount).toBe(1);
|
||||
controller.abort();
|
||||
const result = await flight;
|
||||
expect(result.details.tasks.map(({ status }) => status))
|
||||
.toEqual(['aborted', 'aborted', 'aborted']);
|
||||
expect(processBudget.activeCount).toBe(0);
|
||||
await scheduler.close();
|
||||
});
|
||||
|
||||
it('reclaims an idle parent process instead of deadlocking on a full global budget', async () => {
|
||||
const processBudget = new PiProcessBudget(2);
|
||||
const firstParent = await processBudget.acquire();
|
||||
const idleParent = await processBudget.acquire();
|
||||
let reclaimed = 0;
|
||||
const scheduler = new PiSubagentScheduler({
|
||||
processBudget,
|
||||
reclaimProcessCapacity: async () => {
|
||||
reclaimed += 1;
|
||||
idleParent.release();
|
||||
return true;
|
||||
},
|
||||
openChild: async (input) => ({
|
||||
id: input.taskId,
|
||||
async run() { return { summary: 'done' }; },
|
||||
async stop() {},
|
||||
}),
|
||||
});
|
||||
await expect(scheduler.dispatch({
|
||||
...parent('full-budget'),
|
||||
request: { mode: 'single', tasks: [task('agent-a')] },
|
||||
})).resolves.toMatchObject({
|
||||
details: { tasks: [{ status: 'complete' }] },
|
||||
});
|
||||
expect(reclaimed).toBe(1);
|
||||
expect(processBudget.activeCount).toBe(1);
|
||||
firstParent.release();
|
||||
await scheduler.close();
|
||||
});
|
||||
});
|
||||
Reference in New Issue
Block a user