391 lines
14 KiB
TypeScript
391 lines
14 KiB
TypeScript
// @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, PiWorkerPool } 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 waitingParentSettled = false;
|
|
const waitingParent = processBudget.acquire().then((lease) => {
|
|
waitingParentSettled = true;
|
|
return lease;
|
|
});
|
|
let reclaimed = 0;
|
|
let childRanBeforeWaitingParent = false;
|
|
const scheduler = new PiSubagentScheduler({
|
|
processBudget,
|
|
reclaimProcessCapacity: async () => {
|
|
reclaimed += 1;
|
|
idleParent.release();
|
|
return true;
|
|
},
|
|
openChild: async (input) => ({
|
|
id: input.taskId,
|
|
async run() {
|
|
childRanBeforeWaitingParent = !waitingParentSettled;
|
|
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(childRanBeforeWaitingParent).toBe(true);
|
|
const waitingParentLease = await waitingParent;
|
|
expect(processBudget.activeCount).toBe(2);
|
|
waitingParentLease.release();
|
|
firstParent.release();
|
|
expect(processBudget.activeCount).toBe(0);
|
|
await scheduler.close();
|
|
});
|
|
|
|
it('reclaims a parent that becomes idle after child capacity is already reserved', async () => {
|
|
const parentOpenGate = deferred();
|
|
const processBudget = new PiProcessBudget(1);
|
|
let parentStopped = false;
|
|
let childRan = false;
|
|
const pool = new PiWorkerPool({
|
|
processBudget,
|
|
maxIdle: 1,
|
|
openWorker: async () => {
|
|
await parentOpenGate.promise;
|
|
return {
|
|
worker: {
|
|
id: 'delayed-parent',
|
|
generation: 1,
|
|
async request() {
|
|
return { type: 'response' as const, id: 'parent', success: true as const };
|
|
},
|
|
async send() {},
|
|
subscribe() { return () => undefined; },
|
|
subscribeInvalidation() { return () => undefined; },
|
|
async stop() {
|
|
parentStopped = true;
|
|
return { mode: 'stdin-close' as const, code: 0, signal: null };
|
|
},
|
|
},
|
|
session: { piSessionId: 'parent-session', sessionKey: 'parent-key' },
|
|
};
|
|
},
|
|
});
|
|
const preparingParent = pool.prepare({
|
|
conversationId: 'delayed-parent',
|
|
projectId: 'project-a',
|
|
agentId: 'agent-a',
|
|
title: 'Delayed parent',
|
|
model: {
|
|
model: { accountId: 'account-a', modelId: 'model-a', thinkingLevel: 'medium' },
|
|
modelResolution: 'resolved',
|
|
},
|
|
});
|
|
await expect.poll(() => processBudget.activeCount).toBe(1);
|
|
|
|
const scheduler = new PiSubagentScheduler({
|
|
processBudget,
|
|
reclaimProcessCapacity: (signal) => pool.reclaimIdleWorker(signal),
|
|
openChild: async (input) => ({
|
|
id: input.taskId,
|
|
async run() {
|
|
childRan = true;
|
|
return { summary: 'done' };
|
|
},
|
|
async stop() {},
|
|
}),
|
|
});
|
|
const child = scheduler.dispatch({
|
|
...parent('delayed-idle'),
|
|
request: { mode: 'single', tasks: [task('agent-a')] },
|
|
});
|
|
await expect.poll(() => processBudget.waitingCount).toBe(1);
|
|
|
|
parentOpenGate.resolve();
|
|
await preparingParent;
|
|
await expect(child).resolves.toMatchObject({
|
|
details: { tasks: [{ status: 'complete' }] },
|
|
});
|
|
expect(parentStopped).toBe(true);
|
|
expect(childRan).toBe(true);
|
|
expect(processBudget.activeCount).toBe(0);
|
|
expect(processBudget.waitingCount).toBe(0);
|
|
await scheduler.close();
|
|
await pool.shutdown();
|
|
});
|
|
|
|
it('cancels a queued child reservation when the capacity reclaimer rejects', async () => {
|
|
const processBudget = new PiProcessBudget(1);
|
|
const parentLease = await processBudget.acquire();
|
|
const scheduler = new PiSubagentScheduler({
|
|
processBudget,
|
|
reclaimProcessCapacity: async () => {
|
|
throw new Error('reclaim failed');
|
|
},
|
|
openChild: async () => {
|
|
throw new Error('child must not open');
|
|
},
|
|
});
|
|
|
|
await expect(scheduler.dispatch({
|
|
...parent('reclaim-rejected'),
|
|
request: { mode: 'single', tasks: [task('agent-a')] },
|
|
})).resolves.toMatchObject({
|
|
details: { tasks: [{ status: 'error', errorCode: 'SUBAGENT_CHILD_FAILED' }] },
|
|
});
|
|
expect(processBudget.activeCount).toBe(1);
|
|
expect(processBudget.waitingCount).toBe(0);
|
|
parentLease.release();
|
|
expect(processBudget.activeCount).toBe(0);
|
|
expect(processBudget.waitingCount).toBe(0);
|
|
await scheduler.close();
|
|
});
|
|
|
|
it('cancels both capacity waits when the parent aborts before an idle worker exists', async () => {
|
|
const processBudget = new PiProcessBudget(1);
|
|
const parentLease = await processBudget.acquire();
|
|
let reclaimerCancelled = false;
|
|
const scheduler = new PiSubagentScheduler({
|
|
processBudget,
|
|
reclaimProcessCapacity: async (signal) => await new Promise<boolean>((_resolve, reject) => {
|
|
const cancel = () => {
|
|
reclaimerCancelled = true;
|
|
reject(new Error('reclaim cancelled'));
|
|
};
|
|
if (signal?.aborted) cancel();
|
|
else signal?.addEventListener('abort', cancel, { once: true });
|
|
}),
|
|
openChild: async () => {
|
|
throw new Error('child must not open');
|
|
},
|
|
});
|
|
const controller = new AbortController();
|
|
const child = scheduler.dispatch({
|
|
...parent('reclaim-aborted'),
|
|
request: { mode: 'single', tasks: [task('agent-a')] },
|
|
}, { signal: controller.signal });
|
|
await expect.poll(() => processBudget.waitingCount).toBe(1);
|
|
|
|
controller.abort();
|
|
await expect(child).resolves.toMatchObject({
|
|
details: { tasks: [{ status: 'aborted', errorCode: 'SUBAGENT_ABORTED' }] },
|
|
});
|
|
expect(reclaimerCancelled).toBe(true);
|
|
expect(processBudget.activeCount).toBe(1);
|
|
expect(processBudget.waitingCount).toBe(0);
|
|
parentLease.release();
|
|
expect(processBudget.activeCount).toBe(0);
|
|
await scheduler.close();
|
|
});
|
|
});
|