feat: add Pi conversation worker pool

tests/unit/pi-worker-pool-process-integration.test.ts

@@ -0,0 +1,163 @@
+// @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 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<T = unknown>(command: PiRpcCommand, options?: PiRpcRequestOptions): Promise<PiRpcResponse<T>> {
+    return this.process.request(command, options);
+  }
+
+  subscribe(listener: (event: PiRpcEvent) => void): () => void {
+    return this.process.subscribe(listener);
+  }
+
+  subscribeInvalidation(listener: (error: PiProcessError) => void): () => void {
+    return this.process.subscribeInvalidation(listener);
+  }
+
+  stop(): Promise<PiWorkerStopResult> {
+    return this.process.stop();
+  }
+}
+
+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);
+});