// @vitest-environment node import { mkdir, mkdtemp, rm } from 'node:fs/promises'; import { tmpdir } from 'node:os'; import path from 'node:path'; import { afterEach, describe, expect, it } from 'vitest'; import type { PrepareConversationInput } from '../../electron/coding-runtime/contracts'; import type { PiProcessError } from '../../electron/coding-runtime/pi/process-errors'; import type { PiRpcCommand, PiRpcEvent, PiRpcRequestOptions, PiRpcResponse, } from '../../electron/coding-runtime/pi/rpc-client'; import { PiProcessBudget, PiWorkerPool, type PiConversationWorker, } from '../../electron/coding-runtime/pi/worker-pool'; import { PiWorkerProcess, type PiWorkerStopReason, type PiWorkerStopResult, } from '../../electron/coding-runtime/pi/worker-process'; const scratchRoots: string[] = []; function conversation(conversationId: string): PrepareConversationInput { return { conversationId, projectId: 'project-process', agentId: 'agent-process', title: conversationId, model: { model: { accountId: 'account-process', modelId: 'model-process', thinkingLevel: 'medium' }, modelResolution: 'resolved', }, }; } class ProcessBackedWorker implements PiConversationWorker { constructor( readonly id: string, readonly generation: number, private readonly process: PiWorkerProcess, ) {} request(command: PiRpcCommand, options?: PiRpcRequestOptions): Promise> { return this.process.request(command, options); } send(command: PiRpcCommand): Promise { return this.process.send(command); } subscribe(listener: (event: PiRpcEvent) => void): () => void { return this.process.subscribe(listener); } subscribeInvalidation(listener: (error: PiProcessError) => void): () => void { return this.process.subscribeInvalidation(listener); } stop(reason: PiWorkerStopReason): Promise { return this.process.stop(reason); } } afterEach(async () => { await Promise.all(scratchRoots.splice(0).map((root) => rm(root, { recursive: true, force: true, maxRetries: 3, }))); }); describe('Pi worker pool process integration', () => { it('runs two child processes concurrently and aborts only the addressed Conversation', async () => { const root = await mkdtemp(path.join(tmpdir(), 'makelore-pi-pool-process-')); scratchRoots.push(root); const processBudget = new PiProcessBudget(8); const events: Array<{ conversationId: string; event: PiRpcEvent }> = []; const pool = new PiWorkerPool({ maxRunning: 4, maxIdle: 4, processBudget, openWorker: async ({ conversation: input, generation }) => { const workerRoot = path.join(root, input.conversationId); const configDir = path.join(workerRoot, 'config'); const sessionDir = path.join(workerRoot, 'sessions'); const cwd = path.join(workerRoot, 'project'); await Promise.all([ mkdir(configDir, { recursive: true }), mkdir(sessionDir, { recursive: true }), mkdir(cwd, { recursive: true }), ]); const child = await new PiWorkerProcess({ executablePath: process.execPath, cliPath: path.resolve('tests/fixtures/fake-pi-pool-child.mjs'), cwd, configDir, sessionDir, commandTimeoutMs: 2_000, shutdownGraceMs: 500, }).start(); return { worker: new ProcessBackedWorker( `worker-${input.conversationId}-${generation}`, generation, child, ), session: { piSessionId: `session-${input.conversationId}`, sessionKey: `key-${input.conversationId}`, }, }; }, }); const unsubscribe = pool.subscribe((event) => { if (event.type === 'worker.event') { events.push({ conversationId: event.conversationId, event: event.event }); } }); try { await Promise.all([ pool.prepare(conversation('conversation-left')), pool.prepare(conversation('conversation-right')), ]); const left = pool.startTopLevel({ conversationId: 'conversation-left', runId: 'run-left', command: { type: 'prompt', message: 'left', delayMs: 1_000 }, }); const right = pool.startTopLevel({ conversationId: 'conversation-right', runId: 'run-right', command: { type: 'prompt', message: 'right', delayMs: 150 }, }); await Promise.all([left.accepted, right.accepted]); expect(pool.getState('conversation-left')?.state).toBe('running'); expect(pool.getState('conversation-right')?.state).toBe('running'); await pool.request('conversation-left', { type: 'abort' }); await expect.poll(() => pool.getState('conversation-left')?.state).toBe('idle'); expect(pool.getState('conversation-right')?.state).toBe('running'); await expect.poll(() => pool.getState('conversation-right')?.state).toBe('idle'); expect(events).toEqual(expect.arrayContaining([ expect.objectContaining({ conversationId: 'conversation-left', event: expect.objectContaining({ type: 'agent_settled', marker: 'left', reason: 'aborted' }), }), expect.objectContaining({ conversationId: 'conversation-right', event: expect.objectContaining({ type: 'agent_settled', marker: 'right', reason: 'completed' }), }), ])); expect(events.some(({ conversationId, event }) => ( conversationId === 'conversation-right' && event.reason === 'aborted' ))).toBe(false); } finally { unsubscribe(); await pool.shutdown(); } expect(processBudget.activeCount).toBe(0); }, 10_000); });