makelore/tests/unit/pi-extension-host.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 } from 'vitest';
import { PiManagedExtensionHost } from '../../electron/coding-runtime/pi/extension-host';
import { PiProjectWriteLeaseCoordinator } from '../../electron/coding-runtime/pi/write-lease';

const roots: string[] = [];
const hosts: PiManagedExtensionHost[] = [];

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

async function post(
  registration: Awaited<ReturnType<PiManagedExtensionHost['registerWorker']>>,
  body: Record<string, unknown>,
): Promise<Response> {
  return await fetch(registration.env.MAKELORE_PI_BRIDGE_URL as string, {
    method: 'POST',
    headers: {
      authorization: `Bearer ${registration.env.MAKELORE_PI_WORKER_TOKEN}`,
      'content-type': 'application/json',
    },
    body: JSON.stringify(body),
  });
}

describe('managed Pi extension bridge', () => {
  it('finishes a queued HTTP request while closing a holder and waiter', async () => {
    const root = await mkdtemp(path.join(tmpdir(), 'makelore-pi-extension-close-'));
    roots.push(root);
    const leases = new PiProjectWriteLeaseCoordinator();
    const host = new PiManagedExtensionHost(leases);
    hosts.push(host);
    const holder = await host.registerWorker({
      conversationId: 'conversation-holder', generation: 1, projectId: 'project-a', extensionsDir: root,
    });
    const waiter = await host.registerWorker({
      conversationId: 'conversation-waiter', generation: 1, projectId: 'project-a', extensionsDir: root,
    });
    await Promise.all([
      host.bindRun('conversation-holder', 1, 'run-holder'),
      host.bindRun('conversation-waiter', 1, 'run-waiter'),
    ]);
    const held = await post(holder, {
      action: 'lease.acquire', conversationId: 'conversation-holder', workerGeneration: 1,
      runId: 'run-holder', resourceId: 'held-tool',
    });
    expect(held.status).toBe(200);
    const waiting = post(waiter, {
      action: 'lease.acquire', conversationId: 'conversation-waiter', workerGeneration: 1,
      runId: 'run-waiter', resourceId: 'waiting-tool',
    });
    await expect.poll(() => leases.waitingCount('project-a')).toBe(1);

    await expect(Promise.race([
      host.close().then(() => 'closed'),
      new Promise<string>((resolve) => setTimeout(() => resolve('timeout'), 500)),
    ])).resolves.toBe('closed');
    await expect(waiting).resolves.toMatchObject({ status: 409 });
  });

  it('materializes the active run into a replacement generation before spawn', async () => {
    const root = await mkdtemp(path.join(tmpdir(), 'makelore-pi-extension-rebuild-'));
    roots.push(root);
    const host = new PiManagedExtensionHost();
    hosts.push(host);
    const first = await host.registerWorker({
      conversationId: 'conversation-a', generation: 1, projectId: 'project-a', extensionsDir: root,
    });
    await host.bindRun('conversation-a', 1, 'run-a');
    const replacement = await host.registerWorker({
      conversationId: 'conversation-a', generation: 2, projectId: 'project-a', extensionsDir: root,
    });
    const context = JSON.parse(await readFile(
      replacement.env.MAKELORE_PI_CONTEXT_FILE as string,
      'utf8',
    )) as Record<string, unknown>;
    expect(context).toEqual({
      conversationId: 'conversation-a', workerGeneration: 2, runId: 'run-a',
    });
    await first.dispose();
    const response = await post(replacement, {
      action: 'lease.acquire', conversationId: 'conversation-a', workerGeneration: 2,
      runId: 'run-a', resourceId: 'replacement-tool',
    });
    expect(response.status).toBe(200);
  });

  it('validates worker identity and enforces project-scoped leases over loopback HTTP', async () => {
    const root = await mkdtemp(path.join(tmpdir(), 'makelore-pi-extension-'));
    roots.push(root);
    const host = new PiManagedExtensionHost();
    hosts.push(host);
    const first = await host.registerWorker({
      conversationId: 'conversation-a1', generation: 1, projectId: 'project-a', extensionsDir: root,
    });
    const second = await host.registerWorker({
      conversationId: 'conversation-a2', generation: 1, projectId: 'project-a', extensionsDir: root,
    });
    const other = await host.registerWorker({
      conversationId: 'conversation-b1', generation: 1, projectId: 'project-b', extensionsDir: root,
    });
    await Promise.all([
      host.bindRun('conversation-a1', 1, 'run-a1'),
      host.bindRun('conversation-a2', 1, 'run-a2'),
      host.bindRun('conversation-b1', 1, 'run-b1'),
    ]);
    const identity = (conversationId: string, runId: string, resourceId: string) => ({
      action: 'lease.acquire', conversationId, workerGeneration: 1, runId, resourceId,
    });

    const firstResponse = await post(first, identity('conversation-a1', 'run-a1', 'tool-a1'));
    expect(firstResponse.status).toBe(200);
    const firstLease = await firstResponse.json() as { leaseId: string };
    let sameProjectSettled = false;
    const sameProjectFlight = post(second, identity('conversation-a2', 'run-a2', 'tool-a2'))
      .then((response) => {
        sameProjectSettled = true;
        return response;
      });
    const otherResponse = await post(other, identity('conversation-b1', 'run-b1', 'tool-b1'));
    expect(otherResponse.status).toBe(200);
    await Promise.resolve();
    expect(sameProjectSettled).toBe(false);

    const releaseResponse = await post(first, {
      action: 'lease.release',
      conversationId: 'conversation-a1',
      workerGeneration: 1,
      runId: 'run-a1',
      resourceId: 'tool-a1',
      leaseId: firstLease.leaseId,
    });
    expect(releaseResponse.status).toBe(200);
    expect((await sameProjectFlight).status).toBe(200);

    const forged = await post(second, identity('conversation-a2', 'old-run', 'forged'));
    expect(forged.status).toBe(409);
    await first.dispose();
    const staleToken = await post(first, identity('conversation-a1', 'run-a1', 'stale'));
    expect(staleToken.status).toBe(401);
    const currentWorker = await post(other, {
      action: 'lease.release',
      conversationId: 'conversation-b1', workerGeneration: 1, runId: 'run-b1',
      resourceId: 'tool-b1',
      leaseId: (await otherResponse.clone().json() as { leaseId: string }).leaseId,
    });
    expect(currentWorker.status).toBe(200);
  });
});