369 lines
15 KiB
TypeScript
369 lines
15 KiB
TypeScript
// @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';
|
|
import {
|
|
PiSubagentChildError,
|
|
PiSubagentScheduler,
|
|
} from '../../electron/coding-runtime/pi/subagent';
|
|
import { PiProcessBudget } from '../../electron/coding-runtime/pi/worker-pool';
|
|
|
|
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('single-flights the shared managed extension for concurrent worker registration', async () => {
|
|
const root = await mkdtemp(path.join(tmpdir(), 'makelore-pi-extension-concurrent-'));
|
|
roots.push(root);
|
|
const host = new PiManagedExtensionHost();
|
|
hosts.push(host);
|
|
|
|
const registrations = await Promise.all(Array.from({ length: 8 }, async (_, index) => (
|
|
await host.registerWorker({
|
|
conversationId: `conversation-${index + 1}`,
|
|
generation: 1,
|
|
projectId: `project-${index + 1}`,
|
|
extensionsDir: root,
|
|
})
|
|
)));
|
|
|
|
expect(new Set(registrations.map(({ extensionPath }) => extensionPath)).size).toBe(1);
|
|
expect(host.getDiagnostics().registrations).toEqual({ parent: 8, child: 0 });
|
|
await Promise.all(registrations.map(({ dispose }) => dispose()));
|
|
});
|
|
|
|
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, role: 'parent', 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('streams stable subagent details and rejects recursive child dispatch', async () => {
|
|
const root = await mkdtemp(path.join(tmpdir(), 'makelore-pi-extension-subagent-'));
|
|
roots.push(root);
|
|
const scheduler = new PiSubagentScheduler({
|
|
processBudget: new PiProcessBudget(8),
|
|
openChild: async (input) => ({
|
|
id: input.taskId,
|
|
async run() { return { summary: `done ${input.agentId}` }; },
|
|
async stop() {},
|
|
}),
|
|
});
|
|
const host = new PiManagedExtensionHost();
|
|
const tracked: Array<{ kind: string; id: string }> = [];
|
|
let untracked = 0;
|
|
host.configureSubagents({
|
|
scheduler,
|
|
trackGenerationResource: (input) => {
|
|
tracked.push({ kind: input.kind, id: input.id });
|
|
return () => { untracked += 1; };
|
|
},
|
|
});
|
|
hosts.push(host);
|
|
const parent = await host.registerWorker({
|
|
conversationId: 'conversation-a', generation: 1, projectId: 'project-a', extensionsDir: root,
|
|
});
|
|
const child = await host.registerWorker({
|
|
conversationId: 'conversation-a', generation: 1, projectId: 'project-a', extensionsDir: root,
|
|
role: 'child', runId: 'run-a',
|
|
});
|
|
await host.bindRun('conversation-a', 1, 'run-a');
|
|
const request = {
|
|
action: 'subagent.dispatch', conversationId: 'conversation-a', workerGeneration: 1,
|
|
runId: 'run-a', resourceId: 'subagent-tool',
|
|
request: {
|
|
mode: 'single',
|
|
tasks: [{ agentId: 'agent-a', task: 'Inspect', toolProfile: 'read-only' }],
|
|
},
|
|
};
|
|
const response = await post(parent, request);
|
|
expect(response.status).toBe(200);
|
|
const lines = (await response.text()).trim().split('\n').map((line) => JSON.parse(line));
|
|
expect(lines.at(-1)).toMatchObject({
|
|
done: true,
|
|
details: {
|
|
schema: 'subagent.v1', mode: 'single',
|
|
tasks: [{ agentId: 'agent-a', status: 'complete', summary: 'done agent-a' }],
|
|
},
|
|
});
|
|
expect(tracked).toEqual([{ kind: 'child', id: 'subagent-tool' }]);
|
|
expect(untracked).toBe(1);
|
|
expect((await post(child, request)).status).toBe(403);
|
|
await scheduler.close();
|
|
});
|
|
|
|
it('propagates generation cancellation without an orphan child or leaked budget', async () => {
|
|
const root = await mkdtemp(path.join(tmpdir(), 'makelore-pi-extension-child-cancel-'));
|
|
roots.push(root);
|
|
const processBudget = new PiProcessBudget(8);
|
|
let stopped = 0;
|
|
const scheduler = new PiSubagentScheduler({
|
|
processBudget,
|
|
openChild: async (input) => ({
|
|
id: input.taskId,
|
|
async run(_prompt, signal) {
|
|
await new Promise<void>((_resolve, reject) => {
|
|
const abort = () => reject(new PiSubagentChildError('SUBAGENT_ABORTED'));
|
|
if (signal.aborted) abort();
|
|
else signal.addEventListener('abort', abort, { once: true });
|
|
});
|
|
return { summary: 'unreachable' };
|
|
},
|
|
async stop() { stopped += 1; },
|
|
}),
|
|
});
|
|
let cancelGeneration: (() => void) | undefined;
|
|
let untracked = 0;
|
|
const host = new PiManagedExtensionHost();
|
|
host.configureSubagents({
|
|
scheduler,
|
|
trackGenerationResource: (input) => {
|
|
cancelGeneration = input.cancel;
|
|
return () => { untracked += 1; };
|
|
},
|
|
});
|
|
hosts.push(host);
|
|
const parent = await host.registerWorker({
|
|
conversationId: 'conversation-a', generation: 1, projectId: 'project-a', extensionsDir: root,
|
|
});
|
|
await host.bindRun('conversation-a', 1, 'run-a');
|
|
const flight = post(parent, {
|
|
action: 'subagent.dispatch', conversationId: 'conversation-a', workerGeneration: 1,
|
|
runId: 'run-a', resourceId: 'subagent-tool',
|
|
request: {
|
|
mode: 'single',
|
|
tasks: [{ agentId: 'agent-a', task: 'Wait', toolProfile: 'read-only' }],
|
|
},
|
|
});
|
|
await expect.poll(() => processBudget.activeCount).toBe(1);
|
|
cancelGeneration?.();
|
|
const response = await flight;
|
|
const lines = (await response.text()).trim().split('\n').map((line) => JSON.parse(line));
|
|
expect(lines.at(-1)).toMatchObject({
|
|
done: true,
|
|
details: { tasks: [{ status: 'aborted', errorCode: 'SUBAGENT_ABORTED' }] },
|
|
});
|
|
expect(stopped).toBe(1);
|
|
expect(processBudget.activeCount).toBe(0);
|
|
expect(untracked).toBe(1);
|
|
await scheduler.close();
|
|
});
|
|
|
|
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);
|
|
});
|
|
|
|
it('releases a crashed worker write lease so only its same-project waiter advances', async () => {
|
|
const root = await mkdtemp(path.join(tmpdir(), 'makelore-pi-extension-crash-lease-'));
|
|
roots.push(root);
|
|
const leases = new PiProjectWriteLeaseCoordinator();
|
|
const host = new PiManagedExtensionHost(leases);
|
|
hosts.push(host);
|
|
const crashed = await host.registerWorker({
|
|
conversationId: 'conversation-crashed', generation: 1, projectId: 'project-a', extensionsDir: root,
|
|
});
|
|
const waiter = await host.registerWorker({
|
|
conversationId: 'conversation-waiter', generation: 1, projectId: 'project-a', extensionsDir: root,
|
|
});
|
|
const other = await host.registerWorker({
|
|
conversationId: 'conversation-other', generation: 1, projectId: 'project-b', extensionsDir: root,
|
|
});
|
|
await Promise.all([
|
|
host.bindRun('conversation-crashed', 1, 'run-crashed'),
|
|
host.bindRun('conversation-waiter', 1, 'run-waiter'),
|
|
host.bindRun('conversation-other', 1, 'run-other'),
|
|
]);
|
|
expect((await post(crashed, {
|
|
action: 'lease.acquire', conversationId: 'conversation-crashed', workerGeneration: 1,
|
|
runId: 'run-crashed', resourceId: 'crashed-write',
|
|
})).status).toBe(200);
|
|
const waiting = post(waiter, {
|
|
action: 'lease.acquire', conversationId: 'conversation-waiter', workerGeneration: 1,
|
|
runId: 'run-waiter', resourceId: 'waiting-write',
|
|
});
|
|
expect((await post(other, {
|
|
action: 'lease.acquire', conversationId: 'conversation-other', workerGeneration: 1,
|
|
runId: 'run-other', resourceId: 'other-write',
|
|
})).status).toBe(200);
|
|
await expect.poll(() => leases.waitingCount('project-a')).toBe(1);
|
|
|
|
await crashed.dispose();
|
|
|
|
await expect(waiting).resolves.toMatchObject({ status: 200 });
|
|
expect(leases.waitingCount()).toBe(0);
|
|
expect(leases.activeCount).toBe(2);
|
|
await Promise.all([waiter.dispose(), other.dispose()]);
|
|
expect(leases.activeCount).toBe(0);
|
|
});
|
|
|
|
it('joins a coding child to the same project write lease as its parent', async () => {
|
|
const root = await mkdtemp(path.join(tmpdir(), 'makelore-pi-extension-child-lease-'));
|
|
roots.push(root);
|
|
const leases = new PiProjectWriteLeaseCoordinator();
|
|
const host = new PiManagedExtensionHost(leases);
|
|
hosts.push(host);
|
|
const parent = await host.registerWorker({
|
|
conversationId: 'conversation-a', generation: 1, projectId: 'project-a', extensionsDir: root,
|
|
});
|
|
const child = await host.registerWorker({
|
|
conversationId: 'conversation-a', generation: 1, projectId: 'project-a', extensionsDir: root,
|
|
role: 'child', runId: 'run-a',
|
|
});
|
|
await host.bindRun('conversation-a', 1, 'run-a');
|
|
const identity = { conversationId: 'conversation-a', workerGeneration: 1, runId: 'run-a' };
|
|
const held = await post(parent, {
|
|
...identity, action: 'lease.acquire', resourceId: 'parent-write',
|
|
});
|
|
const parentLease = await held.json() as { leaseId: string };
|
|
let childSettled = false;
|
|
const waiting = post(child, {
|
|
...identity, action: 'lease.acquire', resourceId: 'child-write',
|
|
}).then((response) => {
|
|
childSettled = true;
|
|
return response;
|
|
});
|
|
await expect.poll(() => leases.waitingCount('project-a')).toBe(1);
|
|
expect(childSettled).toBe(false);
|
|
expect((await post(parent, {
|
|
...identity, action: 'lease.release', resourceId: 'parent-write', leaseId: parentLease.leaseId,
|
|
})).status).toBe(200);
|
|
expect((await waiting).status).toBe(200);
|
|
});
|
|
});
|