96 lines
4.0 KiB
TypeScript
96 lines
4.0 KiB
TypeScript
// @vitest-environment node
|
|
|
|
import { mkdtemp, rm } from 'node:fs/promises';
|
|
import { tmpdir } from 'node:os';
|
|
import path from 'node:path';
|
|
import { pathToFileURL } from 'node:url';
|
|
import { afterEach, describe, expect, it } from 'vitest';
|
|
import { PiManagedExtensionHost } from '../../electron/coding-runtime/pi/extension-host';
|
|
|
|
type ExtensionHandler = (...arguments_: unknown[]) => Promise<unknown> | unknown;
|
|
|
|
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 })));
|
|
});
|
|
|
|
describe('Makelore Pi extension bundle', () => {
|
|
it('loads the real bundle and releases its project lease on tool_result', async () => {
|
|
const root = await mkdtemp(path.join(tmpdir(), 'makelore-pi-extension-bundle-'));
|
|
roots.push(root);
|
|
const host = new PiManagedExtensionHost();
|
|
hosts.push(host);
|
|
const extensionWorker = await host.registerWorker({
|
|
conversationId: 'conversation-a1', generation: 1, projectId: 'project-a', extensionsDir: root,
|
|
});
|
|
const waitingWorker = await host.registerWorker({
|
|
conversationId: 'conversation-a2', generation: 1, projectId: 'project-a', extensionsDir: root,
|
|
});
|
|
await Promise.all([
|
|
host.bindRun('conversation-a1', 1, 'run-a1'),
|
|
host.bindRun('conversation-a2', 1, 'run-a2'),
|
|
]);
|
|
|
|
const previousEnvironment = {
|
|
bridge: process.env.MAKELORE_PI_BRIDGE_URL,
|
|
token: process.env.MAKELORE_PI_WORKER_TOKEN,
|
|
context: process.env.MAKELORE_PI_CONTEXT_FILE,
|
|
};
|
|
Object.assign(process.env, extensionWorker.env);
|
|
try {
|
|
const module = await import(/* @vite-ignore */ pathToFileURL(extensionWorker.extensionPath).href) as {
|
|
default(factory: {
|
|
registerTool(tool: { name: string }): void;
|
|
on(event: string, handler: ExtensionHandler): void;
|
|
}): void;
|
|
};
|
|
const handlers = new Map<string, ExtensionHandler>();
|
|
const tools: string[] = [];
|
|
module.default({
|
|
registerTool: (tool) => tools.push(tool.name),
|
|
on: (event, handler) => handlers.set(event, handler),
|
|
});
|
|
expect(tools).toEqual(['ask_user']);
|
|
|
|
const statuses: Array<string | undefined> = [];
|
|
const context = {
|
|
signal: new AbortController().signal,
|
|
ui: { setStatus: (_key: string, text: string | undefined) => statuses.push(text) },
|
|
};
|
|
await handlers.get('tool_call')?.({ toolName: 'read', toolCallId: 'read-1' }, context);
|
|
await handlers.get('tool_call')?.({ toolName: 'write', toolCallId: 'write-1' }, context);
|
|
expect(statuses).toEqual(['等待项目写入', undefined]);
|
|
|
|
let waiterSettled = false;
|
|
const waiting = fetch(waitingWorker.env.MAKELORE_PI_BRIDGE_URL as string, {
|
|
method: 'POST',
|
|
headers: {
|
|
authorization: `Bearer ${waitingWorker.env.MAKELORE_PI_WORKER_TOKEN}`,
|
|
'content-type': 'application/json',
|
|
},
|
|
body: JSON.stringify({
|
|
action: 'lease.acquire', conversationId: 'conversation-a2', workerGeneration: 1,
|
|
runId: 'run-a2', resourceId: 'write-2',
|
|
}),
|
|
}).then((response) => {
|
|
waiterSettled = true;
|
|
return response;
|
|
});
|
|
await Promise.resolve();
|
|
expect(waiterSettled).toBe(false);
|
|
await handlers.get('tool_result')?.({ toolCallId: 'write-1' });
|
|
expect((await waiting).status).toBe(200);
|
|
} finally {
|
|
if (previousEnvironment.bridge === undefined) delete process.env.MAKELORE_PI_BRIDGE_URL;
|
|
else process.env.MAKELORE_PI_BRIDGE_URL = previousEnvironment.bridge;
|
|
if (previousEnvironment.token === undefined) delete process.env.MAKELORE_PI_WORKER_TOKEN;
|
|
else process.env.MAKELORE_PI_WORKER_TOKEN = previousEnvironment.token;
|
|
if (previousEnvironment.context === undefined) delete process.env.MAKELORE_PI_CONTEXT_FILE;
|
|
else process.env.MAKELORE_PI_CONTEXT_FILE = previousEnvironment.context;
|
|
}
|
|
});
|
|
});
|