makelore/tests/unit/pi-subagent-child.test.ts

// @vitest-environment node

import { mkdtemp, readFile, rm } from 'node:fs/promises';
import { tmpdir } from 'node:os';
import path from 'node:path';
import { afterEach, describe, expect, it, vi } from 'vitest';
import { createCodingProjectAgent } from '../../electron/coding-projects/project-config';
import {
  createCodingProjectStore,
  createLocalCodingProject,
  createMemoryCodingProjectStorage,
} from '../../electron/coding-projects/project-store';
import type { ProviderAccount } from '../../electron/shared/providers/types';
import { PiManagedExtensionHost } from '../../electron/coding-runtime/pi/extension-host';
import type { PiProcessError } from '../../electron/coding-runtime/pi/process-errors';
import type {
  PiRpcCommand,
  PiRpcEvent,
  PiRpcRequestOptions,
  PiRpcResponse,
} from '../../electron/coding-runtime/pi/rpc-client';
import {
  createPiManagedSubagentChildOpener,
  type PiSubagentProcessAdapter,
} from '../../electron/coding-runtime/pi/subagent-child';
import type { PiWorkerProcessOptions } from '../../electron/coding-runtime/pi/worker-process';

const roots: string[] = [];
const NOW = '2026-08-23T00:00:00.000Z';

class FakeChildProcess implements PiSubagentProcessAdapter {
  private readonly listeners = new Set<(event: PiRpcEvent) => void>();
  private readonly invalidationListeners = new Set<(error: PiProcessError) => void>();
  stopped = false;

  async start() { return this; }

  async request<T = unknown>(
    command: PiRpcCommand,
    _options?: PiRpcRequestOptions,
  ): Promise<PiRpcResponse<T>> {
    if (command.type === 'prompt') {
      queueMicrotask(() => {
        for (const listener of this.listeners) {
          listener({
            type: 'message_end',
            message: { role: 'assistant', usage: { input: 7, output: 11, cacheRead: 2 } },
          });
          listener({ type: 'agent_settled' });
        }
      });
    }
    if (command.type === 'get_last_assistant_text') {
      return {
        type: 'response', id: 'summary', success: true,
        data: { text: 'managed child summary' } as T,
      };
    }
    return { type: 'response', id: command.type, success: true, data: {} as T };
  }

  subscribe(listener: (event: PiRpcEvent) => void): () => void {
    this.listeners.add(listener);
    return () => this.listeners.delete(listener);
  }

  subscribeInvalidation(listener: (error: PiProcessError) => void): () => void {
    this.invalidationListeners.add(listener);
    return () => this.invalidationListeners.delete(listener);
  }

  async stop() {
    this.stopped = true;
    return { mode: 'stdin-close' as const, code: 0, signal: null };
  }
}

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

describe('managed Pi subagent child opener', () => {
  it('opens an ephemeral managed Agent with the exact tool profile and public result', async () => {
    const root = await mkdtemp(path.join(tmpdir(), 'makelore-pi-subagent-child-'));
    roots.push(root);
    const projectPath = path.join(root, 'project');
    const projectStore = createCodingProjectStore(createMemoryCodingProjectStorage(), {
      createId: () => 'project-a', now: () => NOW,
    });
    await createLocalCodingProject({ projectPath, now: NOW }, projectStore);
    await createCodingProjectAgent(projectPath, {
      id: 'agent-a', avatarId: 'avatar-01', roleName: 'Reviewer', name: 'Agent A',
      model: { accountId: 'account-a', modelId: 'model-a', thinkingLevel: 'high' },
      modelResolution: 'resolved',
      responsibility: { mission: 'Review', owns: [], boundaries: [], collaborators: [], principles: [] },
      prompt: 'PRIVATE CHILD PROMPT', skillIds: ['grilling'],
    }, { now: NOW });
    const account: ProviderAccount = {
      id: 'account-a', vendorId: 'custom', label: 'Account A', authMode: 'api_key',
      apiProtocol: 'openai-completions', baseUrl: 'https://provider.example/v1', model: 'model-a',
      enabled: true, isDefault: true, createdAt: NOW, updatedAt: NOW,
    };
    const host = new PiManagedExtensionHost();
    const processOptions: PiWorkerProcessOptions[] = [];
    const processes: FakeChildProcess[] = [];
    const opener = createPiManagedSubagentChildOpener({
      projectStore,
      executablePath: 'electron.exe',
      cliPath: 'pi-cli.js',
      userDataDir: path.join(root, 'user-data'),
      bundledSkillsDir: path.resolve('resources/coding-skills'),
      extensionHost: host,
      loadProviderInput: async () => ({ accounts: [account], modelSummaries: [] }),
      resolveCredential: async () => 'provider-secret',
      getRevision: () => ({ provider: 3, resources: 4 }),
      createProcess: (options) => {
        processOptions.push(options);
        const child = new FakeChildProcess();
        processes.push(child);
        return child;
      },
    });
    const identity = {
      conversationId: 'conversation-a', workerGeneration: 2, runId: 'run-a', projectId: 'project-a',
      dispatchId: 'dispatch-a', agentId: 'agent-a',
    };
    const readOnly = await opener({
      ...identity, taskId: 'task-read', toolProfile: 'read-only',
    });
    const coding = await opener({
      ...identity, taskId: 'task-code', toolProfile: 'coding',
    });
    expect(processOptions.map(({ tools }) => tools)).toEqual([
      ['read', 'grep', 'find', 'ls'],
      ['read', 'bash', 'edit', 'write', 'grep', 'find', 'ls'],
    ]);
    for (const options of processOptions) {
      expect(options.additionalArgs).toContain('--no-session');
      expect(options.additionalArgs).not.toContain('--session-id');
      expect(options.additionalArgs?.some((argument) => argument.includes('grilling'))).toBe(false);
      expect(JSON.stringify(options.additionalArgs)).not.toContain('PRIVATE CHILD PROMPT');
      expect(JSON.stringify(options.additionalArgs)).not.toContain('provider-secret');
      expect(Object.values(options.env ?? {})).toContain('provider-secret');
      expect(options.env?.MAKELORE_PI_WORKER_ROLE).toBe('child');
    }
    await expect(readOnly.run('Inspect', new AbortController().signal)).resolves.toEqual({
      summary: 'managed child summary',
      usage: { inputTokens: 7, outputTokens: 11, cacheReadTokens: 2 },
    });
    await readOnly.stop('test_injection');
    await coding.stop('test_injection');
    expect(processes.every(({ stopped }) => stopped)).toBe(true);
    await host.close();
  });

  it('reads the current local-proxy credential for every ephemeral child open', async () => {
    const root = await mkdtemp(path.join(tmpdir(), 'makelore-pi-proxy-child-'));
    roots.push(root);
    const projectPath = path.join(root, 'project');
    const userDataDir = path.join(root, 'user-data');
    const projectStore = createCodingProjectStore(createMemoryCodingProjectStorage(), {
      createId: () => 'project-proxy', now: () => NOW,
    });
    await createLocalCodingProject({ projectPath, now: NOW }, projectStore);
    await createCodingProjectAgent(projectPath, {
      id: 'agent-proxy', avatarId: 'avatar-01', roleName: 'Reviewer', name: 'Proxy Agent',
      model: { accountId: 'account-proxy', modelId: 'model-proxy', thinkingLevel: 'medium' },
      modelResolution: 'resolved',
      responsibility: { mission: 'Review', owns: [], boundaries: [], collaborators: [], principles: [] },
    }, { now: NOW });
    const account: ProviderAccount = {
      id: 'account-proxy', vendorId: 'custom', label: 'Proxy account', authMode: 'api_key',
      apiProtocol: 'openai-completions', baseUrl: 'http://127.0.0.1:43123/api/ai-proxy/v1',
      model: 'model-proxy', enabled: true, isDefault: true, createdAt: NOW, updatedAt: NOW,
      metadata: { worksSquareCredentialMode: 'works_square_ai_gateway_proxy' },
    };
    const host = new PiManagedExtensionHost();
    const processOptions: PiWorkerProcessOptions[] = [];
    let currentToken = 'host-token-child-first';
    const getLocalProxyCredential = vi.fn(async () => currentToken);
    const opener = createPiManagedSubagentChildOpener({
      projectStore,
      executablePath: 'electron.exe',
      cliPath: 'pi-cli.js',
      userDataDir,
      bundledSkillsDir: path.resolve('resources/coding-skills'),
      extensionHost: host,
      loadProviderInput: async () => ({ accounts: [account], modelSummaries: [] }),
      resolveCredential: vi.fn(async () => 'stale-secret-store-token'),
      getRevision: () => ({ provider: 1, resources: 1 }),
      getLocalProxyCredential,
      createProcess: (options) => {
        processOptions.push(options);
        return new FakeChildProcess();
      },
    });
    const identity = {
      conversationId: 'conversation-proxy', workerGeneration: 1, runId: 'run-proxy',
      projectId: 'project-proxy', dispatchId: 'dispatch-proxy', agentId: 'agent-proxy',
    };

    const first = await opener({ ...identity, taskId: 'task-child-first', toolProfile: 'read-only' });
    currentToken = 'host-token-child-second';
    const second = await opener({ ...identity, taskId: 'task-child-second', toolProfile: 'coding' });

    expect(getLocalProxyCredential).toHaveBeenCalledTimes(2);
    expect(processOptions.map(({ env }) => Object.values(env ?? {}).find((value) => value.startsWith('host-token-'))))
      .toEqual(['host-token-child-first', 'host-token-child-second']);
    expect(JSON.stringify(processOptions.map(({ additionalArgs }) => additionalArgs))).not.toContain('host-token-');
    expect(await readFile(path.join(userDataDir, 'coding-runtime', 'pi', 'config', 'models.json'), 'utf8'))
      .not.toContain('host-token-');

    await first.stop('test_injection');
    await second.stop('test_injection');
    await host.close();
  });
});