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

// @vitest-environment node

import { mkdtemp, mkdir, rm, writeFile } from 'node:fs/promises';
import { createServer, type ServerResponse } from 'node:http';
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 { PiRpcEvent } from '../../electron/coding-runtime/pi/rpc-client';
import type { PiWorkerProcessOptions } from '../../electron/coding-runtime/pi/worker-process';

const roots: string[] = [];

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('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 invalidJsonExtensionPath = path.join(extensionDir, 'invalid-json-extension.mjs');
    await writeFile(invalidJsonExtensionPath, `
export function createMakeloreRuntime() {
  return async function invalidJsonRuntime(pi) {
    pi.on('message_end', async (event) => ({
      message: { ...event.message, makeloreInvalidJson: 1n },
    }));
  };
}
`);
    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,
      selectedExtensionPath = extensionPath,
    ): 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', selectedExtensionPath,
          '--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', 'broken'].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]);

      const broken = server.createWorker(workerOptions('broken', 1, invalidJsonExtensionPath));
      await broken.start();
      const postAcceptanceFailure = new Promise<PiRpcEvent>((resolve) => {
        const unsubscribe = broken.subscribe((event) => {
          if (event.type !== 'makelore_thread_error') return;
          unsubscribe();
          resolve(event);
        });
      });
      await expect(broken.request({ type: 'prompt', message: 'BREAK_AFTER_ACCEPTANCE' }))
        .resolves.toMatchObject({ success: true });
      await expect(postAcceptanceFailure).resolves.toMatchObject({
        type: 'makelore_thread_error',
        code: 'PROMPT_FAILED_AFTER_ACCEPTANCE',
      });
      await expect(broken.request<{ isStreaming: boolean }>({ type: 'get_state' }))
        .resolves.toMatchObject({ data: { isStreaming: false } });
      await broken.stop('test_injection');

      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);
});