fix: reclaim queued Pi parent capacity
This commit is contained in:
@@ -12,6 +12,7 @@ import type { PiProcessError } from '../../electron/coding-runtime/pi/process-er
|
||||
import { PiProcessError as PiProcessFailure } from '../../electron/coding-runtime/pi/process-errors';
|
||||
import type { PiRpcCommand, PiRpcEvent } from '../../electron/coding-runtime/pi/rpc-client';
|
||||
import type { PiRuntimeTelemetryEvent } from '../../electron/coding-runtime/pi/telemetry';
|
||||
import { PiSubagentScheduler } from '../../electron/coding-runtime/pi/subagent';
|
||||
|
||||
const MODEL = {
|
||||
model: {
|
||||
@@ -254,6 +255,105 @@ describe('Pi worker pool', () => {
|
||||
expect(processBudget.waitingCount).toBe(0);
|
||||
});
|
||||
|
||||
it('suspends and later resumes the oldest queued parent to make child capacity', async () => {
|
||||
const processBudget = new PiProcessBudget(8);
|
||||
const workers = new Map<string, FakeWorker[]>();
|
||||
const pool = new PiWorkerPool({
|
||||
processBudget,
|
||||
maxRunning: 4,
|
||||
maxIdle: 8,
|
||||
openWorker: async ({ conversation: input, generation, existingSession }) => {
|
||||
const worker = new FakeWorker(`worker-${input.conversationId}-${generation}`);
|
||||
workers.set(input.conversationId, [...(workers.get(input.conversationId) ?? []), worker]);
|
||||
return {
|
||||
worker,
|
||||
session: existingSession ?? {
|
||||
piSessionId: `session-${input.conversationId}`,
|
||||
sessionKey: `key-${input.conversationId}`,
|
||||
},
|
||||
};
|
||||
},
|
||||
});
|
||||
const ids = Array.from({ length: 8 }, (_, index) => `conversation-${index + 1}`);
|
||||
for (const id of ids) await pool.prepare(conversation(id));
|
||||
const tickets = ids.map((conversationId, index) => pool.startTopLevel({
|
||||
conversationId,
|
||||
runId: `run-${index + 1}`,
|
||||
command: { type: 'prompt', message: conversationId },
|
||||
}));
|
||||
for (const ticket of tickets) void ticket.accepted.catch(() => undefined);
|
||||
await expect.poll(() => ids.map((id) => workers.get(id)?.[0]?.requests.length ?? 0))
|
||||
.toEqual([1, 1, 1, 1, 0, 0, 0, 0]);
|
||||
expect(ids.slice(4).map((id) => pool.getState(id)?.state))
|
||||
.toEqual(['queued', 'queued', 'queued', 'queued']);
|
||||
expect(processBudget.activeCount).toBe(8);
|
||||
|
||||
const scheduler = new PiSubagentScheduler({
|
||||
processBudget,
|
||||
reclaimProcessCapacity: (signal) => pool.reclaimIdleWorker(signal),
|
||||
openChild: async (input) => ({
|
||||
id: input.taskId,
|
||||
async run() { return { summary: 'child complete' }; },
|
||||
async stop() {},
|
||||
}),
|
||||
});
|
||||
await expect(scheduler.dispatch({
|
||||
conversationId: ids[0]!,
|
||||
workerGeneration: 1,
|
||||
runId: 'run-child',
|
||||
projectId: 'project-a',
|
||||
request: {
|
||||
mode: 'parallel',
|
||||
tasks: Array.from({ length: 4 }, (_, index) => ({
|
||||
agentId: `agent-${index + 1}`,
|
||||
task: `inspect ${index + 1}`,
|
||||
toolProfile: 'read-only',
|
||||
})),
|
||||
},
|
||||
})).resolves.toMatchObject({
|
||||
details: {
|
||||
tasks: Array.from({ length: 4 }, () => ({ status: 'complete' })),
|
||||
},
|
||||
});
|
||||
|
||||
const firstQueuedId = ids[4]!;
|
||||
expect(ids.slice(4).map((id) => workers.get(id)?.[0]?.stopped))
|
||||
.toEqual([true, true, true, true]);
|
||||
expect(pool.getState(firstQueuedId)).toMatchObject({
|
||||
state: 'queued',
|
||||
generation: 1,
|
||||
session: {
|
||||
piSessionId: `session-${firstQueuedId}`,
|
||||
sessionKey: `key-${firstQueuedId}`,
|
||||
},
|
||||
});
|
||||
expect(processBudget.activeCount).toBe(4);
|
||||
expect(processBudget.waitingCount).toBe(0);
|
||||
|
||||
for (let index = 0; index < 4; index += 1) {
|
||||
const queuedId = ids[index + 4]!;
|
||||
workers.get(ids[index]!)?.[0]?.emit({ type: 'agent_settled' });
|
||||
await expect.poll(() => workers.get(queuedId)?.length).toBe(2);
|
||||
await expect.poll(() => workers.get(queuedId)?.[1]?.requests.length).toBe(1);
|
||||
await expect(tickets[index + 4]!.accepted).resolves.toMatchObject({ success: true });
|
||||
}
|
||||
expect(pool.getState(firstQueuedId)).toMatchObject({
|
||||
state: 'running',
|
||||
generation: 2,
|
||||
session: {
|
||||
piSessionId: `session-${firstQueuedId}`,
|
||||
sessionKey: `key-${firstQueuedId}`,
|
||||
},
|
||||
});
|
||||
expect(ids.slice(4).map((id) => workers.get(id)?.length)).toEqual([2, 2, 2, 2]);
|
||||
expect(processBudget.activeCount).toBe(8);
|
||||
|
||||
await scheduler.close();
|
||||
await pool.shutdown();
|
||||
expect(processBudget.activeCount).toBe(0);
|
||||
expect(processBudget.waitingCount).toBe(0);
|
||||
});
|
||||
|
||||
it('never evicts a running worker when the warm-idle LRU exceeds its cap', async () => {
|
||||
const workers = new Map<string, FakeWorker>();
|
||||
const pool = new PiWorkerPool({
|
||||
|
||||
Reference in New Issue
Block a user