makelore/tests/unit/pi-agent-server-process-real.test.ts

// @vitest-environment node

import { copyFile, mkdtemp, mkdir, realpath, rm, symlink, writeFile } from 'node:fs/promises';
import { createServer, type ServerResponse } from 'node:http';
import { createRequire } from 'node:module';
import { tmpdir } from 'node:os';
import path from 'node:path';
import { afterEach, describe, expect, it } from 'vitest';
import { PiAgentServerProcess } from '../../electron/coding-runtime/pi/agent-server-process';
import { materializeMakelorePiExtension } from '../../electron/coding-runtime/pi/extensions/makelore-runtime';
import { MAKELORE_DEFAULT_LANGUAGE_PROMPT } from '../../electron/coding-runtime/pi/resource-loader';
import type { PiWorkerProcessOptions } from '../../electron/coding-runtime/pi/worker-process';

const roots: string[] = [];
const electronExecutable = createRequire(path.resolve('package.json'))('electron') as string;

async function materializePackagedAgentServerLayout(root: string): Promise<{
  configDir: string;
  runtimeRoot: string;
  serverPath: string;
}> {
  const resourcesPath = path.join(root, 'artifact', 'resources');
  const serverPath = path.join(resourcesPath, 'resources', 'pi-agent-server.mjs');
  const runtimeRoot = path.join(resourcesPath, 'pi-runtime');
  const configDir = path.join(root, 'config');
  const piAiRoot = path.join(runtimeRoot, 'node_modules', '@earendil-works', 'pi-ai');
  const [sourceRuntimeRoot, sourcePiAiRoot] = await Promise.all([
    realpath(path.resolve('node_modules/@earendil-works/pi-coding-agent')),
    realpath(path.resolve('node_modules/@earendil-works/pi-ai')),
  ]);
  await Promise.all([
    mkdir(path.dirname(serverPath), { recursive: true }),
    mkdir(path.dirname(piAiRoot), { recursive: true }),
    mkdir(configDir, { recursive: true }),
  ]);
  const directoryLinkType = process.platform === 'win32' ? 'junction' : 'dir';
  await Promise.all([
    copyFile(path.resolve('resources/pi-agent-server.mjs'), serverPath),
    copyFile(path.join(sourceRuntimeRoot, 'package.json'), path.join(runtimeRoot, 'package.json')),
    symlink(path.join(sourceRuntimeRoot, 'dist'), path.join(runtimeRoot, 'dist'), directoryLinkType),
    symlink(sourcePiAiRoot, piAiRoot, directoryLinkType),
  ]);
  return { configDir, runtimeRoot, serverPath };
}

async function startHeldProvider(): Promise<{
  baseUrl: string;
  requests: Array<{ authorization: string; body: string }>;
  activeCount(): number;
  maxActiveCount(): number;
  release(): void;
  close(): Promise<void>;
}> {
  const requests: Array<{ authorization: string; body: string }> = [];
  const held = new Set<ServerResponse>();
  let maxActive = 0;
  const provider = createServer((request, response) => {
    const chunks: Buffer[] = [];
    request.on('data', (chunk: Buffer) => chunks.push(chunk));
    request.on('end', () => {
      requests.push({
        authorization: String(request.headers.authorization ?? ''),
        body: Buffer.concat(chunks).toString('utf8'),
      });
      held.add(response);
      maxActive = Math.max(maxActive, held.size);
      response.once('close', () => held.delete(response));
    });
  });
  await new Promise<void>((resolve, reject) => {
    provider.once('error', reject);
    provider.listen(0, '127.0.0.1', resolve);
  });
  const address = provider.address();
  if (!address || typeof address === 'string') throw new Error('Test Provider did not bind');
  const release = () => {
    for (const response of [...held]) {
      held.delete(response);
      response.writeHead(200, { 'content-type': 'text/event-stream' });
      response.write(`data: ${JSON.stringify({
        id: 'chatcmpl-agent-server-test',
        object: 'chat.completion.chunk',
        created: Math.floor(Date.now() / 1_000),
        model: 'test-model',
        choices: [{ index: 0, delta: { role: 'assistant', content: 'DONE' }, finish_reason: null }],
      })}\n\n`);
      response.write(`data: ${JSON.stringify({
        id: 'chatcmpl-agent-server-test',
        object: 'chat.completion.chunk',
        created: Math.floor(Date.now() / 1_000),
        model: 'test-model',
        choices: [{ index: 0, delta: {}, finish_reason: 'stop' }],
        usage: { prompt_tokens: 1, completion_tokens: 1, total_tokens: 2 },
      })}\n\n`);
      response.end('data: [DONE]\n\n');
    }
  };
  return {
    baseUrl: `http://127.0.0.1:${address.port}/v1`,
    requests,
    activeCount: () => held.size,
    maxActiveCount: () => maxActive,
    release,
    close: async () => {
      release();
      await new Promise<void>((resolve, reject) => {
        provider.close((error) => error ? reject(error) : resolve());
        provider.closeIdleConnections?.();
      });
    },
  };
}

afterEach(async () => {
  await Promise.all(roots.splice(0).map((root) => rm(root, { recursive: true, force: true })));
});

describe('Pi Agent Server real process', () => {
  it('boots from the packaged sibling resource layout', async () => {
    const root = await mkdtemp(path.join(tmpdir(), 'makelore-pi-agent-server-layout-'));
    roots.push(root);
    const layout = await materializePackagedAgentServerLayout(root);
    const server = new PiAgentServerProcess({
      executablePath: electronExecutable,
      ...layout,
    });

    try {
      await expect(server.start()).resolves.toBeUndefined();
      expect(server.processId).toBeTypeOf('number');
    } finally {
      await server.stop();
    }
  }, 10_000);

  it('restarts when start races with a graceful background stop', async () => {
    const root = await mkdtemp(path.join(tmpdir(), 'makelore-pi-agent-server-restart-'));
    roots.push(root);
    const layout = await materializePackagedAgentServerLayout(root);
    const server = new PiAgentServerProcess({
      executablePath: electronExecutable,
      ...layout,
    });

    try {
      await server.start();
      const firstProcessId = server.processId;
      expect(firstProcessId).toBeTypeOf('number');

      const stopping = server.stop();
      const restarting = server.start();
      await Promise.all([stopping, restarting]);
      expect(server.processId).toBeTypeOf('number');
      expect(server.processId).not.toBe(firstProcessId);
    } finally {
      await server.stop();
    }
  }, 20_000);

  it('hosts isolated Conversation threads in one long-lived process', async () => {
    const root = await mkdtemp(path.join(tmpdir(), 'makelore-pi-agent-server-'));
    roots.push(root);
    const projectPath = path.join(root, 'project');
    const configDir = path.join(root, 'config');
    const sessionDir = path.join(root, 'sessions');
    const extensionDir = path.join(root, 'extensions');
    await Promise.all([
      mkdir(projectPath, { recursive: true }),
      mkdir(configDir, { recursive: true }),
      mkdir(sessionDir, { recursive: true }),
      mkdir(extensionDir, { recursive: true }),
    ]);
    const provider = await startHeldProvider();
    await writeFile(path.join(configDir, 'models.json'), JSON.stringify({
      providers: {
        'makelore-test': {
          baseUrl: provider.baseUrl,
          api: 'openai-completions',
          apiKey: '$MAKELORE_TEST_KEY',
          models: [{
            id: 'test-model',
            name: 'Test model',
            reasoning: true,
            input: ['text'],
            contextWindow: 32_000,
            maxTokens: 4_096,
            cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 },
          }],
        },
      },
    }));
    const promptPath = path.join(root, 'system.md');
    const languagePromptPath = path.join(root, 'language.md');
    await Promise.all([
      writeFile(promptPath, 'You are a test Agent.'),
      writeFile(languagePromptPath, MAKELORE_DEFAULT_LANGUAGE_PROMPT),
    ]);
    const extensionPath = await materializeMakelorePiExtension(extensionDir);
    const runtimeRoot = path.resolve('node_modules/@earendil-works/pi-coding-agent');
    const server = new PiAgentServerProcess({
      executablePath: process.execPath,
      serverPath: path.resolve('resources/pi-agent-server.mjs'),
      runtimeRoot,
      configDir,
    });

    const workerOptions = (conversationId: string, generation: number): PiWorkerProcessOptions => {
      const contextFile = path.join(root, `${conversationId}.json`);
      return {
        executablePath: process.execPath,
        cliPath: path.join(runtimeRoot, 'dist', 'cli.js'),
        cwd: projectPath,
        configDir,
        sessionDir,
        conversationId,
        workerGeneration: generation,
        additionalArgs: [
          '--provider', 'makelore-test',
          '--model', 'test-model',
          '--thinking', 'medium',
          '--system-prompt', promptPath,
          '--append-system-prompt', languagePromptPath,
          '--extension', extensionPath,
          '--session-id', `session-${conversationId}`,
        ],
        env: {
          MAKELORE_TEST_KEY: `secret-${conversationId}`,
          MAKELORE_PI_BRIDGE_URL: 'http://127.0.0.1:9/v1/worker',
          MAKELORE_PI_WORKER_TOKEN: `token-${conversationId}`,
          MAKELORE_PI_CONTEXT_FILE: contextFile,
          MAKELORE_PI_WORKER_ROLE: 'parent',
          MAKELORE_PI_PROJECT_PATH: projectPath,
        },
        sensitiveValues: [`secret-${conversationId}`, `token-${conversationId}`],
      };
    };

    await Promise.all(['left', 'right'].map(async (conversationId) => {
      await writeFile(path.join(root, `${conversationId}.json`), JSON.stringify({
        conversationId,
        workerGeneration: 1,
        role: 'parent',
        allowedToolNames: [],
        tools: [],
      }));
    }));

    const left = server.createWorker(workerOptions('left', 1));
    const right = server.createWorker(workerOptions('right', 1));
    try {
      await Promise.all([left.start(), right.start()]);
      const processId = server.processId;
      expect(processId).toBeTypeOf('number');
      expect(server.activeThreadCount).toBe(2);

      await expect(left.request<{ sessionId: string }>({ type: 'get_state' }))
        .resolves.toMatchObject({ data: { sessionId: 'session-left' } });
      await expect(right.request<{ sessionId: string }>({ type: 'get_state' }))
        .resolves.toMatchObject({ data: { sessionId: 'session-right' } });
      await expect(left.request<{ models: Array<{ id: string }> }>({ type: 'get_available_models' }))
        .resolves.toMatchObject({
          data: { models: expect.arrayContaining([expect.objectContaining({ id: 'test-model' })]) },
        });

      const settled = (worker: typeof left) => new Promise<void>((resolve) => {
        const unsubscribe = worker.subscribe((event) => {
          if (event.type !== 'agent_settled') return;
          unsubscribe();
          resolve();
        });
      });
      const leftSettled = settled(left);
      const rightSettled = settled(right);
      await Promise.all([
        left.request({ type: 'prompt', message: 'LEFT_THREAD' }),
        right.request({ type: 'prompt', message: 'RIGHT_THREAD' }),
      ]);
      await expect.poll(() => provider.activeCount()).toBe(2);
      expect(provider.maxActiveCount()).toBe(2);
      expect(provider.requests.map(({ authorization }) => authorization).sort()).toEqual([
        'Bearer secret-left',
        'Bearer secret-right',
      ]);
      expect(provider.requests.some(({ body }) => body.includes('LEFT_THREAD'))).toBe(true);
      expect(provider.requests.some(({ body }) => body.includes('RIGHT_THREAD'))).toBe(true);
      expect(provider.requests.every(({ body }) => body.includes('You are a test Agent.'))).toBe(true);
      expect(provider.requests.every(({ body }) => body.includes('默认使用简体中文'))).toBe(true);
      provider.release();
      await Promise.all([leftSettled, rightSettled]);

      await left.stop('test_injection');
      expect(server.processId).toBe(processId);
      expect(server.activeThreadCount).toBe(1);
      await expect(right.request<{ sessionId: string }>({ type: 'get_state' }))
        .resolves.toMatchObject({ data: { sessionId: 'session-right' } });

      const invalidated = new Promise<string>((resolve) => {
        right.subscribeInvalidation((error) => resolve(error.code));
      });
      process.kill(processId as number, 'SIGKILL');
      await expect(invalidated).resolves.toBe('PI_RPC_EXITED');
      await expect.poll(() => server.processId).toBeUndefined();
      await expect(right.request({ type: 'get_state' })).rejects.toMatchObject({
        code: 'PI_RPC_EXITED',
      });
      await right.stop('test_injection');

      const replacementLeft = server.createWorker(workerOptions('left', 2));
      const replacementRight = server.createWorker(workerOptions('right', 2));
      await Promise.all([replacementLeft.start(), replacementRight.start()]);
      expect(server.processId).toBeTypeOf('number');
      expect(server.processId).not.toBe(processId);
      expect(server.activeThreadCount).toBe(2);
      await expect(replacementRight.request<{ sessionId: string }>({ type: 'get_state' }))
        .resolves.toMatchObject({ data: { sessionId: 'session-right' } });
      await Promise.all([
        replacementLeft.stop('test_injection'),
        replacementRight.stop('test_injection'),
      ]);
    } finally {
      await server.stop();
      await provider.close();
    }
    expect(server.processId).toBeUndefined();
  }, 20_000);
});