feat(coding-runtime): add managed Pi extension host

This commit is contained in:
2026-08-23 09:39:12 +08:00
parent 47159b3cbd
commit 3861c3286a
22 changed files with 1673 additions and 16 deletions

View File

@@ -77,6 +77,10 @@ class RuntimeFakeWorker implements PiConversationWorker {
};
}
async send(command: PiRpcCommand): Promise<void> {
this.requests.push(structuredClone(command));
}
setSessionData(input: { state?: unknown; entries?: unknown; stats?: unknown }): void {
if (input.state !== undefined) this.stateData = structuredClone(input.state);
if (input.entries !== undefined) this.entriesData = structuredClone(input.entries);
@@ -240,6 +244,25 @@ describe('Pi Conversation runtime', () => {
status: 'complete',
}));
expect(JSON.stringify(streamed.nodes)).toContain('Implemented');
workers.get(left.id)!.emit({
type: 'extension_ui_request',
id: 'question-runtime',
method: 'select',
title: 'Choose implementation',
options: ['Small seam', 'Large seam'],
});
await expect.poll(async () => (await runtime.getSnapshot(left.id)).pendingInteractions)
.toContainEqual(expect.objectContaining({ id: 'question-runtime', status: 'pending' }));
await runtime.respondInteraction(left.id, {
interactionId: 'question-runtime',
optionId: 'question-runtime:option:0',
});
expect(workers.get(left.id)!.requests.at(-1)).toEqual({
type: 'extension_ui_response', id: 'question-runtime', value: 'Small seam',
});
expect((await runtime.getSnapshot(left.id)).pendingInteractions).toContainEqual(
expect.objectContaining({ id: 'question-runtime', status: 'answered' }),
);
const durable = {
state: {
sessionId: `session-${left.id}`,

View File

@@ -0,0 +1,95 @@
// @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;
}
});
});

View File

@@ -0,0 +1,120 @@
// @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';
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('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);
});
});

View File

@@ -0,0 +1,44 @@
// @vitest-environment node
import { describe, expect, it } from 'vitest';
import { PiExtensionUiProjector } from '../../electron/coding-runtime/pi/extension-ui-projector';
describe('Pi extension UI projector', () => {
it('projects registered UI safely and refuses stale editor revisions', () => {
let draftRevision = 4;
const projector = new PiExtensionUiProjector({
getDraftRevision: () => draftRevision,
knownWidgetKeys: ['makelore.runtime'],
});
projector.beginRun('conversation-a', 1, 'run-1');
expect(projector.project('conversation-a', 1, 'run-1', {
type: 'extension_ui_request', id: 'status', method: 'setStatus',
statusKey: 'makelore.write-lease', statusText: '等待项目写入',
})).toMatchObject({ kind: 'status', key: 'makelore.write-lease' });
expect(projector.project('conversation-a', 1, 'run-1', {
type: 'extension_ui_request', id: 'draft', method: 'set_editor_text', text: 'new text',
})).toEqual({
kind: 'editor-text', conversationId: 'conversation-a', text: 'new text', draftRevision: 4,
});
draftRevision = 5;
expect(projector.project('conversation-a', 1, 'run-1', {
type: 'extension_ui_request', id: 'stale-draft', method: 'set_editor_text', text: 'overwrite',
})).toBeNull();
expect(projector.getDiagnostics()).toContainEqual({
method: 'set_editor_text', reason: 'stale-draft-revision',
});
});
it('records bounded diagnostics instead of projecting unknown widget payloads', () => {
const projector = new PiExtensionUiProjector({ getDraftRevision: () => 0 });
projector.beginRun('conversation-a', 1, 'run-1');
expect(projector.project('conversation-a', 1, 'run-1', {
type: 'extension_ui_request', id: 'unknown', method: 'setWidget',
widgetKey: 'third-party-widget', widgetLines: ['raw detail'],
})).toBeNull();
expect(projector.getDiagnostics()).toEqual([{
method: 'setWidget', reason: 'unsupported-ui-method',
}]);
});
});

View File

@@ -0,0 +1,78 @@
// @vitest-environment node
import { describe, expect, it } from 'vitest';
import type { ConversationInteraction } from '../../electron/coding-runtime/contracts';
import { PiInteractionStore } from '../../electron/coding-runtime/pi/interaction';
import type { PiRpcCommand } from '../../electron/coding-runtime/pi/rpc-client';
import type { PiGenerationResourceInput } from '../../electron/coding-runtime/pi/worker-pool';
describe('Pi interaction store', () => {
it('responds by exact id and option, then rejects stale generation responses', async () => {
const sent: PiRpcCommand[] = [];
const changes: ConversationInteraction[] = [];
const resources = new Map<string, () => void>();
let generation = 1;
let runId = 'run-1';
const store = new PiInteractionStore({
getState: () => ({
conversationId: 'conversation-a', workerId: 'worker-a', state: 'running', generation,
session: { piSessionId: 'session-a', sessionKey: 'key-a' },
}),
getActiveRun: () => ({ generation, runId }),
send: async (_conversationId, command) => { sent.push(command); },
trackGenerationResource: (input: PiGenerationResourceInput) => {
resources.set(input.id, input.cancel);
return () => resources.delete(input.id);
},
}, (interaction) => changes.push(interaction));
const opened = store.open('conversation-a', 1, 'run-1', {
type: 'extension_ui_request', id: 'question-1', method: 'select', title: 'Choose',
options: ['Alpha', 'Beta'],
});
expect(opened?.status).toBe('pending');
await store.respond('conversation-a', {
interactionId: 'question-1', optionId: 'question-1:option:1',
});
expect(sent).toEqual([{
type: 'extension_ui_response', id: 'question-1', value: 'Beta',
}]);
expect(changes.at(-1)?.status).toBe('answered');
await expect(store.respond('conversation-a', {
interactionId: 'question-1', optionId: 'question-1:option:0',
})).rejects.toThrow('not pending');
store.open('conversation-a', 1, 'run-1', {
type: 'extension_ui_request', id: 'question-2', method: 'confirm', title: 'Continue?', message: 'Proceed',
});
generation = 2;
runId = 'run-2';
await expect(store.respond('conversation-a', {
interactionId: 'question-2', confirmed: true,
})).rejects.toThrow('stale');
expect(changes.at(-1)?.status).toBe('cancelled');
});
it('cancels every pending dialog on abort', async () => {
const sent: PiRpcCommand[] = [];
const changes: ConversationInteraction[] = [];
const store = new PiInteractionStore({
getState: () => ({
conversationId: 'conversation-a', workerId: 'worker-a', state: 'running', generation: 1,
session: { piSessionId: 'session-a', sessionKey: 'key-a' },
}),
getActiveRun: () => ({ generation: 1, runId: 'run-1' }),
send: async (_conversationId, command) => { sent.push(command); },
trackGenerationResource: () => () => undefined,
}, (interaction) => changes.push(interaction));
for (const id of ['question-1', 'question-2']) {
store.open('conversation-a', 1, 'run-1', {
type: 'extension_ui_request', id, method: 'input', title: id,
});
}
await store.cancelRun('conversation-a', 'run-1', true);
expect(sent).toHaveLength(2);
expect(changes.map(({ status }) => status)).toEqual(['cancelled', 'cancelled']);
expect(store.list()).toEqual([]);
});
});

View File

@@ -26,6 +26,7 @@ import type {
} from '../../electron/coding-runtime/pi/rpc-client';
import type { PiWorkerProcessOptions } from '../../electron/coding-runtime/pi/worker-process';
import type { PiRuntimeTelemetryEvent } from '../../electron/coding-runtime/pi/telemetry';
import { PiManagedExtensionHost } from '../../electron/coding-runtime/pi/extension-host';
const roots: string[] = [];
const NOW = '2026-08-22T16:00:00.000Z';
@@ -54,6 +55,8 @@ class OpenerFakeProcess implements PiWorkerProcessAdapter {
};
}
async send(_command: PiRpcCommand): Promise<void> {}
subscribe(_listener: (event: PiRpcEvent) => void): () => void { return () => undefined; }
subscribeInvalidation(_listener: (error: PiProcessError) => void): () => void { return () => undefined; }
async stop() { return { mode: 'stdin-close' as const, code: 0, signal: null }; }
@@ -118,12 +121,14 @@ describe('managed Pi worker opener', () => {
const processOptions: PiWorkerProcessOptions[] = [];
const telemetry: PiRuntimeTelemetryEvent[] = [];
const registry = new PiSessionRegistry({ projectStore });
const extensionHost = new PiManagedExtensionHost();
const opener = createPiManagedWorkerOpener({
registry,
executablePath: 'electron.exe',
cliPath: 'pi-cli.js',
userDataDir,
bundledSkillsDir: path.resolve('resources/coding-skills'),
extensionHost,
loadProviderInput: async () => ({ accounts: [account], modelSummaries: [] }),
resolveCredential: async () => 'provider-secret-value',
createSessionKey: () => 'session-key-a',
@@ -155,10 +160,15 @@ describe('managed Pi worker opener', () => {
expect(argv).toContain('--system-prompt');
expect(argv).toContain('grilling');
expect(argv).toContain('--session-id');
expect(argv).toContain('--extension');
expect(argv).toContain('makelore-runtime-v1.mjs');
expect(options.additionalArgs?.filter((argument) => argument === '--extension')).toHaveLength(1);
expect(argv).not.toContain('PRIVATE MANAGED PROMPT');
expect(argv).not.toContain('provider-secret-value');
expect(Object.values(options.env ?? {})).toContain('provider-secret-value');
expect(options.sensitiveValues).toContain('provider-secret-value');
expect(options.env?.MAKELORE_PI_BRIDGE_URL).toMatch(/^http:\/\/127\.0\.0\.1:/);
expect(options.env?.MAKELORE_PI_CONTEXT_FILE).toContain('worker-');
}
const modelsFile = path.join(userDataDir, 'coding-runtime', 'pi', 'config', 'models.json');
expect(await readFile(modelsFile, 'utf8')).not.toContain('provider-secret-value');
@@ -169,5 +179,8 @@ describe('managed Pi worker opener', () => {
expect(JSON.stringify(telemetry)).not.toContain(created.id);
expect(JSON.stringify(telemetry)).not.toContain('PRIVATE MANAGED PROMPT');
expect(JSON.stringify(telemetry)).not.toContain('provider-secret-value');
await first.worker.stop();
await reopened.worker.stop();
await extensionHost.close();
});
});

View File

@@ -166,7 +166,8 @@ describe('Pi worker process', () => {
'--no-themes',
'--no-context-files',
'--no-approve',
'--no-tools',
'--tools',
'read,bash,edit,write,grep,find,ls,ask_user',
'--model', 'model-a',
]);
});

View File

@@ -53,6 +53,8 @@ class AuthFailureWorker implements PiConversationWorker {
};
}
async send(_command: PiRpcCommand): Promise<void> {}
subscribe(_listener: (event: PiRpcEvent) => void): () => void { return () => undefined; }
subscribeInvalidation(_listener: (error: PiProcessError) => void): () => void { return () => undefined; }
async stop() { return { mode: 'stdin-close' as const, code: 0, signal: null }; }

View File

@@ -48,6 +48,10 @@ class ProcessBackedWorker implements PiConversationWorker {
return this.process.request(command, options);
}
send(command: PiRpcCommand): Promise<void> {
return this.process.send(command);
}
subscribe(listener: (event: PiRpcEvent) => void): () => void {
return this.process.subscribe(listener);
}

View File

@@ -52,6 +52,10 @@ class FakeWorker implements PiConversationWorker {
return { type: 'response' as const, id: 'fake', success: true };
}
async send(command: PiRpcCommand): Promise<void> {
this.requests.push(command);
}
subscribe(listener: (event: PiRpcEvent) => void): () => void {
this.eventListeners.add(listener);
return () => this.eventListeners.delete(listener);

View File

@@ -7,6 +7,7 @@ import { createRequire } from 'node:module';
import { tmpdir } from 'node:os';
import { join, resolve } from 'node:path';
import { PiWorkerProcess } from '../../electron/coding-runtime/pi/worker-process';
import { PiManagedExtensionHost } from '../../electron/coding-runtime/pi/extension-host';
const scratchRoots: string[] = [];
@@ -39,12 +40,23 @@ describe('locked Pi worker process smoke', () => {
const cwd = join(root, 'project');
await Promise.all([mkdir(configDir), mkdir(sessionDir), mkdir(cwd)]);
const extensionHost = new PiManagedExtensionHost();
const extension = await extensionHost.registerWorker({
conversationId: 'real-conversation',
generation: 1,
projectId: 'real-project',
extensionsDir: join(root, 'extensions'),
});
await extensionHost.bindRun('real-conversation', 1, 'real-run');
const worker = await new PiWorkerProcess({
executablePath: electronExecutable,
cliPath: join(packageRoot, 'dist', 'cli.js'),
cwd,
configDir,
sessionDir,
additionalArgs: ['--extension', extension.extensionPath],
env: extension.env,
sensitiveValues: extension.sensitiveValues,
commandTimeoutMs: 5_000,
}).start();
try {
@@ -53,9 +65,12 @@ describe('locked Pi worker process smoke', () => {
command: 'get_state',
success: true,
});
expect(worker.stderrDiagnostic).not.toContain('Failed to load extension');
await expect(worker.stop()).resolves.toMatchObject({ mode: 'stdin-close', code: 0 });
} finally {
await worker.stop().catch(() => undefined);
await extension.dispose();
await extensionHost.close();
}
}, 15_000);
});

View File

@@ -0,0 +1,42 @@
// @vitest-environment node
import { describe, expect, it } from 'vitest';
import { PiProjectWriteLeaseCoordinator } from '../../electron/coding-runtime/pi/write-lease';
describe('Pi project write lease', () => {
it('serializes mutations in one project while allowing other projects to proceed', async () => {
const coordinator = new PiProjectWriteLeaseCoordinator();
const first = await coordinator.acquire('project-a', 'write-a1');
let secondSettled = false;
const secondFlight = coordinator.acquire('project-a', 'write-a2').then((lease) => {
secondSettled = true;
return lease;
});
const otherProject = await coordinator.acquire('project-b', 'write-b1');
await Promise.resolve();
expect(secondSettled).toBe(false);
expect(coordinator.activeCount).toBe(2);
expect(coordinator.waitingCount('project-a')).toBe(1);
first.release();
const second = await secondFlight;
expect(second.holderId).toBe('write-a2');
second.release();
otherProject.release();
expect(coordinator.activeCount).toBe(0);
});
it('removes a cancelled waiter without disturbing the active lease', async () => {
const coordinator = new PiProjectWriteLeaseCoordinator();
const active = await coordinator.acquire('project-a', 'active');
const controller = new AbortController();
const waiting = coordinator.acquire('project-a', 'waiting', controller.signal);
controller.abort();
await expect(waiting).rejects.toThrow('cancelled');
expect(coordinator.activeCount).toBe(1);
expect(coordinator.waitingCount()).toBe(0);
active.release();
});
});