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

// @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<T = unknown>(command: PiRpcCommand, options?: PiRpcRequestOptions): Promise<PiRpcResponse<T>> {
    return this.process.request(command, options);
  }

  send(command: PiRpcCommand): Promise<void> {
    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<PiWorkerStopResult> {
    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);
});