// @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>, body: Record, ): Promise { 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((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; 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); }); });